inspector: reflect storage size used by entrypoints

BUG=tint:919

Change-Id: I80d8ff7a9a56464cd11710e94a3d54c9f2933066
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/56800
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Ken Rockot <rockot@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>

src/inspector/inspector.cc

@@ -14,6 +14,7 @@
 
 #include "src/inspector/inspector.h"
 
+#include <limits>
 #include <utility>
 
 #include "src/ast/bool_literal.h"
@@ -316,6 +317,33 @@ std::map<std::string, uint32_t> Inspector::GetConstantNameToIdMap() {
   return result;
 }
 
+uint32_t Inspector::GetStorageSize(const std::string& entry_point) {
+  auto* func = FindEntryPointByName(entry_point);
+  if (!func) {
+    return 0;
+  }
+
+  size_t size = 0;
+  auto* func_sem = program_->Sem().Get(func);
+  for (auto& ruv : func_sem->ReferencedUniformVariables()) {
+    const sem::Struct* s = ruv.first->Type()->UnwrapRef()->As<sem::Struct>();
+    if (s && s->IsBlockDecorated()) {
+      size += s->Size();
+    }
+  }
+  for (auto& rsv : func_sem->ReferencedStorageBufferVariables()) {
+    const sem::Struct* s = rsv.first->Type()->UnwrapRef()->As<sem::Struct>();
+    if (s) {
+      size += s->Size();
+    }
+  }
+
+  if (size > std::numeric_limits<uint32_t>::max()) {
+    return std::numeric_limits<uint32_t>::max();
+  }
+  return static_cast<uint32_t>(size);
+}
+
 std::vector<ResourceBinding> Inspector::GetResourceBindings(
     const std::string& entry_point) {
   auto* func = FindEntryPointByName(entry_point);

src/inspector/inspector.h

@@ -61,6 +61,11 @@ class Inspector {
   /// @returns map of module-constant name to pipeline constant ID
   std::map<std::string, uint32_t> GetConstantNameToIdMap();
 
+  /// @param entry_point name of the entry point to get information about.
+  /// @returns the total size of shared storage required by an entry point,
+  //           including all uniforms storage buffers.
+  uint32_t GetStorageSize(const std::string& entry_point);
+
   /// @param entry_point name of the entry point to get information about.
   /// @returns vector of all of the resource bindings.
   std::vector<ResourceBinding> GetResourceBindings(

src/inspector/inspector_test.cc

@@ -726,6 +726,8 @@ class InspectorGetConstantIDsTest : public InspectorHelper,
                                     public testing::Test {};
 class InspectorGetConstantNameToIdMapTest : public InspectorHelper,
                                             public testing::Test {};
+class InspectorGetStorageSizeTest : public InspectorHelper,
+                                    public testing::Test {};
 class InspectorGetResourceBindingsTest : public InspectorHelper,
                                          public testing::Test {};
 class InspectorGetUniformBufferResourceBindingsTest : public InspectorHelper,
@@ -1932,6 +1934,40 @@ TEST_F(InspectorGetConstantNameToIdMapTest, WithAndWithoutIds) {
   EXPECT_EQ(result["c"], program_->Sem().Get(c)->ConstantId());
 }
 
+TEST_F(InspectorGetStorageSizeTest, Empty) {
+  MakeEmptyBodyFunction("ep_func",
+                        ast::DecorationList{Stage(ast::PipelineStage::kCompute),
+                                            WorkgroupSize(1)});
+  Inspector& inspector = Build();
+  EXPECT_EQ(0u, inspector.GetStorageSize("ep_func"));
+}
+
+TEST_F(InspectorGetStorageSizeTest, Simple) {
+  ast::Struct* ub_struct_type =
+      MakeUniformBufferType("ub_type", {ty.i32(), ty.i32()});
+  AddUniformBuffer("ub_var", ty.Of(ub_struct_type), 0, 0);
+  MakeStructVariableReferenceBodyFunction("ub_func", "ub_var", {{0, ty.i32()}});
+
+  auto sb = MakeStorageBufferTypes("sb_type", {ty.i32()});
+  AddStorageBuffer("sb_var", sb(), ast::Access::kReadWrite, 1, 0);
+  MakeStructVariableReferenceBodyFunction("sb_func", "sb_var", {{0, ty.i32()}});
+
+  auto ro_sb = MakeStorageBufferTypes("rosb_type", {ty.i32()});
+  AddStorageBuffer("rosb_var", ro_sb(), ast::Access::kRead, 1, 1);
+  MakeStructVariableReferenceBodyFunction("rosb_func", "rosb_var",
+                                          {{0, ty.i32()}});
+
+  MakeCallerBodyFunction("ep_func", {"ub_func", "sb_func", "rosb_func"},
+                         ast::DecorationList{
+                             Stage(ast::PipelineStage::kCompute),
+                             WorkgroupSize(1),
+                         });
+
+  Inspector& inspector = Build();
+
+  EXPECT_EQ(16u, inspector.GetStorageSize("ep_func"));
+}
+
 TEST_F(InspectorGetResourceBindingsTest, Empty) {
   MakeCallerBodyFunction("ep_func", {},
                          ast::DecorationList{