From a2f5e7f09c2fdc3be817144d80e778388c322b4c Mon Sep 17 00:00:00 2001
From: Dan Field <dnfield@google.com>
Date: Wed, 29 Mar 2023 15:10:55 -0700
Subject: [PATCH] [Impeller] Guard against empty grid sizes

---
 .../backend/metal/compute_pass_mtl.mm         |  2 +
 impeller/renderer/compute_pass.cc             |  5 ++
 impeller/renderer/compute_unittests.cc        | 56 +++++++++++++++++++
 3 files changed, 63 insertions(+)

diff --git a/impeller/renderer/backend/metal/compute_pass_mtl.mm b/impeller/renderer/backend/metal/compute_pass_mtl.mm
index dd5e7aee00680..204e2a2c886af 100644
--- a/impeller/renderer/backend/metal/compute_pass_mtl.mm
+++ b/impeller/renderer/backend/metal/compute_pass_mtl.mm
@@ -55,6 +55,8 @@
     return false;
   }
 
+  FML_DCHECK(!grid_size_.IsEmpty() && !thread_group_size_.IsEmpty());
+
   // TODO(dnfield): Support non-serial dispatch type on higher iOS versions.
   auto compute_command_encoder = [buffer_ computeCommandEncoder];
 
diff --git a/impeller/renderer/compute_pass.cc b/impeller/renderer/compute_pass.cc
index 579405831352f..3efda50d4c757 100644
--- a/impeller/renderer/compute_pass.cc
+++ b/impeller/renderer/compute_pass.cc
@@ -47,6 +47,11 @@ bool ComputePass::AddCommand(ComputeCommand command) {
 }
 
 bool ComputePass::EncodeCommands() const {
+  if (grid_size_.IsEmpty() || thread_group_size_.IsEmpty()) {
+    FML_DLOG(WARNING) << "Attempted to encode a compute pass with an empty "
+                         "grid or thread group size.";
+    return false;
+  }
   auto context = context_.lock();
   // The context could have been collected in the meantime.
   if (!context) {
diff --git a/impeller/renderer/compute_unittests.cc b/impeller/renderer/compute_unittests.cc
index 3411103539c04..1ad1dec214924 100644
--- a/impeller/renderer/compute_unittests.cc
+++ b/impeller/renderer/compute_unittests.cc
@@ -285,5 +285,61 @@ TEST_P(ComputeTest, CanCorrectlyDownScaleLargeGridSize) {
   latch.Wait();
 }
 
+TEST_P(ComputeTest, ReturnsEarlyWhenAnyGridDimensionIsZero) {
+  using CS = SampleComputeShader;
+  auto context = GetContext();
+  ASSERT_TRUE(context);
+  ASSERT_TRUE(context->GetCapabilities()->SupportsCompute());
+
+  using SamplePipelineBuilder = ComputePipelineBuilder<CS>;
+  auto pipeline_desc =
+      SamplePipelineBuilder::MakeDefaultPipelineDescriptor(*context);
+  ASSERT_TRUE(pipeline_desc.has_value());
+  auto compute_pipeline =
+      context->GetPipelineLibrary()->GetPipeline(pipeline_desc).Get();
+  ASSERT_TRUE(compute_pipeline);
+
+  auto cmd_buffer = context->CreateCommandBuffer();
+  auto pass = cmd_buffer->CreateComputePass();
+  ASSERT_TRUE(pass && pass->IsValid());
+
+  static constexpr size_t kCount = 5;
+
+  // Intentionally making the grid size obscenely large. No GPU will tolerate
+  // this.
+  pass->SetGridSize(ISize(0, 1));
+  pass->SetThreadGroupSize(ISize(0, 1));
+
+  ComputeCommand cmd;
+  cmd.label = "Compute";
+  cmd.pipeline = compute_pipeline;
+
+  CS::Info info{.count = kCount};
+  CS::Input0<kCount> input_0;
+  CS::Input1<kCount> input_1;
+  for (size_t i = 0; i < kCount; i++) {
+    input_0.elements[i] = Vector4(2.0 + i, 3.0 + i, 4.0 + i, 5.0 * i);
+    input_1.elements[i] = Vector4(6.0, 7.0, 8.0, 9.0);
+  }
+
+  input_0.fixed_array[1] = IPoint32(2, 2);
+  input_1.fixed_array[0] = UintPoint32(3, 3);
+  input_0.some_int = 5;
+  input_1.some_struct = CS::SomeStruct{.vf = Point(3, 4), .i = 42};
+
+  auto output_buffer = CreateHostVisibleDeviceBuffer<CS::Output<kCount>>(
+      context, "Output Buffer");
+
+  CS::BindInfo(cmd, pass->GetTransientsBuffer().EmplaceUniform(info));
+  CS::BindInput0(cmd,
+                 pass->GetTransientsBuffer().EmplaceStorageBuffer(input_0));
+  CS::BindInput1(cmd,
+                 pass->GetTransientsBuffer().EmplaceStorageBuffer(input_1));
+  CS::BindOutput(cmd, output_buffer->AsBufferView());
+
+  ASSERT_TRUE(pass->AddCommand(std::move(cmd)));
+  ASSERT_FALSE(pass->EncodeCommands());
+}
+
 }  // namespace testing
 }  // namespace impeller