Add support for push constants generated by Clspv

src/clspv_helper.cc

@@ -94,6 +94,24 @@ Result Compile(Pipeline::ShaderInfo* shader_info,
 
       shader_info->AddDescriptorEntry(entry.kernel_arg_data.kernel_name,
                                       std::move(descriptor_entry));
+    } else if (entry.kind ==
+               clspv::version0::DescriptorMapEntry::PushConstant) {
+      Pipeline::ShaderInfo::PushConstant push_constant;
+      switch (entry.push_constant_data.pc) {
+        case clspv::PushConstant::Dimensions:
+          push_constant.type =
+              Pipeline::ShaderInfo::PushConstant::PushConstantType::kDimensions;
+          break;
+        case clspv::PushConstant::GlobalOffset:
+          push_constant.type = Pipeline::ShaderInfo::PushConstant::
+              PushConstantType::kGlobalOffset;
+          break;
+        default:
+          return Result("unhandled push constant generated by Clspv");
+      }
+      push_constant.offset = entry.push_constant_data.offset;
+      push_constant.size = entry.push_constant_data.size;
+      shader_info->AddPushConstant(std::move(push_constant));
     } else {
       assert(entry.kind == clspv::version0::DescriptorMapEntry::Sampler);
       // Create a new sampler info.

src/executor.cc

@@ -71,6 +71,9 @@ Result Executor::Execute(Engine* engine,
       r = pipeline->GenerateOpenCLLiteralSamplers();
       if (!r.IsSuccess())
         return r;
+      r = pipeline->GenerateOpenCLPushConstants();
+      if (!r.IsSuccess())
+        return r;
     }
 
     for (auto& pipeline : script->GetPipelines()) {

src/pipeline.cc

@@ -15,6 +15,7 @@
 #include "src/pipeline.h"
 
 #include <algorithm>
+#include <cstring>
 #include <limits>
 #include <set>
 
@@ -45,6 +46,7 @@ const uint32_t kOpenCLFilterModeLinearBit = 0x20;
 
 const char* Pipeline::kGeneratedColorBuffer = "framebuffer";
 const char* Pipeline::kGeneratedDepthBuffer = "depth_buffer";
+const char* Pipeline::kGeneratedPushConstantBuffer = "push_constant_buffer";
 
 Pipeline::ShaderInfo::ShaderInfo(Shader* shader, ShaderType type)
     : shader_(shader), shader_type_(type), entry_point_("main") {}
@@ -852,4 +854,62 @@ Result Pipeline::GenerateOpenCLLiteralSamplers() {
   return {};
 }
 
+Result Pipeline::GenerateOpenCLPushConstants() {
+  if (!IsCompute() || GetShaders().empty() ||
+      GetShaders()[0].GetShader()->GetFormat() != kShaderFormatOpenCLC) {
+    return {};
+  }
+
+  const auto& shader_info = GetShaders()[0];
+  if (shader_info.GetPushConstants().empty())
+    return {};
+
+  // Determine size and contents of the push constant buffer.
+  std::vector<uint32_t> bytes;
+  for (const auto& pc : shader_info.GetPushConstants()) {
+    assert(pc.size % sizeof(uint32_t) == 0);
+    assert(pc.offset % sizeof(uint32_t) == 0);
+    uint32_t elements = (pc.offset + pc.size) / sizeof(uint32_t);
+    if (elements > bytes.size()) {
+      bytes.resize(elements);
+    }
+
+    uint32_t base = pc.offset / sizeof(uint32_t);
+    switch (pc.type) {
+      case Pipeline::ShaderInfo::PushConstant::PushConstantType::kDimensions:
+        // All compute kernel launches are 3D.
+        bytes[base] = 3;
+        break;
+      case Pipeline::ShaderInfo::PushConstant::PushConstantType::kGlobalOffset:
+        // Global offsets are not currently supported.
+        bytes[base] = 0;
+        bytes[base + 1] = 0;
+        bytes[base + 2] = 0;
+        break;
+    }
+  }
+
+  TypeParser parser;
+  auto type = parser.Parse("R8_UINT");
+  auto fmt = MakeUnique<Format>(type.get());
+
+  // Create buffer containing push constant data.
+  std::unique_ptr<Buffer> buf = MakeUnique<Buffer>();
+  buf->SetName(kGeneratedPushConstantBuffer);
+  buf->SetFormat(fmt.get());
+  buf->SetSizeInBytes(static_cast<uint32_t>(bytes.size() * sizeof(uint32_t)));
+  std::memcpy(buf->ValuePtr()->data(), bytes.data(),
+              bytes.size() * sizeof(uint32_t));
+
+  Result r = SetPushConstantBuffer(buf.get());
+  if (!r.IsSuccess())
+    return r;
+
+  formats_.push_back(std::move(fmt));
+  types_.push_back(std::move(type));
+  opencl_push_constants_ = std::move(buf);
+
+  return {};
+}
+
 }  // namespace amber

src/pipeline.h

@@ -112,6 +112,24 @@ class Pipeline {
       return descriptor_map_;
     }
 
+    /// Push constant information for an OpenCL-C shader.
+    struct PushConstant {
+      enum class PushConstantType {
+        kDimensions = 0,
+        kGlobalOffset,
+      };
+      PushConstantType type;
+      uint32_t offset = 0;
+      uint32_t size = 0;
+    };
+
+    void AddPushConstant(PushConstant&& pc) {
+      push_constants_.emplace_back(std::move(pc));
+    }
+    const std::vector<PushConstant>& GetPushConstants() const {
+      return push_constants_;
+    }
+
    private:
     Shader* shader_ = nullptr;
     ShaderType shader_type_;
@@ -121,6 +139,7 @@ class Pipeline {
     std::map<uint32_t, uint32_t> specialization_;
     std::unordered_map<std::string, std::vector<DescriptorMapEntry>>
         descriptor_map_;
+    std::vector<PushConstant> push_constants_;
     std::vector<std::string> compile_options_;
   };
 
@@ -157,6 +176,7 @@ class Pipeline {
 
   static const char* kGeneratedColorBuffer;
   static const char* kGeneratedDepthBuffer;
+  static const char* kGeneratedPushConstantBuffer;
 
   explicit Pipeline(PipelineType type);
   ~Pipeline();
@@ -318,6 +338,9 @@ class Pipeline {
   /// descriptor map. This should be called after all other samplers are bound.
   Result GenerateOpenCLLiteralSamplers();
 
+  /// Generate the push constant buffers necessary for OpenCL kernels.
+  Result GenerateOpenCLPushConstants();
+
  private:
   void UpdateFramebufferSizes();
 
@@ -346,6 +369,7 @@ class Pipeline {
   /// Maps (descriptor set, binding) to the buffer for that binding pair.
   std::map<std::pair<uint32_t, uint32_t>, Buffer*> opencl_pod_buffer_map_;
   std::vector<std::unique_ptr<Sampler>> opencl_literal_samplers_;
+  std::unique_ptr<Buffer> opencl_push_constants_;
 };
 
 }  // namespace amber

src/pipeline_test.cc

@@ -821,4 +821,39 @@ TEST_F(PipelineTest, OpenCLGenerateLiteralSamplers) {
   }
 }
 
+TEST_F(PipelineTest, OpenCLGeneratePushConstants) {
+  Pipeline p(PipelineType::kCompute);
+  p.SetName("my_pipeline");
+
+  Shader cs(kShaderTypeCompute);
+  cs.SetFormat(kShaderFormatOpenCLC);
+  p.AddShader(&cs, kShaderTypeCompute);
+  p.SetShaderEntryPoint(&cs, "my_main");
+
+  Pipeline::ShaderInfo::PushConstant pc1;
+  pc1.type = Pipeline::ShaderInfo::PushConstant::PushConstantType::kDimensions;
+  pc1.offset = 0;
+  pc1.size = 4;
+  p.GetShaders()[0].AddPushConstant(std::move(pc1));
+
+  Pipeline::ShaderInfo::PushConstant pc2;
+  pc2.type =
+      Pipeline::ShaderInfo::PushConstant::PushConstantType::kGlobalOffset;
+  pc2.offset = 16;
+  pc2.size = 12;
+  p.GetShaders()[0].AddPushConstant(std::move(pc2));
+
+  auto r = p.GenerateOpenCLPushConstants();
+  ASSERT_TRUE(r.IsSuccess());
+
+  const auto& buf = p.GetPushConstantBuffer();
+  EXPECT_EQ(28U, buf.buffer->GetSizeInBytes());
+
+  const uint32_t* bytes = buf.buffer->GetValues<uint32_t>();
+  EXPECT_EQ(3U, bytes[0]);
+  EXPECT_EQ(0U, bytes[4]);
+  EXPECT_EQ(0U, bytes[5]);
+  EXPECT_EQ(0U, bytes[6]);
+}
+
 }  // namespace amber

tests/cases/opencl_generated_push_constants.amber

@@ -0,0 +1,79 @@
+#!amber
+# Copyright 2020 The Amber Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+SHADER compute work_dim_shader OPENCL-C
+kernel void foo(global int* out) {
+  *out = get_work_dim();
+}
+END
+
+SHADER compute global_offset_shader OPENCL-C
+kernel void foo(global int* out) {
+  out[0] = get_global_offset(0);
+  out[1] = get_global_offset(1);
+  out[2] = get_global_offset(2);
+}
+END
+
+SHADER compute all_constants_shader OPENCL-C
+kernel void foo(global int* out) {
+  out[0] = get_global_offset(0);
+  out[1] = get_global_offset(1);
+  out[2] = get_global_offset(2);
+  out[3] = get_work_dim();
+}
+END
+
+BUFFER work_dim_buf DATA_TYPE uint32 SIZE 1 FILL -1
+BUFFER global_offset_buf DATA_TYPE uint32 SIZE 3 FILL -1
+BUFFER all_constants_buf DATA_TYPE uint32 SIZE 4 FILL -1
+
+PIPELINE compute work_dim_pipeline
+  ATTACH work_dim_shader ENTRY_POINT foo
+  COMPILE_OPTIONS work_dim_shader
+     -work-dim
+  END
+  BIND BUFFER work_dim_buf AS storage DESCRIPTOR_SET 0 BINDING 0
+END
+
+PIPELINE compute global_offset_pipeline
+  ATTACH global_offset_shader ENTRY_POINT foo
+  COMPILE_OPTIONS global_offset_shader
+     -global-offset
+  END
+  BIND BUFFER global_offset_buf AS storage DESCRIPTOR_SET 0 BINDING 0
+END
+
+PIPELINE compute all_constants_pipeline
+  ATTACH all_constants_shader ENTRY_POINT foo
+  COMPILE_OPTIONS all_constants_shader
+     -work-dim -global-offset
+  END
+  BIND BUFFER all_constants_buf AS storage DESCRIPTOR_SET 0 BINDING 0
+END
+
+RUN work_dim_pipeline 1 1 1
+EXPECT work_dim_buf IDX 0 EQ 3
+
+RUN global_offset_pipeline 1 1 1
+EXPECT global_offset_buf IDX 0 EQ 0
+EXPECT global_offset_buf IDX 4 EQ 0
+EXPECT global_offset_buf IDX 8 EQ 0
+
+RUN all_constants_pipeline 1 1 1
+EXPECT all_constants_buf IDX 0 EQ 0
+EXPECT all_constants_buf IDX 4 EQ 0
+EXPECT all_constants_buf IDX 8 EQ 0
+EXPECT all_constants_buf IDX 12 EQ 3

tests/run_tests.py

@@ -95,6 +95,7 @@
 OPENCL_CASES = [
   "opencl_bind_buffer.amber",
   "opencl_c_copy.amber",
+  "opencl_generated_push_constants.amber",
   "opencl_read_and_write_image3d_rgba32i.amber",
   "opencl_read_image.amber",
   "opencl_read_image_literal_sampler.amber",