writers: Add flag to disable workgroup memory init

Bug: tint:1003
Change-Id: Ia30a2c51b5d3f8c6a01bed5299eac51dc3ad6337
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/58843
Kokoro: Kokoro <noreply+kokoro@google.com>
Commit-Queue: Ben Clayton <bclayton@google.com>
Auto-Submit: Ben Clayton <bclayton@google.com>
Reviewed-by: James Price <jrprice@google.com>

src/transform/hlsl.cc

@@ -31,6 +31,7 @@
 #include "src/transform/zero_init_workgroup_memory.h"
 
 TINT_INSTANTIATE_TYPEINFO(tint::transform::Hlsl);
+TINT_INSTANTIATE_TYPEINFO(tint::transform::Hlsl::Config);
 
 namespace tint {
 namespace transform {
@@ -38,18 +39,22 @@ namespace transform {
 Hlsl::Hlsl() = default;
 Hlsl::~Hlsl() = default;
 
-Output Hlsl::Run(const Program* in, const DataMap&) {
+Output Hlsl::Run(const Program* in, const DataMap& inputs) {
   Manager manager;
   DataMap data;
 
+  auto* cfg = inputs.Get<Config>();
+
   // Attempt to convert `loop`s into for-loops. This is to try and massage the
   // output into something that will not cause FXC to choke or misbehave.
   manager.Add<FoldTrivialSingleUseLets>();
   manager.Add<LoopToForLoop>();
 
-  // ZeroInitWorkgroupMemory must come before CanonicalizeEntryPointIO as
-  // ZeroInitWorkgroupMemory may inject new builtin parameters.
-  manager.Add<ZeroInitWorkgroupMemory>();
+  if (!cfg || !cfg->disable_workgroup_init) {
+    // ZeroInitWorkgroupMemory must come before CanonicalizeEntryPointIO as
+    // ZeroInitWorkgroupMemory may inject new builtin parameters.
+    manager.Add<ZeroInitWorkgroupMemory>();
+  }
   manager.Add<CanonicalizeEntryPointIO>();
   manager.Add<InlinePointerLets>();
   // Simplify cleans up messy `*(&(expr))` expressions from InlinePointerLets.
@@ -97,5 +102,9 @@ void Hlsl::AddEmptyEntryPoint(CloneContext& ctx) const {
                  ctx.dst->WorkgroupSize(1)});
 }
 
+Hlsl::Config::Config(bool disable_wi) : disable_workgroup_init(disable_wi) {}
+Hlsl::Config::Config(const Config&) = default;
+Hlsl::Config::~Config() = default;
+
 }  // namespace transform
 }  // namespace tint

src/transform/hlsl.h

@@ -29,6 +29,23 @@ namespace transform {
 /// behavior.
 class Hlsl : public Castable<Hlsl, Transform> {
  public:
+  /// Configuration options for the Hlsl sanitizer transform.
+  struct Config : public Castable<Data, transform::Data> {
+    /// Constructor
+    /// @param disable_workgroup_init `true` to disable workgroup memory zero
+    ///        initialization
+    explicit Config(bool disable_workgroup_init = false);
+
+    /// Copy constructor
+    Config(const Config&);
+
+    /// Destructor
+    ~Config() override;
+
+    /// Set to `true` to disable workgroup memory zero initialization
+    bool disable_workgroup_init = false;
+  };
+
   /// Constructor
   Hlsl();
   ~Hlsl() override;

src/transform/msl.cc

@@ -73,9 +73,11 @@ Output Msl::Run(const Program* in, const DataMap& inputs) {
     }
   }
 
-  // ZeroInitWorkgroupMemory must come before CanonicalizeEntryPointIO as
-  // ZeroInitWorkgroupMemory may inject new builtin parameters.
-  manager.Add<ZeroInitWorkgroupMemory>();
+  if (!cfg || !cfg->disable_workgroup_init) {
+    // ZeroInitWorkgroupMemory must come before CanonicalizeEntryPointIO as
+    // ZeroInitWorkgroupMemory may inject new builtin parameters.
+    manager.Add<ZeroInitWorkgroupMemory>();
+  }
   manager.Add<CanonicalizeEntryPointIO>();
   manager.Add<ExternalTextureTransform>();
   manager.Add<PromoteInitializersToConstVar>();
@@ -280,9 +282,12 @@ void Msl::HandleModuleScopeVariables(CloneContext& ctx) const {
   }
 }
 
-Msl::Config::Config(uint32_t buffer_size_ubo_idx, uint32_t sample_mask)
+Msl::Config::Config(uint32_t buffer_size_ubo_idx,
+                    uint32_t sample_mask,
+                    bool disable_wi)
     : buffer_size_ubo_index(buffer_size_ubo_idx),
-      fixed_sample_mask(sample_mask) {}
+      fixed_sample_mask(sample_mask),
+      disable_workgroup_init(disable_wi) {}
 Msl::Config::Config(const Config&) = default;
 Msl::Config::~Config() = default;
 

src/transform/msl.h

@@ -33,8 +33,11 @@ class Msl : public Castable<Msl, Transform> {
     /// Constructor
     /// @param buffer_size_ubo_idx the index to use for the buffer size UBO
     /// @param sample_mask the fixed sample mask to use for fragment shaders
-    explicit Config(uint32_t buffer_size_ubo_idx,
-                    uint32_t sample_mask = 0xFFFFFFFF);
+    /// @param disable_workgroup_init `true` to disable workgroup memory zero
+    ///        initialization
+    Config(uint32_t buffer_size_ubo_idx,
+           uint32_t sample_mask = 0xFFFFFFFF,
+           bool disable_workgroup_init = false);
 
     /// Copy constructor
     Config(const Config&);
@@ -43,10 +46,13 @@ class Msl : public Castable<Msl, Transform> {
     ~Config() override;
 
     /// The index to use when generating a UBO to receive storage buffer sizes.
-    uint32_t buffer_size_ubo_index;
+    uint32_t buffer_size_ubo_index = 0;
 
     /// The fixed sample mask to combine with fragment shader outputs.
-    uint32_t fixed_sample_mask;
+    uint32_t fixed_sample_mask = 0xFFFFFFFF;
+
+    /// Set to `true` to disable workgroup memory zero initialization
+    bool disable_workgroup_init = false;
   };
 
   /// Information produced by the sanitizer that users may need to act on.

src/transform/spirv.cc

@@ -45,8 +45,12 @@ Spirv::Spirv() = default;
 Spirv::~Spirv() = default;
 
 Output Spirv::Run(const Program* in, const DataMap& data) {
+  auto* cfg = data.Get<Config>();
+
   Manager manager;
-  manager.Add<ZeroInitWorkgroupMemory>();
+  if (!cfg || !cfg->disable_workgroup_init) {
+    manager.Add<ZeroInitWorkgroupMemory>();
+  }
   manager.Add<InlinePointerLets>();  // Required for arrayLength()
   manager.Add<Simplify>();           // Required for arrayLength()
   manager.Add<FoldConstants>();
@@ -58,8 +62,6 @@ Output Spirv::Run(const Program* in, const DataMap& data) {
     return transformedInput;
   }
 
-  auto* cfg = data.Get<Config>();
-
   ProgramBuilder out;
   CloneContext ctx(&out, &transformedInput.program);
   HandleEntryPointIOTypes(ctx);
@@ -427,7 +429,8 @@ void Spirv::HoistToOutputVariables(CloneContext& ctx,
   }
 }
 
-Spirv::Config::Config(bool emit_vps) : emit_vertex_point_size(emit_vps) {}
+Spirv::Config::Config(bool emit_vps, bool disable_wi)
+    : emit_vertex_point_size(emit_vps), disable_workgroup_init(disable_wi) {}
 
 Spirv::Config::Config(const Config&) = default;
 Spirv::Config::~Config() = default;

src/transform/spirv.h

@@ -35,7 +35,10 @@ class Spirv : public Castable<Spirv, Transform> {
   struct Config : public Castable<Config, Data> {
     /// Constructor
     /// @param emit_vertex_point_size `true` to generate a PointSize builtin
-    explicit Config(bool emit_vertex_point_size = false);
+    /// @param disable_workgroup_init `true` to disable workgroup memory zero
+    ///        initialization
+    Config(bool emit_vertex_point_size = false,
+           bool disable_workgroup_init = false);
 
     /// Copy constructor.
     Config(const Config&);
@@ -49,7 +52,10 @@ class Spirv : public Castable<Spirv, Transform> {
 
     /// Set to `true` to generate a PointSize builtin and have it set to 1.0
     /// from all vertex shaders in the module.
-    bool emit_vertex_point_size;
+    bool emit_vertex_point_size = false;
+
+    /// Set to `true` to disable workgroup memory zero initialization
+    bool disable_workgroup_init = false;
   };
 
   /// Constructor