[ir] Add function attributes

This CL adds the pipeline_stage and workgroup_size attributes into the
IR function.

Bug: tint:1915
Change-Id: I245dbf0104a1784cff364535106b3e520322ac73
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/130920
Commit-Queue: Dan Sinclair <dsinclair@chromium.org>
Reviewed-by: Ben Clayton <bclayton@google.com>
Reviewed-by: James Price <jrprice@google.com>
Kokoro: Kokoro <noreply+kokoro@google.com>

src/tint/ir/builder_impl.cc

@@ -65,6 +65,7 @@
 #include "src/tint/program.h"
 #include "src/tint/sem/builtin.h"
 #include "src/tint/sem/call.h"
+#include "src/tint/sem/function.h"
 #include "src/tint/sem/materialize.h"
 #include "src/tint/sem/module.h"
 #include "src/tint/sem/switch_statement.h"
@@ -212,6 +213,39 @@ void BuilderImpl::EmitFunction(const ast::Function* ast_func) {
 
     if (ast_func->IsEntryPoint()) {
         builder.ir.entry_points.Push(ir_func);
+
+        switch (ast_func->PipelineStage()) {
+            case ast::PipelineStage::kVertex:
+                ir_func->pipeline_stage = Function::PipelineStage::kVertex;
+                break;
+            case ast::PipelineStage::kFragment:
+                ir_func->pipeline_stage = Function::PipelineStage::kFragment;
+                break;
+            case ast::PipelineStage::kCompute: {
+                ir_func->pipeline_stage = Function::PipelineStage::kCompute;
+
+                const auto* sem = program_->Sem().Get(ast_func);
+                auto wg_size = sem->WorkgroupSize();
+
+                uint32_t x = wg_size[0].value();
+                uint32_t y = 1;
+                uint32_t z = 1;
+                if (wg_size[1].has_value()) {
+                    y = wg_size[1].value();
+
+                    if (wg_size[2].has_value()) {
+                        z = wg_size[2].value();
+                    }
+                }
+
+                ir_func->workgroup_size = {x, y, z};
+                break;
+            }
+            default: {
+                TINT_ICE(IR, diagnostics_) << "Invalid pipeline stage";
+                return;
+            }
+        }
     }
 
     {
@@ -222,7 +256,6 @@ void BuilderImpl::EmitFunction(const ast::Function* ast_func) {
 
         // TODO(dsinclair): Store return type and attributes
         // TODO(dsinclair): Store parameters
-        // TODO(dsinclair): Store attributes
 
         // If the branch target has already been set then a `return` was called. Only set in the
         // case where `return` wasn't called.

src/tint/ir/builder_impl_test.cc

@@ -88,7 +88,7 @@ TEST_F(IR_BuilderImplTest, IfStatement) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -136,7 +136,7 @@ TEST_F(IR_BuilderImplTest, IfStatement_TrueReturns) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(2u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -183,7 +183,7 @@ TEST_F(IR_BuilderImplTest, IfStatement_FalseReturns) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(2u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -230,7 +230,7 @@ TEST_F(IR_BuilderImplTest, IfStatement_BothReturn) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(2u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -273,7 +273,7 @@ TEST_F(IR_BuilderImplTest, IfStatement_JumpChainToMerge) {
     ASSERT_NE(loop_flow->continuing.target, nullptr);
     ASSERT_NE(loop_flow->merge.target, nullptr);
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -330,7 +330,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithBreak) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -388,7 +388,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithContinue) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -463,7 +463,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithContinuing_BreakIf) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -538,7 +538,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithReturn) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -595,7 +595,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithOnlyReturn) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -657,7 +657,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithOnlyReturn_ContinuingBreakIf) {
     // This is 1 because only the loop branch happens. The subsequent if return is dead code.
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -711,7 +711,7 @@ TEST_F(IR_BuilderImplTest, Loop_WithIf_BothBranchesBreak) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -857,7 +857,7 @@ TEST_F(IR_BuilderImplTest, Loop_Nested) {
     EXPECT_EQ(1u, func->start_target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -1001,7 +1001,7 @@ TEST_F(IR_BuilderImplTest, While) {
     EXPECT_EQ(1u, if_flow->false_.target->inbound_branches.Length());
     EXPECT_EQ(1u, if_flow->merge.target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -1071,7 +1071,7 @@ TEST_F(IR_BuilderImplTest, While_Return) {
     EXPECT_EQ(1u, if_flow->false_.target->inbound_branches.Length());
     EXPECT_EQ(1u, if_flow->merge.target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -1178,7 +1178,7 @@ TEST_F(IR_BuilderImplTest, For_NoInitCondOrContinuing) {
     EXPECT_EQ(1u, flow->merge.target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -1237,7 +1237,7 @@ TEST_F(IR_BuilderImplTest, Switch) {
     EXPECT_EQ(3u, flow->merge.target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -1301,7 +1301,7 @@ TEST_F(IR_BuilderImplTest, Switch_MultiSelector) {
     EXPECT_EQ(1u, flow->merge.target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -1345,7 +1345,7 @@ TEST_F(IR_BuilderImplTest, Switch_OnlyDefault) {
     EXPECT_EQ(1u, flow->merge.target->inbound_branches.Length());
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -1398,7 +1398,7 @@ TEST_F(IR_BuilderImplTest, Switch_WithBreak) {
     // This is 1 because the if is dead-code eliminated and the return doesn't happen.
     EXPECT_EQ(1u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -1457,7 +1457,7 @@ TEST_F(IR_BuilderImplTest, Switch_AllReturn) {
     EXPECT_EQ(0u, flow->merge.target->inbound_branches.Length());
     EXPECT_EQ(2u, func->end_target->inbound_branches.Length());
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   branch %fn2
 
@@ -1596,7 +1596,7 @@ TEST_F(IR_BuilderImplTest, Emit_Var_NoInit) {
     ASSERT_TRUE(r) << Error();
     auto m = r.Move();
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   %1(ref<function, u32, read_write>) = var function read_write
   ret
@@ -1614,7 +1614,7 @@ TEST_F(IR_BuilderImplTest, Emit_Var_Init) {
     ASSERT_TRUE(r) << Error();
     auto m = r.Move();
 
-    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn1 = block
   %1(ref<function, u32, read_write>) = var function read_write
   store %1(ref<function, u32, read_write>), 2u
@@ -1782,7 +1782,7 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Binary_LogicalAnd) {
   ret true
 func_end
 
-%fn2 = func test_function
+%fn2 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn3 = block
   %1(bool) = call my_func
   %2(bool) = var function read_write
@@ -1817,7 +1817,7 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Binary_LogicalOr) {
   ret true
 func_end
 
-%fn2 = func test_function
+%fn2 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn3 = block
   %1(bool) = call my_func
   %2(bool) = var function read_write
@@ -1999,7 +1999,7 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Binary_Compound) {
   ret 0.0f
 func_end
 
-%fn2 = func test_function
+%fn2 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn3 = block
   %1(f32) = call my_func
   %2(bool) = lt %1(f32), 2.0f
@@ -2041,7 +2041,7 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Binary_Compound_WithConstEval) {
   ret true
 func_end
 
-%fn2 = func test_function
+%fn2 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn3 = block
   %1(bool) = call my_func, false
   ret
@@ -2137,7 +2137,7 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Construct) {
 store %1(ref<private, f32, read_write>), 1.0f
 ret
 
-%fn1 = func test_function
+%fn1 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn2 = block
   %2(vec3<f32>) = construct 2.0f, 3.0f, %1(ref<private, f32, read_write>)
   ret
@@ -2161,7 +2161,7 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Convert) {
 store %1(ref<private, i32, read_write>), 1i
 ret
 
-%fn1 = func test_function
+%fn1 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn2 = block
   %2(f32) = convert i32, %1(ref<private, i32, read_write>)
   ret
@@ -2201,7 +2201,7 @@ TEST_F(IR_BuilderImplTest, EmitExpression_Builtin) {
 store %1(ref<private, f32, read_write>), 1.0f
 ret
 
-%fn1 = func test_function
+%fn1 = func test_function [@compute @workgroup_size(1, 1, 1)]
   %fn2 = block
   %2(f32) = asin %1(ref<private, f32, read_write>)
   ret
@@ -2210,5 +2210,54 @@ func_end
 )");
 }
 
+TEST_F(IR_BuilderImplTest, EmitFunction_Vertex) {
+    Func("test", utils::Empty, ty.vec4<f32>(), utils::Vector{Return(vec4<f32>(0_f, 0_f, 0_f, 0_f))},
+         utils::Vector{Stage(ast::PipelineStage::kVertex)},
+         utils::Vector{Builtin(builtin::BuiltinValue::kPosition)});
+
+    auto r = Build();
+    ASSERT_TRUE(r) << Error();
+    auto m = r.Move();
+
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test [@vertex]
+  %fn1 = block
+  ret vec4<f32> 0.0f
+func_end
+
+)");
+}
+
+TEST_F(IR_BuilderImplTest, EmitFunction_Fragment) {
+    Func("test", utils::Empty, ty.void_(), utils::Empty,
+         utils::Vector{Stage(ast::PipelineStage::kFragment)});
+
+    auto r = Build();
+    ASSERT_TRUE(r) << Error();
+    auto m = r.Move();
+
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test [@fragment]
+  %fn1 = block
+  ret
+func_end
+
+)");
+}
+
+TEST_F(IR_BuilderImplTest, EmitFunction_Compute) {
+    Func("test", utils::Empty, ty.void_(), utils::Empty,
+         utils::Vector{Stage(ast::PipelineStage::kCompute), WorkgroupSize(8_i, 4_i, 2_i)});
+
+    auto r = Build();
+    ASSERT_TRUE(r) << Error();
+    auto m = r.Move();
+
+    EXPECT_EQ(Disassemble(m), R"(%fn0 = func test [@compute @workgroup_size(8, 4, 2)]
+  %fn1 = block
+  ret
+func_end
+
+)");
+}
+
 }  // namespace
 }  // namespace tint::ir

src/tint/ir/disassembler.cc

@@ -92,7 +92,18 @@ void Disassembler::Walk(const FlowNode* node) {
         [&](const ir::Function* f) {
             TINT_SCOPED_ASSIGNMENT(in_function_, true);
 
-            Indent() << "%fn" << GetIdForNode(f) << " = func " << f->name.Name() << std::endl;
+            Indent() << "%fn" << GetIdForNode(f) << " = func " << f->name.Name();
+            if (f->pipeline_stage != Function::PipelineStage::kUndefined) {
+                out_ << " [@" << f->pipeline_stage;
+
+                if (f->workgroup_size) {
+                    auto arr = f->workgroup_size.value();
+                    out_ << " @workgroup_size(" << arr[0] << ", " << arr[1] << ", " << arr[2]
+                         << ")";
+                }
+                out_ << "]";
+            }
+            out_ << std::endl;
 
             {
                 ScopedIndent func_indent(&indent_size_);

src/tint/ir/function.cc

@@ -22,4 +22,18 @@ Function::Function() : Base() {}
 
 Function::~Function() = default;
 
+utils::StringStream& operator<<(utils::StringStream& out, Function::PipelineStage value) {
+    switch (value) {
+        case Function::PipelineStage::kVertex:
+            return out << "vertex";
+        case Function::PipelineStage::kFragment:
+            return out << "fragment";
+        case Function::PipelineStage::kCompute:
+            return out << "compute";
+        default:
+            break;
+    }
+    return out << "<unknown>";
+}
+
 }  // namespace tint::ir

src/tint/ir/function.h

@@ -15,6 +15,8 @@
 #ifndef SRC_TINT_IR_FUNCTION_H_
 #define SRC_TINT_IR_FUNCTION_H_
 
+#include <optional>
+
 #include "src/tint/ir/flow_node.h"
 #include "src/tint/symbol.h"
 
@@ -29,6 +31,18 @@ namespace tint::ir {
 /// An IR representation of a function
 class Function : public utils::Castable<Function, FlowNode> {
   public:
+    /// The pipeline stage for an entry point
+    enum class PipelineStage {
+        /// Not a pipeline entry point
+        kUndefined,
+        /// Vertex
+        kCompute,
+        /// Fragment
+        kFragment,
+        /// Vertex
+        kVertex,
+    };
+
     /// Constructor
     Function();
     ~Function() override;
@@ -36,6 +50,12 @@ class Function : public utils::Castable<Function, FlowNode> {
     /// The function name
     Symbol name;
 
+    /// The pipeline stage for the function, `kUndefined` if the function is not an entry point
+    PipelineStage pipeline_stage = PipelineStage::kUndefined;
+
+    /// If this is a `compute` entry point, holds the workgroup size information
+    std::optional<std::array<uint32_t, 3>> workgroup_size;
+
     /// The start target is the first block in a function.
     Block* start_target = nullptr;
     /// The end target is the end of the function. It is used as the branch target if a return is
@@ -43,6 +63,8 @@ class Function : public utils::Castable<Function, FlowNode> {
     Terminator* end_target = nullptr;
 };
 
+utils::StringStream& operator<<(utils::StringStream& out, Function::PipelineStage value);
+
 }  // namespace tint::ir
 
 #endif  // SRC_TINT_IR_FUNCTION_H_