Dawn/Native: Implement SetImmediateData() in D3D12 Backend

This CL implement SetImmediateData() API in D3D12 backend. User
could use SetImmediateData() in renderPass/ComputePass to upload
small constants.

Bug: 366291600
Change-Id: I5c52533f7b3b37292969b816f46d77b8585a834f
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/250974
Reviewed-by: Rafael Cintron <rafael.cintron@microsoft.com>
Commit-Queue: Shaobo Yan <shaoboyan@microsoft.com>
Reviewed-by: Corentin Wallez <cwallez@chromium.org>
diff --git a/src/dawn/native/ImmediateConstantsTracker.h b/src/dawn/native/ImmediateConstantsTracker.h
index adcba55..f8da372 100644
--- a/src/dawn/native/ImmediateConstantsTracker.h
+++ b/src/dawn/native/ImmediateConstantsTracker.h
@@ -71,6 +71,8 @@
     alignas(T) unsigned char mData[sizeof(T)] = {0};
 };
 
+// TODO(crbug.com/366291600): Add inheritance ability(like BindGroupTracker) so that it can inherit
+// immediate constants in native backend if supported.
 template <typename T, typename PipelineType>
 class UserImmediateConstantsTrackerBase {
   public:
diff --git a/src/dawn/native/d3d12/CommandBufferD3D12.cpp b/src/dawn/native/d3d12/CommandBufferD3D12.cpp
index 6ee2f7d..e20009e 100644
--- a/src/dawn/native/d3d12/CommandBufferD3D12.cpp
+++ b/src/dawn/native/d3d12/CommandBufferD3D12.cpp
@@ -37,6 +37,7 @@
 #include "dawn/native/CommandValidation.h"
 #include "dawn/native/DynamicUploader.h"
 #include "dawn/native/Error.h"
+#include "dawn/native/ImmediateConstantsTracker.h"
 #include "dawn/native/Queue.h"
 #include "dawn/native/RenderBundle.h"
 #include "dawn/native/d3d12/BindGroupD3D12.h"
@@ -403,6 +404,56 @@
     return {};
 }
 
+template <typename T>
+class ImmediateConstantTracker : public T {
+  public:
+    ImmediateConstantTracker() = default;
+
+    // Calling this after BindGroupTrackerBase::Apply() to update root signature.
+    void Apply(CommandRecordingContext* commandContext) {
+        auto* lastPipeline = this->mLastPipeline;
+        DAWN_ASSERT(lastPipeline != nullptr);
+
+        ImmediateConstantMask pipelineMask = lastPipeline->GetImmediateMask();
+        ImmediateConstantMask uploadBits = this->mDirty & pipelineMask;
+        for (auto&& [offset, size] : IterateRanges(uploadBits)) {
+            uint32_t immediateContentStartOffset =
+                static_cast<uint32_t>(offset) * kImmediateConstantElementByteSize;
+            uint32_t immediateRangeStartOffset =
+                GetImmediateIndexInPipeline(static_cast<uint32_t>(offset), pipelineMask);
+            SetRootConstant(commandContext->GetCommandList(),
+                            ToBackend(lastPipeline->GetLayout())->GetImmediatesParameterIndex(),
+                            size,
+                            this->mContent.template Get<uint32_t>(immediateContentStartOffset),
+                            immediateRangeStartOffset);
+        }
+
+        // Reset all dirty bits after uploading.
+        this->mDirty.reset();
+    }
+
+  private:
+    static constexpr bool kIsRenderImmediateConstants =
+        std::is_same_v<T, RenderImmediateConstantsTrackerBase>;
+    static constexpr bool kIsComputeImmediateConstants =
+        std::is_same_v<T, ComputeImmediateConstantsTrackerBase>;
+
+    void SetRootConstant(ID3D12GraphicsCommandList* commandList,
+                         uint32_t parameterIndex,
+                         uint32_t rootConstantsLength,
+                         const void* rootConstantsData,
+                         uint32_t registerOffset) const {
+        if constexpr (kIsRenderImmediateConstants) {
+            commandList->SetGraphicsRoot32BitConstants(parameterIndex, rootConstantsLength,
+                                                       rootConstantsData, registerOffset);
+        } else {
+            static_assert(kIsComputeImmediateConstants);
+            commandList->SetComputeRoot32BitConstants(parameterIndex, rootConstantsLength,
+                                                      rootConstantsData, registerOffset);
+        }
+    }
+};
+
 }  // anonymous namespace
 
 class DescriptorHeapState;
@@ -1225,6 +1276,7 @@
 
     Command type;
     ComputePipeline* lastPipeline = nullptr;
+    ImmediateConstantTracker<ComputeImmediateConstantsTrackerBase> immediates = {};
     while (mCommands.NextCommandId(&type)) {
         switch (type) {
             case Command::Dispatch: {
@@ -1239,6 +1291,7 @@
                 DAWN_TRY(TransitionAndClearForSyncScope(
                     commandContext, resourceUsages.dispatchUsages[currentDispatch]));
                 DAWN_TRY(bindingTracker->Apply(commandContext));
+                immediates.Apply(commandContext);
 
                 RecordNumWorkgroupsForDispatch(commandList, lastPipeline, dispatch);
                 commandList->Dispatch(dispatch->x, dispatch->y, dispatch->z);
@@ -1252,6 +1305,7 @@
                 DAWN_TRY(TransitionAndClearForSyncScope(
                     commandContext, resourceUsages.dispatchUsages[currentDispatch]));
                 DAWN_TRY(bindingTracker->Apply(commandContext));
+                immediates.Apply(commandContext);
 
                 ComPtr<ID3D12CommandSignature> signature =
                     lastPipeline->GetDispatchIndirectCommandSignature();
@@ -1282,6 +1336,7 @@
                 commandList->SetPipelineState(pipeline->GetPipelineState());
 
                 bindingTracker->OnSetPipeline(pipeline);
+                immediates.OnSetPipeline(pipeline);
                 lastPipeline = pipeline;
                 break;
             }
@@ -1300,6 +1355,15 @@
                 break;
             }
 
+            case Command::SetImmediateData: {
+                SetImmediateDataCmd* cmd = mCommands.NextCommand<SetImmediateDataCmd>();
+                DAWN_ASSERT(cmd->size > 0);
+                uint8_t* value = nullptr;
+                value = mCommands.NextData<uint8_t>(cmd->size);
+                immediates.SetImmediateData(cmd->offset, value, cmd->size);
+                break;
+            }
+
             case Command::InsertDebugMarker: {
                 InsertDebugMarkerCmd* cmd = mCommands.NextCommand<InsertDebugMarkerCmd>();
                 const char* label = mCommands.NextData<char>(cmd->length + 1);
@@ -1344,9 +1408,6 @@
                 break;
             }
 
-            case Command::SetImmediateData:
-                return DAWN_UNIMPLEMENTED_ERROR("SetImmediateData unimplemented");
-
             default:
                 DAWN_UNREACHABLE();
         }
@@ -1572,6 +1633,7 @@
 
     RenderPipeline* lastPipeline = nullptr;
     VertexBufferTracker vertexBufferTracker = {};
+    ImmediateConstantTracker<RenderImmediateConstantsTrackerBase> immediates = {};
 
     auto EncodeRenderBundleCommand = [&](CommandIterator* iter, Command type) -> MaybeError {
         switch (type) {
@@ -1582,6 +1644,7 @@
                 vertexBufferTracker.Apply(commandList, lastPipeline);
                 RecordFirstIndexOffset(commandList, lastPipeline, draw->firstVertex,
                                        draw->firstInstance);
+                immediates.Apply(commandContext);
                 commandList->DrawInstanced(draw->vertexCount, draw->instanceCount,
                                            draw->firstVertex, draw->firstInstance);
                 break;
@@ -1594,6 +1657,7 @@
                 vertexBufferTracker.Apply(commandList, lastPipeline);
                 RecordFirstIndexOffset(commandList, lastPipeline, draw->baseVertex,
                                        draw->firstInstance);
+                immediates.Apply(commandContext);
                 commandList->DrawIndexedInstanced(draw->indexCount, draw->instanceCount,
                                                   draw->firstIndex, draw->baseVertex,
                                                   draw->firstInstance);
@@ -1605,6 +1669,7 @@
 
                 DAWN_TRY(bindingTracker->Apply(commandContext));
                 vertexBufferTracker.Apply(commandList, lastPipeline);
+                immediates.Apply(commandContext);
 
                 Buffer* buffer = ToBackend(draw->indirectBuffer.Get());
                 ComPtr<ID3D12CommandSignature> signature =
@@ -1619,6 +1684,7 @@
 
                 DAWN_TRY(bindingTracker->Apply(commandContext));
                 vertexBufferTracker.Apply(commandList, lastPipeline);
+                immediates.Apply(commandContext);
 
                 Buffer* buffer = ToBackend(draw->indirectBuffer.Get());
                 DAWN_ASSERT(buffer != nullptr);
@@ -1635,6 +1701,7 @@
 
                 DAWN_TRY(bindingTracker->Apply(commandContext));
                 vertexBufferTracker.Apply(commandList, lastPipeline);
+                immediates.Apply(commandContext);
 
                 Buffer* indirectBuffer = ToBackend(draw->indirectBuffer.Get());
                 DAWN_ASSERT(indirectBuffer != nullptr);
@@ -1661,6 +1728,7 @@
 
                 DAWN_TRY(bindingTracker->Apply(commandContext));
                 vertexBufferTracker.Apply(commandList, lastPipeline);
+                immediates.Apply(commandContext);
 
                 Buffer* indirectBuffer = ToBackend(draw->indirectBuffer.Get());
                 DAWN_ASSERT(indirectBuffer != nullptr);
@@ -1725,6 +1793,7 @@
                 commandList->IASetPrimitiveTopology(pipeline->GetD3D12PrimitiveTopology());
 
                 bindingTracker->OnSetPipeline(pipeline);
+                immediates.OnSetPipeline(pipeline);
 
                 lastPipeline = pipeline;
                 break;
@@ -1744,6 +1813,15 @@
                 break;
             }
 
+            case Command::SetImmediateData: {
+                SetImmediateDataCmd* cmd = mCommands.NextCommand<SetImmediateDataCmd>();
+                DAWN_ASSERT(cmd->size > 0);
+                uint8_t* value = nullptr;
+                value = mCommands.NextData<uint8_t>(cmd->size);
+                immediates.SetImmediateData(cmd->offset, value, cmd->size);
+                break;
+            }
+
             case Command::SetIndexBuffer: {
                 SetIndexBufferCmd* cmd = iter->NextCommand<SetIndexBufferCmd>();
 
@@ -1873,9 +1951,6 @@
                 break;
             }
 
-            case Command::SetImmediateData:
-                return DAWN_UNIMPLEMENTED_ERROR("SetImmediateData unimplemented");
-
             default: {
                 DAWN_TRY(EncodeRenderBundleCommand(&mCommands, type));
                 break;
diff --git a/src/dawn/native/d3d12/PhysicalDeviceD3D12.cpp b/src/dawn/native/d3d12/PhysicalDeviceD3D12.cpp
index 7a9818a..9989ba9 100644
--- a/src/dawn/native/d3d12/PhysicalDeviceD3D12.cpp
+++ b/src/dawn/native/d3d12/PhysicalDeviceD3D12.cpp
@@ -339,6 +339,15 @@
         2 * limits->v1.maxDynamicUniformBuffersPerPipelineLayout -
         3 * limits->v1.maxDynamicStorageBuffersPerPipelineLayout;
 
+    // Report kMaxSupportedImmediateDataBytes if availableRootSignatureSlots is enough.
+    // Otherwise, reserve all available slots for immediates.
+    constexpr uint32_t kMaxSupportedImmediateDataSlots =
+        kMaxSupportedImmediateDataBytes / kImmediateConstantElementByteSize;
+    uint32_t maxImmediateDataSlots =
+        std::min(availableRootSignatureSlots, kMaxSupportedImmediateDataSlots);
+    availableRootSignatureSlots -= maxImmediateDataSlots;
+    limits->v1.maxImmediateSize = maxImmediateDataSlots * kImmediateConstantElementByteSize;
+
     while (availableRootSignatureSlots >= 2) {
         // Start by incrementing maxDynamicStorageBuffersPerPipelineLayout since the
         // default is just 4 and developers likely want more. This scheme currently
@@ -357,10 +366,11 @@
         }
     }
 
-    DAWN_ASSERT(2 * limits->v1.maxBindGroups +
-                    2 * limits->v1.maxDynamicUniformBuffersPerPipelineLayout +
-                    3 * limits->v1.maxDynamicStorageBuffersPerPipelineLayout <=
-                kMaxRootSignatureSize - kReservedSlots);
+    DAWN_ASSERT(
+        2 * limits->v1.maxBindGroups + 2 * limits->v1.maxDynamicUniformBuffersPerPipelineLayout +
+            3 * limits->v1.maxDynamicStorageBuffersPerPipelineLayout +
+            limits->v1.maxImmediateSize / kImmediateConstantElementByteSize + kReservedSlots ==
+        kMaxRootSignatureSize);
 
     // https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/sm5-attributes-numthreads
     limits->v1.maxComputeWorkgroupSizeX = D3D12_CS_THREAD_GROUP_MAX_X;
diff --git a/src/dawn/native/d3d12/PipelineLayoutD3D12.cpp b/src/dawn/native/d3d12/PipelineLayoutD3D12.cpp
index 304e7dd..4ab1670a 100644
--- a/src/dawn/native/d3d12/PipelineLayoutD3D12.cpp
+++ b/src/dawn/native/d3d12/PipelineLayoutD3D12.cpp
@@ -46,14 +46,19 @@
 // but are not directly related to allocation of the root signature.
 // In the root signature, it the index of the root parameter where these registers are
 // used that determines the layout of the root signature.
+// TODO(crbug.com/366291600): Use Immediates to support internal constants.
 static constexpr uint32_t kRenderOrComputeInternalRegisterSpace = kMaxBindGroups + 1;
 static constexpr uint32_t kRenderOrComputeInternalBaseRegister = 0;
 
 static constexpr uint32_t kDynamicStorageBufferLengthsRegisterSpace = kMaxBindGroups + 2;
 static constexpr uint32_t kDynamicStorageBufferLengthsBaseRegister = 0;
 
+static constexpr uint32_t kImmediatesRegisterSpace = kMaxBindGroups + 3;
+static constexpr uint32_t kImmediatesBaseRegister = 0;
+
 static constexpr uint32_t kInvalidDynamicStorageBufferLengthsParameterIndex =
     std::numeric_limits<uint32_t>::max();
+static constexpr uint32_t kInvalidImmediatesParameterIndex = std::numeric_limits<uint32_t>::max();
 
 D3D12_ROOT_PARAMETER_TYPE RootParameterType(wgpu::BufferBindingType type) {
     switch (type) {
@@ -328,6 +333,20 @@
             kInvalidDynamicStorageBufferLengthsParameterIndex;
     }
 
+    if (GetImmediateDataRangeByteSize() > 0) {
+        D3D12_ROOT_PARAMETER1 immediateConstants{};
+        immediateConstants.ShaderVisibility = D3D12_SHADER_VISIBILITY_ALL;
+        immediateConstants.ParameterType = D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS;
+        immediateConstants.Constants.Num32BitValues =
+            GetImmediateDataRangeByteSize() / sizeof(uint32_t);
+        immediateConstants.Constants.RegisterSpace = kImmediatesRegisterSpace;
+        immediateConstants.Constants.ShaderRegister = kImmediatesBaseRegister;
+        mImmediatesParameterIndex = rootParameters.size();
+        rootParameters.emplace_back(immediateConstants);
+    } else {
+        mImmediatesParameterIndex = kInvalidImmediatesParameterIndex;
+    }
+
     D3D12_VERSIONED_ROOT_SIGNATURE_DESC versionedRootSignatureDescriptor = {};
     versionedRootSignatureDescriptor.Version = D3D_ROOT_SIGNATURE_VERSION_1_1;
     versionedRootSignatureDescriptor.Desc_1_1.NumParameters = rootParameters.size();
@@ -468,6 +487,19 @@
     return mDynamicStorageBufferLengthsParameterIndex;
 }
 
+uint32_t PipelineLayout::GetImmediatesRegisterSpace() const {
+    return kImmediatesRegisterSpace;
+}
+
+uint32_t PipelineLayout::GetImmediatesShaderRegister() const {
+    return kImmediatesBaseRegister;
+}
+
+uint32_t PipelineLayout::GetImmediatesParameterIndex() const {
+    DAWN_ASSERT(mImmediatesParameterIndex != kInvalidImmediatesParameterIndex);
+    return mImmediatesParameterIndex;
+}
+
 ID3D12CommandSignature* PipelineLayout::GetDispatchIndirectCommandSignatureWithNumWorkgroups() {
     // mDispatchIndirectCommandSignatureWithNumWorkgroups won't be created until it is needed.
     if (mDispatchIndirectCommandSignatureWithNumWorkgroups.Get() != nullptr) {
diff --git a/src/dawn/native/d3d12/PipelineLayoutD3D12.h b/src/dawn/native/d3d12/PipelineLayoutD3D12.h
index 0ab5cf9..1e25ea8 100644
--- a/src/dawn/native/d3d12/PipelineLayoutD3D12.h
+++ b/src/dawn/native/d3d12/PipelineLayoutD3D12.h
@@ -65,6 +65,10 @@
     uint32_t GetDynamicStorageBufferLengthsShaderRegister() const;
     uint32_t GetDynamicStorageBufferLengthsParameterIndex() const;
 
+    uint32_t GetImmediatesRegisterSpace() const;
+    uint32_t GetImmediatesShaderRegister() const;
+    uint32_t GetImmediatesParameterIndex() const;
+
     ID3D12RootSignature* GetRootSignature() const;
 
     ID3DBlob* GetRootSignatureBlob() const;
@@ -110,6 +114,7 @@
     uint32_t mFirstIndexOffsetParameterIndex;
     uint32_t mNumWorkgroupsParameterIndex;
     uint32_t mDynamicStorageBufferLengthsParameterIndex;
+    uint32_t mImmediatesParameterIndex;
     ComPtr<ID3D12RootSignature> mRootSignature;
     // Store the root signature blob to put in pipeline cachekey
     ComPtr<ID3DBlob> mRootSignatureBlob;
diff --git a/src/dawn/native/d3d12/ShaderModuleD3D12.cpp b/src/dawn/native/d3d12/ShaderModuleD3D12.cpp
index 644951e..b86b894 100644
--- a/src/dawn/native/d3d12/ShaderModuleD3D12.cpp
+++ b/src/dawn/native/d3d12/ShaderModuleD3D12.cpp
@@ -350,6 +350,9 @@
     // read by the shader.
     req.hlsl.tintOptions.array_length_from_uniform = std::move(arrayLengthFromUniform);
 
+    req.hlsl.tintOptions.immediate_binding_point = tint::BindingPoint{
+        layout->GetImmediatesRegisterSpace(), layout->GetImmediatesShaderRegister()};
+
     if (stage == SingleShaderStage::Vertex) {
         // Now that only vertex shader can have interstage outputs.
         // Pass in the actually used interstage locations for tint to potentially truncate unused
diff --git a/src/dawn/tests/end2end/ImmediateDataTests.cpp b/src/dawn/tests/end2end/ImmediateDataTests.cpp
index c320111..2c8a188 100644
--- a/src/dawn/tests/end2end/ImmediateDataTests.cpp
+++ b/src/dawn/tests/end2end/ImmediateDataTests.cpp
@@ -485,7 +485,7 @@
     }
 }
 
-DAWN_INSTANTIATE_TEST(ImmediateDataTests, D3D11Backend(), VulkanBackend());
+DAWN_INSTANTIATE_TEST(ImmediateDataTests, D3D11Backend(), D3D12Backend(), VulkanBackend());
 
 }  // anonymous namespace
 }  // namespace dawn