Shader specialization

docs/amber_script.md

@@ -197,6 +197,12 @@ The following commands are all specified within the `PIPELINE` command.
   # Attach a 'multi' shader to the pipeline of |shader_type| and use the entry
   # point with |name|. The provided shader _must_ be a 'multi' shader.
   ATTACH {name_of_multi_shader} TYPE {shader_type} ENTRY_POINT {name}
+
+  # Attach specialized shader. Specialization can be specified multiple times.
+  # Specialization values must be a 32-bit type. Shader type and entry point
+  # must be specified prior to specializing the shader.
+  ATTACH {name_of_shader} SPECIALIZE 1 AS uint32 4
+  ATTACH {name_of_shader} SPECIALIZE 1 AS uint32 4 SPECIALIZE 4 AS float 1.0
 ```
 
 ```groovy

src/amberscript/parser.cc

@@ -472,7 +472,7 @@ Result Parser::ParsePipelineAttach(Pipeline* pipeline) {
 
     type = token->AsString();
   }
-  if (type != "ENTRY_POINT")
+  if (set_shader_type && type != "ENTRY_POINT")
     return Result("Unknown ATTACH parameter: " + type);
 
   if (shader->GetType() == ShaderType::kShaderTypeMulti && !set_shader_type)
@@ -482,15 +482,81 @@ Result Parser::ParsePipelineAttach(Pipeline* pipeline) {
   if (!r.IsSuccess())
     return r;
 
+  if (type == "ENTRY_POINT") {
+    token = tokenizer_->NextToken();
+    if (!token->IsString())
+      return Result("missing shader name in ATTACH ENTRY_POINT command");
+
+    r = pipeline->SetShaderEntryPoint(shader, token->AsString());
+    if (!r.IsSuccess())
+      return r;
+
+    token = tokenizer_->NextToken();
+  }
+
+  while (true) {
+    if (token->IsString() && token->AsString() == "SPECIALIZE") {
+      r = ParseShaderSpecialization(pipeline);
+      if (!r.IsSuccess())
+        return r;
+
+      token = tokenizer_->NextToken();
+    } else {
+      if (token->IsEOL() || token->IsEOS())
+        return {};
+      if (token->IsString())
+        return Result("Unknown ATTACH parameter: " + token->AsString());
+      return Result("extra parameters after ATTACH command");
+    }
+  }
+}
+
+Result Parser::ParseShaderSpecialization(Pipeline* pipeline) {
+  auto token = tokenizer_->NextToken();
+  if (!token->IsInteger())
+    return Result("specialization ID must be an integer");
+
+  auto spec_id = token->AsUint32();
+
+  token = tokenizer_->NextToken();
+  if (!token->IsString() || token->AsString() != "AS")
+    return Result("expected AS as next token");
+
   token = tokenizer_->NextToken();
   if (!token->IsString())
-    return Result("missing shader name in ATTACH ENTRY_POINT command");
+    return Result("expected data type in SPECIALIZE subcommand");
 
-  r = pipeline->SetShaderEntryPoint(shader, token->AsString());
+  DatumType type;
+  auto r = ToDatumType(token->AsString(), &type);
   if (!r.IsSuccess())
     return r;
 
-  return ValidateEndOfStatement("ATTACH command");
+  token = tokenizer_->NextToken();
+  uint32_t value = 0;
+  switch (type.GetType()) {
+    case DataType::kUint32:
+    case DataType::kInt32:
+      value = token->AsUint32();
+      break;
+    case DataType::kFloat: {
+      r = token->ConvertToDouble();
+      if (!r.IsSuccess())
+        return Result("value is not a floating point value");
+      union {
+        uint32_t u;
+        float f;
+      } u;
+      u.f = token->AsFloat();
+      value = u.u;
+      break;
+    }
+    default:
+      return Result(
+          "only 32-bit types are currently accepted for specialization values");
+  }
+  auto& shader = pipeline->GetShaders()[pipeline->GetShaders().size() - 1];
+  shader.AddSpecialization(spec_id, value);
+  return {};
 }
 
 Result Parser::ParsePipelineShaderOptimizations(Pipeline* pipeline) {

src/amberscript/parser.h

@@ -77,6 +77,7 @@ class Parser : public amber::Parser {
   Result ParseDerivePipelineBlock();
   Result ParsePipelineBody(const std::string& cmd_name,
                            std::unique_ptr<Pipeline> pipeline);
+  Result ParseShaderSpecialization(Pipeline* pipeline);
 
   /// Parses a set of values out of the token stream. |name| is the name of the
   /// current command we're parsing for error purposes. The |type| is the type

src/amberscript/parser_attach_test.cc

@@ -215,7 +215,7 @@ END)";
   Parser parser;
   Result r = parser.Parse(in);
   ASSERT_FALSE(r.IsSuccess());
-  EXPECT_EQ("6: extra parameters after ATTACH command", r.Error());
+  EXPECT_EQ("6: Unknown ATTACH parameter: INVALID", r.Error());
 }
 
 TEST_F(AmberScriptParserTest, PiplineMultiShaderAttach) {
@@ -307,5 +307,198 @@ END)";
   EXPECT_EQ("6: ATTACH missing TYPE for multi shader", r.Error());
 }
 
+TEST_F(AmberScriptParserTest, PipelineSpecializationUint32) {
+  std::string in = R"(
+SHADER compute my_shader GLSL
+#shaders
+END
+PIPELINE compute my_pipeline
+  ATTACH my_shader TYPE compute ENTRY_POINT my_ep SPECIALIZE 1 AS uint32 4
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  EXPECT_EQ(r.Error(), "");
+  ASSERT_TRUE(r.IsSuccess());
+
+  auto script = parser.GetScript();
+  const auto& pipelines = script->GetPipelines();
+  ASSERT_EQ(1U, pipelines.size());
+
+  const auto* pipeline = pipelines[0].get();
+  const auto& shaders = pipeline->GetShaders();
+  ASSERT_EQ(1U, shaders.size());
+
+  EXPECT_EQ(1, shaders[0].GetSpecialization().size());
+  EXPECT_EQ(4, shaders[0].GetSpecialization().at(1));
+}
+
+TEST_F(AmberScriptParserTest, PipelineSpecializationInt32) {
+  std::string in = R"(
+SHADER compute my_shader GLSL
+#shaders
+END
+PIPELINE compute my_pipeline
+  ATTACH my_shader TYPE compute ENTRY_POINT my_ep SPECIALIZE 2 AS int32 -1
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_TRUE(r.IsSuccess());
+
+  auto script = parser.GetScript();
+  const auto& pipelines = script->GetPipelines();
+  ASSERT_EQ(1U, pipelines.size());
+
+  const auto* pipeline = pipelines[0].get();
+  const auto& shaders = pipeline->GetShaders();
+  ASSERT_EQ(1U, shaders.size());
+
+  EXPECT_EQ(1, shaders[0].GetSpecialization().size());
+  EXPECT_EQ(0xffffffff, shaders[0].GetSpecialization().at(2));
+}
+
+TEST_F(AmberScriptParserTest, PipelineSpecializationFloat) {
+  std::string in = R"(
+SHADER compute my_shader GLSL
+#shaders
+END
+PIPELINE compute my_pipeline
+  ATTACH my_shader TYPE compute ENTRY_POINT my_ep SPECIALIZE 3 AS float 1.1
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_TRUE(r.IsSuccess());
+
+  auto script = parser.GetScript();
+  const auto& pipelines = script->GetPipelines();
+  ASSERT_EQ(1U, pipelines.size());
+
+  const auto* pipeline = pipelines[0].get();
+  const auto& shaders = pipeline->GetShaders();
+  ASSERT_EQ(1U, shaders.size());
+
+  EXPECT_EQ(1, shaders[0].GetSpecialization().size());
+  EXPECT_EQ(0x3f8ccccd, shaders[0].GetSpecialization().at(3));
+}
+
+TEST_F(AmberScriptParserTest, PipelineSpecializationIDIsString) {
+  std::string in = R"(
+SHADER compute my_shader GLSL
+#shaders
+END
+PIPELINE compute my_pipeline
+  ATTACH my_shader TYPE compute ENTRY_POINT my_ep SPECIALIZE s3 AS float 1.1
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("6: specialization ID must be an integer", r.Error());
+}
+
+TEST_F(AmberScriptParserTest, PipelineSpecializationNoAS) {
+  std::string in = R"(
+SHADER compute my_shader GLSL
+#shaders
+END
+PIPELINE compute my_pipeline
+  ATTACH my_shader TYPE compute ENTRY_POINT my_ep SPECIALIZE 1 ASa float 1.1
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("6: expected AS as next token", r.Error());
+}
+
+TEST_F(AmberScriptParserTest, PipelineSpecializationNotDataType) {
+  std::string in = R"(
+SHADER compute my_shader GLSL
+#shaders
+END
+PIPELINE compute my_pipeline
+  ATTACH my_shader TYPE compute ENTRY_POINT my_ep SPECIALIZE 1 AS uint 1.1
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ("6: invalid data_type provided", r.Error());
+}
+
+TEST_F(AmberScriptParserTest, PipelineSpecializationBadDataType) {
+  std::string in = R"(
+SHADER compute my_shader GLSL
+#shaders
+END
+PIPELINE compute my_pipeline
+  ATTACH my_shader ENTRY_POINT my_ep SPECIALIZE 1 AS uint8 1.1
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_FALSE(r.IsSuccess());
+  EXPECT_EQ(
+      "6: only 32-bit types are currently accepted for specialization values",
+      r.Error());
+}
+
+TEST_F(AmberScriptParserTest, PipelineSpecializationMultipleSpecializations) {
+  std::string in = R"(
+SHADER compute my_shader GLSL
+#shaders
+END
+PIPELINE compute my_pipeline
+  ATTACH my_shader TYPE compute ENTRY_POINT my_ep \
+      SPECIALIZE 1 AS uint32 4 \
+      SPECIALIZE 2 AS uint32 5 \
+      SPECIALIZE 5 AS uint32 1
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_TRUE(r.IsSuccess());
+
+  auto script = parser.GetScript();
+  const auto& pipelines = script->GetPipelines();
+  ASSERT_EQ(1U, pipelines.size());
+
+  const auto* pipeline = pipelines[0].get();
+  const auto& shaders = pipeline->GetShaders();
+  ASSERT_EQ(1U, shaders.size());
+
+  EXPECT_EQ(3, shaders[0].GetSpecialization().size());
+  EXPECT_EQ(4, shaders[0].GetSpecialization().at(1));
+  EXPECT_EQ(5, shaders[0].GetSpecialization().at(2));
+  EXPECT_EQ(1, shaders[0].GetSpecialization().at(5));
+}
+
+TEST_F(AmberScriptParserTest, PipelineSpecializationNoType) {
+  std::string in = R"(
+SHADER compute my_shader GLSL
+#shaders
+END
+PIPELINE compute my_pipeline
+  ATTACH my_shader SPECIALIZE 1 AS uint32 4
+END)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  ASSERT_TRUE(r.IsSuccess());
+
+  auto script = parser.GetScript();
+  const auto& pipelines = script->GetPipelines();
+  ASSERT_EQ(1U, pipelines.size());
+
+  const auto* pipeline = pipelines[0].get();
+  const auto& shaders = pipeline->GetShaders();
+  ASSERT_EQ(1U, shaders.size());
+
+  EXPECT_EQ(1, shaders[0].GetSpecialization().size());
+  EXPECT_EQ(4, shaders[0].GetSpecialization().at(1));
+}
+
 }  // namespace amberscript
 }  // namespace amber

src/amberscript/parser_pipeline_test.cc

@@ -414,5 +414,41 @@ END
   EXPECT_EQ("3: missing pipeline name for DERIVE_PIPELINE command", r.Error());
 }
 
+TEST_F(AmberScriptParserTest, DerivePipelineSpecialized) {
+  std::string in = R"(
+SHADER compute my_shader GLSL
+#shaders
+END
+PIPELINE compute p1
+  ATTACH my_shader SPECIALIZE 3 AS uint32 4
+END
+DERIVE_PIPELINE p2 FROM p1
+END
+)";
+
+  Parser parser;
+  Result r = parser.Parse(in);
+  EXPECT_EQ("", r.Error());
+  ASSERT_TRUE(r.IsSuccess());
+
+  auto script = parser.GetScript();
+  const auto& pipelines = script->GetPipelines();
+  ASSERT_EQ(2U, pipelines.size());
+
+  const auto* p1 = pipelines[0].get();
+  const auto& s1 = p1->GetShaders();
+  ASSERT_EQ(1U, s1.size());
+
+  EXPECT_EQ(1, s1[0].GetSpecialization().size());
+  EXPECT_EQ(4, s1[0].GetSpecialization().at(3));
+
+  const auto* p2 = pipelines[1].get();
+  const auto& s2 = p2->GetShaders();
+  ASSERT_EQ(1U, s2.size());
+
+  EXPECT_EQ(1, s2[0].GetSpecialization().size());
+  EXPECT_EQ(4, s2[0].GetSpecialization().at(3));
+}
+
 }  // namespace amberscript
 }  // namespace amber

src/pipeline.h

@@ -15,6 +15,7 @@
 #ifndef SRC_PIPELINE_H_
 #define SRC_PIPELINE_H_
 
+#include <map>
 #include <memory>
 #include <string>
 #include <utility>
@@ -59,12 +60,20 @@ class Pipeline {
     const std::vector<uint32_t> GetData() const { return data_; }
     void SetData(std::vector<uint32_t>&& data) { data_ = std::move(data); }
 
+    const std::map<uint32_t, uint32_t>& GetSpecialization() const {
+      return specialization_;
+    }
+    void AddSpecialization(uint32_t spec_id, uint32_t value) {
+      specialization_[spec_id] = value;
+    }
+
    private:
     Shader* shader_ = nullptr;
     ShaderType shader_type_;
     std::vector<std::string> shader_optimizations_;
     std::string entry_point_;
     std::vector<uint32_t> data_;
+    std::map<uint32_t, uint32_t> specialization_;
   };
 
   /// Information on a buffer attached to the pipeline.