[amberscript] Add DATA to IMAGE command (#934)

This CL adds a DATA initializer to the IMAGE command. This allows
setting the data to be used in the image as part of the initializer.

4 files changed, 175 insertions, 9 deletions

diff --git a/docs/amber_script.md b/docs/amber_script.md
index c4d249f..57e0632 100644
--- a/docs/amber_script.md
+++ b/docs/amber_script.md
@@ -215,12 +215,9 @@ either image buffers or, what the target API would refer to as a buffer.
 Sized arrays and structures are not currently representable.
 
 ```groovy
-# Filling the buffer with a given set of data. The values must be
-# of |type| data. The data can be provided as the type or as a hex value.
-# Buffers are STD430 by default.
-BUFFER {name} DATA_TYPE {type} {STD140 | STD430} DATA
-_value_+
-END
+# Filling the buffer with a given initializer. Initializer data must be
+# of |type|. Buffers are STD430 by default.
+BUFFER {name} DATA_TYPE {type} {STD140 | STD430} {initializer}
 
 # Defines a buffer which is filled with data as specified by the `initializer`.
 BUFFER {name} DATA_TYPE {type} {STD140 | STD430} SIZE _size_in_items_ \
@@ -287,6 +284,14 @@ IMAGE {name} DATA_TYPE {type} {dimensionality} \
 #### Buffer Initializers
 
 ```groovy
+# Filling the buffer with a given set of data. The values must be
+# of the correct type. The data can be provided as the type or as a hex
+# value.
+DATA
+_value_+
+END
+
+```groovy
 # Fill the buffer with a single value.
 FILL _value_
 
@@ -488,7 +493,7 @@ The following commands are all specified within the `PIPELINE` command.
   #  - a power-of-two integer that _must_ be greater or equal to minSubgroupSize
   #    and be less than or equal to maxSubgroupSize
   # - MIN to set the required subgroup size to the minSubgroupSize
-  # - MAX to set the required subgroup size to the maxSubgroupSize  
+  # - MAX to set the required subgroup size to the maxSubgroupSize
   SUBROUP {name_of_shader}
     FULLY_POPULATED {fully_populated_enable}
     VARYING_SIZE {varying_size_enable}
diff --git a/src/amberscript/parser.cc b/src/amberscript/parser.cc
index 778a659..fc5c948 100644
--- a/src/amberscript/parser.cc
+++ b/src/amberscript/parser.cc
@@ -1937,8 +1937,10 @@ Result Parser::ParseImage() {
 
   token = tokenizer_->PeekNextToken();
   while (token->IsIdentifier()) {
-    if (token->AsString() == "FILL" || token->AsString() == "SERIES_FROM")
+    if (token->AsString() == "FILL" || token->AsString() == "SERIES_FROM" ||
+        token->AsString() == "DATA") {
       break;
+    }
 
     tokenizer_->NextToken();
 
@@ -2045,7 +2047,18 @@ Result Parser::ParseImage() {
   // Parse initializers.
   token = tokenizer_->NextToken();
   if (token->IsIdentifier()) {
-    if (token->AsString() == "FILL") {
+    if (token->AsString() == "DATA") {
+      Result r = ParseBufferInitializerData(buffer.get());
+      if (!r.IsSuccess())
+        return r;
+
+      if (size_in_items != buffer->ElementCount()) {
+        return Result(
+            "Elements provided in data does not match size specified: " +
+            std::to_string(size_in_items) + " specified vs " +
+            std::to_string(buffer->ElementCount()) + " provided");
+      }
+    } else if (token->AsString() == "FILL") {
       Result r = ParseBufferInitializerFill(buffer.get(), size_in_items);
       if (!r.IsSuccess())
         return r;
diff --git a/src/amberscript/parser_image_test.cc b/src/amberscript/parser_image_test.cc
index bfa3fb6..53c8ad7 100644
--- a/src/amberscript/parser_image_test.cc
+++ b/src/amberscript/parser_image_test.cc
@@ -309,5 +309,61 @@ IMAGE image DATA_TYPE uint32 DIM_3D WIDTH 3 HEIGHT 4 DEPTH 5
   EXPECT_EQ(60u, buffer->ElementCount());
 }
 
+TEST_F(AmberScriptParserTest, ImageWithData) {
+  std::string in = R"(
+IMAGE image DATA_TYPE float DIM_3D HEIGHT 2 WIDTH 2 DEPTH 2 DATA
+ 0.11 0.12
+ 0.21 0.22
+
+ 0.31 0.32
+ 0.41 0.42
+END
+)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_TRUE(r.IsSuccess()) << r.Error();
+  auto script = parser.GetScript();
+  const auto& buffers = script->GetBuffers();
+  ASSERT_EQ(1U, buffers.size());
+
+  ASSERT_TRUE(buffers[0] != nullptr);
+  EXPECT_EQ("image", buffers[0]->GetName());
+
+  auto* buffer = buffers[0].get();
+  EXPECT_TRUE(buffer->GetFormat()->IsFloat32());
+  EXPECT_EQ(ImageDimension::k3D, buffer->GetImageDimension());
+  EXPECT_EQ(2u, buffer->GetWidth());
+  EXPECT_EQ(2u, buffer->GetHeight());
+  EXPECT_EQ(2u, buffer->GetDepth());
+  EXPECT_EQ(8u, buffer->ElementCount());
+
+  auto* values = buffer->GetValues<float>();
+  std::vector<float> result = {0.11f, 0.12f, 0.21f, 0.22f,
+                               0.31f, 0.32f, 0.41f, 0.42f};
+
+  EXPECT_EQ((*buffer->ValuePtr()).size(), 8u * sizeof(float));
+  for (size_t i = 0; i < result.size(); ++i) {
+    EXPECT_FLOAT_EQ(result[i], values[i]);
+  }
+}
+
+TEST_F(AmberScriptParserTest, ImageDataSizeIncorrect) {
+  std::string in = R"(
+IMAGE image DATA_TYPE float DIM_3D HEIGHT 2 WIDTH 2 DEPTH 2 DATA
+ 0.11 0.12
+ 0.21 0.22
+END
+)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ(
+      "6: Elements provided in data does not match size specified: 8 specified "
+      "vs 4 provided",
+      r.Error());
+}
+
 }  // namespace amberscript
 }  // namespace amber
diff --git a/tests/cases/image_data.amber b/tests/cases/image_data.amber
new file mode 100644
index 0000000..3d1f79f
--- /dev/null
+++ b/tests/cases/image_data.amber
@@ -0,0 +1,92 @@
+#!amber
+# Copyright 2021 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 compute_shader GLSL
+#version 460
+
+layout(local_size_x = 4, local_size_y = 4, local_size_z = 1) in;
+
+layout(binding = 0, r32f) uniform readonly image3D inImg;
+layout(binding = 1, r32f) uniform image3D outImg;
+
+void main() {
+  // Get current pixel
+  const int current_x = int(gl_GlobalInvocationID.x);
+  const int current_y = int(gl_GlobalInvocationID.y);
+
+  for (int idx = 0; idx < 4; ++idx) {
+    vec4 result = imageLoad(inImg, ivec3(current_x, current_y, 3 - idx));
+    imageStore(outImg, ivec3(current_x, current_y, idx), result);
+  }
+}
+END
+
+IMAGE outputImage DATA_TYPE float DIM_3D WIDTH 4 HEIGHT 4 DEPTH 4 FILL 0.0
+
+IMAGE inputImage DATA_TYPE float DIM_3D WIDTH 4 HEIGHT 4 DEPTH 4 DATA
+  0.110 0.111 0.112 0.113
+  0.120 0.121 0.122 0.123
+  0.130 0.131 0.132 0.133
+  0.140 0.141 0.142 0.143
+
+  0.210 0.211 0.212 0.213
+  0.220 0.221 0.222 0.223
+  0.230 0.231 0.232 0.233
+  0.240 0.241 0.242 0.243
+
+  0.310 0.311 0.312 0.313
+  0.320 0.321 0.322 0.323
+  0.330 0.331 0.332 0.333
+  0.340 0.341 0.342 0.343
+
+  0.410 0.411 0.412 0.413
+  0.420 0.421 0.422 0.423
+  0.430 0.431 0.432 0.433
+  0.440 0.441 0.442 0.443
+END
+
+IMAGE expectedImage DATA_TYPE float DIM_3D WIDTH 4 HEIGHT 4 DEPTH 4 DATA
+  0.410 0.411 0.412 0.413
+  0.420 0.421 0.422 0.423
+  0.430 0.431 0.432 0.433
+  0.440 0.441 0.442 0.443
+
+  0.310 0.311 0.312 0.313
+  0.320 0.321 0.322 0.323
+  0.330 0.331 0.332 0.333
+  0.340 0.341 0.342 0.343
+
+  0.210 0.211 0.212 0.213
+  0.220 0.221 0.222 0.223
+  0.230 0.231 0.232 0.233
+  0.240 0.241 0.242 0.243
+
+  0.110 0.111 0.112 0.113
+  0.120 0.121 0.122 0.123
+  0.130 0.131 0.132 0.133
+  0.140 0.141 0.142 0.143
+END
+
+PIPELINE compute pipeline
+  ATTACH compute_shader
+
+  BIND BUFFER inputImage AS storage_image DESCRIPTOR_SET 0 BINDING 0
+  BIND BUFFER outputImage AS storage_image DESCRIPTOR_SET 0 BINDING 1
+END
+
+RUN pipeline 1 1 1
+
+EXPECT outputImage EQ_BUFFER expectedImage
+