D3D11: Manage builtin variables with ImmediateConstantTracker

This CL replaces write to uniform buffer mechanism with
ImmediateConstantTracker to manage builtin variables in D3D11 backend.

It is the base to switch to use ImmediateConstants to support internal
immediate constants.

Bug:366291600

Change-Id: I2db6560cbcae29acd9fd088fcd8503d842433bf8
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/238474
Reviewed-by: Quyen Le <lehoangquyen@chromium.org>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
Commit-Queue: Shaobo Yan <shaoboyan@microsoft.com>
diff --git a/src/dawn/common/ityp_bitset.h b/src/dawn/common/ityp_bitset.h
index 74bdcfb..7938bc02 100644
--- a/src/dawn/common/ityp_bitset.h
+++ b/src/dawn/common/ityp_bitset.h
@@ -33,6 +33,7 @@
 #include <limits>
 
 #include "dawn/common/Assert.h"
+#include "dawn/common/BitSetRangeIterator.h"
 #include "dawn/common/Math.h"
 #include "dawn/common/Platform.h"
 #include "dawn/common/TypedInteger.h"
@@ -242,6 +243,10 @@
         return bitset(static_cast<const Base&>(lhs) ^ static_cast<const Base&>(rhs));
     }
 
+    friend BitSetRangeIterator<N, Index> IterateRanges(const bitset& bitset) {
+        return BitSetRangeIterator<N, Index>(static_cast<const Base&>(bitset));
+    }
+
     friend struct std::hash<bitset>;
 };
 
diff --git a/src/dawn/native/ImmediateConstantsLayout.h b/src/dawn/native/ImmediateConstantsLayout.h
index 7116001..bb6d682 100644
--- a/src/dawn/native/ImmediateConstantsLayout.h
+++ b/src/dawn/native/ImmediateConstantsLayout.h
@@ -115,6 +115,17 @@
            GetImmediateConstantBlockBits(offset, size).to_ulong();
 }
 
+template <typename Object, typename Member>
+std::optional<uint32_t> GetImmediateByteOffsetInPipelineIfAny(
+    Member Object::* ptr,
+    const ImmediateConstantMask& pipelineImmediateMask) {
+    if (!HasImmediateConstants(ptr, pipelineImmediateMask)) {
+        return std::nullopt;
+    }
+
+    return GetImmediateByteOffsetInPipeline(ptr, pipelineImmediateMask);
+}
+
 }  // namespace dawn::native
 
 #endif  // SRC_DAWN_NATIVE_IMMEDIATECONSTANTSLAYOUT_H_
diff --git a/src/dawn/native/d3d11/CommandBufferD3D11.cpp b/src/dawn/native/d3d11/CommandBufferD3D11.cpp
index e3a13a3..593834a 100644
--- a/src/dawn/native/d3d11/CommandBufferD3D11.cpp
+++ b/src/dawn/native/d3d11/CommandBufferD3D11.cpp
@@ -33,6 +33,7 @@
 #include <utility>
 #include <vector>
 
+#include "dawn/common/BitSetRangeIterator.h"
 #include "dawn/common/WindowsUtils.h"
 #include "dawn/native/ApplyClearColorValueWithDrawHelper.h"
 #include "dawn/native/ChainUtils.h"
@@ -40,6 +41,7 @@
 #include "dawn/native/CommandValidation.h"
 #include "dawn/native/Commands.h"
 #include "dawn/native/ExternalTexture.h"
+#include "dawn/native/ImmediateConstantsTracker.h"
 #include "dawn/native/RenderBundle.h"
 #include "dawn/native/d3d/D3DError.h"
 #include "dawn/native/d3d11/BindGroupTrackerD3D11.h"
@@ -216,6 +218,44 @@
     return pixelLocalStorageUAVs;
 }
 
+template <typename T>
+class ImmediateConstantTracker : public T {
+  public:
+    ImmediateConstantTracker() = default;
+
+    MaybeError Apply(const ScopedSwapStateCommandRecordingContext* commandContext) {
+        auto* lastPipeline = this->mLastPipeline;
+        if (!lastPipeline) {
+            return {};
+        }
+
+        ImmediateConstantMask uploadBits = this->mDirty & lastPipeline->GetImmediateMask();
+        uint32_t immediateRangeStartOffset = 0;
+        uint32_t immediateContentStartOffset = 0;
+        for (auto&& [offset, size] : IterateRanges(uploadBits)) {
+            immediateContentStartOffset =
+                static_cast<uint32_t>(offset) * kImmediateConstantElementByteSize;
+            commandContext->WriteUniformBufferRange(
+                immediateRangeStartOffset,
+                this->mContent.template Get<uint32_t>(immediateContentStartOffset),
+                size * kImmediateConstantElementByteSize);
+            immediateRangeStartOffset += size * kImmediateConstantElementByteSize;
+        }
+
+        // Reset all dirty bits after uploading.
+        this->mDirty.reset();
+
+        return commandContext->FlushUniformBuffer();
+    }
+
+    uint32_t GetFirstIndexContentStartOffset() {
+        uint32_t startIndex =
+            offsetof(RenderImmediateConstants, firstVertex) / kImmediateConstantElementByteSize;
+        ImmediateConstantMask prefixBits = ImmediateConstantMask((1u << startIndex) - 1u);
+        return (prefixBits & this->mDirty).count() * kImmediateConstantElementByteSize;
+    }
+};
+
 }  // namespace
 
 // Create CommandBuffer
@@ -496,6 +536,8 @@
     ComputePipeline* lastPipeline = nullptr;
     ComputePassBindGroupTracker bindGroupTracker(commandContext);
 
+    ImmediateConstantTracker<ComputeImmediateConstantsTrackerBase> immediates = {};
+
     Command type;
     while (mCommands.NextCommandId(&type)) {
         switch (type) {
@@ -508,8 +550,8 @@
                 DispatchCmd* dispatch = mCommands.NextCommand<DispatchCmd>();
 
                 DAWN_TRY(bindGroupTracker.Apply());
-
-                DAWN_TRY(RecordNumWorkgroupsForDispatch(lastPipeline, commandContext, dispatch));
+                immediates.SetNumWorkgroups(dispatch->x, dispatch->y, dispatch->z);
+                DAWN_TRY(immediates.Apply(commandContext));
                 commandContext->GetD3D11DeviceContext3()->Dispatch(dispatch->x, dispatch->y,
                                                                    dispatch->z);
 
@@ -520,8 +562,8 @@
                 DispatchIndirectCmd* dispatch = mCommands.NextCommand<DispatchIndirectCmd>();
 
                 DAWN_TRY(bindGroupTracker.Apply());
-
                 auto* indirectBuffer = ToGPUUsableBuffer(dispatch->indirectBuffer.Get());
+                DAWN_TRY(immediates.Apply(commandContext));
 
                 if (lastPipeline->UsesNumWorkgroups()) {
                     // Copy indirect args into the uniform buffer for built-in workgroup variables.
@@ -543,6 +585,7 @@
                 lastPipeline = ToBackend(cmd->pipeline).Get();
                 lastPipeline->ApplyNow(commandContext);
                 bindGroupTracker.OnSetPipeline(lastPipeline);
+                immediates.OnSetPipeline(lastPipeline);
                 break;
             }
 
@@ -676,6 +719,8 @@
     std::array<float, 4> blendColor = {0.0f, 0.0f, 0.0f, 0.0f};
     uint32_t stencilReference = 0;
 
+    ImmediateConstantTracker<RenderImmediateConstantsTrackerBase> immediates = {};
+
     auto DoRenderBundleCommand = [&](CommandIterator* iter, Command type) -> MaybeError {
         switch (type) {
             case Command::Draw: {
@@ -683,8 +728,8 @@
 
                 DAWN_TRY(bindGroupTracker.Apply());
                 vertexBufferTracker.Apply(lastPipeline);
-                DAWN_TRY(RecordFirstIndexOffset(lastPipeline, commandContext, draw->firstVertex,
-                                                draw->firstInstance));
+                immediates.SetFirstIndexOffset(draw->firstVertex, draw->firstInstance);
+                DAWN_TRY(immediates.Apply(commandContext));
                 commandContext->GetD3D11DeviceContext3()->DrawInstanced(
                     draw->vertexCount, draw->instanceCount, draw->firstVertex, draw->firstInstance);
 
@@ -696,8 +741,8 @@
 
                 DAWN_TRY(bindGroupTracker.Apply());
                 vertexBufferTracker.Apply(lastPipeline);
-                DAWN_TRY(RecordFirstIndexOffset(lastPipeline, commandContext, draw->baseVertex,
-                                                draw->firstInstance));
+                immediates.SetFirstIndexOffset(draw->baseVertex, draw->firstInstance);
+                DAWN_TRY(immediates.Apply(commandContext));
                 commandContext->GetD3D11DeviceContext3()->DrawIndexedInstanced(
                     draw->indexCount, draw->instanceCount, draw->firstIndex, draw->baseVertex,
                     draw->firstInstance);
@@ -713,6 +758,7 @@
 
                 DAWN_TRY(bindGroupTracker.Apply());
                 vertexBufferTracker.Apply(lastPipeline);
+                DAWN_TRY(immediates.Apply(commandContext));
 
                 if (lastPipeline->UsesVertexIndex() || lastPipeline->UsesInstanceIndex()) {
                     // Copy StartVertexLocation and StartInstanceLocation into the uniform buffer
@@ -722,7 +768,8 @@
                         offsetof(D3D11_DRAW_INSTANCED_INDIRECT_ARGS, StartVertexLocation);
                     DAWN_TRY(Buffer::Copy(commandContext, indirectBuffer, offset,
                                           sizeof(uint32_t) * 2,
-                                          commandContext->GetInternalUniformBuffer(), 0));
+                                          commandContext->GetInternalUniformBuffer(),
+                                          immediates.GetFirstIndexContentStartOffset()));
                 }
 
                 ID3D11Buffer* d3dBuffer;
@@ -742,6 +789,7 @@
 
                 DAWN_TRY(bindGroupTracker.Apply());
                 vertexBufferTracker.Apply(lastPipeline);
+                DAWN_TRY(immediates.Apply(commandContext));
 
                 if (lastPipeline->UsesVertexIndex() || lastPipeline->UsesInstanceIndex()) {
                     // Copy StartVertexLocation and StartInstanceLocation into the uniform buffer
@@ -751,7 +799,8 @@
                         offsetof(D3D11_DRAW_INDEXED_INSTANCED_INDIRECT_ARGS, BaseVertexLocation);
                     DAWN_TRY(Buffer::Copy(commandContext, indirectBuffer, offset,
                                           sizeof(uint32_t) * 2,
-                                          commandContext->GetInternalUniformBuffer(), 0));
+                                          commandContext->GetInternalUniformBuffer(),
+                                          immediates.GetFirstIndexContentStartOffset()));
                 }
 
                 ID3D11Buffer* d3dBuffer;
@@ -769,6 +818,7 @@
                 lastPipeline = ToBackend(cmd->pipeline.Get());
                 lastPipeline->ApplyNow(commandContext, blendColor, stencilReference);
                 bindGroupTracker.OnSetPipeline(lastPipeline);
+                immediates.OnSetPipeline(lastPipeline);
 
                 break;
             }
@@ -978,38 +1028,4 @@
     }
 }
 
-MaybeError CommandBuffer::RecordFirstIndexOffset(
-    RenderPipeline* renderPipeline,
-    const ScopedSwapStateCommandRecordingContext* commandContext,
-    uint32_t firstVertex,
-    uint32_t firstInstance) {
-    constexpr uint32_t kFirstVertexOffset = 0;
-    constexpr uint32_t kFirstInstanceOffset = 1;
-
-    if (renderPipeline->UsesVertexIndex()) {
-        commandContext->WriteUniformBuffer(kFirstVertexOffset, firstVertex);
-    }
-    if (renderPipeline->UsesInstanceIndex()) {
-        commandContext->WriteUniformBuffer(kFirstInstanceOffset, firstInstance);
-    }
-
-    return commandContext->FlushUniformBuffer();
-}
-
-MaybeError CommandBuffer::RecordNumWorkgroupsForDispatch(
-    ComputePipeline* computePipeline,
-    const ScopedSwapStateCommandRecordingContext* commandContext,
-    DispatchCmd* dispatchCmd) {
-    if (!computePipeline->UsesNumWorkgroups()) {
-        // Workgroup size is not used in shader, so we don't need to update the uniform buffer. The
-        // original value in the uniform buffer will not be used, so we don't need to clear it.
-        return {};
-    }
-
-    commandContext->WriteUniformBuffer(/*offset=*/0, dispatchCmd->x);
-    commandContext->WriteUniformBuffer(/*offset=*/1, dispatchCmd->y);
-    commandContext->WriteUniformBuffer(/*offset=*/2, dispatchCmd->z);
-    return commandContext->FlushUniformBuffer();
-}
-
 }  // namespace dawn::native::d3d11
diff --git a/src/dawn/native/d3d11/CommandBufferD3D11.h b/src/dawn/native/d3d11/CommandBufferD3D11.h
index beb201a..7975bf2 100644
--- a/src/dawn/native/d3d11/CommandBufferD3D11.h
+++ b/src/dawn/native/d3d11/CommandBufferD3D11.h
@@ -57,15 +57,6 @@
     void HandleDebugCommands(const ScopedSwapStateCommandRecordingContext* commandContext,
                              CommandIterator* iter,
                              Command command);
-
-    MaybeError RecordFirstIndexOffset(RenderPipeline* renderPipeline,
-                                      const ScopedSwapStateCommandRecordingContext* commandContext,
-                                      uint32_t firstVertex,
-                                      uint32_t firstInstance);
-    MaybeError RecordNumWorkgroupsForDispatch(
-        ComputePipeline* computePipeline,
-        const ScopedSwapStateCommandRecordingContext* commandContext,
-        DispatchCmd* dispatchCmd);
 };
 
 }  // namespace dawn::native::d3d11
diff --git a/src/dawn/native/d3d11/CommandRecordingContextD3D11.cpp b/src/dawn/native/d3d11/CommandRecordingContextD3D11.cpp
index aa7125e..26019d0 100644
--- a/src/dawn/native/d3d11/CommandRecordingContextD3D11.cpp
+++ b/src/dawn/native/d3d11/CommandRecordingContextD3D11.cpp
@@ -134,12 +134,13 @@
     return Get()->mD3D11DeviceContext3->Flush1(ContextType, hEvent);
 }
 
-void ScopedCommandRecordingContext::WriteUniformBuffer(uint32_t offset, uint32_t element) const {
-    DAWN_ASSERT(offset < CommandRecordingContext::kMaxNumBuiltinElements);
-    if (Get()->mUniformBufferData[offset] != element) {
-        Get()->mUniformBufferData[offset] = element;
-        Get()->mUniformBufferDirty = true;
-    }
+void ScopedCommandRecordingContext::WriteUniformBufferRange(uint32_t offset,
+                                                            const void* data,
+                                                            size_t size) const {
+    DAWN_ASSERT(offset < kMaxImmediateConstantsPerPipeline);
+    DAWN_ASSERT(size <= sizeof(uint32_t) * (kMaxImmediateConstantsPerPipeline - offset));
+    std::memcpy(&Get()->mUniformBufferData[offset], data, size);
+    Get()->mUniformBufferDirty = true;
 }
 
 MaybeError ScopedCommandRecordingContext::FlushUniformBuffer() const {
@@ -308,12 +309,12 @@
 // static
 ResultOrError<Ref<BufferBase>> CommandRecordingContext::CreateInternalUniformBuffer(
     DeviceBase* device) {
-    // Create a uniform buffer for built in variables.
+    // Create a uniform buffer for user and internal ImmediateConstants.
     BufferDescriptor descriptor;
-    descriptor.size = sizeof(uint32_t) * kMaxNumBuiltinElements;
+    descriptor.size = sizeof(uint32_t) * kMaxImmediateConstantsPerPipeline;
     descriptor.usage = wgpu::BufferUsage::Uniform | wgpu::BufferUsage::CopyDst;
     descriptor.mappedAtCreation = false;
-    descriptor.label = "BuiltinUniform";
+    descriptor.label = "ImmediateConstantsInternalBuffer";
 
     Ref<BufferBase> uniformBuffer;
     // Lock the device to protect the clearing of the built-in uniform buffer.
diff --git a/src/dawn/native/d3d11/CommandRecordingContextD3D11.h b/src/dawn/native/d3d11/CommandRecordingContextD3D11.h
index 3a8c5ff..97b5e47 100644
--- a/src/dawn/native/d3d11/CommandRecordingContextD3D11.h
+++ b/src/dawn/native/d3d11/CommandRecordingContextD3D11.h
@@ -30,6 +30,7 @@
 
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/inlined_vector.h"
+#include "dawn/common/Constants.h"
 #include "dawn/common/MutexProtected.h"
 #include "dawn/common/NonCopyable.h"
 #include "dawn/common/Ref.h"
@@ -107,11 +108,9 @@
     ComPtr<ID3D11Multithread> mD3D11Multithread;
     ComPtr<ID3DUserDefinedAnnotation> mD3DUserDefinedAnnotation;
 
-    // The maximum number of builtin elements is 4 (vec4). It must be multiple of 4.
-    static constexpr size_t kMaxNumBuiltinElements = 4;
     // The uniform buffer for built-in variables.
     Ref<GPUUsableBuffer> mUniformBuffer;
-    std::array<uint32_t, kMaxNumBuiltinElements> mUniformBufferData;
+    std::array<uint32_t, kMaxImmediateConstantsPerPipeline> mUniformBufferData{};
     bool mUniformBufferDirty = true;
 
     absl::flat_hash_set<Ref<d3d::KeyedMutex>> mAcquiredKeyedMutexes;
@@ -168,8 +167,8 @@
     void Flush() const;
     void Flush1(D3D11_CONTEXT_TYPE ContextType, HANDLE hEvent) const;
 
-    // Write the built-in variable value to the uniform buffer.
-    void WriteUniformBuffer(uint32_t offset, uint32_t element) const;
+    // Write immediate data to the uniform buffer.
+    void WriteUniformBufferRange(uint32_t offset, const void* data, size_t size) const;
     MaybeError FlushUniformBuffer() const;
 
     MaybeError AcquireKeyedMutex(Ref<d3d::KeyedMutex> keyedMutex) const;
diff --git a/src/dawn/native/d3d11/ComputePipelineD3D11.cpp b/src/dawn/native/d3d11/ComputePipelineD3D11.cpp
index f34a3e4..038c754 100644
--- a/src/dawn/native/d3d11/ComputePipelineD3D11.cpp
+++ b/src/dawn/native/d3d11/ComputePipelineD3D11.cpp
@@ -31,6 +31,7 @@
 #include <utility>
 
 #include "dawn/native/CreatePipelineAsyncEvent.h"
+#include "dawn/native/ImmediateConstantsLayout.h"
 #include "dawn/native/d3d/D3DError.h"
 #include "dawn/native/d3d11/DeviceD3D11.h"
 #include "dawn/native/d3d11/ShaderModuleD3D11.h"
@@ -51,6 +52,11 @@
     Device* device = ToBackend(GetDevice());
     uint32_t compileFlags = 0;
 
+    if (UsesNumWorkgroups()) {
+        mImmediateMask |= GetImmediateConstantBlockBits(
+            offsetof(ComputeImmediateConstants, numWorkgroups), sizeof(NumWorkgroupsDimensions));
+    }
+
     if (!device->IsToggleEnabled(Toggle::UseDXC) &&
         !device->IsToggleEnabled(Toggle::FxcOptimizations)) {
         compileFlags |= D3DCOMPILE_OPTIMIZATION_LEVEL0;
@@ -70,9 +76,10 @@
     }
 
     d3d::CompiledShader compiledShader;
-    DAWN_TRY_ASSIGN(compiledShader, ToBackend(programmableStage.module)
-                                        ->Compile(programmableStage, SingleShaderStage::Compute,
-                                                  ToBackend(GetLayout()), compileFlags));
+    DAWN_TRY_ASSIGN(compiledShader,
+                    ToBackend(programmableStage.module)
+                        ->Compile(programmableStage, SingleShaderStage::Compute,
+                                  ToBackend(GetLayout()), compileFlags, GetImmediateMask()));
     DAWN_TRY(CheckHRESULT(device->GetD3D11Device()->CreateComputeShader(
                               compiledShader.shaderBlob.Data(), compiledShader.shaderBlob.Size(),
                               nullptr, &mComputeShader),
diff --git a/src/dawn/native/d3d11/RenderPipelineD3D11.cpp b/src/dawn/native/d3d11/RenderPipelineD3D11.cpp
index afc1a43..fca67bd 100644
--- a/src/dawn/native/d3d11/RenderPipelineD3D11.cpp
+++ b/src/dawn/native/d3d11/RenderPipelineD3D11.cpp
@@ -36,6 +36,7 @@
 
 #include "dawn/common/Range.h"
 #include "dawn/native/CreatePipelineAsyncEvent.h"
+#include "dawn/native/ImmediateConstantsLayout.h"
 #include "dawn/native/d3d/D3DError.h"
 #include "dawn/native/d3d/ShaderUtils.h"
 #include "dawn/native/d3d11/DeviceD3D11.h"
@@ -234,6 +235,16 @@
       mD3DPrimitiveTopology(D3DPrimitiveTopology(GetPrimitiveTopology())) {}
 
 MaybeError RenderPipeline::InitializeImpl() {
+    // Set firstVertex and firstInstance bits together to ensure non immediate case has correct
+    // offset.
+    // TODO(crbug.com/366291600): Setting these bits respectively after immediate covers all cases.
+    if (UsesVertexIndex() || UsesInstanceIndex()) {
+        mImmediateMask |= GetImmediateConstantBlockBits(
+            offsetof(RenderImmediateConstants, firstVertex), kImmediateConstantElementByteSize);
+        mImmediateMask |= GetImmediateConstantBlockBits(
+            offsetof(RenderImmediateConstants, firstInstance), kImmediateConstantElementByteSize);
+    }
+
     DAWN_TRY(InitializeRasterizerState());
     DAWN_TRY(InitializeBlendState());
     DAWN_TRY(InitializeShaders());
@@ -462,18 +473,16 @@
             additionalCompileFlags |= D3DCOMPILE_IEEE_STRICTNESS;
         }
 
-        DAWN_TRY_ASSIGN(
-            compiledShader[SingleShaderStage::Vertex],
-            ToBackend(programmableStage.module)
-                ->Compile(programmableStage, SingleShaderStage::Vertex, ToBackend(GetLayout()),
-                          compileFlags | additionalCompileFlags, usedInterstageVariables));
+        DAWN_TRY_ASSIGN(compiledShader[SingleShaderStage::Vertex],
+                        ToBackend(programmableStage.module)
+                            ->Compile(programmableStage, SingleShaderStage::Vertex,
+                                      ToBackend(GetLayout()), compileFlags | additionalCompileFlags,
+                                      GetImmediateMask(), usedInterstageVariables));
         const Blob& shaderBlob = compiledShader[SingleShaderStage::Vertex].shaderBlob;
         DAWN_TRY(CheckHRESULT(device->GetD3D11Device()->CreateVertexShader(
                                   shaderBlob.Data(), shaderBlob.Size(), nullptr, &mVertexShader),
                               "D3D11 create vertex shader"));
         DAWN_TRY(InitializeInputLayout(shaderBlob));
-        mUsesVertexIndex = compiledShader[SingleShaderStage::Vertex].usesVertexIndex;
-        mUsesInstanceIndex = compiledShader[SingleShaderStage::Vertex].usesInstanceIndex;
     }
 
     std::optional<tint::hlsl::writer::PixelLocalOptions> pixelLocalOptions;
@@ -530,11 +539,12 @@
             additionalCompileFlags |= D3DCOMPILE_IEEE_STRICTNESS;
         }
 
-        DAWN_TRY_ASSIGN(compiledShader[SingleShaderStage::Fragment],
-                        ToBackend(programmableStage.module)
-                            ->Compile(programmableStage, SingleShaderStage::Fragment,
-                                      ToBackend(GetLayout()), compileFlags | additionalCompileFlags,
-                                      usedInterstageVariables, pixelLocalOptions));
+        DAWN_TRY_ASSIGN(
+            compiledShader[SingleShaderStage::Fragment],
+            ToBackend(programmableStage.module)
+                ->Compile(programmableStage, SingleShaderStage::Fragment, ToBackend(GetLayout()),
+                          compileFlags | additionalCompileFlags, GetImmediateMask(),
+                          usedInterstageVariables, pixelLocalOptions));
         DAWN_TRY(CheckHRESULT(device->GetD3D11Device()->CreatePixelShader(
                                   compiledShader[SingleShaderStage::Fragment].shaderBlob.Data(),
                                   compiledShader[SingleShaderStage::Fragment].shaderBlob.Size(),
diff --git a/src/dawn/native/d3d11/RenderPipelineD3D11.h b/src/dawn/native/d3d11/RenderPipelineD3D11.h
index 0ff9900..b2f5e32 100644
--- a/src/dawn/native/d3d11/RenderPipelineD3D11.h
+++ b/src/dawn/native/d3d11/RenderPipelineD3D11.h
@@ -54,9 +54,6 @@
     void ApplyDepthStencilState(const ScopedSwapStateCommandRecordingContext* commandContext,
                                 uint32_t stencilReference);
 
-    bool UsesVertexIndex() const { return mUsesVertexIndex; }
-    bool UsesInstanceIndex() const { return mUsesInstanceIndex; }
-
   private:
     RenderPipeline(Device* device, const UnpackedPtr<RenderPipelineDescriptor>& descriptor);
     ~RenderPipeline() override;
@@ -77,8 +74,6 @@
     ComPtr<ID3D11PixelShader> mPixelShader;
     ComPtr<ID3D11BlendState> mBlendState;
     ComPtr<ID3D11DepthStencilState> mDepthStencilState;
-    bool mUsesVertexIndex = false;
-    bool mUsesInstanceIndex = false;
 };
 
 }  // namespace dawn::native::d3d11
diff --git a/src/dawn/native/d3d11/ShaderModuleD3D11.cpp b/src/dawn/native/d3d11/ShaderModuleD3D11.cpp
index 4554437..d5c7073 100644
--- a/src/dawn/native/d3d11/ShaderModuleD3D11.cpp
+++ b/src/dawn/native/d3d11/ShaderModuleD3D11.cpp
@@ -33,6 +33,7 @@
 
 #include "dawn/common/Assert.h"
 #include "dawn/common/Log.h"
+#include "dawn/native/ImmediateConstantsLayout.h"
 #include "dawn/native/Pipeline.h"
 #include "dawn/native/TintUtils.h"
 #include "dawn/native/d3d/D3DCompilationRequest.h"
@@ -80,6 +81,7 @@
     SingleShaderStage stage,
     const PipelineLayout* layout,
     uint32_t compileFlags,
+    const ImmediateConstantMask& pipelineImmediateMask,
     const std::optional<dawn::native::d3d::InterStageShaderVariablesMask>& usedInterstageVariables,
     const std::optional<tint::hlsl::writer::PixelLocalOptions>& pixelLocalOptions) {
     Device* device = ToBackend(GetDevice());
@@ -193,23 +195,34 @@
     req.hlsl.inputProgram = UseTintProgram();
     req.hlsl.entryPointName = programmableStage.entryPoint.c_str();
     req.hlsl.stage = stage;
-    // Put the firstIndex into the internally reserved group and binding to avoid conflicting with
-    // any existing bindings.
-    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,
+    if (!useTintIR) {
+        if (stage == SingleShaderStage::Vertex) {
+            // 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 (entryPoint.usesNumWorkgroups) {
+            DAWN_ASSERT(stage == SingleShaderStage::Compute);
+            req.hlsl.tintOptions.root_constant_binding_point =
+                tint::BindingPoint{0, PipelineLayout::kNumWorkgroupsConstantBufferSlot};
         }
     }
 
@@ -224,14 +237,20 @@
         device->IsToggleEnabled(Toggle::DisableWorkgroupInit);
     req.hlsl.tintOptions.bindings = std::move(bindings);
 
-    if (entryPoint.usesNumWorkgroups) {
-        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};
+    // Immediate data available in TintIR only.
+    if (useTintIR) {
+        req.hlsl.tintOptions.immediate_binding_point =
+            tint::BindingPoint{0, PipelineLayout::kReservedConstantBufferSlot};
+        if (stage == SingleShaderStage::Compute) {
+            req.hlsl.tintOptions.num_workgroups_start_offset =
+                GetImmediateByteOffsetInPipelineIfAny(&ComputeImmediateConstants::numWorkgroups,
+                                                      pipelineImmediateMask);
+        } else {
+            req.hlsl.tintOptions.first_index_offset = GetImmediateByteOffsetInPipelineIfAny(
+                &RenderImmediateConstants::firstVertex, pipelineImmediateMask);
+            req.hlsl.tintOptions.first_instance_offset = GetImmediateByteOffsetInPipelineIfAny(
+                &RenderImmediateConstants::firstInstance, pipelineImmediateMask);
+        }
     }
 
     if (stage == SingleShaderStage::Vertex) {
diff --git a/src/dawn/native/d3d11/ShaderModuleD3D11.h b/src/dawn/native/d3d11/ShaderModuleD3D11.h
index 30ee9f3..eb0d88f 100644
--- a/src/dawn/native/d3d11/ShaderModuleD3D11.h
+++ b/src/dawn/native/d3d11/ShaderModuleD3D11.h
@@ -62,6 +62,7 @@
         SingleShaderStage stage,
         const PipelineLayout* layout,
         uint32_t compileFlags,
+        const ImmediateConstantMask& pipelineImmediateMask,
         const std::optional<dawn::native::d3d::InterStageShaderVariablesMask>&
             usedInterstageVariables = {},
         const std::optional<tint::hlsl::writer::PixelLocalOptions>& pixelLocalOptions = {});