writer/spirv: Support declaring workgroup variables with 0 initializer

This patch adds an option to declare the workgroup variables with zero
initializer in Build() instead of transform::ZeroInitWorkgroupMemory
in Sanitize(). This option will be enabled when the Vulkan extension
VK_KHR_zero_initialize_workgroup_memory is enabled on the API side.

BUG=dawn:1302

Change-Id: Ia580df98ec161ec6f2d3099a01dbedb8bf848bf2
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/82580
Kokoro: Kokoro <noreply+kokoro@google.com>
Reviewed-by: Ben Clayton <bclayton@chromium.org>
Commit-Queue: Ben Clayton <bclayton@chromium.org>

src/tint/writer/spirv/builder.cc

@@ -300,8 +300,10 @@ Builder::AccessorInfo::AccessorInfo() : source_id(0), source_type(nullptr) {}
 
 Builder::AccessorInfo::~AccessorInfo() {}
 
-Builder::Builder(const Program* program)
-    : builder_(ProgramBuilder::Wrap(program)), scope_stack_({}) {}
+Builder::Builder(const Program* program, bool zero_initialize_workgroup_memory)
+    : builder_(ProgramBuilder::Wrap(program)),
+      scope_stack_({}),
+      zero_initialize_workgroup_memory_(zero_initialize_workgroup_memory) {}
 
 Builder::~Builder() = default;
 
@@ -861,8 +863,13 @@ bool Builder::GenerateGlobalVariable(const ast::Variable* var) {
     if (!type->Is<sem::Sampler>()) {
       // If we don't have a constructor and we're an Output or Private
       // variable, then WGSL requires that we zero-initialize.
+      // If we're a Workgroup variable, and the
+      // VK_KHR_zero_initialize_workgroup_memory extension is enabled, we should
+      // also zero-initialize.
       if (sem->StorageClass() == ast::StorageClass::kPrivate ||
-          sem->StorageClass() == ast::StorageClass::kOutput) {
+          sem->StorageClass() == ast::StorageClass::kOutput ||
+          (zero_initialize_workgroup_memory_ &&
+           sem->StorageClass() == ast::StorageClass::kWorkgroup)) {
         init_id = GenerateConstantNullIfNeeded(type);
         if (init_id == 0) {
           return 0;

src/tint/writer/spirv/builder.h

@@ -91,7 +91,10 @@ class Builder {
 
   /// Constructor
   /// @param program the program
-  explicit Builder(const Program* program);
+  /// @param zero_initialize_workgroup_memory `true` to initialize all the
+  /// variables in the Workgroup storage class with OpConstantNull
+  Builder(const Program* program,
+          bool zero_initialize_workgroup_memory = false);
   ~Builder();
 
   /// Generates the SPIR-V instructions for the given program
@@ -624,6 +627,7 @@ class Builder {
   std::vector<uint32_t> continue_stack_;
   std::unordered_set<uint32_t> capability_set_;
   bool has_overridable_workgroup_size_ = false;
+  bool zero_initialize_workgroup_memory_ = false;
 
   struct ContinuingInfo {
     ContinuingInfo(const ast::Statement* last_statement,

src/tint/writer/spirv/builder_global_variable_test.cc

@@ -622,6 +622,51 @@ OpDecorate %5 DescriptorSet 0
 )");
 }
 
+TEST_F(BuilderTest, GlobalVar_WorkgroupWithZeroInit) {
+  auto* type_scalar = ty.i32();
+  auto* var_scalar = Global("a", type_scalar, ast::StorageClass::kWorkgroup);
+
+  auto* type_array = ty.array<f32, 16>();
+  auto* var_array = Global("b", type_array, ast::StorageClass::kWorkgroup);
+
+  auto* type_struct = Structure("C",
+                                {
+                                    Member("a", ty.i32()),
+                                    Member("b", ty.i32()),
+                                },
+                                {create<ast::StructBlockAttribute>()});
+  auto* var_struct =
+      Global("c", ty.Of(type_struct), ast::StorageClass::kWorkgroup);
+
+  program = std::make_unique<Program>(std::move(*this));
+
+  constexpr bool kZeroInitializeWorkgroupMemory = true;
+  std::unique_ptr<spirv::Builder> b = std::make_unique<spirv::Builder>(
+      program.get(), kZeroInitializeWorkgroupMemory);
+
+  EXPECT_TRUE(b->GenerateGlobalVariable(var_scalar)) << b->error();
+  EXPECT_TRUE(b->GenerateGlobalVariable(var_array)) << b->error();
+  EXPECT_TRUE(b->GenerateGlobalVariable(var_struct)) << b->error();
+  ASSERT_FALSE(b->has_error()) << b->error();
+
+  EXPECT_EQ(DumpInstructions(b->types()), R"(%3 = OpTypeInt 32 1
+%2 = OpTypePointer Workgroup %3
+%4 = OpConstantNull %3
+%1 = OpVariable %2 Workgroup %4
+%8 = OpTypeFloat 32
+%9 = OpTypeInt 32 0
+%10 = OpConstant %9 16
+%7 = OpTypeArray %8 %10
+%6 = OpTypePointer Workgroup %7
+%11 = OpConstantNull %7
+%5 = OpVariable %6 Workgroup %11
+%14 = OpTypeStruct %3 %3
+%13 = OpTypePointer Workgroup %14
+%15 = OpConstantNull %14
+%12 = OpVariable %13 Workgroup %15
+)");
+}
+
 }  // namespace
 }  // namespace spirv
 }  // namespace writer

src/tint/writer/spirv/generator.cc

@@ -28,8 +28,11 @@ Result Generate(const Program* program, const Options& options) {
   Result result;
 
   // Sanitize the program.
+  bool disable_workgroup_init_in_sanitizer =
+      options.disable_workgroup_init ||
+      options.use_zero_initialize_workgroup_memory_extension;
   auto sanitized_result = Sanitize(program, options.emit_vertex_point_size,
-                                   options.disable_workgroup_init);
+                                   disable_workgroup_init_in_sanitizer);
   if (!sanitized_result.program.IsValid()) {
     result.success = false;
     result.error = sanitized_result.program.Diagnostics().str();
@@ -37,7 +40,11 @@ Result Generate(const Program* program, const Options& options) {
   }
 
   // Generate the SPIR-V code.
-  auto builder = std::make_unique<Builder>(&sanitized_result.program);
+  bool zero_initialize_workgroup_memory =
+      !options.disable_workgroup_init &&
+      options.use_zero_initialize_workgroup_memory_extension;
+  auto builder = std::make_unique<Builder>(&sanitized_result.program,
+                                           zero_initialize_workgroup_memory);
   auto writer = std::make_unique<BinaryWriter>();
   if (!builder->Build()) {
     result.success = false;

src/tint/writer/spirv/generator.h

@@ -41,6 +41,10 @@ struct Options {
 
   /// Set to `true` to disable workgroup memory zero initialization
   bool disable_workgroup_init = false;
+
+  /// Set to `true` to initialize workgroup memory with OpConstantNull when
+  /// VK_KHR_zero_initialize_workgroup_memory is enabled.
+  bool use_zero_initialize_workgroup_memory_extension = false;
 };
 
 /// The result produced when generating SPIR-V.