ast: Support non-literal workgroup_size parameters

Change the type of the values in an ast::WorkgroupDecoration to be
ast::Expression nodes, so that they can represent both
ast::ScalarExpression (literal) and ast::IdentifierExpression
(module-scope constant).

The Resolver processes these nodes to produce a uint32_t for the
default value on each dimension, and captures a reference to the
module-scope constant if it is overridable (which will soon be used by
the inspector and backends).

The WGSL parser now uses `primary_expression` to parse arguments to
workgroup_size.

Also added some WorkgroupSize() helpers to ProgramBuilder.

Bug: tint:713
Change-Id: I44b7b0021b925c84f25f65e26dc7da6b19ede508
Reviewed-on: https://dawn-review.googlesource.com/c/tint/+/51262
Commit-Queue: James Price <jrprice@google.com>
Auto-Submit: James Price <jrprice@google.com>
Reviewed-by: Ben Clayton <bclayton@google.com>

src/ast/function_test.cc

@@ -107,7 +107,7 @@ TEST_F(FunctionTest, Assert_DifferentProgramID_Deco) {
         ProgramBuilder b2;
         b1.Func("func", VariableList{}, b1.ty.void_(), StatementList{},
                 DecorationList{
-                    b2.create<WorkgroupDecoration>(2, 4, 6),
+                    b2.WorkgroupSize(2, 4, 6),
                 });
       },
       "internal compiler error");
@@ -121,7 +121,7 @@ TEST_F(FunctionTest, Assert_DifferentProgramID_ReturnDeco) {
         b1.Func("func", VariableList{}, b1.ty.void_(), StatementList{},
                 DecorationList{},
                 DecorationList{
-                    b2.create<WorkgroupDecoration>(2, 4, 6),
+                    b2.WorkgroupSize(2, 4, 6),
                 });
       },
       "internal compiler error");
@@ -159,10 +159,14 @@ TEST_F(FunctionTest, ToStr_WithDecoration) {
                  StatementList{
                      create<DiscardStatement>(),
                  },
-                 DecorationList{create<WorkgroupDecoration>(2, 4, 6)});
+                 DecorationList{WorkgroupSize(2, 4, 6)});
 
   EXPECT_EQ(str(f), R"(Function func -> __void
-WorkgroupDecoration{2 4 6}
+WorkgroupDecoration{
+  ScalarConstructor[not set]{2}
+  ScalarConstructor[not set]{4}
+  ScalarConstructor[not set]{6}
+}
 ()
 {
   Discard{}

src/ast/int_literal.h

@@ -25,6 +25,9 @@ class IntLiteral : public Castable<IntLiteral, Literal> {
  public:
   ~IntLiteral() override;
 
+  /// @returns the literal value as an i32
+  int32_t value_as_i32() const { return static_cast<int32_t>(value_); }
+
   /// @returns the literal value as a u32
   uint32_t value_as_u32() const { return value_; }
 

src/ast/workgroup_decoration.cc

@@ -23,36 +23,36 @@ namespace ast {
 
 WorkgroupDecoration::WorkgroupDecoration(ProgramID program_id,
                                          const Source& source,
-                                         uint32_t x)
-    : WorkgroupDecoration(program_id, source, x, 1, 1) {}
-
-WorkgroupDecoration::WorkgroupDecoration(ProgramID program_id,
-                                         const Source& source,
-                                         uint32_t x,
-                                         uint32_t y)
-    : WorkgroupDecoration(program_id, source, x, y, 1) {}
-
-WorkgroupDecoration::WorkgroupDecoration(ProgramID program_id,
-                                         const Source& source,
-                                         uint32_t x,
-                                         uint32_t y,
-                                         uint32_t z)
+                                         ast::Expression* x,
+                                         ast::Expression* y,
+                                         ast::Expression* z)
     : Base(program_id, source), x_(x), y_(y), z_(z) {}
 
 WorkgroupDecoration::~WorkgroupDecoration() = default;
 
-void WorkgroupDecoration::to_str(const sem::Info&,
+void WorkgroupDecoration::to_str(const sem::Info& sem,
                                  std::ostream& out,
                                  size_t indent) const {
   make_indent(out, indent);
-  out << "WorkgroupDecoration{" << x_ << " " << y_ << " " << z_ << "}"
-      << std::endl;
+  out << "WorkgroupDecoration{" << std::endl;
+  x_->to_str(sem, out, indent + 2);
+  if (y_) {
+    y_->to_str(sem, out, indent + 2);
+    if (z_) {
+      z_->to_str(sem, out, indent + 2);
+    }
+  }
+  make_indent(out, indent);
+  out << "}" << std::endl;
 }
 
 WorkgroupDecoration* WorkgroupDecoration::Clone(CloneContext* ctx) const {
   // Clone arguments outside of create() call to have deterministic ordering
   auto src = ctx->Clone(source());
-  return ctx->dst->create<WorkgroupDecoration>(src, x_, y_, z_);
+  auto* x = ctx->Clone(x_);
+  auto* y = ctx->Clone(y_);
+  auto* z = ctx->Clone(z_);
+  return ctx->dst->create<WorkgroupDecoration>(src, x, y, z);
 }
 
 }  // namespace ast

src/ast/workgroup_decoration.h

@@ -15,47 +15,35 @@
 #ifndef SRC_AST_WORKGROUP_DECORATION_H_
 #define SRC_AST_WORKGROUP_DECORATION_H_
 
-#include <tuple>
+#include <array>
 
 #include "src/ast/decoration.h"
 
 namespace tint {
 namespace ast {
 
+// Forward declaration
+class Expression;
+
 /// A workgroup decoration
 class WorkgroupDecoration : public Castable<WorkgroupDecoration, Decoration> {
  public:
   /// constructor
   /// @param program_id the identifier of the program that owns this node
   /// @param source the source of this decoration
-  /// @param x the workgroup x dimension size
-  WorkgroupDecoration(ProgramID program_id, const Source& source, uint32_t x);
-  /// constructor
-  /// @param program_id the identifier of the program that owns this node
-  /// @param source the source of this decoration
-  /// @param x the workgroup x dimension size
-  /// @param y the workgroup x dimension size
+  /// @param x the workgroup x dimension expression
+  /// @param y the optional workgroup y dimension expression
+  /// @param z the optional workgroup z dimension expression
   WorkgroupDecoration(ProgramID program_id,
                       const Source& source,
-                      uint32_t x,
-                      uint32_t y);
-  /// constructor
-  /// @param program_id the identifier of the program that owns this node
-  /// @param source the source of this decoration
-  /// @param x the workgroup x dimension size
-  /// @param y the workgroup x dimension size
-  /// @param z the workgroup x dimension size
-  WorkgroupDecoration(ProgramID program_id,
-                      const Source& source,
-                      uint32_t x,
-                      uint32_t y,
-                      uint32_t z);
+                      ast::Expression* x,
+                      ast::Expression* y = nullptr,
+                      ast::Expression* z = nullptr);
+
   ~WorkgroupDecoration() override;
 
   /// @returns the workgroup dimensions
-  std::tuple<uint32_t, uint32_t, uint32_t> values() const {
-    return {x_, y_, z_};
-  }
+  std::array<ast::Expression*, 3> values() const { return {x_, y_, z_}; }
 
   /// Outputs the decoration to the given stream
   /// @param sem the semantic info for the program
@@ -72,9 +60,9 @@ class WorkgroupDecoration : public Castable<WorkgroupDecoration, Decoration> {
   WorkgroupDecoration* Clone(CloneContext* ctx) const override;
 
  private:
-  uint32_t const x_;
-  uint32_t const y_;
-  uint32_t const z_;
+  ast::Expression* x_ = nullptr;
+  ast::Expression* y_ = nullptr;
+  ast::Expression* z_ = nullptr;
 };
 
 }  // namespace ast

src/ast/workgroup_decoration_test.cc

@@ -24,40 +24,84 @@ namespace {
 using WorkgroupDecorationTest = TestHelper;
 
 TEST_F(WorkgroupDecorationTest, Creation_1param) {
-  auto* d = create<WorkgroupDecoration>(2);
-  uint32_t x = 0;
-  uint32_t y = 0;
-  uint32_t z = 0;
-  std::tie(x, y, z) = d->values();
-  EXPECT_EQ(x, 2u);
-  EXPECT_EQ(y, 1u);
-  EXPECT_EQ(z, 1u);
+  auto* d = WorkgroupSize(2);
+  auto values = d->values();
+  ASSERT_NE(values[0], nullptr);
+  auto* x_scalar = values[0]->As<ast::ScalarConstructorExpression>();
+  ASSERT_TRUE(x_scalar);
+  ASSERT_TRUE(x_scalar->literal()->Is<ast::IntLiteral>());
+  EXPECT_EQ(x_scalar->literal()->As<ast::IntLiteral>()->value_as_u32(), 2u);
+
+  EXPECT_EQ(values[1], nullptr);
+  EXPECT_EQ(values[2], nullptr);
 }
 TEST_F(WorkgroupDecorationTest, Creation_2param) {
-  auto* d = create<WorkgroupDecoration>(2, 4);
-  uint32_t x = 0;
-  uint32_t y = 0;
-  uint32_t z = 0;
-  std::tie(x, y, z) = d->values();
-  EXPECT_EQ(x, 2u);
-  EXPECT_EQ(y, 4u);
-  EXPECT_EQ(z, 1u);
+  auto* d = WorkgroupSize(2, 4);
+  auto values = d->values();
+  ASSERT_NE(values[0], nullptr);
+  auto* x_scalar = values[0]->As<ast::ScalarConstructorExpression>();
+  ASSERT_TRUE(x_scalar);
+  ASSERT_TRUE(x_scalar->literal()->Is<ast::IntLiteral>());
+  EXPECT_EQ(x_scalar->literal()->As<ast::IntLiteral>()->value_as_u32(), 2u);
+
+  ASSERT_NE(values[1], nullptr);
+  auto* y_scalar = values[1]->As<ast::ScalarConstructorExpression>();
+  ASSERT_TRUE(y_scalar);
+  ASSERT_TRUE(y_scalar->literal()->Is<ast::IntLiteral>());
+  EXPECT_EQ(y_scalar->literal()->As<ast::IntLiteral>()->value_as_u32(), 4u);
+
+  EXPECT_EQ(values[2], nullptr);
 }
 
 TEST_F(WorkgroupDecorationTest, Creation_3param) {
-  auto* d = create<WorkgroupDecoration>(2, 4, 6);
-  uint32_t x = 0;
-  uint32_t y = 0;
-  uint32_t z = 0;
-  std::tie(x, y, z) = d->values();
-  EXPECT_EQ(x, 2u);
-  EXPECT_EQ(y, 4u);
-  EXPECT_EQ(z, 6u);
+  auto* d = WorkgroupSize(2, 4, 6);
+  auto values = d->values();
+  ASSERT_NE(values[0], nullptr);
+  auto* x_scalar = values[0]->As<ast::ScalarConstructorExpression>();
+  ASSERT_TRUE(x_scalar);
+  ASSERT_TRUE(x_scalar->literal()->Is<ast::IntLiteral>());
+  EXPECT_EQ(x_scalar->literal()->As<ast::IntLiteral>()->value_as_u32(), 2u);
+
+  ASSERT_NE(values[1], nullptr);
+  auto* y_scalar = values[1]->As<ast::ScalarConstructorExpression>();
+  ASSERT_TRUE(y_scalar);
+  ASSERT_TRUE(y_scalar->literal()->Is<ast::IntLiteral>());
+  EXPECT_EQ(y_scalar->literal()->As<ast::IntLiteral>()->value_as_u32(), 4u);
+
+  ASSERT_NE(values[2], nullptr);
+  auto* z_scalar = values[2]->As<ast::ScalarConstructorExpression>();
+  ASSERT_TRUE(z_scalar);
+  ASSERT_TRUE(z_scalar->literal()->Is<ast::IntLiteral>());
+  EXPECT_EQ(z_scalar->literal()->As<ast::IntLiteral>()->value_as_u32(), 6u);
+}
+
+TEST_F(WorkgroupDecorationTest, Creation_WithIdentifier) {
+  auto* d = WorkgroupSize(2, 4, "depth");
+  auto values = d->values();
+  ASSERT_NE(values[0], nullptr);
+  auto* x_scalar = values[0]->As<ast::ScalarConstructorExpression>();
+  ASSERT_TRUE(x_scalar);
+  ASSERT_TRUE(x_scalar->literal()->Is<ast::IntLiteral>());
+  EXPECT_EQ(x_scalar->literal()->As<ast::IntLiteral>()->value_as_u32(), 2u);
+
+  ASSERT_NE(values[1], nullptr);
+  auto* y_scalar = values[1]->As<ast::ScalarConstructorExpression>();
+  ASSERT_TRUE(y_scalar);
+  ASSERT_TRUE(y_scalar->literal()->Is<ast::IntLiteral>());
+  EXPECT_EQ(y_scalar->literal()->As<ast::IntLiteral>()->value_as_u32(), 4u);
+
+  ASSERT_NE(values[2], nullptr);
+  auto* z_ident = values[2]->As<ast::IdentifierExpression>();
+  ASSERT_TRUE(z_ident);
+  EXPECT_EQ(Symbols().NameFor(z_ident->symbol()), "depth");
 }
 
 TEST_F(WorkgroupDecorationTest, ToStr) {
-  auto* d = create<WorkgroupDecoration>(2, 4, 6);
-  EXPECT_EQ(str(d), R"(WorkgroupDecoration{2 4 6}
+  auto* d = WorkgroupSize(2, "height");
+  EXPECT_EQ(str(d), R"(WorkgroupDecoration{
+  ScalarConstructor[not set]{2}
+  Identifier[not set]{height}
+}
 )");
 }