[inspector] Implement GetImmediateBlockInfo

This CL implements a helper function GetImmediateBlockInfo in the
inspector. It returns a Bitset<16> indicating which 4-byte slots
of the immediate block are accessible (used) by the entry point.
This handles padding in structs and matrices correctly.

Bug:366291600
Change-Id: I5b9576e2b31b7597634e71bf8584afbed7e58ac8
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/281495
Reviewed-by: James Price <jrprice@google.com>
Commit-Queue: Shaobo Yan <shaoboyan@microsoft.com>
diff --git a/src/tint/lang/wgsl/inspector/inspector.cc b/src/tint/lang/wgsl/inspector/inspector.cc
index 654bbf7..137eacb 100644
--- a/src/tint/lang/wgsl/inspector/inspector.cc
+++ b/src/tint/lang/wgsl/inspector/inspector.cc
@@ -27,6 +27,7 @@
 
 #include "src/tint/lang/wgsl/inspector/inspector.h"
 
+#include <functional>
 #include <unordered_set>
 #include <utility>
 
@@ -42,6 +43,7 @@
 #include "src/tint/lang/core/type/i32.h"
 #include "src/tint/lang/core/type/input_attachment.h"
 #include "src/tint/lang/core/type/matrix.h"
+#include "src/tint/lang/core/type/memory_view.h"
 #include "src/tint/lang/core/type/multisampled_texture.h"
 #include "src/tint/lang/core/type/resource_type.h"
 #include "src/tint/lang/core/type/sampled_texture.h"
@@ -724,6 +726,63 @@
     return metadata;
 }
 
+std::bitset<kImmediateSlotCount> Inspector::GetImmediateBlockInfo(const std::string& entry_point) {
+    auto* func = FindEntryPointByName(entry_point);
+    if (!func) {
+        return {};
+    }
+
+    auto* func_sem = program_.Sem().Get(func);
+
+    const sem::GlobalVariable* immediate_var = nullptr;
+    for (const sem::Variable* var : func_sem->TransitivelyReferencedGlobals()) {
+        if (var->AddressSpace() == core::AddressSpace::kImmediate) {
+            immediate_var = var->As<sem::GlobalVariable>();
+            break;
+        }
+    }
+
+    if (!immediate_var) {
+        return {};
+    }
+
+    auto* mv = immediate_var->Type()->As<core::type::MemoryView>();
+    auto* type = mv->StoreType();
+
+    std::bitset<kImmediateSlotCount> accessible_slots;
+
+    std::function<void(const core::type::Type*, uint32_t)> mark_slots =
+        [&](const core::type::Type* t, uint32_t offset) {
+            tint::Switch(
+                t,
+                [&](const sem::Struct* str) {
+                    for (auto* member : str->Members()) {
+                        mark_slots(member->Type(), offset + member->Offset());
+                    }
+                },
+                [&](const core::type::Matrix* mat) {
+                    uint32_t col_stride = mat->ColumnStride();
+                    for (uint32_t i = 0; i < mat->Columns(); i++) {
+                        mark_slots(mat->ColumnType(), offset + i * col_stride);
+                    }
+                },
+                [&](const core::type::Type* other) {
+                    uint32_t s = other->Size();
+                    uint32_t start_slot = offset / kImmediateSlotSize;
+                    uint32_t end_byte = offset + s;
+                    uint32_t end_slot = (end_byte - 1) / kImmediateSlotSize;
+                    for (uint32_t i = start_slot; i <= end_slot; i++) {
+                        TINT_ASSERT(i < accessible_slots.size());
+                        accessible_slots[i] = true;
+                    }
+                });
+        };
+
+    mark_slots(type, 0);
+
+    return accessible_slots;
+}
+
 std::vector<std::string> Inspector::GetUsedExtensionNames() {
     auto& extensions = program_.Sem().Module()->Extensions();
     std::vector<std::string> out;
diff --git a/src/tint/lang/wgsl/inspector/inspector.h b/src/tint/lang/wgsl/inspector/inspector.h
index 8a3a3f9..af8fcf0 100644
--- a/src/tint/lang/wgsl/inspector/inspector.h
+++ b/src/tint/lang/wgsl/inspector/inspector.h
@@ -28,6 +28,8 @@
 #ifndef SRC_TINT_LANG_WGSL_INSPECTOR_INSPECTOR_H_
 #define SRC_TINT_LANG_WGSL_INSPECTOR_INSPECTOR_H_
 
+#include <array>
+#include <bitset>
 #include <map>
 #include <string>
 #include <tuple>
@@ -49,6 +51,13 @@
 /// A temporary alias to sem::SamplerTexturePair. [DEPRECATED]
 using SamplerTexturePair = sem::SamplerTexturePair;
 
+/// The maximum size of the immediate buffer in bytes.
+constexpr size_t kMaxImmediateSize = 64;
+/// The size of a single slot in the immediate buffer in bytes.
+constexpr size_t kImmediateSlotSize = 4;
+/// The number of slots in the immediate buffer.
+constexpr size_t kImmediateSlotCount = kMaxImmediateSize / kImmediateSlotSize;
+
 /// Extracts information from a program
 class Inspector {
   public:
@@ -104,6 +113,12 @@
         const std::string& entry_point,
         const BindingPoint& non_sampler_placeholder);
 
+    /// @param entry_point name of the entry point to get information about
+    /// @returns a bitset where the nth bit is true if the nth 4-byte slot of the immediate block
+    /// may be used by the entry point. Returns an empty bitset if the immediate block is not
+    /// referenced.
+    std::bitset<kImmediateSlotCount> GetImmediateBlockInfo(const std::string& entry_point);
+
     /// @returns vector of all valid extension names used by the program. There
     /// will be no duplicated names in the returned vector even if an extension
     /// is enabled multiple times.
diff --git a/src/tint/lang/wgsl/inspector/inspector_test.cc b/src/tint/lang/wgsl/inspector/inspector_test.cc
index 957406c..1843a2d 100644
--- a/src/tint/lang/wgsl/inspector/inspector_test.cc
+++ b/src/tint/lang/wgsl/inspector/inspector_test.cc
@@ -4355,4 +4355,205 @@
     return out;
 }
 
+using InspectorImmediateBlockTest = InspectorTest;
+
+TEST_F(InspectorImmediateBlockTest, NoImmediateVariable) {
+    auto* src = R"(
+@compute @workgroup_size(1)
+fn main() {}
+)";
+    Inspector& inspector = Initialize(src);
+    auto info = inspector.GetImmediateBlockInfo("main");
+    EXPECT_TRUE(info.none());
+}
+
+TEST_F(InspectorImmediateBlockTest, UnusedImmediateVariable) {
+    auto* src = R"(
+requires immediate_address_space;
+var<immediate> a : u32;
+@compute @workgroup_size(1)
+fn main() {}
+)";
+    Inspector& inspector = Initialize(src);
+    auto info = inspector.GetImmediateBlockInfo("main");
+    EXPECT_TRUE(info.none());
+}
+
+TEST_F(InspectorImmediateBlockTest, SimpleScalar) {
+    auto* src = R"(
+requires immediate_address_space;
+var<immediate> a : u32;
+@compute @workgroup_size(1)
+fn main() {
+  _ = a;
+}
+)";
+    Inspector& inspector = Initialize(src);
+    auto info = inspector.GetImmediateBlockInfo("main");
+    EXPECT_TRUE(info[0]);
+    for (size_t i = 1; i < info.size(); ++i) {
+        EXPECT_FALSE(info[i]);
+    }
+}
+
+TEST_F(InspectorImmediateBlockTest, SimpleVector) {
+    auto* src = R"(
+requires immediate_address_space;
+var<immediate> a : vec4<f32>;
+@compute @workgroup_size(1)
+fn main() {
+  _ = a;
+}
+)";
+    Inspector& inspector = Initialize(src);
+    auto info = inspector.GetImmediateBlockInfo("main");
+    for (size_t i = 0; i < 4; ++i) {
+        EXPECT_TRUE(info[i]);
+    }
+    for (size_t i = 4; i < info.size(); ++i) {
+        EXPECT_FALSE(info[i]);
+    }
+}
+
+TEST_F(InspectorImmediateBlockTest, MatrixWithPadding) {
+    auto* src = R"(
+requires immediate_address_space;
+var<immediate> a : mat3x3<f32>;
+@compute @workgroup_size(1)
+fn main() {
+  _ = a;
+}
+)";
+    Inspector& inspector = Initialize(src);
+    auto info = inspector.GetImmediateBlockInfo("main");
+    // mat3x3<f32> has 3 columns of vec3<f32>.
+    // vec3<f32> is 12 bytes, aligned to 16.
+    // Size is 48 bytes.
+
+    // Col 0: slots 0, 1, 2. Slot 3 is padding.
+    EXPECT_TRUE(info[0]);
+    EXPECT_TRUE(info[1]);
+    EXPECT_TRUE(info[2]);
+    EXPECT_FALSE(info[3]);
+
+    // Col 1: slots 4, 5, 6. Slot 7 is padding.
+    EXPECT_TRUE(info[4]);
+    EXPECT_TRUE(info[5]);
+    EXPECT_TRUE(info[6]);
+    EXPECT_FALSE(info[7]);
+
+    // Col 2: slots 8, 9, 10. Slot 11 is padding.
+    EXPECT_TRUE(info[8]);
+    EXPECT_TRUE(info[9]);
+    EXPECT_TRUE(info[10]);
+    EXPECT_FALSE(info[11]);
+
+    for (size_t i = 12; i < info.size(); ++i) {
+        EXPECT_FALSE(info[i]);
+    }
+}
+
+TEST_F(InspectorImmediateBlockTest, StructWithPadding) {
+    auto* src = R"(
+requires immediate_address_space;
+struct S {
+  a : u32,
+  b : vec3<f32>,
+}
+var<immediate> u : S;
+@compute @workgroup_size(1)
+fn main() {
+  _ = u;
+}
+)";
+    Inspector& inspector = Initialize(src);
+    auto info = inspector.GetImmediateBlockInfo("main");
+    // S alignment is 16 (vec3<f32>).
+    // a: offset 0, size 4.
+    // padding: 12 bytes.
+    // b: offset 16, size 12.
+    // padding: 4 bytes (to align struct size to 16).
+    // Total size: 32 bytes.
+
+    // a
+    EXPECT_TRUE(info[0]);
+    // padding
+    EXPECT_FALSE(info[1]);
+    EXPECT_FALSE(info[2]);
+    EXPECT_FALSE(info[3]);
+    // b
+    EXPECT_TRUE(info[4]);
+    EXPECT_TRUE(info[5]);
+    EXPECT_TRUE(info[6]);
+    // padding
+    EXPECT_FALSE(info[7]);
+
+    for (size_t i = 8; i < info.size(); ++i) {
+        EXPECT_FALSE(info[i]);
+    }
+}
+
+TEST_F(InspectorImmediateBlockTest, IndirectUse) {
+    auto* src = R"(
+requires immediate_address_space;
+var<immediate> a : u32;
+fn foo() {
+  _ = a;
+}
+@compute @workgroup_size(1)
+fn main() {
+  foo();
+}
+)";
+    Inspector& inspector = Initialize(src);
+    auto info = inspector.GetImmediateBlockInfo("main");
+    EXPECT_TRUE(info[0]);
+    for (size_t i = 1; i < info.size(); ++i) {
+        EXPECT_FALSE(info[i]);
+    }
+}
+
+TEST_F(InspectorImmediateBlockTest, PartialAccess) {
+    auto* src = R"(
+requires immediate_address_space;
+struct S {
+  a : u32,
+  b : u32,
+}
+var<immediate> u : S;
+@compute @workgroup_size(1)
+fn main() {
+  _ = u.b;
+}
+)";
+    Inspector& inspector = Initialize(src);
+    auto info = inspector.GetImmediateBlockInfo("main");
+    // Even with partial access, we mark all non-padding bytes.
+    EXPECT_TRUE(info[0]);
+    EXPECT_TRUE(info[1]);
+    for (size_t i = 2; i < info.size(); ++i) {
+        EXPECT_FALSE(info[i]);
+    }
+}
+
+TEST_F(InspectorImmediateBlockTest, FunctionParameter) {
+    auto* src = R"(
+requires immediate_address_space;
+var<immediate> a : u32;
+fn foo(val : u32) {
+  _ = val;
+}
+@compute @workgroup_size(1)
+fn main() {
+  foo(a);
+}
+)";
+    Inspector& inspector = Initialize(src);
+    auto info = inspector.GetImmediateBlockInfo("main");
+    EXPECT_TRUE(info[0]);
+    for (size_t i = 1; i < info.size(); ++i) {
+        EXPECT_FALSE(info[i]);
+    }
+}
+
 }  // namespace tint::inspector