HLSL-IR: implement FirstIndexOffset

The AST path implemented this as a separate transform, but for IR, this
is handled by the ShaderIO transform, which simplifies things as we only
need to add the offset once to the vertex_index and instance_index
parameters passed into the "inner" entry point function.

Bug: 380044486
Change-Id: Ie53edaa98ce3f3cae6d73bae7d1930b5aed7490d
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/219954
Commit-Queue: Antonio Maiorano <[email protected]>
Reviewed-by: James Price <[email protected]>
Reviewed-by: Kai Ninomiya <[email protected]>

src/dawn/native/d3d/ShaderUtils.cpp

@@ -208,9 +208,6 @@ ResultOrError<ComPtr<ID3DBlob>> CompileShaderFXC(const d3d::D3DBytecodeCompilati
 MaybeError TranslateToHLSL(d3d::HlslCompilationRequest r,
                            CacheKey::UnsafeUnkeyedValue<dawn::platform::Platform*> tracePlatform,
                            CompiledShader* compiledShader) {
-    std::ostringstream errorStream;
-    errorStream << "Tint HLSL failure:\n";
-
     tint::ast::transform::Manager transformManager;
     tint::ast::transform::DataMap transformInputs;
 
@@ -232,7 +229,7 @@ MaybeError TranslateToHLSL(d3d::HlslCompilationRequest r,
             std::move(requestedNames));
     }
 
-    if (r.stage == SingleShaderStage::Vertex) {
+    if (!r.useTintIR && r.stage == SingleShaderStage::Vertex) {
         transformManager.Add<tint::ast::transform::FirstIndexOffset>();
         transformInputs.Add<tint::ast::transform::FirstIndexOffset::BindingPoint>(
             r.firstIndexOffsetShaderRegister, r.firstIndexOffsetRegisterSpace);
@@ -257,7 +254,7 @@ MaybeError TranslateToHLSL(d3d::HlslCompilationRequest r,
 
     bool usesVertexIndex = false;
     bool usesInstanceIndex = false;
-    if (r.stage == SingleShaderStage::Vertex) {
+    if (!r.useTintIR && r.stage == SingleShaderStage::Vertex) {
         if (auto* data = transformOutputs.Get<tint::ast::transform::FirstIndexOffset::Data>()) {
             usesVertexIndex = data->has_vertex_index;
             usesInstanceIndex = data->has_instance_index;
@@ -301,6 +298,11 @@ MaybeError TranslateToHLSL(d3d::HlslCompilationRequest r,
     DAWN_INVALID_IF(result != tint::Success, "An error occurred while generating HLSL:\n%s",
                     result.Failure().reason.Str());
 
+    if (r.useTintIR && r.stage == SingleShaderStage::Vertex) {
+        usesVertexIndex = result->has_vertex_index;
+        usesInstanceIndex = result->has_instance_index;
+    }
+
     compiledShader->usesVertexIndex = usesVertexIndex;
     compiledShader->usesInstanceIndex = usesInstanceIndex;
     compiledShader->hlslSource = std::move(result->hlsl);

src/dawn/native/d3d11/ShaderModuleD3D11.cpp

@@ -88,13 +88,14 @@ ResultOrError<d3d::CompiledShader> ShaderModule::Compile(
 
     ScopedTintICEHandler scopedICEHandler(device);
     const EntryPointMetadata& entryPoint = GetEntryPoint(programmableStage.entryPoint);
+    const bool useTintIR = device->IsToggleEnabled(Toggle::UseTintIR);
 
     d3d::D3DCompilationRequest req = {};
     req.tracePlatform = UnsafeUnkeyedValue(device->GetPlatform());
     req.hlsl.shaderModel = 50;
     req.hlsl.disableSymbolRenaming = device->IsToggleEnabled(Toggle::DisableSymbolRenaming);
     req.hlsl.dumpShaders = device->IsToggleEnabled(Toggle::DumpShaders);
-    req.hlsl.useTintIR = device->IsToggleEnabled(Toggle::UseTintIR);
+    req.hlsl.useTintIR = useTintIR;
 
     req.bytecode.hasShaderF16Feature = false;
     req.bytecode.compileFlags = compileFlags;
@@ -195,19 +196,22 @@ ResultOrError<d3d::CompiledShader> ShaderModule::Compile(
     req.hlsl.stage = stage;
     // Put the firstIndex into the internally reserved group and binding to avoid conflicting with
     // any existing bindings.
-    req.hlsl.firstIndexOffsetRegisterSpace = PipelineLayout::kReservedConstantsBindGroupIndex;
-    req.hlsl.firstIndexOffsetShaderRegister = PipelineLayout::kFirstIndexOffsetBindingNumber;
-    // Remap to the desired space and binding, [0, kFirstIndexOffsetConstantBufferSlot].
-    {
-        tint::BindingPoint srcBindingPoint{req.hlsl.firstIndexOffsetRegisterSpace,
-                                           req.hlsl.firstIndexOffsetShaderRegister};
-        // D3D11 (HLSL SM5.0) doesn't support spaces, so we have to put the firstIndex in the
-        // default space(0)
-        tint::BindingPoint dstBindingPoint{0u, PipelineLayout::kFirstIndexOffsetConstantBufferSlot};
-
-        bindings.uniform.emplace(
-            srcBindingPoint,
-            tint::hlsl::writer::binding::Uniform{dstBindingPoint.group, dstBindingPoint.binding});
+    if (!useTintIR) {
+        req.hlsl.firstIndexOffsetRegisterSpace = PipelineLayout::kReservedConstantsBindGroupIndex;
+        req.hlsl.firstIndexOffsetShaderRegister = PipelineLayout::kFirstIndexOffsetBindingNumber;
+        // Remap to the desired space and binding, [0, kFirstIndexOffsetConstantBufferSlot].
+        {
+            tint::BindingPoint srcBindingPoint{req.hlsl.firstIndexOffsetRegisterSpace,
+                                               req.hlsl.firstIndexOffsetShaderRegister};
+            // D3D11 (HLSL SM5.0) doesn't support spaces, so we have to put the firstIndex in the
+            // default space(0)
+            tint::BindingPoint dstBindingPoint{0u,
+                                               PipelineLayout::kFirstIndexOffsetConstantBufferSlot};
+
+            bindings.uniform.emplace(srcBindingPoint,
+                                     tint::hlsl::writer::binding::Uniform{dstBindingPoint.group,
+                                                                          dstBindingPoint.binding});
+        }
     }
 
     req.hlsl.substituteOverrideConfig = std::move(substituteOverrideConfig);
@@ -222,10 +226,13 @@ ResultOrError<d3d::CompiledShader> ShaderModule::Compile(
     req.hlsl.tintOptions.bindings = std::move(bindings);
 
     if (entryPoint.usesNumWorkgroups) {
-        // D3D11 (HLSL SM5.0) doesn't support spaces, so we have to put the numWorkgroups in the
-        // default space(0)
+        DAWN_ASSERT(stage == SingleShaderStage::Compute);
         req.hlsl.tintOptions.root_constant_binding_point =
             tint::BindingPoint{0, PipelineLayout::kNumWorkgroupsConstantBufferSlot};
+    } else if (useTintIR && stage == SingleShaderStage::Vertex) {
+        // For vertex shaders, use root constant to add FirstIndexOffset, if needed
+        req.hlsl.tintOptions.root_constant_binding_point =
+            tint::BindingPoint{0, PipelineLayout::kFirstIndexOffsetConstantBufferSlot};
     }
 
     if (stage == SingleShaderStage::Vertex) {

src/dawn/native/d3d12/ShaderModuleD3D12.cpp

@@ -133,14 +133,15 @@ ResultOrError<d3d::CompiledShader> ShaderModule::Compile(
 
     ScopedTintICEHandler scopedICEHandler(device);
     const EntryPointMetadata& entryPoint = GetEntryPoint(programmableStage.entryPoint);
+    const bool useTintIR = device->IsToggleEnabled(Toggle::UseTintIR);
 
     d3d::D3DCompilationRequest req = {};
     req.tracePlatform = UnsafeUnkeyedValue(device->GetPlatform());
     req.hlsl.shaderModel = ToBackend(device->GetPhysicalDevice())
                                ->GetAppliedShaderModelUnderToggles(device->GetTogglesState());
     req.hlsl.disableSymbolRenaming = device->IsToggleEnabled(Toggle::DisableSymbolRenaming);
     req.hlsl.dumpShaders = device->IsToggleEnabled(Toggle::DumpShaders);
-    req.hlsl.useTintIR = device->IsToggleEnabled(Toggle::UseTintIR);
+    req.hlsl.useTintIR = useTintIR;
 
     req.bytecode.hasShaderF16Feature = device->HasFeature(Feature::ShaderF16);
     req.bytecode.compileFlags = compileFlags;
@@ -327,8 +328,10 @@ ResultOrError<d3d::CompiledShader> ShaderModule::Compile(
     req.hlsl.inputProgram = &(tintProgram->program);
     req.hlsl.entryPointName = programmableStage.entryPoint.c_str();
     req.hlsl.stage = stage;
-    req.hlsl.firstIndexOffsetShaderRegister = layout->GetFirstIndexOffsetShaderRegister();
-    req.hlsl.firstIndexOffsetRegisterSpace = layout->GetFirstIndexOffsetRegisterSpace();
+    if (!useTintIR) {
+        req.hlsl.firstIndexOffsetRegisterSpace = layout->GetFirstIndexOffsetRegisterSpace();
+        req.hlsl.firstIndexOffsetShaderRegister = layout->GetFirstIndexOffsetShaderRegister();
+    }
     req.hlsl.substituteOverrideConfig = std::move(substituteOverrideConfig);
 
     req.hlsl.tintOptions.disable_robustness = !device->IsRobustnessEnabled();
@@ -341,8 +344,14 @@ ResultOrError<d3d::CompiledShader> ShaderModule::Compile(
                                         : tint::hlsl::writer::Options::Compiler::kDXC;
 
     if (entryPoint.usesNumWorkgroups) {
+        DAWN_ASSERT(stage == SingleShaderStage::Compute);
         req.hlsl.tintOptions.root_constant_binding_point = tint::BindingPoint{
             layout->GetNumWorkgroupsRegisterSpace(), layout->GetNumWorkgroupsShaderRegister()};
+    } else if (useTintIR && stage == SingleShaderStage::Vertex) {
+        // For vertex shaders, use root constant to add FirstIndexOffset, if needed
+        req.hlsl.tintOptions.root_constant_binding_point =
+            tint::BindingPoint{layout->GetFirstIndexOffsetRegisterSpace(),
+                               layout->GetFirstIndexOffsetShaderRegister()};
     }
 
     // TODO(dawn:549): HLSL generation outputs the indices into the

src/tint/lang/hlsl/writer/common/output.h

@@ -76,6 +76,12 @@ struct Output {
 
     /// The workgroup size information, if the entry point was a compute shader
     WorkgroupInfo workgroup_info{};
+
+    /// True if the shader uses vertex_index
+    bool has_vertex_index = false;
+
+    /// True if the shader uses instance_index
+    bool has_instance_index = false;
 };
 
 }  // namespace tint::hlsl::writer

src/tint/lang/hlsl/writer/printer/printer.cc

@@ -1569,6 +1569,17 @@ class Printer : public tint::TextGenerator {
                         ++which_clip_distance;
                     } else {
                         name = builtin_to_attribute(builtin.value());
+
+                        switch (builtin.value()) {
+                            case core::BuiltinValue::kVertexIndex:
+                                result_.has_vertex_index = true;
+                                break;
+                            case core::BuiltinValue::kInstanceIndex:
+                                result_.has_instance_index = true;
+                                break;
+                            default:
+                                break;
+                        }
                     }
                     TINT_ASSERT(!name.empty());
 

src/tint/lang/hlsl/writer/raise/raise.cc

@@ -173,6 +173,7 @@ Result<SuccessType> Raise(core::ir::Module& module, const Options& options) {
     {
         raise::ShaderIOConfig config;
         config.num_workgroups_binding = options.root_constant_binding_point;
+        config.first_index_offset_binding = options.root_constant_binding_point;
         config.add_input_position_member = pixel_local_enabled;
         config.truncate_interstage_variables = options.truncate_interstage_variables;
         config.interstage_locations = std::move(options.interstage_locations);

src/tint/lang/hlsl/writer/raise/shader_io.cc

@@ -73,6 +73,10 @@ struct StateImpl : core::ir::transform::ShaderIOBackendState {
     std::optional<uint32_t> second_clip_distance_index;
     Hashset<uint32_t, 4> truncated_indices;
 
+    // If set, points to a var of type struct with fields for offsets to apply to vertex_index and
+    // instance_index
+    core::ir::Var* tint_first_index_offset = nullptr;
+
     /// Constructor
     StateImpl(core::ir::Module& mod, core::ir::Function* f, const ShaderIOConfig& c)
         : ShaderIOBackendState(mod, f), config(c) {}
@@ -199,6 +203,7 @@ struct StateImpl : core::ir::transform::ShaderIOBackendState {
         }
 
         Vector<MemberInfo, 4> input_data;
+        bool has_vertex_or_instance_index = false;
         for (uint32_t i = 0; i < inputs.Length(); ++i) {
             // Save the index of certain builtins for GetIndex. Although struct members will not be
             // added for these inputs, we still add entries to input_data so that other inputs with
@@ -210,6 +215,10 @@ struct StateImpl : core::ir::transform::ShaderIOBackendState {
                     subgroup_size_index = i;
                 } else if (*builtin == core::BuiltinValue::kNumWorkgroups) {
                     num_workgroups_index = i;
+                } else if (*builtin == core::BuiltinValue::kVertexIndex) {
+                    has_vertex_or_instance_index = true;
+                } else if (*builtin == core::BuiltinValue::kInstanceIndex) {
+                    has_vertex_or_instance_index = true;
                 }
             }
 
@@ -218,6 +227,22 @@ struct StateImpl : core::ir::transform::ShaderIOBackendState {
 
         input_indices.Resize(input_data.Length());
 
+        if (config.first_index_offset_binding.has_value() && has_vertex_or_instance_index) {
+            // Create a FirstIndexOffset uniform buffer. GetInput will use this to offset the
+            // vertex/instance index.
+            TINT_ASSERT(func->IsVertex());
+            tint::Vector<tint::core::type::Manager::StructMemberDesc, 2> members;
+            auto* str = ty.Struct(ir.symbols.New("tint_first_index_offset_struct"),
+                                  {
+                                      {ir.symbols.New("vertex_index"), ty.u32(), {}},
+                                      {ir.symbols.New("instance_index"), ty.u32(), {}},
+                                  });
+            tint_first_index_offset = b.Var("tint_first_index_offset", uniform, str);
+            tint_first_index_offset->SetBindingPoint(config.first_index_offset_binding->group,
+                                                     config.first_index_offset_binding->binding);
+            ir.root_block->Append(tint_first_index_offset);
+        }
+
         // Sort the struct members to satisfy HLSL interfacing matching rules.
         // We use stable_sort so that two members with the same attributes maintain their relative
         // ordering (e.g. kClipDistance).
@@ -406,12 +431,27 @@ struct StateImpl : core::ir::transform::ShaderIOBackendState {
 
         core::ir::Value* v = builder.Access(inputs[idx].type, input_param, u32(index))->Result(0);
 
-        // If this is an input position builtin we need to invert the 'w' component of the vector.
         if (inputs[idx].attributes.builtin == core::BuiltinValue::kPosition) {
+            // If this is an input position builtin we need to invert the 'w' component of the
+            // vector.
             auto* w = builder.Access(ty.f32(), v, 3_u);
             auto* div = builder.Divide(ty.f32(), 1.0_f, w);
             auto* swizzle = builder.Swizzle(ty.vec3<f32>(), v, {0, 1, 2});
-            v = builder.Construct(ty.vec4<f32>(), swizzle, div)->Results()[0];
+            v = builder.Construct(ty.vec4<f32>(), swizzle, div)->Result(0);
+        } else if (config.first_index_offset_binding.has_value() &&
+                   inputs[idx].attributes.builtin == core::BuiltinValue::kVertexIndex) {
+            // Apply vertex_index offset
+            TINT_ASSERT(tint_first_index_offset);
+            auto* vertex_index_offset =
+                builder.Access(ty.ptr<uniform, u32>(), tint_first_index_offset, 0_u);
+            v = builder.Add<u32>(v, builder.Load(vertex_index_offset))->Result(0);
+        } else if (config.first_index_offset_binding.has_value() &&
+                   inputs[idx].attributes.builtin == core::BuiltinValue::kInstanceIndex) {
+            // Apply instance_index offset
+            TINT_ASSERT(tint_first_index_offset);
+            auto* instance_index_offset =
+                builder.Access(ty.ptr<uniform, u32>(), tint_first_index_offset, 1_u);
+            v = builder.Add<u32>(v, builder.Load(instance_index_offset))->Result(0);
         }
 
         return v;

src/tint/lang/hlsl/writer/raise/shader_io.h

@@ -47,6 +47,12 @@ struct ShaderIOConfig {
     /// group plus 1 is used if at least one resource is bound, otherwise group 0 binding 0 is used.
     std::optional<BindingPoint> num_workgroups_binding;
 
+    /// The binding point to use for the first_index_offset uniform buffer. If set, and if a vertex
+    /// entry point contains a vertex_index or instance_index input parameter (or both), this
+    /// transform will add a uniform buffer with both indices, and will add the offsets to the input
+    /// variables, respectively.
+    std::optional<BindingPoint> first_index_offset_binding;
+
     /// If one doesn't exist, adds a @position member to the input struct as the last member.
     /// This is used for PixelLocal, for which Dawn requires such a member in the final HLSL shader.
     bool add_input_position_member = false;