From 651383231ca26c24b46f095cce3ccaf9be86f0c6 Mon Sep 17 00:00:00 2001
From: Alan Baker <alanbaker@google.com>
Date: Tue, 17 Mar 2020 14:01:39 -0400
Subject: [PATCH 1/2] Set the propoer bytes for OpenCL SET commands

Fixes #814

* Convert the SET values into bytes before filling in the buffer
* Add a parser check that the type is a scalar
  * add a test
---
 src/amberscript/parser.cc                   |  3 +++
 src/amberscript/parser_pipeline_set_test.cc | 17 ++++++++++++
 src/pipeline.cc                             | 30 ++++++++++++++++++++-
 3 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/amberscript/parser.cc b/src/amberscript/parser.cc
index 8d124836d..953fc12b5 100644
--- a/src/amberscript/parser.cc
+++ b/src/amberscript/parser.cc
@@ -1024,6 +1024,9 @@ Result Parser::ParsePipelineSet(Pipeline* pipeline) {
   if (!type)
     return Result("invalid data type '" + token->AsString() + "' provided");
 
+  if (type->IsVec() || type->IsMatrix() || type->IsArray() || type->IsStruct())
+    return Result("data type must be a scalar type");
+
   token = tokenizer_->NextToken();
   if (!token->IsInteger() && !token->IsDouble())
     return Result("expected data value");
diff --git a/src/amberscript/parser_pipeline_set_test.cc b/src/amberscript/parser_pipeline_set_test.cc
index 7d78ca739..46f0dd45f 100644
--- a/src/amberscript/parser_pipeline_set_test.cc
+++ b/src/amberscript/parser_pipeline_set_test.cc
@@ -241,5 +241,22 @@ END
   EXPECT_EQ("7: SET can only be used with OPENCL-C shaders", r.Error());
 }
 
+TEST_F(AmberScriptParserTest, OpenCLSetNonScalarDataType) {
+  std::string in = R"(
+SHADER compute my_shader OPENCL-C
+#shader
+END
+PIPELINE compute my_pipeline
+  ATTACH my_shader
+  SET KERNEL ARG_NAME arg_a AS vec4<uint32> 0 0 0 0
+END
+)";
+
+  Parser parser;
+  auto r = parser.Parse(in);
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("7: data type must be a scalar type", r.Error());
+}
+
 }  // namespace amberscript
 }  // namespace amber
diff --git a/src/pipeline.cc b/src/pipeline.cc
index eca6326b7..ad1f0894b 100644
--- a/src/pipeline.cc
+++ b/src/pipeline.cc
@@ -755,7 +755,35 @@ Result Pipeline::GenerateOpenCLPodBuffers() {
       return Result(message);
     }
 
-    Result r = buffer->SetDataWithOffset({arg_info.value}, offset);
+    // Convert the argument value into bytes. Currently, only scalar arguments
+    // are supported.
+    const auto buffer_byte_size = buffer->GetFormat()->SizeInBytes();
+    const auto arg_byte_size = arg_info.fmt->SizeInBytes();
+    std::vector<Value> data_bytes;
+    for (auto i = 0; i < arg_byte_size; ++i) {
+      Value v;
+      if (arg_info.value.IsFloat()) {
+        if (arg_byte_size == sizeof(double)) {
+          union {
+            uint64_t u;
+            double d;
+          } u;
+          u.d = arg_info.value.AsDouble();
+          v.SetIntValue((u.u >> (i * 8)) & 0xff);
+        } else {
+          union {
+            uint32_t u;
+            float f;
+          } u;
+          u.f = arg_info.value.AsFloat();
+          v.SetIntValue((u.u >> (i * 8)) & 0xff);
+        }
+      } else {
+        v.SetIntValue((arg_info.value.AsUint64() >> (i * 8)) & 0xff);
+      }
+      data_bytes.push_back(v);
+    }
+    Result r = buffer->SetDataWithOffset(data_bytes, offset);
     if (!r.IsSuccess())
       return r;
   }

From e521f6c3694fecaf1ccce5d8a74fa0e2aa59f12c Mon Sep 17 00:00:00 2001
From: Alan Baker <alanbaker@google.com>
Date: Tue, 17 Mar 2020 14:54:55 -0400
Subject: [PATCH 2/2] Fixes

* Remove unused variable
* Fix sign compare mismatch
---
 src/pipeline.cc | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/pipeline.cc b/src/pipeline.cc
index ad1f0894b..478e6e179 100644
--- a/src/pipeline.cc
+++ b/src/pipeline.cc
@@ -757,10 +757,9 @@ Result Pipeline::GenerateOpenCLPodBuffers() {
 
     // Convert the argument value into bytes. Currently, only scalar arguments
     // are supported.
-    const auto buffer_byte_size = buffer->GetFormat()->SizeInBytes();
     const auto arg_byte_size = arg_info.fmt->SizeInBytes();
     std::vector<Value> data_bytes;
-    for (auto i = 0; i < arg_byte_size; ++i) {
+    for (uint32_t i = 0; i < arg_byte_size; ++i) {
       Value v;
       if (arg_info.value.IsFloat()) {
         if (arg_byte_size == sizeof(double)) {