[glsl] Add support for linear_indexing language feature

Fix: 482840477
Change-Id: I8a7197df562584a60ffcd6f2ed92ab9439cffa3e
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/293115
Reviewed-by: James Price <jrprice@google.com>
Commit-Queue: Alan Baker <alanbaker@google.com>
diff --git a/src/tint/lang/glsl/writer/raise/shader_io.cc b/src/tint/lang/glsl/writer/raise/shader_io.cc
index e8b14ae..96a6165 100644
--- a/src/tint/lang/glsl/writer/raise/shader_io.cc
+++ b/src/tint/lang/glsl/writer/raise/shader_io.cc
@@ -51,15 +51,27 @@
     /// The output variables.
     Vector<core::ir::Var*, 4> output_vars;
 
+    Vector<uint32_t, 4> input_indices;
+
     /// The original type of vertex input variables that we are bgra-swizzling. Keyed by location.
     Hashmap<uint32_t, const core::type::Type*, 1> bgra_swizzle_original_types;
 
     /// The configuration options.
     const ShaderIOConfig& config;
 
+    std::optional<uint32_t> global_invocation_index_index;
+    std::optional<uint32_t> global_invocation_id_index;
+    std::optional<uint32_t> workgroup_index_index;
+    std::optional<uint32_t> workgroup_id_index;
+    std::optional<uint32_t> num_workgroups_index;
+
     /// Constructor
     StateImpl(core::ir::Module& mod, core::ir::Function* f, const ShaderIOConfig& cfg)
-        : ShaderIOBackendState(mod, f), config(cfg) {}
+        : ShaderIOBackendState(mod, f), config(cfg) {
+        if (auto wgsize = func->WorkgroupSizeAsConst()) {
+            workgroup_size = wgsize;
+        }
+    }
 
     /// Destructor
     ~StateImpl() override {}
@@ -79,6 +91,7 @@
             StringStream name;
             name << ir.NameOf(func).Name();
 
+            uint32_t index = static_cast<uint32_t>(input_indices.Length());
             const core::type::MemoryView* ptr = nullptr;
             if (io.attributes.builtin) {
                 switch (io.attributes.builtin.value()) {
@@ -91,6 +104,28 @@
                     case core::BuiltinValue::kSampleIndex:
                         ptr = ty.ptr(addrspace, ty.i32(), access);
                         break;
+                    // Record an index for polyfilled inputs.
+                    case core::BuiltinValue::kGlobalInvocationIndex:
+                        global_invocation_index_index = index;
+                        input_indices.Push(index);
+                        continue;
+                    case core::BuiltinValue::kWorkgroupIndex:
+                        workgroup_index_index = index;
+                        input_indices.Push(index);
+                        continue;
+                    // Save the indices of the builtins below for use in polyfills.
+                    case core::BuiltinValue::kGlobalInvocationId:
+                        global_invocation_id_index = index;
+                        ptr = ty.ptr(addrspace, io.type, access);
+                        break;
+                    case core::BuiltinValue::kWorkgroupId:
+                        workgroup_id_index = index;
+                        ptr = ty.ptr(addrspace, io.type, access);
+                        break;
+                    case core::BuiltinValue::kNumWorkgroups:
+                        num_workgroups_index = index;
+                        ptr = ty.ptr(addrspace, io.type, access);
+                        break;
                     default:
                         ptr = ty.ptr(addrspace, io.type, access);
                         break;
@@ -123,12 +158,33 @@
             auto* var = b.Var(name.str(), ptr);
             var->SetAttributes(io.attributes);
             ir.root_block->Append(var);
+            input_indices.Push(static_cast<uint32_t>(vars.Length()));
             vars.Push(var);
         }
     }
 
     /// @copydoc ShaderIO::BackendState::FinalizeInputs
     Vector<core::ir::FunctionParam*, 4> FinalizeInputs() override {
+        // The following builtin values are polyfilled using other builtin values:
+        // * workgroup_index - workgroup_id and num_workgroups
+        // * global_invocation_index - global_invocation_id, num_workgroups (and workgroup size)
+        const bool has_global_invocation_index =
+            HasBuiltinInput(core::BuiltinValue::kGlobalInvocationIndex);
+        const bool has_workgroup_index = HasBuiltinInput(core::BuiltinValue::kWorkgroupIndex);
+        const bool needs_workgroup_id = has_workgroup_index;
+        if (needs_workgroup_id) {
+            RequireBuiltinInput(core::BuiltinValue::kWorkgroupId, ty.vec3u(), "workgroup_id");
+        }
+        const bool needs_num_workgroups = has_workgroup_index || has_global_invocation_index;
+        if (needs_num_workgroups) {
+            RequireBuiltinInput(core::BuiltinValue::kNumWorkgroups, ty.vec3u(), "num_workgroups");
+        }
+        const bool needs_global_invocation_id = has_global_invocation_index;
+        if (needs_global_invocation_id) {
+            RequireBuiltinInput(core::BuiltinValue::kGlobalInvocationId, ty.vec3u(),
+                                "global_invocation_id");
+        }
+
         MakeVars(input_vars, inputs, core::AddressSpace::kIn, core::Access::kRead, "_Input");
         return tint::Empty;
     }
@@ -141,7 +197,16 @@
 
     /// @copydoc ShaderIO::BackendState::GetInput
     core::ir::Value* GetInput(core::ir::Builder& builder, uint32_t idx) override {
-        auto* from = input_vars[idx]->Result();
+        if (idx == global_invocation_index_index) {
+            return PolyfillGlobalInvocationIndex(builder, global_invocation_id_index.value(),
+                                                 num_workgroups_index.value());
+        }
+        if (idx == workgroup_index_index) {
+            return PolyfillWorkgroupIndex(builder, workgroup_id_index.value(),
+                                          num_workgroups_index.value());
+        }
+        auto input_index = input_indices[idx];
+        auto* from = input_vars[input_index]->Result();
         auto* value = builder.Load(from)->Result();
 
         auto& builtin = inputs[idx].attributes.builtin;
diff --git a/src/tint/lang/glsl/writer/raise/shader_io_test.cc b/src/tint/lang/glsl/writer/raise/shader_io_test.cc
index 9a75674..472338a 100644
--- a/src/tint/lang/glsl/writer/raise/shader_io_test.cc
+++ b/src/tint/lang/glsl/writer/raise/shader_io_test.cc
@@ -1338,5 +1338,269 @@
     EXPECT_EQ(expect, str());
 }
 
+TEST_F(GlslWriter_ShaderIOTest, WorkgroupIndex_ReuseExistingBuiltins) {
+    auto* workgroup_id = b.FunctionParam("wgid", ty.vec3u());
+    workgroup_id->SetBuiltin(core::BuiltinValue::kWorkgroupId);
+
+    auto* num_workgroups = b.FunctionParam("numwgs", ty.vec3u());
+    num_workgroups->SetBuiltin(core::BuiltinValue::kNumWorkgroups);
+
+    auto* workgroup_index = b.FunctionParam("wgindex", ty.u32());
+    workgroup_index->SetBuiltin(core::BuiltinValue::kWorkgroupIndex);
+
+    auto* ep = b.ComputeFunction("foo", 3_u, 2_u, 1_u);
+    ep->SetParams({workgroup_id, num_workgroups, workgroup_index});
+    b.Append(ep->Block(), [&] {
+        b.Let("x", b.Add(workgroup_index, 0_u));
+        b.Return(ep);
+    });
+
+    auto* src = R"(
+%foo = @compute @workgroup_size(3u, 2u, 1u) func(%wgid:vec3<u32> [@workgroup_id], %numwgs:vec3<u32> [@num_workgroups], %wgindex:u32 [@workgroup_index]):void {
+  $B1: {
+    %5:u32 = add %wgindex, 0u
+    %x:u32 = let %5
+    ret
+  }
+}
+)";
+    EXPECT_EQ(src, str());
+
+    auto* expect = R"(
+$B1: {  # root
+  %foo_workgroup_id:ptr<__in, vec3<u32>, read> = var undef @builtin(workgroup_id)
+  %foo_num_workgroups:ptr<__in, vec3<u32>, read> = var undef @builtin(num_workgroups)
+}
+
+%foo_inner = func(%wgid:vec3<u32>, %numwgs:vec3<u32>, %wgindex:u32):void {
+  $B2: {
+    %7:u32 = add %wgindex, 0u
+    %x:u32 = let %7
+    ret
+  }
+}
+%foo = @compute @workgroup_size(3u, 2u, 1u) func():void {
+  $B3: {
+    %10:vec3<u32> = load %foo_workgroup_id
+    %11:vec3<u32> = load %foo_num_workgroups
+    %12:vec3<u32> = load %foo_workgroup_id
+    %13:vec3<u32> = load %foo_num_workgroups
+    %14:u32 = access %13, 0u
+    %15:u32 = access %13, 1u
+    %16:u32 = mul %14, %15
+    %17:u32 = access %12, 2u
+    %18:u32 = mul %17, %16
+    %19:u32 = access %12, 1u
+    %20:u32 = mul %19, %14
+    %21:u32 = access %12, 0u
+    %22:u32 = add %21, %20
+    %23:u32 = add %22, %18
+    %24:void = call %foo_inner, %10, %11, %23
+    ret
+  }
+}
+)";
+
+    core::ir::transform::ImmediateDataLayout immediate_data;
+    ShaderIOConfig config{immediate_data};
+    Run(ShaderIO, config);
+
+    EXPECT_EQ(expect, str());
+}
+
+TEST_F(GlslWriter_ShaderIOTest, WorkgroupIndex_AddMissingBuiltins) {
+    auto* workgroup_index = b.FunctionParam("wgindex", ty.u32());
+    workgroup_index->SetBuiltin(core::BuiltinValue::kWorkgroupIndex);
+
+    auto* ep = b.ComputeFunction("foo", 3_u, 2_u, 1_u);
+    ep->SetParams({workgroup_index});
+    b.Append(ep->Block(), [&] {
+        b.Let("x", b.Add(workgroup_index, 0_u));
+        b.Return(ep);
+    });
+
+    auto* src = R"(
+%foo = @compute @workgroup_size(3u, 2u, 1u) func(%wgindex:u32 [@workgroup_index]):void {
+  $B1: {
+    %3:u32 = add %wgindex, 0u
+    %x:u32 = let %3
+    ret
+  }
+}
+)";
+    EXPECT_EQ(src, str());
+
+    auto* expect = R"(
+$B1: {  # root
+  %foo_workgroup_id:ptr<__in, vec3<u32>, read> = var undef @builtin(workgroup_id)
+  %foo_num_workgroups:ptr<__in, vec3<u32>, read> = var undef @builtin(num_workgroups)
+}
+
+%foo_inner = func(%wgindex:u32):void {
+  $B2: {
+    %5:u32 = add %wgindex, 0u
+    %x:u32 = let %5
+    ret
+  }
+}
+%foo = @compute @workgroup_size(3u, 2u, 1u) func():void {
+  $B3: {
+    %8:vec3<u32> = load %foo_workgroup_id
+    %9:vec3<u32> = load %foo_num_workgroups
+    %10:u32 = access %9, 0u
+    %11:u32 = access %9, 1u
+    %12:u32 = mul %10, %11
+    %13:u32 = access %8, 2u
+    %14:u32 = mul %13, %12
+    %15:u32 = access %8, 1u
+    %16:u32 = mul %15, %10
+    %17:u32 = access %8, 0u
+    %18:u32 = add %17, %16
+    %19:u32 = add %18, %14
+    %20:void = call %foo_inner, %19
+    ret
+  }
+}
+)";
+
+    core::ir::transform::ImmediateDataLayout immediate_data;
+    ShaderIOConfig config{immediate_data};
+    Run(ShaderIO, config);
+
+    EXPECT_EQ(expect, str());
+}
+
+TEST_F(GlslWriter_ShaderIOTest, GlobalInvocationIndex_ReuseExistingBuiltins) {
+    auto* num_workgroups = b.FunctionParam("numwgs", ty.vec3u());
+    num_workgroups->SetBuiltin(core::BuiltinValue::kNumWorkgroups);
+
+    auto* global_index = b.FunctionParam("gindex", ty.u32());
+    global_index->SetBuiltin(core::BuiltinValue::kGlobalInvocationIndex);
+
+    auto* ep = b.ComputeFunction("foo", 3_u, 2_u, 1_u);
+    ep->SetParams({num_workgroups, global_index});
+    b.Append(ep->Block(), [&] {
+        b.Let("x", b.Add(global_index, 0_u));
+        b.Return(ep);
+    });
+
+    auto* src = R"(
+%foo = @compute @workgroup_size(3u, 2u, 1u) func(%numwgs:vec3<u32> [@num_workgroups], %gindex:u32 [@global_invocation_index]):void {
+  $B1: {
+    %4:u32 = add %gindex, 0u
+    %x:u32 = let %4
+    ret
+  }
+}
+)";
+    EXPECT_EQ(src, str());
+
+    auto* expect = R"(
+$B1: {  # root
+  %foo_num_workgroups:ptr<__in, vec3<u32>, read> = var undef @builtin(num_workgroups)
+  %foo_global_invocation_id:ptr<__in, vec3<u32>, read> = var undef @builtin(global_invocation_id)
+}
+
+%foo_inner = func(%numwgs:vec3<u32>, %gindex:u32):void {
+  $B2: {
+    %6:u32 = add %gindex, 0u
+    %x:u32 = let %6
+    ret
+  }
+}
+%foo = @compute @workgroup_size(3u, 2u, 1u) func():void {
+  $B3: {
+    %9:vec3<u32> = load %foo_num_workgroups
+    %10:vec3<u32> = load %foo_num_workgroups
+    %11:vec3<u32> = load %foo_global_invocation_id
+    %12:u32 = access %11, 0u
+    %13:u32 = access %11, 1u
+    %14:u32 = access %11, 2u
+    %15:u32 = access %10, 0u
+    %16:u32 = access %10, 1u
+    %17:u32 = mul %15, 3u
+    %18:u32 = mul %16, 2u
+    %19:u32 = mul %17, %18
+    %20:u32 = mul %14, %19
+    %21:u32 = mul %13, %17
+    %22:u32 = add %12, %21
+    %23:u32 = add %22, %20
+    %24:void = call %foo_inner, %9, %23
+    ret
+  }
+}
+)";
+
+    core::ir::transform::ImmediateDataLayout immediate_data;
+    ShaderIOConfig config{immediate_data};
+    Run(ShaderIO, config);
+
+    EXPECT_EQ(expect, str());
+}
+
+TEST_F(GlslWriter_ShaderIOTest, GlobalInvocationIndex_AddMissingBuiltins) {
+    auto* global_index = b.FunctionParam("gindex", ty.u32());
+    global_index->SetBuiltin(core::BuiltinValue::kGlobalInvocationIndex);
+
+    auto* ep = b.ComputeFunction("foo", 3_u, 2_u, 1_u);
+    ep->SetParams({global_index});
+    b.Append(ep->Block(), [&] {
+        b.Let("x", b.Add(global_index, 0_u));
+        b.Return(ep);
+    });
+
+    auto* src = R"(
+%foo = @compute @workgroup_size(3u, 2u, 1u) func(%gindex:u32 [@global_invocation_index]):void {
+  $B1: {
+    %3:u32 = add %gindex, 0u
+    %x:u32 = let %3
+    ret
+  }
+}
+)";
+    EXPECT_EQ(src, str());
+
+    auto* expect = R"(
+$B1: {  # root
+  %foo_num_workgroups:ptr<__in, vec3<u32>, read> = var undef @builtin(num_workgroups)
+  %foo_global_invocation_id:ptr<__in, vec3<u32>, read> = var undef @builtin(global_invocation_id)
+}
+
+%foo_inner = func(%gindex:u32):void {
+  $B2: {
+    %5:u32 = add %gindex, 0u
+    %x:u32 = let %5
+    ret
+  }
+}
+%foo = @compute @workgroup_size(3u, 2u, 1u) func():void {
+  $B3: {
+    %8:vec3<u32> = load %foo_num_workgroups
+    %9:vec3<u32> = load %foo_global_invocation_id
+    %10:u32 = access %9, 0u
+    %11:u32 = access %9, 1u
+    %12:u32 = access %9, 2u
+    %13:u32 = access %8, 0u
+    %14:u32 = access %8, 1u
+    %15:u32 = mul %13, 3u
+    %16:u32 = mul %14, 2u
+    %17:u32 = mul %15, %16
+    %18:u32 = mul %12, %17
+    %19:u32 = mul %11, %15
+    %20:u32 = add %10, %19
+    %21:u32 = add %20, %18
+    %22:void = call %foo_inner, %21
+    ret
+  }
+}
+)";
+
+    core::ir::transform::ImmediateDataLayout immediate_data;
+    ShaderIOConfig config{immediate_data};
+    Run(ShaderIO, config);
+
+    EXPECT_EQ(expect, str());
+}
+
 }  // namespace
 }  // namespace tint::glsl::writer::raise
diff --git a/src/tint/lang/glsl/writer/writer.cc b/src/tint/lang/glsl/writer/writer.cc
index d1eb0c7..c91bf33 100644
--- a/src/tint/lang/glsl/writer/writer.cc
+++ b/src/tint/lang/glsl/writer/writer.cc
@@ -198,13 +198,6 @@
         if (attributes.builtin == core::BuiltinValue::kCullDistance) {
             return Failure("cull_distance is not supported by the GLSL backend");
         }
-        if (attributes.builtin == core::BuiltinValue::kGlobalInvocationIndex) {
-            return Failure(
-                "@builtin(global_invocation_index) is not supported by the GLSL backend");
-        }
-        if (attributes.builtin == core::BuiltinValue::kWorkgroupIndex) {
-            return Failure("@builtin(workgroup_index) is not supported by the GLSL backend");
-        }
         if (attributes.color.has_value()) {
             return Failure("@color attribute is not supported by the GLSL backend");
         }