diff --git a/include/API/Buffer.h b/include/API/Buffer.h index 43499bdd6..539bd3d9d 100644 --- a/include/API/Buffer.h +++ b/include/API/Buffer.h @@ -70,7 +70,7 @@ struct BufferCreateDesc { false}; } - static BufferCreateDesc scratchBuffer() { + static BufferCreateDesc gpuOnlyStorage() { return BufferCreateDesc{MemoryLocation::GpuOnly, MemoryBacking::Automatic, BufferUsage::Storage, @@ -78,6 +78,8 @@ struct BufferCreateDesc { {}, false}; } + + static BufferCreateDesc scratchBuffer() { return gpuOnlyStorage(); } }; class Buffer { diff --git a/include/API/CommandBuffer.h b/include/API/CommandBuffer.h index db7988af6..761d53825 100644 --- a/include/API/CommandBuffer.h +++ b/include/API/CommandBuffer.h @@ -26,6 +26,7 @@ namespace offloadtest { +class DescriptorPool; class RenderPass; class Texture; @@ -46,6 +47,8 @@ class CommandBuffer { GPUAPI getKind() const { return Kind; } + virtual void bindPool(const DescriptorPool &Pool) {}; + /// Create a compute command encoder for recording dispatch commands. /// Barriers are automatically inserted between commands. virtual llvm::Expected> diff --git a/include/API/DX/AccelerationStructure.h b/include/API/DX/AccelerationStructure.h new file mode 100644 index 000000000..6d05ce6b4 --- /dev/null +++ b/include/API/DX/AccelerationStructure.h @@ -0,0 +1,49 @@ +//===- DX/AccelerationStructure.h - Offload API DX AccelStruct API --------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_ACCELERATIONSTRUCTURE_H +#define OFFLOADTEST_API_DX_ACCELERATIONSTRUCTURE_H + +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +#include "API/AccelerationStructure.h" + +#include +#include + +namespace offloadtest { +using Microsoft::WRL::ComPtr; + +class DXAccelerationStructure : public offloadtest::AccelerationStructure { +public: + ComPtr Resource; + D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle; + + DXAccelerationStructure(ComPtr Resource, + const AccelerationStructureSizes &Sizes) + : offloadtest::AccelerationStructure(GPUAPI::DirectX, Sizes), + Resource(Resource) {} + + D3D12_GPU_VIRTUAL_ADDRESS getGPUVirtualAddress() const { + return Resource->GetGPUVirtualAddress(); + } + + static bool classof(const offloadtest::AccelerationStructure *AS) { + return AS->getAPI() == GPUAPI::DirectX; + } +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_ACCELERATIONSTRUCTURE_H diff --git a/include/API/DX/Buffer.h b/include/API/DX/Buffer.h new file mode 100644 index 000000000..cc009d9a5 --- /dev/null +++ b/include/API/DX/Buffer.h @@ -0,0 +1,73 @@ +//===- DX/Buffer.h - Offload API DX Buffer API ----------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_BUFFER_H +#define OFFLOADTEST_API_DX_BUFFER_H + +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +#include "API/Buffer.h" + +#include +#include + +namespace offloadtest { +using Microsoft::WRL::ComPtr; + +class DXBuffer : public offloadtest::Buffer { +public: + ComPtr Buffer; + std::string Name; + BufferCreateDesc Desc; + size_t SizeInBytes; + uint64_t CounterOffsetInBytes; + + // Contract: If a command on the command buffer needs a resource to be in a + // different state it should always transition it back to the PreferredState + // afterwards. The PreferredState is the state of the most common use case for + // that resource. This allows us to do state transitions without state + // tracking. + D3D12_RESOURCE_STATES PreferredState; + + D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle; + D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle; + D3D12_CPU_DESCRIPTOR_HANDLE CBVHandle; + + DXBuffer(ComPtr Buffer, llvm::StringRef Name, + BufferCreateDesc Desc, size_t SizeInBytes, + uint64_t CounterOffsetInBytes, D3D12_RESOURCE_STATES PreferredState, + D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle, + D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle, + D3D12_CPU_DESCRIPTOR_HANDLE CBVHandle); + DXBuffer(const DXBuffer &) = delete; + DXBuffer(DXBuffer &&) = delete; + DXBuffer &operator=(const DXBuffer &) = delete; + DXBuffer &operator=(DXBuffer &&) = delete; + + size_t getSizeInBytes() const override; + + size_t querySparseTileSizeInBytes(const Device & /*Dev*/) const override; + + llvm::Expected map() override; + + void unmap() override; + + const BufferCreateDesc &getDesc() const override; + + static bool classof(const offloadtest::Buffer *B); +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_BUFFER_H \ No newline at end of file diff --git a/include/API/DX/CommandBuffer.h b/include/API/DX/CommandBuffer.h new file mode 100644 index 000000000..f430a1ddc --- /dev/null +++ b/include/API/DX/CommandBuffer.h @@ -0,0 +1,77 @@ +//===- DX/CommandBuffer.h - Offload API DX CommandBuffer API --------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_COMMANDBUFFER_H +#define OFFLOADTEST_API_DX_COMMANDBUFFER_H + +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +#include "API/Buffer.h" +#include "API/CommandBuffer.h" +#include "API/DX/Common.h" +#include "API/DX/Encoder.h" + +#include "llvm/ADT/SmallVector.h" + +#include +#include + +namespace offloadtest { +using Microsoft::WRL::ComPtr; + +class DXDevice; + +class DXCommandBuffer : public offloadtest::CommandBuffer { +public: + ComPtr Allocator; + ComPtr CmdList; + /// Back-pointer to the owning device. Used by encoders that need access to + /// device-level resources (e.g. allocating AS scratch buffers). + DXDevice *Dev = nullptr; + /// Whether a UAV barrier is pending from a prior compute command. + bool PendingUAVBarrier = false; + llvm::SmallVector PendingTransitions; + /// Buffers that must outlive command-buffer submission (e.g. AS scratch + /// and TLAS instance buffers used during builds). + llvm::SmallVector> KeepAliveOwned; + + static llvm::Expected> + create(ComPtr Device); + + ~DXCommandBuffer() override = default; + + static bool classof(const CommandBuffer *CB); + + void addPendingUAVBarrier(); + void addResourceTransition(ID3D12Resource *Resource, + D3D12_RESOURCE_STATES StateBefore, + D3D12_RESOURCE_STATES StateAfter); + + void flushBarrier(); + + void bindPool(const DescriptorPool &Pool) override; + + llvm::Expected> + createComputeEncoder() override; + + llvm::Expected> + createRenderEncoder(const offloadtest::RenderPassBeginDesc &Desc) override; + +private: + DXCommandBuffer() : CommandBuffer(GPUAPI::DirectX) {} +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_COMMANDBUFFER_H diff --git a/include/API/DX/Common.h b/include/API/DX/Common.h new file mode 100644 index 000000000..7241925ba --- /dev/null +++ b/include/API/DX/Common.h @@ -0,0 +1,27 @@ +//===- DX/Common.h - Offload API DX API -----------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_COMMON_H +#define OFFLOADTEST_API_DX_COMMON_H + +#include +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +namespace offloadtest { +using ID3D12DeviceX = ID3D12Device5; +using ID3D12GraphicsCommandListX = ID3D12GraphicsCommandList6; +using Microsoft::WRL::ComPtr; +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_COMMON_H diff --git a/include/API/DX/Descriptors.h b/include/API/DX/Descriptors.h new file mode 100644 index 000000000..44aa15d28 --- /dev/null +++ b/include/API/DX/Descriptors.h @@ -0,0 +1,130 @@ +//===- DX/Descriptors.h - Offload API DX Descriptors API ------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_DESCRIPTORS_H +#define OFFLOADTEST_API_DX_DESCRIPTORS_H + +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +#include "API/DX/Common.h" +#include "API/Descriptors.h" + +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Error.h" + +#include + +namespace offloadtest { + +using Microsoft::WRL::ComPtr; + +struct DXDescriptorSet { + D3D12_CPU_DESCRIPTOR_HANDLE CSUHandle; + D3D12_GPU_DESCRIPTOR_HANDLE CSUHandleGPU; + + D3D12_CPU_DESCRIPTOR_HANDLE SamplerHandle; + D3D12_GPU_DESCRIPTOR_HANDLE SamplerHandleGPU; +}; + +class DXDescriptorSets : public DescriptorSets { +public: + llvm::SmallVector Sets; + + DXDescriptorSets(llvm::SmallVector Sets); + + static bool classof(const DescriptorSets *S); +}; + +class DXDescriptorPool : public DescriptorPool { +public: + ComPtr CSUHeap; + ComPtr SamplerHeap; + + uint32_t CSUIncSize; + uint32_t SamplerIncSize; + + D3D12_CPU_DESCRIPTOR_HANDLE CSUHandle; + D3D12_GPU_DESCRIPTOR_HANDLE CSUHandleGPU; + + D3D12_CPU_DESCRIPTOR_HANDLE SamplerHandle; + D3D12_GPU_DESCRIPTOR_HANDLE SamplerHandleGPU; + + std::atomic CSUAllocator = 0; + std::atomic SamplerAllocator = 0; + + DXDescriptorPool(ComPtr CSUHeap, + ComPtr SamplerHeap, + uint32_t CSUIncSize, uint32_t SamplerIncSize); + + static llvm::Expected> + create(ComPtr Dev); + + void allocateDescriptors(uint32_t Count, D3D12_CPU_DESCRIPTOR_HANDLE &CPU, + D3D12_GPU_DESCRIPTOR_HANDLE &GPU); + + void allocateSamplers(uint32_t Count, D3D12_CPU_DESCRIPTOR_HANDLE &CPU, + D3D12_GPU_DESCRIPTOR_HANDLE &GPU); + + void reset() override; + + static bool classof(const DescriptorPool *P); +}; + +class DXDescriptorSetsBuilder : public DescriptorSetsBuilder { +public: + struct SetState { + D3D12_CPU_DESCRIPTOR_HANDLE CSUHandle; + D3D12_CPU_DESCRIPTOR_HANDLE SamplerHandle; + }; + + ComPtr Dev; + llvm::SmallVector Sets; + llvm::SmallVector SetStates; + uint32_t CSUIncSize; + uint32_t SamplerIncSize; + + DXDescriptorSetsBuilder(ComPtr Dev, + llvm::SmallVector Sets, + uint32_t CSUIncSize, uint32_t SamplerIncSize); + + DescriptorSetsBuilder &constant(uint32_t SetIndex, + llvm::ArrayRef B, + VKBind) override; + + DescriptorSetsBuilder & + read(uint32_t SetIndex, llvm::ArrayRef B, VKBind) override; + DescriptorSetsBuilder &read(uint32_t SetIndex, + llvm::ArrayRef T, + llvm::ArrayRef S, + VKBind) override; + DescriptorSetsBuilder &read(uint32_t SetIndex, + llvm::ArrayRef A, + VKBind) override; + + DescriptorSetsBuilder & + write(uint32_t SetIndex, llvm::ArrayRef B, VKBind) override; + DescriptorSetsBuilder & + write(uint32_t SetIndex, llvm::ArrayRef T, VKBind) override; + + DescriptorSetsBuilder &sampler(uint32_t SetIndex, + llvm::ArrayRef S, + VKBind) override; + + std::unique_ptr build() override; + + static bool classof(const DescriptorSetsBuilder *B); +}; +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_DESCRIPTORS_H diff --git a/include/API/DX/Device.h b/include/API/DX/Device.h new file mode 100644 index 000000000..fcfa026af --- /dev/null +++ b/include/API/DX/Device.h @@ -0,0 +1,184 @@ +//===- DX/Device.h - Offload API DX Device API ----------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_DEVICE_H +#define OFFLOADTEST_API_DX_DEVICE_H + +#include "API/DX/Common.h" +#include "API/DX/Queue.h" +#include "API/Device.h" + +#include +#include + +namespace offloadtest { + +struct DescriptorSetsLayout; + +struct DescriptorAllocator { + ComPtr Heap; + std::atomic NextIndex{0}; + uint32_t DescIncSize; + uint32_t Capacity; + + static llvm::Expected + create(ID3D12DeviceX *Device, D3D12_DESCRIPTOR_HEAP_TYPE Type, + uint32_t Capacity); + + llvm::Expected allocate(); + + DescriptorAllocator(DescriptorAllocator &&Other) + : Heap(std::move(Other.Heap)), + NextIndex(Other.NextIndex.load(std::memory_order_relaxed)), + DescIncSize(Other.DescIncSize), Capacity(Other.Capacity) {} + DescriptorAllocator(const DescriptorAllocator &) = delete; + DescriptorAllocator &operator=(const DescriptorAllocator &) = delete; + DescriptorAllocator &operator=(DescriptorAllocator &&) = delete; + + DescriptorAllocator(ComPtr Heap, uint32_t DescIncSize, + uint32_t Capacity) + : Heap(Heap), DescIncSize(DescIncSize), Capacity(Capacity) {} +}; + +class DXDevice : public offloadtest::Device { +public: + ComPtr Adapter; + ComPtr Device; + DXQueue GraphicsQueue; + Capabilities Caps; + DescriptorAllocator RTVAllocator; + DescriptorAllocator DSVAllocator; + DescriptorAllocator CSUAllocator; + DescriptorAllocator SamplerAllocator; + + DXDevice(ComPtr A, ComPtr D, DXQueue Q, + DescriptorAllocator RTVAllocator, DescriptorAllocator DSVAllocator, + DescriptorAllocator CSUAllocator, + DescriptorAllocator SamplerAllocator, std::string Desc, + std::string DriverVer); + DXDevice(const DXDevice &) = delete; + DXDevice &operator=(const DXDevice &) = delete; + + ~DXDevice() override; + + llvm::StringRef getAPIName() const override; + GPUAPI getAPI() const override; + + static bool classof(const offloadtest::Device *D); + + Queue &getGraphicsQueue() override; + + llvm::Error + createRootSignatureFromShader(llvm::StringRef Name, + const ShaderContainer &Shader, + ComPtr &OutRootSignature, + DescriptorSetsLayout &Layout); + + llvm::Error createRootSignatureFromBindingsDesc( + llvm::StringRef Name, const BindingsDesc &BndDesc, bool IsGraphics, + ComPtr &OutRootSignature, + DescriptorSetsLayout &Layout); + + llvm::Error createRootSignature(llvm::StringRef Name, + const BindingsDesc &BndDesc, + const ShaderContainer &Shader, + bool IsGraphics, + ComPtr &OutRootSignature, + DescriptorSetsLayout &Layout); + + llvm::Expected> + createPipelineCs(llvm::StringRef Name, const BindingsDesc &BndDesc, + ShaderContainer CS) override; + + llvm::Expected> + createTraditionalRasterPipeline( + llvm::StringRef Name, const BindingsDesc &BndDesc, + const TraditionalRasterPipelineCreateDesc &Desc) override; + + llvm::Expected> createMeshShaderRasterPipeline( + llvm::StringRef Name, const BindingsDesc &BindingsDesc, + const MeshShaderRasterPipelineCreateDesc &Desc) override; + + llvm::Expected> + createPipelineRT(llvm::StringRef Name, const BindingsDesc &BndDesc, + const RayTracingPipelineCreateDesc &Desc) override; + + llvm::Expected> + createShaderBindingTable(const PipelineState &PSO, + const ShaderBindingTableDesc &Desc) override; + + llvm::Expected> + createFence(llvm::StringRef Name) override; + + llvm::Expected> + createMemoryHeap(std::string Name, size_t SizeInBytes) override; + + llvm::Expected> + createBuffer(std::string Name, const BufferCreateDesc &Desc, + size_t SizeInBytes) override; + + llvm::Expected> + createTexture(std::string Name, const TextureCreateDesc &Desc) override; + + llvm::Expected> + createSampler(std::string Name, const SamplerCreateDesc &Desc) override; + + uint32_t getTextureUploadRowStrideInBytes( + const TextureCreateDesc &Desc) const override; + + TextureUploadLayout + getTextureUploadLayout(const TextureCreateDesc &Desc) const override; + + static llvm::Expected> + create(ComPtr Adapter, const DeviceConfig &Config); + + const Capabilities &getCapabilities() override; + + void queryCapabilities(); + + llvm::Expected> + createCommandBuffer() override; + + llvm::Expected> + createRenderPass(const offloadtest::RenderPassDesc &Desc) override; + + llvm::Expected + getBLASBuildSizes(llvm::ArrayRef Triangles) override; + + llvm::Expected + getBLASBuildSizes(llvm::ArrayRef AABBs) override; + + AccelerationStructureSizes queryBLASPrebuildSize( + llvm::ArrayRef GeomDescs); + + llvm::Expected> + allocateAS(const AccelerationStructureSizes &Sizes, const char *Kind); + + llvm::Expected + getTLASBuildSizes(uint32_t InstanceCount) override; + + llvm::Expected> + createBLAS(const AccelerationStructureSizes &Sizes) override; + + llvm::Expected> + createTLAS(const AccelerationStructureSizes &Sizes, + uint32_t /*InstanceCount*/) override; + + llvm::Expected> + createDescriptorPool() override; + + llvm::Expected> + createDescriptorSetsBuilder(DescriptorPool &Pool, + const PipelineState &Pipeline) override; + + llvm::Error executeProgram(Pipeline &P) override; +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_DEVICE_H diff --git a/include/API/DX/Encoder.h b/include/API/DX/Encoder.h new file mode 100644 index 000000000..f87959d79 --- /dev/null +++ b/include/API/DX/Encoder.h @@ -0,0 +1,139 @@ +//===- DX/Encoder.h - Offload API DX Encoder API --------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_ENCODER_H +#define OFFLOADTEST_API_DX_ENCODER_H + +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +#include "API/CommandBuffer.h" +#include "API/Encoder.h" + +#include +#include + +namespace offloadtest { +using Microsoft::WRL::ComPtr; + +class DXCommandBuffer; + +class DXComputeEncoder : public offloadtest::ComputeEncoder { + DXCommandBuffer &CB; + + void addUAVBarrier(); + +public: + DXComputeEncoder(DXCommandBuffer &CB) + : ComputeEncoder(GPUAPI::DirectX), CB(CB) {} + + ~DXComputeEncoder() override; + + static bool classof(const CommandEncoder *E) { + return E->getAPI() == GPUAPI::DirectX; + } + + // D3D12 debug labels require WinPixEventRuntime for the proper event + // encoding. Without it, BeginEvent/EndEvent/SetMarker with metadata type 0 + // crash the D3D12 debug layer, so leave these as no-ops for now. + void pushDebugGroup(llvm::StringRef Label) override {} + void popDebugGroup() override {} + void insertDebugSignpost(llvm::StringRef Label) override {} + + void bindDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets); + + void bindComputeDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) override; + + void bindRayTracingDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) override; + + llvm::Error dispatch(const offloadtest::PipelineState &PSO, + uint32_t GroupCountX, uint32_t GroupCountY, + uint32_t GroupCountZ) override; + + llvm::Error copyBufferToBuffer(offloadtest::Buffer &Src, size_t SrcOffset, + offloadtest::Buffer &Dst, size_t DstOffset, + size_t Size) override; + + llvm::Error copyBufferToTexture(Buffer &Src, Texture &Dst) override; + + llvm::Error copyCounterToBuffer(offloadtest::Buffer &Src, + offloadtest::Buffer &Dst) override; + + llvm::Error copyTextureToBuffer(Texture &Src, Buffer &Dst) override; + + // Defined out-of-line below — needs DXDevice's full type for access to the + // ID3D12Device5 entry point and helper allocators. + llvm::Error batchBuildAS(llvm::ArrayRef Items) override; + + // Defined out-of-line below — needs DXDevice's full type for access to + // Device5 and the DXRayTracingPipelineState definition. + llvm::Error dispatchRays(const PipelineState &PSO, + const ShaderBindingTable &SBT, uint32_t Width, + uint32_t Height, uint32_t Depth) override; + + void endEncodingImpl() override; +}; + +class DXRenderEncoder : public offloadtest::RenderEncoder { + DXCommandBuffer &CB; + offloadtest::RenderPassBeginDesc Desc; + + // Encoder contract: viewport and scissor must both be set before draw(). + bool ViewportSet = false; + bool ScissorSet = false; + + llvm::Error bindCommonDrawState(const offloadtest::PipelineState &PSO); + +public: + DXRenderEncoder(DXCommandBuffer &CB, + const offloadtest::RenderPassBeginDesc &Desc); + DXRenderEncoder(const DXRenderEncoder &CB) = delete; + DXRenderEncoder(DXRenderEncoder &&CB) = delete; + DXRenderEncoder &operator=(DXRenderEncoder &CB) = delete; + DXRenderEncoder &operator=(const DXRenderEncoder &&CB) = delete; + + ~DXRenderEncoder() override; + + static bool classof(const CommandEncoder *E); + + // See DXComputeEncoder for why these are no-ops. + void pushDebugGroup(llvm::StringRef Label) override; + void popDebugGroup() override; + void insertDebugSignpost(llvm::StringRef Label) override; + + void setViewport(const offloadtest::Viewport &VP) override; + + void setScissor(const offloadtest::ScissorRect &Rect) override; + + void setVertexBuffer(uint32_t Slot, offloadtest::Buffer *VB, size_t Offset, + uint32_t Stride) override; + + void bindDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) override; + + llvm::Error drawInstanced(const offloadtest::PipelineState &PSO, + uint32_t VertexCount, uint32_t InstanceCount, + uint32_t FirstVertex, + uint32_t FirstInstance) override; + llvm::Error dispatchMesh(const offloadtest::PipelineState &PSO, + uint32_t GroupCountX, uint32_t GroupCountY, + uint32_t GroupCountZ) override; + void endEncodingImpl() override; +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_ENCODER_H diff --git a/include/API/DX/MemoryHeap.h b/include/API/DX/MemoryHeap.h new file mode 100644 index 000000000..4648a0ecf --- /dev/null +++ b/include/API/DX/MemoryHeap.h @@ -0,0 +1,44 @@ +//===- DX/MemoryHeap.h - Offload API DX MemoryHeap API --------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_MEMORYHEAP_H +#define OFFLOADTEST_API_DX_MEMORYHEAP_H + +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +#include "API/DX/Common.h" +#include "API/Device.h" + +#include +#include + +namespace offloadtest { +using Microsoft::WRL::ComPtr; + +class DXMemoryHeap : public offloadtest::MemoryHeap { +public: + std::string Name; + ComPtr Heap; + + DXMemoryHeap(ComPtr Heap, llvm::StringRef Name); + + static bool classof(const offloadtest::MemoryHeap *H); + + static llvm::Expected> + create(ID3D12DeviceX *Device, llvm::StringRef Name, size_t SizeInBytes); +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_MEMORYHEAP_H diff --git a/include/API/DX/PipelineState.h b/include/API/DX/PipelineState.h new file mode 100644 index 000000000..108b83123 --- /dev/null +++ b/include/API/DX/PipelineState.h @@ -0,0 +1,99 @@ +//===- DX/PipelineState.h - Offload API DX PipelineState API --------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_PIPELINESTATE_H +#define OFFLOADTEST_API_DX_PIPELINESTATE_H + +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +#include "API/Device.h" + +#include +#include + +namespace offloadtest { +using Microsoft::WRL::ComPtr; + +enum class RootParameterType : uint32_t { + DescriptorTable = 0, + SamplerTable, + Constant, + CBV, + SRV, + UAV, +}; + +struct RootSignatureLayout { + RootParameterType ParameterType : 3; + uint32_t Count : 29; + + RootSignatureLayout(RootParameterType ParameterType, uint32_t Count) + : ParameterType(ParameterType), Count(Count) {} + RootSignatureLayout() = delete; +}; + +struct DescriptorCountPair { + uint32_t DescriptorCount; + uint32_t SamplerCount; +}; + +struct DescriptorSetsLayout { + llvm::SmallVector RSigLayout; + llvm::SmallVector Sets; +}; + +class DXPipelineState : public offloadtest::PipelineState { +public: + std::string Name; + ComPtr RootSig; + DescriptorSetsLayout Layout; + + ComPtr PSO; + // Only set for graphics pipelines. + std::optional Topology; + // True for pipelines created via createPipelineRT — used by SBT / dispatch + // code to safely downcast to DXRayTracingPipelineState (parallel to + // VulkanPipelineState::IsRayTracing). + bool IsRayTracing = false; + + DXPipelineState(llvm::StringRef Name, ComPtr RootSig, + DescriptorSetsLayout Layout, ComPtr PSO, + std::optional Topology, + bool IsRT = false); + + static bool classof(const offloadtest::PipelineState *B); +}; + +/// RT pipeline state: holds the ID3D12StateObject + cached +/// ID3D12StateObjectProperties for SBT identifier queries plus a +/// shader-name → identifier-pointer map. The `void *` identifiers are +/// owned by Properties — keep it alive for the SBT's lifetime. +class DXRayTracingPipelineState : public DXPipelineState { +public: + ComPtr StateObject; + ComPtr Properties; + llvm::StringMap ShaderIdentifiers; + + DXRayTracingPipelineState(llvm::StringRef Name, + ComPtr RootSig, + DescriptorSetsLayout Layout, + ComPtr SO, + ComPtr Props); + + static bool classof(const offloadtest::PipelineState *B); +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_PIPELINESTATE_H \ No newline at end of file diff --git a/include/API/DX/Queue.h b/include/API/DX/Queue.h new file mode 100644 index 000000000..4f2c143f8 --- /dev/null +++ b/include/API/DX/Queue.h @@ -0,0 +1,97 @@ +//===- DX/Queue.h - Offload API DX Queue API ----------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_QUEUE_H +#define OFFLOADTEST_API_DX_QUEUE_H + +#include "API/DX/Common.h" +#include "API/Device.h" + +#include +#include + +namespace offloadtest { + +class DXFence : public offloadtest::Fence { +public: +#ifdef _WIN32 + DXFence(ComPtr Fence, HANDLE Event, llvm::StringRef Name) +#else // WSL + DXFence(ComPtr Fence, int Event, llvm::StringRef Name) +#endif + : Name(Name), Fence(Fence), Event(Event) { + } + + ~DXFence() override; + + std::string Name; + ComPtr Fence; +#ifdef _WIN32 + HANDLE Event; +#else // WSL + int Event; +#endif + + static llvm::Expected> create(ID3D12DeviceX *Device, + llvm::StringRef Name); + + uint64_t getFenceValue() override; + + llvm::Error waitForCompletion(uint64_t SignalValue) override; +}; + +class DXQueue : public offloadtest::Queue { +public: + using Queue::submit; + + ComPtr Queue; + std::unique_ptr SubmitFence; + uint64_t FenceCounter = 0; + // Batches of command buffers submitted to the GPU that may still be + // in-flight. The ID3D12CommandAllocator owns the backing memory for + // recorded commands, so it must outlive GPU execution. Each batch + // records the fence value it signals so we can non-blockingly query + // progress and release completed batches. + struct InFlightBatch { + uint64_t FenceValue; + llvm::SmallVector> CBs; + }; + llvm::SmallVector InFlightBatches; + + DXQueue(ComPtr Queue, + std::unique_ptr SubmitFence) + : Queue(Queue), SubmitFence(std::move(SubmitFence)) {} + DXQueue(DXQueue &&) = default; + ~DXQueue() override; + + static llvm::Expected + createGraphicsQueue(ComPtr Device); + + llvm::Expected + submit(llvm::SmallVector> CBs) + override; + + llvm::Expected + updateTileMappings(offloadtest::Buffer &Resource, + llvm::ArrayRef Mappings) override; + + llvm::Expected + updateTileMappings(offloadtest::Texture &Resource, + llvm::ArrayRef Mappings) override; + + // Shared tile-mapping path for buffers and textures: they differ only in how + // a TileRegion maps onto the underlying ID3D12Resource, so both variants + // resolve to a native resource and funnel through here. + llvm::Expected + updateTileMappingsImpl(ID3D12Resource *Resource, + llvm::ArrayRef Mappings); +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_QUEUE_H \ No newline at end of file diff --git a/include/API/DX/RenderPass.h b/include/API/DX/RenderPass.h new file mode 100644 index 000000000..654616f24 --- /dev/null +++ b/include/API/DX/RenderPass.h @@ -0,0 +1,30 @@ +//===- DX/RenderPass.h - Offload API DX Render Pass API -------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_RENDERPASS_H +#define OFFLOADTEST_API_DX_RENDERPASS_H + +#include "API/RenderPass.h" + +namespace offloadtest { + +class DXRenderPass final : public offloadtest::RenderPass { +public: + offloadtest::RenderPassDesc Desc; + + explicit DXRenderPass(offloadtest::RenderPassDesc Desc) + : RenderPass(GPUAPI::DirectX), Desc(std::move(Desc)) {} + + static bool classof(const offloadtest::RenderPass *RP) { + return RP->getAPI() == GPUAPI::DirectX; + } +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_RENDERPASS_H diff --git a/include/API/DX/SBT.h b/include/API/DX/SBT.h new file mode 100644 index 000000000..d75b4b6bc --- /dev/null +++ b/include/API/DX/SBT.h @@ -0,0 +1,54 @@ +//===- DX/SBT.h - Offload API DX Shader Binding Table API -----------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_SBT_H +#define OFFLOADTEST_API_DX_SBT_H + +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +#include "API/Device.h" + +#include +#include + +namespace offloadtest { +using Microsoft::WRL::ComPtr; + +class DXShaderBindingTable : public offloadtest::ShaderBindingTable { +public: + ComPtr Buffer; + // Pre-built ranges for D3D12_DISPATCH_RAYS_DESC. Sizes are zero for + // empty regions; raygen is a single record so it uses the no-stride + // variant. + D3D12_GPU_VIRTUAL_ADDRESS_RANGE RayGenRange{}; + D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE MissRange{}; + D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE HitGroupRange{}; + D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE CallableRange{}; + + DXShaderBindingTable(ComPtr Buf, + D3D12_GPU_VIRTUAL_ADDRESS_RANGE RG, + D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE MS, + D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE HG, + D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE CL) + : offloadtest::ShaderBindingTable(GPUAPI::DirectX), Buffer(Buf), + RayGenRange(RG), MissRange(MS), HitGroupRange(HG), CallableRange(CL) {} + + static bool classof(const offloadtest::ShaderBindingTable *S) { + return S->getAPI() == GPUAPI::DirectX; + } +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_SBT_H \ No newline at end of file diff --git a/include/API/DX/Sampler.h b/include/API/DX/Sampler.h new file mode 100644 index 000000000..9fec6a733 --- /dev/null +++ b/include/API/DX/Sampler.h @@ -0,0 +1,44 @@ +//===- DX/Sampler.h - Offload API DX Sampler API --------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_SAMPLER_H +#define OFFLOADTEST_API_DX_SAMPLER_H + +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +#include "API/Sampler.h" + +#include +#include + +namespace offloadtest { +using Microsoft::WRL::ComPtr; + +class DXSampler : public offloadtest::Sampler { +public: + D3D12_CPU_DESCRIPTOR_HANDLE Handle = {}; + std::string Name; + SamplerCreateDesc Desc; + + DXSampler(llvm::StringRef Name, SamplerCreateDesc Desc, + D3D12_CPU_DESCRIPTOR_HANDLE Handle); + + const SamplerCreateDesc &getDesc() const override; + + static bool classof(const offloadtest::Sampler *S); +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_SAMPLER_H \ No newline at end of file diff --git a/include/API/DX/Texture.h b/include/API/DX/Texture.h new file mode 100644 index 000000000..51a90b0ce --- /dev/null +++ b/include/API/DX/Texture.h @@ -0,0 +1,63 @@ +//===- DX/Texture.h - Offload API DX Texture API --------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DX_TEXTURE_H +#define OFFLOADTEST_API_DX_TEXTURE_H + +#include +#include + +// The windows headers define these macros which conflict with the C++ standard +// library. Undefining them before including any LLVM C++ code prevents errors. +#undef max +#undef min + +#include "API/Texture.h" + +#include +#include + +namespace offloadtest { +using Microsoft::WRL::ComPtr; + +class DXTexture : public offloadtest::Texture { +public: + ComPtr Resource; + + // Contract: If a command on the command buffer needs a resource to be in a + // different state it should always transition it back to the PreferredState + // afterwards. The PreferredState is the state of the most common use case for + // that resource. This allows us to do state transitions without state + // tracking. + D3D12_RESOURCE_STATES PreferredState; + + // Optional descriptors, depending on Desc.Usage. + // A zero ptr means no descriptor was created for that view type. + D3D12_CPU_DESCRIPTOR_HANDLE RTVHandle = {}; + D3D12_CPU_DESCRIPTOR_HANDLE DSVHandle = {}; + D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle = {}; + D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle = {}; + + std::string Name; + TextureCreateDesc Desc; + + DXTexture(ComPtr Resource, llvm::StringRef Name, + TextureCreateDesc Desc, D3D12_RESOURCE_STATES PreferredState, + D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle, + D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle); + + TileShape querySparseTileShape(const Device &Dev) const override; + + const TextureCreateDesc &getDesc() const override; + + static bool classof(const offloadtest::Texture *T); +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DX_TEXTURE_H \ No newline at end of file diff --git a/include/API/Descriptors.h b/include/API/Descriptors.h new file mode 100644 index 000000000..47c2368bd --- /dev/null +++ b/include/API/Descriptors.h @@ -0,0 +1,92 @@ +//===- Descriptors.h - Offload API Buffer ---------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_DESCRIPTORS_H +#define OFFLOADTEST_API_DESCRIPTORS_H + +#include "API/API.h" + +#include "llvm/ADT/ArrayRef.h" + +namespace offloadtest { + +class AccelerationStructure; +class Buffer; +class Sampler; +class Texture; + +struct VKBind { + uint32_t Binding; + uint32_t CounterBinding; +}; + +class DescriptorPool { + GPUAPI API; + +public: + virtual ~DescriptorPool() {} // TOOD(manon): Move this def to a .cpp file + GPUAPI getAPI() const { return API; } + + virtual void reset() = 0; + +protected: + explicit DescriptorPool(GPUAPI API) : API(API) {} +}; + +class DescriptorSets { + GPUAPI API; + +public: + virtual ~DescriptorSets() {} // TOOD(manon): Move this def to a .cpp file + + GPUAPI getAPI() const { return API; } + +protected: + explicit DescriptorSets(GPUAPI API) : API(API) {} +}; + +class DescriptorSetsBuilder { + GPUAPI API; + +public: + virtual ~DescriptorSetsBuilder() {} // TODO(manon): move to .cpp file + + GPUAPI getAPI() const { return API; } + + virtual DescriptorSetsBuilder & + constant(uint32_t SetIndex, llvm::ArrayRef, VKBind) = 0; + + virtual DescriptorSetsBuilder & + read(uint32_t SetIndex, llvm::ArrayRef, VKBind) = 0; + virtual DescriptorSetsBuilder &read(uint32_t SetIndex, + llvm::ArrayRef, + llvm::ArrayRef, + VKBind) = 0; + virtual DescriptorSetsBuilder & + read(uint32_t SetIndex, llvm::ArrayRef, + VKBind) = 0; + + virtual DescriptorSetsBuilder & + write(uint32_t SetIndex, llvm::ArrayRef, VKBind) = 0; + virtual DescriptorSetsBuilder & + write(uint32_t SetIndex, llvm::ArrayRef, VKBind) = 0; + + virtual DescriptorSetsBuilder & + sampler(uint32_t SetIndex, llvm::ArrayRef, VKBind) = 0; + + virtual std::unique_ptr build() = 0; + +protected: + explicit DescriptorSetsBuilder(GPUAPI API) : API(API) {} +}; +} // namespace offloadtest + +#endif // OFFLOADTEST_API_DESCRIPTORS_H diff --git a/include/API/Device.h b/include/API/Device.h index 089416ed3..f7e263ab8 100644 --- a/include/API/Device.h +++ b/include/API/Device.h @@ -19,7 +19,9 @@ #include "API/Buffer.h" #include "API/Capabilities.h" #include "API/CommandBuffer.h" +#include "API/Descriptors.h" #include "API/RenderPass.h" +#include "API/Sampler.h" #include "API/ShaderBindingTable.h" #include "API/Texture.h" @@ -299,14 +301,13 @@ class Queue { }; class Device { -protected: +public: std::string Description; std::string DriverName; std::string DriverVersion; std::string GPUGeneration; uint16_t FamilyPrefix = 0; -public: virtual const Capabilities &getCapabilities() = 0; virtual llvm::StringRef getAPIName() const = 0; virtual GPUAPI getAPI() const = 0; @@ -346,6 +347,9 @@ class Device { virtual llvm::Expected> createTexture(std::string Name, const TextureCreateDesc &Desc) = 0; + virtual llvm::Expected> + createSampler(std::string Name, const SamplerCreateDesc &Desc) = 0; + virtual llvm::Expected> createMemoryHeap(std::string Name, size_t SizeInBytes) = 0; @@ -354,6 +358,12 @@ class Device { virtual uint32_t getTextureUploadRowStrideInBytes(const TextureCreateDesc &Desc) const = 0; + // The layout an upload buffer must have to feed createTextureWithData / + // copyBufferToTexture for the given texture description. Encodes per-mip + // offsets, row pitch, and total size in the backend's required alignment. + virtual TextureUploadLayout + getTextureUploadLayout(const TextureCreateDesc &Desc) const = 0; + virtual llvm::Expected> createRenderPass(const RenderPassDesc &Desc) = 0; @@ -382,7 +392,14 @@ class Device { createBLAS(const AccelerationStructureSizes &Sizes) = 0; virtual llvm::Expected> - createTLAS(const AccelerationStructureSizes &Sizes) = 0; + createTLAS(const AccelerationStructureSizes &Sizes, + uint32_t InstanceCount) = 0; + + virtual llvm::Expected> + createDescriptorPool() = 0; + virtual llvm::Expected> + createDescriptorSetsBuilder(DescriptorPool &Pool, + const PipelineState &Pipeline) = 0; virtual ~Device() = 0; @@ -421,6 +438,9 @@ createBufferWithData(Device &Dev, std::string Name, size_t SizeInBytes, ComputeEncoder *Encoder, std::unique_ptr *OutUploadBuffer); +// Create a texture and upload `Data` (tightly-packed across mip levels) into +// it via a staging buffer recorded on `Encoder`. The staging buffer is handed +// back through `OutUploadBuffer` and must outlive command-buffer submission. llvm::Expected> createTextureWithData(Device &Dev, std::string Name, const TextureCreateDesc &Desc, const void *Data, diff --git a/include/API/Encoder.h b/include/API/Encoder.h index 8bea72d08..a5bcd74ff 100644 --- a/include/API/Encoder.h +++ b/include/API/Encoder.h @@ -21,12 +21,13 @@ namespace offloadtest { +class AccelerationStructure; class Buffer; -class Texture; +class DescriptorSets; class PipelineState; -class Texture; -class AccelerationStructure; class ShaderBindingTable; +class Texture; + struct BLASBuildRequest; struct TLASBuildRequest; @@ -87,6 +88,12 @@ class ComputeEncoder : public CommandEncoder { public: using CommandEncoder::CommandEncoder; + virtual void bindComputeDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) = 0; + + virtual void bindRayTracingDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) = 0; + /// Dispatch a compute grid. GroupCount specifies how many workgroups to /// launch in each dimension. The workgroup size is derived from \p PSO /// (e.g. the shader's numthreads attribute), which is also bound for the @@ -150,6 +157,9 @@ class RenderEncoder : public CommandEncoder { virtual void setVertexBuffer(uint32_t Slot, Buffer *VB, size_t Offset, uint32_t Stride) = 0; + virtual void bindDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) = 0; + virtual llvm::Error drawInstanced(const PipelineState &PSO, uint32_t VertexCount, uint32_t InstanceCount, diff --git a/include/API/Sampler.h b/include/API/Sampler.h new file mode 100644 index 000000000..121c3470c --- /dev/null +++ b/include/API/Sampler.h @@ -0,0 +1,67 @@ +//===- Sampler.h - Offload API Texture ------------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_SAMPLER_H +#define OFFLOADTEST_API_SAMPLER_H + +#include "API/API.h" +#include "API/Resources.h" + +namespace offloadtest { + +enum class FilterMode { Nearest, Linear }; + +enum class AddressMode { Clamp, Repeat, Mirror, Border, MirrorOnce }; + +enum class CompareFunction { + Never, + Less, + Equal, + LessEqual, + Greater, + NotEqual, + GreaterEqual, + Always +}; + +enum class SamplerKind { Sampler, SamplerComparison }; + +struct SamplerCreateDesc { + FilterMode MinFilter = FilterMode::Linear; + FilterMode MagFilter = FilterMode::Linear; + AddressMode Address = AddressMode::Clamp; + float MinLOD = 0.0f; + float MaxLOD = std::numeric_limits::max(); + float MipLODBias = 0.0f; + CompareFunction ComparisonOp = CompareFunction::Never; + SamplerKind Kind = SamplerKind::Sampler; +}; + +class Sampler { + GPUAPI API; + +public: + virtual ~Sampler(); + Sampler(const Sampler &) = delete; + // Sampler(Sampler &&) = delete; + Sampler &operator=(const Sampler &) = delete; + // Sampler &operator=(Sampler &&) = delete; + + GPUAPI getAPI() const { return API; } + virtual const SamplerCreateDesc &getDesc() const = 0; + +protected: + explicit Sampler(GPUAPI API) : API(API) {} +}; + +} // namespace offloadtest + +#endif // OFFLOADTEST_API_SAMPLER_H diff --git a/include/API/Texture.h b/include/API/Texture.h index 1b851393b..ea82fe62b 100644 --- a/include/API/Texture.h +++ b/include/API/Texture.h @@ -16,6 +16,7 @@ #include "API/Resources.h" #include "llvm/ADT/BitmaskEnum.h" +#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/Casting.h" #include "llvm/Support/Error.h" @@ -105,12 +106,6 @@ inline llvm::Error validateTextureCreateDesc(const TextureCreateDesc &Desc) { std::errc::not_supported, "DepthStencil combined with Storage is not yet supported."); - // Depth formats require DepthStencil usage; non-depth formats forbid it. - if (IsDepth && !IsDS) - return llvm::createStringError( - std::errc::invalid_argument, - "Depth format '%s' requires DepthStencil usage.", - getFormatName(Desc.Fmt).data()); if (!IsDepth && IsDS) return llvm::createStringError( std::errc::invalid_argument, @@ -154,6 +149,24 @@ struct TileShape { uint32_t Depth = 1; }; +struct SubresourceFootprint { + uint64_t Offset = 0; // Byte offset of this subresource in the buffer. + uint32_t RowPitchInBytes = 0; // Destination row stride (may include padding). + uint32_t RowSizeInBytes = 0; // Tightly-packed bytes per row to copy. + uint32_t NumRows = 0; // Number of rows in this subresource. +}; + +struct TextureUploadLayout { + llvm::SmallVector Subresources; // One entry per mip. + uint64_t TotalSizeInBytes = 0; +}; + +// Compute a tightly-packed upload layout (no row or subresource padding) for +// the given texture description. Suitable for backends whose buffer-to-texture +// copy consumes a tightly-packed staging buffer (e.g. Vulkan, Metal). +TextureUploadLayout +computeTightTextureUploadLayout(const TextureCreateDesc &Desc); + class Texture { GPUAPI API; diff --git a/lib/API/Util.h b/include/API/Util.h similarity index 100% rename from lib/API/Util.h rename to include/API/Util.h diff --git a/include/Support/Pipeline.h b/include/Support/Pipeline.h index f8dc1785d..e62681f46 100644 --- a/include/Support/Pipeline.h +++ b/include/Support/Pipeline.h @@ -16,6 +16,7 @@ #include "API/AccelerationStructure.h" #include "API/Enums.h" #include "API/Resources.h" +#include "API/Sampler.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/Error.h" @@ -140,24 +141,7 @@ static inline DescriptorKind getDescriptorKind(ResourceKind RK) { llvm_unreachable("All cases handled"); } -enum class FilterMode { Nearest, Linear }; - -enum class AddressMode { Clamp, Repeat, Mirror, Border, MirrorOnce }; - -enum class CompareFunction { - Never, - Less, - Equal, - LessEqual, - Greater, - NotEqual, - GreaterEqual, - Always -}; - -enum class SamplerKind { Sampler, SamplerComparison }; - -struct Sampler { +struct YAMLSampler { std::string Name; FilterMode MinFilter = FilterMode::Linear; FilterMode MagFilter = FilterMode::Linear; @@ -270,7 +254,7 @@ struct Resource { DirectXBinding DXBinding; std::optional VKBinding; CPUBuffer *BufferPtr = nullptr; - Sampler *SamplerPtr = nullptr; + YAMLSampler *SamplerPtr = nullptr; bool HasCounter = false; std::optional TilesMapped; bool IsReserved = false; @@ -630,7 +614,7 @@ struct Pipeline { IOBindings Bindings; llvm::SmallVector PushConstants; llvm::SmallVector Buffers; - llvm::SmallVector Samplers; + llvm::SmallVector Samplers; llvm::SmallVector Results; llvm::SmallVector Sets; DispatchParametersSet DispatchParameters; @@ -671,7 +655,7 @@ struct Pipeline { return nullptr; } - Sampler *getSampler(llvm::StringRef Name) { + YAMLSampler *getSampler(llvm::StringRef Name) { for (auto &S : Samplers) if (Name == S.Name) return &S; @@ -712,7 +696,7 @@ struct Pipeline { LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::DescriptorSet) LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::Resource) LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::CPUBuffer) -LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::Sampler) +LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::YAMLSampler) LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::Shader) LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::dx::RootParameter) LLVM_YAML_IS_SEQUENCE_VECTOR(offloadtest::Result) @@ -744,8 +728,8 @@ template <> struct MappingTraits { static void mapping(IO &I, offloadtest::CPUBuffer &R); }; -template <> struct MappingTraits { - static void mapping(IO &I, offloadtest::Sampler &S); +template <> struct MappingTraits { + static void mapping(IO &I, offloadtest::YAMLSampler &S); }; template <> struct MappingTraits { diff --git a/lib/API/CMakeLists.txt b/lib/API/CMakeLists.txt index 9f3836e83..237ad19ba 100644 --- a/lib/API/CMakeLists.txt +++ b/lib/API/CMakeLists.txt @@ -5,7 +5,7 @@ if (OFFLOADTEST_ENABLE_VULKAN) endif() if (OFFLOADTEST_ENABLE_D3D12) - list(APPEND api_sources DX/Device.cpp DX/DXFeatures.cpp) + list(APPEND api_sources DX/Buffer.cpp DX/CommandBuffer.cpp DX/Descriptors.cpp DX/Device.cpp DX/DXFeatures.cpp DX/Encoder.cpp DX/MemoryHeap.cpp DX/PipelineState.cpp DX/Queue.cpp DX/Sampler.cpp DX/Texture.cpp) list(APPEND api_headers PRIVATE ${DIRECTX_HEADERS_PATH}/include/directx ${DIRECTX_HEADERS_PATH}/include/dxguids @@ -21,7 +21,8 @@ endif() if (APPLE) list(APPEND api_sources MTL/MTLDevice.cpp MTL/MTLDescriptorHeap.cpp - MTL/MTLTopLevelArgumentBuffer.cpp) + MTL/MTLTopLevelArgumentBuffer.cpp + MTL/ResidencyTracker.cpp) list(APPEND api_libraries "-framework Metal" "-framework MetalKit" "-framework AppKit" diff --git a/lib/API/DX/Buffer.cpp b/lib/API/DX/Buffer.cpp new file mode 100644 index 000000000..0e2ed188b --- /dev/null +++ b/lib/API/DX/Buffer.cpp @@ -0,0 +1,53 @@ +//===- DX/Buffer.cpp - DirectX Buffer API ---------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#include "API/DX/Buffer.h" +#include "Support/WinError.h" + +using namespace offloadtest; + +DXBuffer::DXBuffer(ComPtr Buffer, llvm::StringRef Name, + BufferCreateDesc Desc, size_t SizeInBytes, + uint64_t CounterOffsetInBytes, + D3D12_RESOURCE_STATES PreferredState, + D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle, + D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle, + D3D12_CPU_DESCRIPTOR_HANDLE CBVHandle) + : offloadtest::Buffer(GPUAPI::DirectX), Buffer(Buffer), Name(Name), + Desc(Desc), SizeInBytes(SizeInBytes), + CounterOffsetInBytes(CounterOffsetInBytes), + PreferredState(PreferredState), SRVHandle(SRVHandle), + UAVHandle(UAVHandle), CBVHandle(CBVHandle) {} + +size_t DXBuffer::getSizeInBytes() const { return SizeInBytes; } + +size_t DXBuffer::querySparseTileSizeInBytes(const Device & /*Dev*/) const { + return D3D12_TILED_RESOURCE_TILE_SIZE_IN_BYTES; +} + +llvm::Expected DXBuffer::map() { + if (Desc.Location == MemoryLocation::GpuOnly) + return llvm::createStringError(std::errc::invalid_argument, + "Cannot map a GpuOnly buffer."); + void *Ptr = nullptr; + if (auto Err = + HR::toError(Buffer->Map(0, nullptr, &Ptr), "Failed to map buffer.")) + return std::move(Err); + return Ptr; +} + +void DXBuffer::unmap() { Buffer->Unmap(0, nullptr); } + +const BufferCreateDesc &DXBuffer::getDesc() const { return Desc; } + +bool DXBuffer::classof(const offloadtest::Buffer *B) { + return B->getAPI() == GPUAPI::DirectX; +} \ No newline at end of file diff --git a/lib/API/DX/CommandBuffer.cpp b/lib/API/DX/CommandBuffer.cpp new file mode 100644 index 000000000..4cfa57d96 --- /dev/null +++ b/lib/API/DX/CommandBuffer.cpp @@ -0,0 +1,228 @@ +//===- DX/CommandBuffer.cpp - DirectX CommandBuffer API -------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#include + +#include "API/DX/CommandBuffer.h" +#include "API/DX/Descriptors.h" +#include "API/DX/Encoder.h" +#include "API/DX/RenderPass.h" +#include "API/DX/Texture.h" +#include "API/Util.h" + +#include "Support/WinError.h" + +#include "llvm/Support/Casting.h" + +using namespace offloadtest; + +llvm::Expected> +DXCommandBuffer::create(ComPtr Device) { + auto CB = std::unique_ptr(new DXCommandBuffer()); + if (auto Err = HR::toError( + Device->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT, + IID_PPV_ARGS(&CB->Allocator)), + "Failed to create command allocator.")) + return Err; + if (auto Err = HR::toError( + Device->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT, + CB->Allocator.Get(), nullptr, + IID_PPV_ARGS(&CB->CmdList)), + "Failed to create command list.")) + return Err; + return CB; +} + +bool DXCommandBuffer::classof(const CommandBuffer *CB) { + return CB->getKind() == GPUAPI::DirectX; +} + +void DXCommandBuffer::addPendingUAVBarrier() { PendingUAVBarrier = true; } +void DXCommandBuffer::addResourceTransition(ID3D12Resource *Resource, + D3D12_RESOURCE_STATES StateBefore, + D3D12_RESOURCE_STATES StateAfter) { + + for (auto TransIt = PendingTransitions.begin(), + End = PendingTransitions.end(); + TransIt != End; ++TransIt) { + if (TransIt->Transition.pResource != Resource) + continue; + + assert(StateBefore == TransIt->Transition.StateAfter); + + if (StateAfter == TransIt->Transition.StateBefore) { + // We actually need the resource to be in the original state! + // Remove the transition from the list of Pending Transitions. + PendingTransitions.erase(TransIt); + + // NOTE: This could cause an execution barrier issue since legacy + // barriers don't allow us to express execution barriers, cache + // flushes, and layout transitions separately. + // As a workaround we add a UAV barrier so an execution barrier and + // cache flush is inserted. + addPendingUAVBarrier(); + } else { + TransIt->Transition.StateAfter = StateAfter; + } + + return; + } + + PendingTransitions.push_back( + CD3DX12_RESOURCE_BARRIER::Transition(Resource, StateBefore, StateAfter)); +} + +void DXCommandBuffer::flushBarrier() { + + if (PendingUAVBarrier) { + PendingTransitions.push_back(CD3DX12_RESOURCE_BARRIER::UAV(nullptr)); + PendingUAVBarrier = false; + } + + if (!PendingTransitions.empty()) { + CmdList->ResourceBarrier(PendingTransitions.size(), + PendingTransitions.data()); + PendingTransitions.clear(); + } +} + +void DXCommandBuffer::bindPool(const DescriptorPool &Pool) { + const DXDescriptorPool &PoolDX = llvm::cast(Pool); + ID3D12DescriptorHeap *const Heaps[] = {PoolDX.CSUHeap.Get(), + PoolDX.SamplerHeap.Get()}; + CmdList->SetDescriptorHeaps(2, Heaps); +} + +llvm::Expected> +DXCommandBuffer::createComputeEncoder() { + auto Enc = std::make_unique(*this); + Enc->pushDebugGroup("ComputeEncoder"); + return Enc; +} + +llvm::Expected> +DXCommandBuffer::createRenderEncoder( + const offloadtest::RenderPassBeginDesc &Desc) { + // The pass carries format / load / store policy; the begin desc supplies + // the actual textures. Walk both in lockstep. + if (!Desc.Pass) + return llvm::createStringError( + std::errc::invalid_argument, + "RenderPassBeginDesc is missing its RenderPass."); + auto &DXPass = llvm::cast(*Desc.Pass); + const offloadtest::RenderPassDesc &PassDesc = DXPass.Desc; + + if (Desc.ColorAttachments.size() != PassDesc.ColorAttachments.size()) + return llvm::createStringError( + std::errc::invalid_argument, + "RenderPassBeginDesc color attachment count does not match its " + "RenderPass."); + if (PassDesc.DepthStencil.has_value() != (Desc.DepthStencil != nullptr)) + return llvm::createStringError(std::errc::invalid_argument, + "RenderPassBeginDesc depth-stencil " + "presence does not match its RenderPass."); + + if (auto Err = findAndValidateRenderPassTextureSize(Desc, nullptr, nullptr)) + return Err; + + // Validate attachments and gather the RTV / DSV CPU handles. RT and DSV + // descriptors are owned by the textures themselves; this just collects + // them for OMSetRenderTargets. + llvm::SmallVector RTTextures; + llvm::SmallVector RTVHandles; + RTTextures.reserve(Desc.ColorAttachments.size()); + RTVHandles.reserve(Desc.ColorAttachments.size()); + for (offloadtest::Texture *Tex : Desc.ColorAttachments) { + if (!Tex) + return llvm::createStringError( + std::errc::invalid_argument, + "RenderPassBeginDesc has a null color attachment texture."); + auto &DXTex = llvm::cast(*Tex); + if (DXTex.RTVHandle.ptr == 0) + return llvm::createStringError( + std::errc::invalid_argument, + "Color attachment texture was not created with RenderTarget usage."); + RTTextures.push_back(&DXTex); + RTVHandles.push_back(DXTex.RTVHandle); + } + + DXTexture *DSTexture = nullptr; + D3D12_CPU_DESCRIPTOR_HANDLE DSVHandle = {}; + if (Desc.DepthStencil) { + auto &DXDS = llvm::cast(*Desc.DepthStencil); + if (DXDS.DSVHandle.ptr == 0) + return llvm::createStringError( + std::errc::invalid_argument, + "Depth-stencil texture was not created with DepthStencil usage."); + DSTexture = &DXDS; + DSVHandle = DXDS.DSVHandle; + } + + // State transitions + for (offloadtest::Texture *Tex : Desc.ColorAttachments) { + auto &DXTex = llvm::cast(*Tex); + if (DXTex.PreferredState != D3D12_RESOURCE_STATE_RENDER_TARGET) + this->addResourceTransition(DXTex.Resource.Get(), DXTex.PreferredState, + D3D12_RESOURCE_STATE_RENDER_TARGET); + } + if (Desc.DepthStencil) { + auto &DXTex = llvm::cast(*Desc.DepthStencil); + if (DXTex.PreferredState != D3D12_RESOURCE_STATE_DEPTH_WRITE) + this->addResourceTransition(DXTex.Resource.Get(), DXTex.PreferredState, + D3D12_RESOURCE_STATE_DEPTH_WRITE); + } + + this->flushBarrier(); + + CmdList->OMSetRenderTargets(static_cast(RTVHandles.size()), + RTVHandles.data(), + /*RTsSingleHandleToDescriptorRange=*/false, + Desc.DepthStencil ? &DSVHandle : nullptr); + + for (size_t I = 0; I < PassDesc.ColorAttachments.size(); ++I) { + if (PassDesc.ColorAttachments[I].Load != offloadtest::LoadAction::Clear) + continue; + if (!RTTextures[I]->Desc.OptimizedClearValue) + return llvm::createStringError( + std::errc::invalid_argument, + "LoadAction::Clear requires the render target to have been " + "created with an OptimizedClearValue."); + const auto *CV = + std::get_if(&*RTTextures[I]->Desc.OptimizedClearValue); + assert(CV && "RenderTarget OptimizedClearValue must be a ClearColor"); + const float ClearArr[4] = {CV->R, CV->G, CV->B, CV->A}; + CmdList->ClearRenderTargetView(RTVHandles[I], ClearArr, 0, nullptr); + } + if (PassDesc.DepthStencil) { + D3D12_CLEAR_FLAGS Flags = static_cast(0); + if (PassDesc.DepthStencil->DepthLoad == offloadtest::LoadAction::Clear) + Flags |= D3D12_CLEAR_FLAG_DEPTH; + if (PassDesc.DepthStencil->StencilLoad == offloadtest::LoadAction::Clear) + Flags |= D3D12_CLEAR_FLAG_STENCIL; + if (Flags != 0) { + if (!DSTexture->Desc.OptimizedClearValue) + return llvm::createStringError( + std::errc::invalid_argument, + "LoadAction::Clear requires the depth-stencil texture to have " + "been created with an OptimizedClearValue."); + const auto *CV = + std::get_if(&*DSTexture->Desc.OptimizedClearValue); + assert(CV && + "DepthStencil OptimizedClearValue must be a ClearDepthStencil"); + CmdList->ClearDepthStencilView(DSVHandle, Flags, CV->Depth, CV->Stencil, + 0, nullptr); + } + } + + auto Enc = std::make_unique(*this, Desc); + Enc->pushDebugGroup("RenderEncoder"); + return Enc; +} diff --git a/lib/API/DX/DXFeatures.cpp b/lib/API/DX/DXFeatures.cpp index 7cc59ac2b..e3b52aca4 100644 --- a/lib/API/DX/DXFeatures.cpp +++ b/lib/API/DX/DXFeatures.cpp @@ -10,73 +10,55 @@ //===----------------------------------------------------------------------===// #include "DXFeatures.h" -#include "llvm/Support/ScopedPrinter.h" +#include "llvm/ADT/Enum.h" using namespace offloadtest; using namespace offloadtest::directx; using namespace llvm; -#define SHADER_MODEL_ENUM(NewCase, Str, Value) {#Str, NewCase}, -static const EnumEntry ShaderModelNames[]{ +#define SHADER_MODEL_ENUM(NewCase, Str, Value) {{#Str}, NewCase}, +static constexpr EnumStringDef ShaderModelDefs[]{ #include "DXFeatures.def" }; +static constexpr auto ShaderModelNames = BUILD_ENUM_STRINGS(ShaderModelDefs); -static ArrayRef> getShaderModels() { - return ArrayRef(ShaderModelNames); -} - -#define ROOT_SIGNATURE_ENUM(NewCase, Str, Value) {#Str, NewCase}, -static const EnumEntry RootSignatureNames[]{ +#define ROOT_SIGNATURE_ENUM(NewCase, Str, Value) {{#Str}, NewCase}, +static constexpr EnumStringDef RootSignatureDefs[]{ #include "DXFeatures.def" }; +static constexpr auto RootSignatureNames = + BUILD_ENUM_STRINGS(RootSignatureDefs); -static ArrayRef> getRootSignatures() { - return ArrayRef(RootSignatureNames); -} - -#define MESH_SHADER_TIER_ENUM(NewCase, Str, Value) {#Str, NewCase}, -static const EnumEntry MeshShaderTierNames[]{ +#define MESH_SHADER_TIER_ENUM(NewCase, Str, Value) {{#Str}, NewCase}, +static constexpr EnumStringDef MeshShaderTierDefs[]{ #include "DXFeatures.def" }; +static constexpr auto MeshShaderTierNames = + BUILD_ENUM_STRINGS(MeshShaderTierDefs); -static ArrayRef> getMeshShaderTiers() { - return ArrayRef(MeshShaderTierNames); -} - -#define RAYTRACING_TIER_ENUM(NewCase, Str, Value) {#Str, NewCase}, -static const EnumEntry RaytracingTierNames[]{ +#define RAYTRACING_TIER_ENUM(NewCase, Str, Value) {{#Str}, NewCase}, +static constexpr EnumStringDef RaytracingTierDefs[]{ #include "DXFeatures.def" }; - -static ArrayRef> getRaytracingTiers() { - return ArrayRef(RaytracingTierNames); -} - -template -static std::string enumEntryToString(ArrayRef> EnumValues, - T Value) { - for (const EnumEntry &I : EnumValues) - if (I.Value == Value) - return I.Name.str(); - llvm_unreachable("All cases must be covered"); -} +static constexpr auto RaytracingTierNames = + BUILD_ENUM_STRINGS(RaytracingTierDefs); std::string CapabilityPrinter::toString( const directx::ShaderModel &V) { - return enumEntryToString(getShaderModels(), V); + return EnumStrings(ShaderModelNames).toString(V).str(); } std::string CapabilityPrinter::toString( const directx::RootSignature &V) { - return enumEntryToString(getRootSignatures(), V); + return EnumStrings(RootSignatureNames).toString(V).str(); } std::string CapabilityPrinter::toString( const directx::MeshShaderTier &V) { - return enumEntryToString(getMeshShaderTiers(), V); + return EnumStrings(MeshShaderTierNames).toString(V).str(); } std::string CapabilityPrinter::toString( const directx::RaytracingTier &V) { - return enumEntryToString(getRaytracingTiers(), V); -} + return EnumStrings(RaytracingTierNames).toString(V).str(); +} \ No newline at end of file diff --git a/lib/API/DX/DXFeatures.h b/lib/API/DX/DXFeatures.h index fff4cb19a..db0b53424 100644 --- a/lib/API/DX/DXFeatures.h +++ b/lib/API/DX/DXFeatures.h @@ -13,11 +13,6 @@ #define OFFLOADTEST_API_DXFEATURES_H #include "API/Capabilities.h" -#include "llvm/ADT/ArrayRef.h" - -namespace llvm { -template struct EnumEntry; -} // namespace llvm namespace offloadtest { namespace directx { diff --git a/lib/API/DX/DXResources.h b/lib/API/DX/DXResources.h index aaa911794..c5224f14e 100644 --- a/lib/API/DX/DXResources.h +++ b/lib/API/DX/DXResources.h @@ -115,6 +115,15 @@ inline DXGI_FORMAT getDXGIIndexFormat(IndexFormat Fmt) { llvm_unreachable("All IndexFormat cases handled"); } +// D3D12 requires the RowPitch in a placed subresource footprint (used for +// texture <-> buffer copies via CopyTextureRegion) to be a multiple of +// D3D12_TEXTURE_DATA_PITCH_ALIGNMENT (256 bytes). For textures whose natural +// row size (Width * elementSize) is already a multiple of 256, this is a +// no-op; for smaller rows it pads up. +inline uint32_t getAlignedTexturePitch(uint32_t Width, uint32_t ElementSize) { + return llvm::alignTo(Width * ElementSize, D3D12_TEXTURE_DATA_PITCH_ALIGNMENT); +} + } // namespace offloadtest #endif // OFFLOADTEST_API_DXRESOURCES_H diff --git a/lib/API/DX/Descriptors.cpp b/lib/API/DX/Descriptors.cpp new file mode 100644 index 000000000..4e56ad573 --- /dev/null +++ b/lib/API/DX/Descriptors.cpp @@ -0,0 +1,227 @@ +//===- DX/Descriptors.cpp - DirectX Desriptors API ------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "API/DX/Descriptors.h" +#include "API/DX/AccelerationStructure.h" +#include "API/DX/Buffer.h" +#include "API/DX/Sampler.h" +#include "API/DX/Texture.h" + +#include "Support/WinError.h" + +using namespace offloadtest; + +DXDescriptorSets::DXDescriptorSets(llvm::SmallVector Sets) + : DescriptorSets(GPUAPI::DirectX), Sets(std::move(Sets)) {} + +bool DXDescriptorSets::classof(const DescriptorSets *S) { + return S->getAPI() == GPUAPI::DirectX; +} + +DXDescriptorPool::DXDescriptorPool(ComPtr CSUHeap, + ComPtr SamplerHeap, + uint32_t CSUIncSize, uint32_t SamplerIncSize) + : DescriptorPool(GPUAPI::DirectX), CSUHeap(CSUHeap), + SamplerHeap(SamplerHeap), CSUIncSize(CSUIncSize), + SamplerIncSize(SamplerIncSize) { + CSUHandle = this->CSUHeap->GetCPUDescriptorHandleForHeapStart(); + CSUHandleGPU = this->CSUHeap->GetGPUDescriptorHandleForHeapStart(); + + SamplerHandle = this->SamplerHeap->GetCPUDescriptorHandleForHeapStart(); + SamplerHandleGPU = this->SamplerHeap->GetGPUDescriptorHandleForHeapStart(); +} + +llvm::Expected> +DXDescriptorPool::create(ComPtr Dev) { + const uint32_t DescriptorCount = 4096; + const uint32_t SamplerCount = 512; + + const D3D12_DESCRIPTOR_HEAP_DESC HeapDesc = { + D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV, DescriptorCount, + D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE, 0}; + + ComPtr CSUHeap; + if (auto Err = HR::toError( + Dev->CreateDescriptorHeap(&HeapDesc, IID_PPV_ARGS(&CSUHeap)), + "Failed to create descriptor heap.")) + return Err; + + const D3D12_DESCRIPTOR_HEAP_DESC SamplerHeapDesc = { + D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER, SamplerCount, + D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE, 0}; + ComPtr SamplerHeap; + if (auto Err = HR::toError(Dev->CreateDescriptorHeap( + &SamplerHeapDesc, IID_PPV_ARGS(&SamplerHeap)), + "Failed to create sampler descriptor heap.")) + return Err; + + const uint32_t CSUIncSize = Dev->GetDescriptorHandleIncrementSize( + D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); + const uint32_t SamplerIncSize = + Dev->GetDescriptorHandleIncrementSize(D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER); + + return std::make_unique(CSUHeap, SamplerHeap, CSUIncSize, + SamplerIncSize); +} + +void DXDescriptorPool::allocateDescriptors(uint32_t Count, + D3D12_CPU_DESCRIPTOR_HANDLE &CPU, + D3D12_GPU_DESCRIPTOR_HANDLE &GPU) { + const uint32_t Idx = CSUAllocator.fetch_add(Count); + const uint32_t Offset = Idx * CSUIncSize; + CPU = {CSUHandle.ptr + Offset}; + GPU = {CSUHandleGPU.ptr + Offset}; +} + +void DXDescriptorPool::allocateSamplers(uint32_t Count, + D3D12_CPU_DESCRIPTOR_HANDLE &CPU, + D3D12_GPU_DESCRIPTOR_HANDLE &GPU) { + const uint32_t Idx = SamplerAllocator.fetch_add(Count); + const uint32_t Offset = Idx * SamplerIncSize; + CPU = {SamplerHandle.ptr + Offset}; + GPU = {SamplerHandleGPU.ptr + Offset}; +} + +void DXDescriptorPool::reset() { + CSUAllocator.store(0); + SamplerAllocator.store(0); +} + +bool DXDescriptorPool::classof(const DescriptorPool *P) { + return P->getAPI() == GPUAPI::DirectX; +} + +DXDescriptorSetsBuilder::DXDescriptorSetsBuilder( + ComPtr Dev, llvm::SmallVector Sets, + uint32_t CSUIncSize, uint32_t SamplerIncSize) + : DescriptorSetsBuilder(GPUAPI::DirectX), Dev(Dev), Sets(std::move(Sets)), + CSUIncSize(CSUIncSize), SamplerIncSize(SamplerIncSize) { + SetStates.reserve(this->Sets.size()); + for (const auto &Set : this->Sets) { + SetStates.push_back({Set.CSUHandle, Set.SamplerHandle}); + } +} + +DescriptorSetsBuilder & +DXDescriptorSetsBuilder::constant(uint32_t SetIndex, + llvm::ArrayRef B, VKBind) { + if (SetIndex >= SetStates.size()) + return *this; + SetState &State = SetStates[SetIndex]; + for (const Buffer *Buf : B) { + const DXBuffer &BufferDX = llvm::cast(*Buf); + Dev->CopyDescriptorsSimple(1, State.CSUHandle, BufferDX.CBVHandle, + D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); + State.CSUHandle.ptr += CSUIncSize; + } + return *this; +} + +DescriptorSetsBuilder & +DXDescriptorSetsBuilder::read(uint32_t SetIndex, + llvm::ArrayRef B, VKBind) { + if (SetIndex >= SetStates.size()) + return *this; + SetState &State = SetStates[SetIndex]; + for (const Buffer *Buf : B) { + const DXBuffer &BufferDX = llvm::cast(*Buf); + Dev->CopyDescriptorsSimple(1, State.CSUHandle, BufferDX.SRVHandle, + D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); + State.CSUHandle.ptr += CSUIncSize; + } + return *this; +} +DescriptorSetsBuilder & +DXDescriptorSetsBuilder::read(uint32_t SetIndex, + llvm::ArrayRef T, + llvm::ArrayRef S, VKBind) { + assert((S.empty() || S.size() == T.size()) && + "Sampler list must either be empty or match " + "texture list when binding descriptors."); + + if (SetIndex >= SetStates.size()) + return *this; + SetState &State = SetStates[SetIndex]; + for (size_t I = 0, N = T.size(); I < N; ++I) { + const DXTexture &TextureDX = llvm::cast(*T[I]); + assert((S.empty() || S[I] == nullptr) && + "DirectX 12 does not support combined image samplers."); + Dev->CopyDescriptorsSimple(1, State.CSUHandle, TextureDX.SRVHandle, + D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); + State.CSUHandle.ptr += CSUIncSize; + } + return *this; +} +DescriptorSetsBuilder & +DXDescriptorSetsBuilder::read(uint32_t SetIndex, + llvm::ArrayRef A, + VKBind) { + if (SetIndex >= SetStates.size()) + return *this; + SetState &State = SetStates[SetIndex]; + for (const AccelerationStructure *AS : A) { + const DXAccelerationStructure &AccelStructDX = + llvm::cast(*AS); + Dev->CopyDescriptorsSimple(1, State.CSUHandle, AccelStructDX.SRVHandle, + D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); + State.CSUHandle.ptr += CSUIncSize; + } + return *this; +} + +DescriptorSetsBuilder & +DXDescriptorSetsBuilder::write(uint32_t SetIndex, + llvm::ArrayRef B, VKBind) { + if (SetIndex >= SetStates.size()) + return *this; + SetState &State = SetStates[SetIndex]; + for (const Buffer *Buf : B) { + const DXBuffer &BufferDX = llvm::cast(*Buf); + Dev->CopyDescriptorsSimple(1, State.CSUHandle, BufferDX.UAVHandle, + D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); + State.CSUHandle.ptr += CSUIncSize; + } + return *this; +} +DescriptorSetsBuilder & +DXDescriptorSetsBuilder::write(uint32_t SetIndex, + llvm::ArrayRef T, VKBind) { + if (SetIndex >= SetStates.size()) + return *this; + SetState &State = SetStates[SetIndex]; + for (const Texture *Tex : T) { + const DXTexture &TextureDX = llvm::cast(*Tex); + Dev->CopyDescriptorsSimple(1, State.CSUHandle, TextureDX.UAVHandle, + D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); + State.CSUHandle.ptr += CSUIncSize; + } + return *this; +} + +DescriptorSetsBuilder & +DXDescriptorSetsBuilder::sampler(uint32_t SetIndex, + llvm::ArrayRef S, VKBind) { + if (SetIndex >= SetStates.size()) + return *this; + SetState &State = SetStates[SetIndex]; + for (const Sampler *Sampl : S) { + const DXSampler &SamplerDX = llvm::cast(*Sampl); + Dev->CopyDescriptorsSimple(1, State.SamplerHandle, SamplerDX.Handle, + D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER); + State.SamplerHandle.ptr += SamplerIncSize; + } + return *this; +} + +std::unique_ptr DXDescriptorSetsBuilder::build() { + return std::make_unique(std::move(Sets)); +} + +bool DXDescriptorSetsBuilder::classof(const DescriptorSetsBuilder *B) { + return B->getAPI() == GPUAPI::DirectX; +} diff --git a/lib/API/DX/Device.cpp b/lib/API/DX/Device.cpp index afce98aa6..506d70064 100644 --- a/lib/API/DX/Device.cpp +++ b/lib/API/DX/Device.cpp @@ -32,9 +32,21 @@ #undef min #include "API/Capabilities.h" -#include "API/Device.h" -#include "API/Encoder.h" +#include "API/DX/AccelerationStructure.h" +#include "API/DX/Buffer.h" +#include "API/DX/CommandBuffer.h" +#include "API/DX/Descriptors.h" +#include "API/DX/Device.h" +#include "API/DX/Encoder.h" +#include "API/DX/MemoryHeap.h" +#include "API/DX/PipelineState.h" +#include "API/DX/Queue.h" +#include "API/DX/RenderPass.h" +#include "API/DX/SBT.h" +#include "API/DX/Sampler.h" +#include "API/DX/Texture.h" #include "API/FormatConversion.h" +#include "API/Util.h" #include "DXFeatures.h" #include "Support/Pipeline.h" #include "Support/WinError.h" @@ -48,7 +60,7 @@ #include "llvm/Support/FormatVariadic.h" #include "llvm/Support/Signals.h" -#include "../Util.h" +#include "../Support/OffloadMigration.h" #include #include @@ -58,9 +70,6 @@ using namespace offloadtest; using Microsoft::WRL::ComPtr; -using ID3D12DeviceX = ID3D12Device5; -using ID3D12GraphicsCommandListX = ID3D12GraphicsCommandList6; - template <> char CapabilityValueEnum::ID = 0; template <> char CapabilityValueEnum::ID = 0; template <> char CapabilityValueEnum::ID = 0; @@ -69,6 +78,37 @@ template <> char CapabilityValueEnum::ID = 0; static std::mutex SignalHandlerMutex; static llvm::SmallVector SignalHandlerDevices; +static size_t getCBVSize(size_t Sz) { return (Sz + 255u) & 0xFFFFFFFFFFFFFF00; } + +static D3D12_HIT_GROUP_TYPE getDXHitGroupType(HitGroupType T) { + switch (T) { + case HitGroupType::Triangles: + return D3D12_HIT_GROUP_TYPE_TRIANGLES; + case HitGroupType::Procedural: + return D3D12_HIT_GROUP_TYPE_PROCEDURAL_PRIMITIVE; + } + llvm_unreachable("All HitGroupType cases handled"); +} + +static std::wstring widen(llvm::StringRef S) { + // Entry-point names and hit-group names are ASCII; a straight 1:1 widen + // is sufficient. + return std::wstring(S.begin(), S.end()); +} + +static llvm::Error configureInfoQueue(ID3D12DeviceX *Device) { +#ifdef _WIN32 + ComPtr InfoQueue; + if (auto Err = HR::toError(Device->QueryInterface(InfoQueue.GetAddressOf()), + "Error initializing info queue")) + return Err; + InfoQueue->SetBreakOnSeverity(D3D12_MESSAGE_SEVERITY_CORRUPTION, TRUE); + InfoQueue->SetBreakOnSeverity(D3D12_MESSAGE_SEVERITY_ERROR, TRUE); + InfoQueue->SetBreakOnSeverity(D3D12_MESSAGE_SEVERITY_WARNING, TRUE); +#endif + return llvm::Error::success(); +} + static void dumpD3DInfoQueues(void *) { const std::lock_guard Lock(SignalHandlerMutex); for (ID3D12DeviceX *Device : SignalHandlerDevices) { @@ -94,15 +134,6 @@ static void dumpD3DInfoQueues(void *) { } } -// D3D12 requires the RowPitch in a placed subresource footprint (used for -// texture <-> buffer copies via CopyTextureRegion) to be a multiple of -// D3D12_TEXTURE_DATA_PITCH_ALIGNMENT (256 bytes). For textures whose natural -// row size (Width * elementSize) is already a multiple of 256, this is a -// no-op; for smaller rows it pads up. -static uint32_t getAlignedTexturePitch(uint32_t Width, uint32_t ElementSize) { - return llvm::alignTo(Width * ElementSize, D3D12_TEXTURE_DATA_PITCH_ALIGNMENT); -} - static D3D12_PRIMITIVE_TOPOLOGY_TYPE getDXPrimitiveTopologyType(PrimitiveTopology Topology) { switch (Topology) { @@ -139,2140 +170,1235 @@ getDXPrimitiveTopology(PrimitiveTopology Topology, llvm_unreachable("All PrimitiveTopology cases handled"); } -static uint64_t getAlignedTextureBufferSize(const CPUBuffer &B) { - const uint64_t AlignedPitch = - getAlignedTexturePitch(B.OutputProps.Width, B.getElementSize()); - const uint64_t LastRowSize = - uint64_t(B.OutputProps.Width) * B.getElementSize(); - return uint64_t(B.OutputProps.Height - 1) * AlignedPitch + LastRowSize; +static D3D12_FILTER getDXFilterMode(FilterMode MinFilter, FilterMode MagFilter, + bool IsComparison) { + if (IsComparison) { + if (MinFilter == FilterMode::Nearest) + return MagFilter == FilterMode::Nearest + ? D3D12_FILTER_COMPARISON_MIN_MAG_MIP_POINT + : D3D12_FILTER_COMPARISON_MIN_POINT_MAG_LINEAR_MIP_POINT; + else + return MagFilter == FilterMode::Nearest + ? D3D12_FILTER_COMPARISON_MIN_LINEAR_MAG_MIP_POINT + : D3D12_FILTER_COMPARISON_MIN_MAG_LINEAR_MIP_POINT; + } else { + if (MinFilter == FilterMode::Nearest) + return MagFilter == FilterMode::Nearest + ? D3D12_FILTER_MIN_MAG_MIP_POINT + : D3D12_FILTER_MIN_POINT_MAG_LINEAR_MIP_POINT; + else + return MagFilter == FilterMode::Nearest + ? D3D12_FILTER_MIN_LINEAR_MAG_MIP_POINT + : D3D12_FILTER_MIN_MAG_LINEAR_MIP_POINT; + } } -static BufferUsage bufferUsageFromResourceKind(ResourceKind Kind) { - // Determine Buffer Usage - switch (Kind) { - case ResourceKind::Buffer: - case ResourceKind::StructuredBuffer: - case ResourceKind::ByteAddressBuffer: - case ResourceKind::RWBuffer: - case ResourceKind::RWStructuredBuffer: - case ResourceKind::RWByteAddressBuffer: - return BufferUsage::Storage; - case ResourceKind::ConstantBuffer: - return BufferUsage::ConstantBuffer; - case ResourceKind::Texture2D: - case ResourceKind::RWTexture2D: - case ResourceKind::Sampler: - case ResourceKind::SampledTexture2D: - case ResourceKind::AccelerationStructure: - llvm_unreachable("Invalid case, ResourceKind is not a buffer."); - } - llvm_unreachable("All ResourceKind cases handled"); +static D3D12_TEXTURE_ADDRESS_MODE getDXTextureAddressMode(AddressMode Mode) { + switch (Mode) { + case AddressMode::Clamp: + return D3D12_TEXTURE_ADDRESS_MODE_CLAMP; + case AddressMode::Repeat: + return D3D12_TEXTURE_ADDRESS_MODE_WRAP; + case AddressMode::Mirror: + return D3D12_TEXTURE_ADDRESS_MODE_MIRROR; + case AddressMode::Border: + return D3D12_TEXTURE_ADDRESS_MODE_BORDER; + case AddressMode::MirrorOnce: + return D3D12_TEXTURE_ADDRESS_MODE_MIRROR_ONCE; + } + llvm_unreachable("All cases handled."); } -static BufferShaderAccessType bufferShaderAccessTypeFromResourceKind( - const Resource &Resource, BufferShaderAccessTypeParams &OutParams) { - // Determine Buffer Access Type - switch (Resource.Kind) { - case ResourceKind::Buffer: - case ResourceKind::RWBuffer: { - auto FmtOrErr = - toFormat(Resource.BufferPtr->Format, Resource.BufferPtr->Channels); - if (!FmtOrErr) { - printf("Invalid format! FMT: %d, CHANNELS: %d\n", - Resource.BufferPtr->Format, Resource.BufferPtr->Channels); - assert(false && "Invalid format."); - } - OutParams.Fmt = *FmtOrErr; - return BufferShaderAccessType::Typed; - } - case ResourceKind::StructuredBuffer: - case ResourceKind::RWStructuredBuffer: - OutParams.StructureStride = Resource.BufferPtr->getElementSize(); - return BufferShaderAccessType::Structured; - case ResourceKind::ByteAddressBuffer: - case ResourceKind::RWByteAddressBuffer: - case ResourceKind::ConstantBuffer: - return BufferShaderAccessType::Raw; - case ResourceKind::Texture2D: - case ResourceKind::RWTexture2D: - case ResourceKind::Sampler: - case ResourceKind::SampledTexture2D: - case ResourceKind::AccelerationStructure: - llvm_unreachable( - "Invalid case, non-buffers should have been filtered out."); - } - llvm_unreachable("All ResourceKind cases handled"); +static D3D12_COMPARISON_FUNC getDXComparisonFunc(CompareFunction ComparisonOp) { + switch (ComparisonOp) { + case CompareFunction::Never: + return D3D12_COMPARISON_FUNC_NEVER; + case CompareFunction::Less: + return D3D12_COMPARISON_FUNC_LESS; + case CompareFunction::Equal: + return D3D12_COMPARISON_FUNC_EQUAL; + case CompareFunction::LessEqual: + return D3D12_COMPARISON_FUNC_LESS_EQUAL; + case CompareFunction::Greater: + return D3D12_COMPARISON_FUNC_GREATER; + case CompareFunction::NotEqual: + return D3D12_COMPARISON_FUNC_NOT_EQUAL; + case CompareFunction::GreaterEqual: + return D3D12_COMPARISON_FUNC_GREATER_EQUAL; + case CompareFunction::Always: + return D3D12_COMPARISON_FUNC_ALWAYS; + } + llvm_unreachable("All cases handled."); } -namespace { -class DXDevice; +llvm::Expected DescriptorAllocator::create( + ID3D12DeviceX *Device, D3D12_DESCRIPTOR_HEAP_TYPE Type, uint32_t Capacity) { + ComPtr Heap; + const D3D12_DESCRIPTOR_HEAP_DESC HeapDesc = { + Type, Capacity, D3D12_DESCRIPTOR_HEAP_FLAG_NONE, 0}; + if (auto Err = HR::toError( + Device->CreateDescriptorHeap(&HeapDesc, IID_PPV_ARGS(&Heap)), + "Failed to create descriptor heap for DescriptorAllocator.")) + return Err; + const uint32_t DescIncSize = Device->GetDescriptorHandleIncrementSize(Type); + return DescriptorAllocator(Heap, DescIncSize, Capacity); +} -class DXBuffer : public offloadtest::Buffer { -public: - ComPtr Buffer; - std::string Name; - BufferCreateDesc Desc; - size_t SizeInBytes; - uint64_t CounterOffsetInBytes; - - // Contract: If a command on the command buffer needs a resource to be in a - // different state it should always transition it back to the PreferredState - // afterwards. The PreferredState is the state of the most common use case for - // that resource. This allows us to do state transitions without state - // tracking. - D3D12_RESOURCE_STATES PreferredState; - - D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle; - D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle; - D3D12_CPU_DESCRIPTOR_HANDLE CBVHandle; - - DXBuffer(ComPtr Buffer, llvm::StringRef Name, - BufferCreateDesc Desc, size_t SizeInBytes, - uint64_t CounterOffsetInBytes, D3D12_RESOURCE_STATES PreferredState, - D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle, - D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle, - D3D12_CPU_DESCRIPTOR_HANDLE CBVHandle) - : offloadtest::Buffer(GPUAPI::DirectX), Buffer(Buffer), Name(Name), - Desc(Desc), SizeInBytes(SizeInBytes), - CounterOffsetInBytes(CounterOffsetInBytes), - PreferredState(PreferredState), SRVHandle(SRVHandle), - UAVHandle(UAVHandle), CBVHandle(CBVHandle) {} - DXBuffer(const DXBuffer &) = delete; - DXBuffer(DXBuffer &&) = delete; - DXBuffer &operator=(const DXBuffer &) = delete; - DXBuffer &operator=(DXBuffer &&) = delete; - - size_t getSizeInBytes() const override { return SizeInBytes; } - - size_t querySparseTileSizeInBytes(const Device & /*Dev*/) const override { - return D3D12_TILED_RESOURCE_TILE_SIZE_IN_BYTES; - } +llvm::Expected DescriptorAllocator::allocate() { + // TODO(manon): Use a better allocator that can also free descriptors. + const uint32_t Index = NextIndex.fetch_add(1, std::memory_order_relaxed); + if (Index >= Capacity) + return llvm::createStringError(std::errc::not_enough_memory, + "Descriptor heap allocator exhausted."); + D3D12_CPU_DESCRIPTOR_HANDLE Handle = + Heap->GetCPUDescriptorHandleForHeapStart(); + Handle.ptr += Index * DescIncSize; + return Handle; +} - llvm::Expected map() override { - if (Desc.Location == MemoryLocation::GpuOnly) - return llvm::createStringError(std::errc::invalid_argument, - "Cannot map a GpuOnly buffer."); - void *Ptr = nullptr; - if (auto Err = - HR::toError(Buffer->Map(0, nullptr, &Ptr), "Failed to map buffer.")) - return std::move(Err); - return Ptr; +DXDevice::DXDevice(ComPtr A, ComPtr D, DXQueue Q, + DescriptorAllocator RTVAllocator, + DescriptorAllocator DSVAllocator, + DescriptorAllocator CSUAllocator, + DescriptorAllocator SamplerAllocator, std::string Desc, + std::string DriverVer) + : Adapter(A), Device(D), GraphicsQueue(std::move(Q)), + RTVAllocator(std::move(RTVAllocator)), + DSVAllocator(std::move(DSVAllocator)), + CSUAllocator(std::move(CSUAllocator)), + SamplerAllocator(std::move(SamplerAllocator)) { + Description = std::move(Desc); + DriverVersion = std::move(DriverVer); + DriverName = "DirectX"; + + DXCoreHardwareID HardwareID; + if (SUCCEEDED(Adapter->GetProperty(DXCoreAdapterProperty::HardwareID, + &HardwareID))) { + // 0x8086 is the Vendor ID for Intel + if (HardwareID.vendorID == 0x8086) { + FamilyPrefix = static_cast(HardwareID.deviceID) & 0xFF00; + const IntelGpuEra Era = + getIntelGpuEra(static_cast(HardwareID.deviceID)); + if (Era == IntelGpuEra::Gen7_to_10) + GPUGeneration = "Intel Gen7-10"; + else if (Era == IntelGpuEra::Gen11_to_14_and_Xe) + GPUGeneration = "Intel Gen11-14/Xe"; + else + GPUGeneration = "Intel Unknown"; + } else { + // We don't have a need yet to identify other GPU vendors. + GPUGeneration = "Unknown"; + } } +} - void unmap() override { Buffer->Unmap(0, nullptr); } +DXDevice::~DXDevice() { + const std::lock_guard Lock(SignalHandlerMutex); + llvm::erase(SignalHandlerDevices, Device.Get()); +} - const BufferCreateDesc &getDesc() const override { return Desc; } +llvm::StringRef DXDevice::getAPIName() const { return "DirectX"; } +GPUAPI DXDevice::getAPI() const { return GPUAPI::DirectX; } - static bool classof(const offloadtest::Buffer *B) { - return B->getAPI() == GPUAPI::DirectX; - } -}; - -class DXTexture : public offloadtest::Texture { -public: - ComPtr Resource; - - // Contract: If a command on the command buffer needs a resource to be in a - // different state it should always transition it back to the PreferredState - // afterwards. The PreferredState is the state of the most common use case for - // that resource. This allows us to do state transitions without state - // tracking. - D3D12_RESOURCE_STATES PreferredState; - - // Optional descriptors, depending on Desc.Usage. - // A zero ptr means no descriptor was created for that view type. - D3D12_CPU_DESCRIPTOR_HANDLE RTVHandle = {}; - D3D12_CPU_DESCRIPTOR_HANDLE DSVHandle = {}; - D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle = {}; - D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle = {}; +bool DXDevice::classof(const offloadtest::Device *D) { + return D->getAPI() == GPUAPI::DirectX; +} - std::string Name; - TextureCreateDesc Desc; +Queue &DXDevice::getGraphicsQueue() { return GraphicsQueue; } + +llvm::Error DXDevice::createRootSignatureFromShader( + llvm::StringRef Name, const ShaderContainer &Shader, + ComPtr &OutRootSignature, + DescriptorSetsLayout &Layout) { + // Try pulling a root signature from the DXIL first + auto ExContainer = + llvm::object::DXContainer::create(Shader.Shader->getMemBufferRef()); + // If this fails we really have a problem... + if (!ExContainer) + return ExContainer.takeError(); + + bool HasRootSigPart = false; + for (const auto &Part : *ExContainer) + if (memcmp(Part.Part.Name, "RTS0", 4) == 0) + HasRootSigPart = true; + + if (HasRootSigPart) { + const llvm::StringRef Binary = Shader.Shader->getBuffer(); + if (auto Err = HR::toError( + Device->CreateRootSignature(0, Binary.data(), Binary.size(), + IID_PPV_ARGS(&OutRootSignature)), + "Failed to create root signature.")) + return Err; - DXTexture(ComPtr Resource, llvm::StringRef Name, - TextureCreateDesc Desc, D3D12_RESOURCE_STATES PreferredState, - D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle, - D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle) - : offloadtest::Texture(GPUAPI::DirectX), Resource(Resource), - PreferredState(PreferredState), SRVHandle(SRVHandle), - UAVHandle(UAVHandle), Name(Name), Desc(Desc) {} + const std::wstring WStr(Name.begin(), Name.end()); + OutRootSignature->SetName(WStr.c_str()); - TileShape querySparseTileShape(const Device &Dev) const override; + // Deserialize the root signature to determine how we need to bind + // descriptor tables + ComPtr Deserializer; + if (auto Err = HR::toError( + D3D12CreateVersionedRootSignatureDeserializer( + Binary.data(), Binary.size(), IID_PPV_ARGS(&Deserializer)), + "Failed to create Root Signature Deserializer")) + return Err; - const TextureCreateDesc &getDesc() const override { return Desc; } + const D3D12_VERSIONED_ROOT_SIGNATURE_DESC *RootSigDesc = nullptr; + if (auto Err = HR::toError(Deserializer->GetRootSignatureDescAtVersion( + D3D_ROOT_SIGNATURE_VERSION_1, &RootSigDesc), + "Failed to deseralize root signature")) + return Err; - static bool classof(const offloadtest::Texture *T) { - return T->getAPI() == GPUAPI::DirectX; - } -}; + for (uint32_t I = 0; I < RootSigDesc->Desc_1_0.NumParameters; ++I) { + const D3D12_ROOT_PARAMETER &Parameter = + RootSigDesc->Desc_1_0.pParameters[I]; + switch (Parameter.ParameterType) { + case D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE: { + uint32_t DescriptorCount = 0; + for (uint32_t I = 0; I < Parameter.DescriptorTable.NumDescriptorRanges; + ++I) { + const uint32_t RangeSize = + Parameter.DescriptorTable.pDescriptorRanges[I].NumDescriptors; + if (RangeSize == UINT_MAX) + DescriptorCount += 1024; // Force unbounded arrays to 1024 entries + // in this framework + else + DescriptorCount += RangeSize; + } -class DXPipelineState : public offloadtest::PipelineState { -public: - std::string Name; - ComPtr RootSig; - ComPtr PSO; - // Only set for graphics pipelines. - std::optional Topology; - // True for pipelines created via createPipelineRT — used by SBT / dispatch - // code to safely downcast to DXRayTracingPipelineState (parallel to - // VulkanPipelineState::IsRayTracing). - bool IsRayTracing = false; - - DXPipelineState(llvm::StringRef Name, ComPtr RootSig, - ComPtr PSO, - std::optional Topology, - bool IsRT = false) - : offloadtest::PipelineState(GPUAPI::DirectX), Name(Name), - RootSig(RootSig), PSO(PSO), Topology(Topology), IsRayTracing(IsRT) {} - - static bool classof(const offloadtest::PipelineState *B) { - return B->getAPI() == GPUAPI::DirectX; - } -}; - -/// RT pipeline state: holds the ID3D12StateObject + cached -/// ID3D12StateObjectProperties for SBT identifier queries plus a -/// shader-name → identifier-pointer map. The `void *` identifiers are -/// owned by Properties — keep it alive for the SBT's lifetime. -class DXRayTracingPipelineState : public DXPipelineState { -public: - ComPtr StateObject; - ComPtr Properties; - llvm::StringMap ShaderIdentifiers; - - DXRayTracingPipelineState(llvm::StringRef Name, - ComPtr RootSig, - ComPtr SO, - ComPtr Props) - : DXPipelineState(Name, RootSig, /*PSO=*/nullptr, std::nullopt, - /*IsRT=*/true), - StateObject(SO), Properties(Props) {} - - static bool classof(const offloadtest::PipelineState *B) { - if (B->getAPI() != GPUAPI::DirectX) - return false; - return static_cast(B)->IsRayTracing; + if (Parameter.DescriptorTable.NumDescriptorRanges > 0 && + Parameter.DescriptorTable.pDescriptorRanges[0].RangeType == + D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER) { + Layout.RSigLayout.push_back(RootSignatureLayout( + RootParameterType::SamplerTable, DescriptorCount)); + Layout.Sets.push_back({0, DescriptorCount}); + } else { + Layout.RSigLayout.push_back(RootSignatureLayout( + RootParameterType::DescriptorTable, DescriptorCount)); + Layout.Sets.push_back({DescriptorCount, 0}); + } + break; + } + case D3D12_ROOT_PARAMETER_TYPE_32BIT_CONSTANTS: + Layout.RSigLayout.push_back(RootSignatureLayout( + RootParameterType::Constant, Parameter.Constants.Num32BitValues)); + break; + case D3D12_ROOT_PARAMETER_TYPE_CBV: + Layout.RSigLayout.push_back( + RootSignatureLayout(RootParameterType::CBV, 1)); + break; + case D3D12_ROOT_PARAMETER_TYPE_SRV: + Layout.RSigLayout.push_back( + RootSignatureLayout(RootParameterType::SRV, 1)); + break; + case D3D12_ROOT_PARAMETER_TYPE_UAV: + Layout.RSigLayout.push_back( + RootSignatureLayout(RootParameterType::UAV, 1)); + break; + } + } } -}; -class DXShaderBindingTable : public offloadtest::ShaderBindingTable { -public: - ComPtr Buffer; - // Pre-built ranges for D3D12_DISPATCH_RAYS_DESC. Sizes are zero for - // empty regions; raygen is a single record so it uses the no-stride - // variant. - D3D12_GPU_VIRTUAL_ADDRESS_RANGE RayGenRange{}; - D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE MissRange{}; - D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE HitGroupRange{}; - D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE CallableRange{}; - - DXShaderBindingTable(ComPtr Buf, - D3D12_GPU_VIRTUAL_ADDRESS_RANGE RG, - D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE MS, - D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE HG, - D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE CL) - : offloadtest::ShaderBindingTable(GPUAPI::DirectX), Buffer(Buf), - RayGenRange(RG), MissRange(MS), HitGroupRange(HG), CallableRange(CL) {} - - static bool classof(const offloadtest::ShaderBindingTable *S) { - return S->getAPI() == GPUAPI::DirectX; - } -}; + return llvm::Error::success(); +} -class DXAccelerationStructure : public offloadtest::AccelerationStructure { -public: - ComPtr Resource; - D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle; +llvm::Error DXDevice::createRootSignatureFromBindingsDesc( + llvm::StringRef Name, const BindingsDesc &BndDesc, bool IsGraphics, + ComPtr &OutRootSignature, + DescriptorSetsLayout &Layout) { + uint32_t DescriptorCount = 0; + for (auto &D : BndDesc.DescriptorSetDescs) + DescriptorCount += D.ResourceBindings.size(); + + std::vector RootParams; + const std::unique_ptr Ranges( + new D3D12_DESCRIPTOR_RANGE[DescriptorCount]); + uint32_t RangeIdx = 0; + for (const auto &Set : BndDesc.DescriptorSetDescs) { + uint32_t DescriptorIdx = 0; + const uint32_t StartRangeIdx = RangeIdx; + for (const auto &Binding : Set.ResourceBindings) { + const DescriptorKind Kind = getDescriptorKind(Binding.Kind); + if (Kind == DescriptorKind::SAMPLER) + continue; - DXAccelerationStructure(ComPtr Resource, - const AccelerationStructureSizes &Sizes) - : offloadtest::AccelerationStructure(GPUAPI::DirectX, Sizes), - Resource(Resource) {} + switch (Kind) { + case DescriptorKind::SRV: + Ranges.get()[RangeIdx].RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SRV; + break; + case DescriptorKind::UAV: + Ranges.get()[RangeIdx].RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_UAV; + break; + case DescriptorKind::CBV: + Ranges.get()[RangeIdx].RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_CBV; + break; + case DescriptorKind::SAMPLER: + llvm_unreachable("Sampler should have been filtered out."); + break; + } + Ranges.get()[RangeIdx].NumDescriptors = Binding.DescriptorCount; + Ranges.get()[RangeIdx].BaseShaderRegister = Binding.DXBinding.Register; + Ranges.get()[RangeIdx].RegisterSpace = Binding.DXBinding.Space; + Ranges.get()[RangeIdx].OffsetInDescriptorsFromTableStart = DescriptorIdx; + RangeIdx++; + DescriptorIdx += Binding.DescriptorCount; + } + const uint32_t RangeCount = static_cast(RangeIdx - StartRangeIdx); + if (RangeCount > 0) { + RootParams.push_back( + D3D12_ROOT_PARAMETER{D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE, + {D3D12_ROOT_DESCRIPTOR_TABLE{ + RangeCount, &Ranges.get()[StartRangeIdx]}}, + D3D12_SHADER_VISIBILITY_ALL}); + Layout.RSigLayout.push_back(RootSignatureLayout( + RootParameterType::DescriptorTable, DescriptorIdx)); + } - D3D12_GPU_VIRTUAL_ADDRESS getGPUVirtualAddress() const { - return Resource->GetGPUVirtualAddress(); - } + uint32_t SamplerDescriptorIdx = 0; + const uint32_t SamplerStartRangeIdx = RangeIdx; + for (const auto &Binding : Set.ResourceBindings) { + const DescriptorKind Kind = getDescriptorKind(Binding.Kind); + if (Kind != DescriptorKind::SAMPLER) + continue; - static bool classof(const offloadtest::AccelerationStructure *AS) { - return AS->getAPI() == GPUAPI::DirectX; - } -}; + Ranges.get()[RangeIdx].RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SAMPLER; + Ranges.get()[RangeIdx].NumDescriptors = Binding.DescriptorCount; + Ranges.get()[RangeIdx].BaseShaderRegister = Binding.DXBinding.Register; + Ranges.get()[RangeIdx].RegisterSpace = Binding.DXBinding.Space; + Ranges.get()[RangeIdx].OffsetInDescriptorsFromTableStart = + SamplerDescriptorIdx; -class DXFence : public offloadtest::Fence { -public: -#ifdef _WIN32 - DXFence(ComPtr Fence, HANDLE Event, llvm::StringRef Name) -#else // WSL - DXFence(ComPtr Fence, int Event, llvm::StringRef Name) -#endif - : Name(Name), Fence(Fence), Event(Event) { - } + assert(Binding.DescriptorCount == 1 && "Manon expected this to be 1."); + RangeIdx++; + SamplerDescriptorIdx += Binding.DescriptorCount; + } - std::string Name; - ComPtr Fence; -#ifdef _WIN32 - HANDLE Event; -#else // WSL - int Event; -#endif + const uint32_t SamplerRangeCount = + static_cast(RangeIdx - SamplerStartRangeIdx); + if (SamplerRangeCount > 0) { + RootParams.push_back(D3D12_ROOT_PARAMETER{ + D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE, + {D3D12_ROOT_DESCRIPTOR_TABLE{ + static_cast(RangeIdx - SamplerStartRangeIdx), + &Ranges.get()[SamplerStartRangeIdx]}}, + D3D12_SHADER_VISIBILITY_ALL}); + Layout.RSigLayout.push_back(RootSignatureLayout( + RootParameterType::SamplerTable, SamplerDescriptorIdx)); + } - static llvm::Expected> create(ID3D12DeviceX *Device, - llvm::StringRef Name) { - ComPtr Fence; - if (auto Err = HR::toError( - Device->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&Fence)), - "Failed to create Fence.")) - return Err; + Layout.Sets.push_back({ + DescriptorIdx, // DescriptorCount + SamplerDescriptorIdx // SamplerCount + }); + } -#ifdef _WIN32 - HANDLE Event = CreateEventA(nullptr, false, false, nullptr); - if (!Event) -#else // WSL - int Event = eventfd(0, EFD_CLOEXEC | EFD_NONBLOCK); - if (Event == -1) -#endif - return llvm::createStringError(std::errc::device_or_resource_busy, - "Failed to create event."); + CD3DX12_ROOT_SIGNATURE_DESC Desc; + Desc.Init( + static_cast(RootParams.size()), RootParams.data(), 0, nullptr, + IsGraphics ? D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT + : D3D12_ROOT_SIGNATURE_FLAG_NONE); + + ComPtr Signature; + ComPtr Error; + if (auto Err = HR::toError( + D3D12SerializeRootSignature(&Desc, D3D_ROOT_SIGNATURE_VERSION_1, + &Signature, &Error), + "Failed to seialize root signature.")) { + const std::string Msg = + std::string(reinterpret_cast(Error->GetBufferPointer()), + Error->GetBufferSize() / sizeof(char)); + return joinErrors( + std::move(Err), + llvm::createStringError(std::errc::protocol_error, Msg.c_str())); + } + + if (auto Err = HR::toError( + Device->CreateRootSignature(0, Signature->GetBufferPointer(), + Signature->GetBufferSize(), + IID_PPV_ARGS(&OutRootSignature)), + "Failed to create root signature.")) + return Err; - return std::make_unique(Fence, Event, Name); - } + const std::wstring WStr(Name.begin(), Name.end()); + OutRootSignature->SetName(WStr.c_str()); - ~DXFence() { -#ifdef _WIN32 - CloseHandle(Event); -#else // WSL - close(Event); -#endif - } + return llvm::Error::success(); +} - uint64_t getFenceValue() override { return Fence->GetCompletedValue(); } +llvm::Error +DXDevice::createRootSignature(llvm::StringRef Name, const BindingsDesc &BndDesc, + const ShaderContainer &Shader, bool IsGraphics, + ComPtr &OutRootSignature, + DescriptorSetsLayout &Layout) { + assert(OutRootSignature.Get() == nullptr); - llvm::Error waitForCompletion(uint64_t SignalValue) override { - if (Fence->GetCompletedValue() >= SignalValue) - return llvm::Error::success(); + if (auto Err = + createRootSignatureFromShader(Name, Shader, OutRootSignature, Layout)) + return Err; -#ifdef _WIN32 - if (auto Err = HR::toError(Fence->SetEventOnCompletion(SignalValue, Event), - "Failed to register end event.")) - return Err; - WaitForSingleObject(Event, INFINITE); -#else // WSL - if (auto Err = - HR::toError(Fence->SetEventOnCompletion( - SignalValue, reinterpret_cast(Event)), - "Failed to register end event.")) - return Err; - pollfd PollEvent; - PollEvent.fd = Event; - PollEvent.events = POLLIN; - PollEvent.revents = 0; - if (poll(&PollEvent, 1, -1) == -1) - return llvm::createStringError( - std::error_code(errno, std::system_category()), strerror(errno)); -#endif + if (OutRootSignature.Get() != nullptr) return llvm::Error::success(); - } -}; -class DXMemoryHeap : public offloadtest::MemoryHeap { -public: - DXMemoryHeap(ComPtr Heap, llvm::StringRef Name) - : offloadtest::MemoryHeap(GPUAPI::DirectX), Name(Name), Heap(Heap) {} + return createRootSignatureFromBindingsDesc(Name, BndDesc, IsGraphics, + OutRootSignature, Layout); +} - std::string Name; - ComPtr Heap; +llvm::Expected> +DXDevice::createPipelineCs(llvm::StringRef Name, const BindingsDesc &BndDesc, + ShaderContainer CS) { + ComPtr RootSig; + DescriptorSetsLayout Layout; + if (auto Err = createRootSignature(Name, BndDesc, CS, + /*IsGraphics=*/false, RootSig, Layout)) + return Err; - static bool classof(const offloadtest::MemoryHeap *H) { - return H->getAPI() == GPUAPI::DirectX; - } + auto DXIL = CS.Shader->getBuffer(); + const D3D12_COMPUTE_PIPELINE_STATE_DESC Desc = { + RootSig.Get(), + {DXIL.data(), DXIL.size()}, + 0, + { + nullptr, + 0, + }, + D3D12_PIPELINE_STATE_FLAG_NONE}; - static llvm::Expected> - create(ID3D12DeviceX *Device, llvm::StringRef Name, size_t SizeInBytes) { - D3D12_HEAP_DESC HeapDesc = {}; - HeapDesc.Properties = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_DEFAULT); - HeapDesc.Alignment = D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT; - // Heap sizes must be a multiple of the placement alignment (64KB), which - // is also the tiled-resource tile size. - HeapDesc.SizeInBytes = - llvm::alignTo(SizeInBytes, D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT); - HeapDesc.Flags = D3D12_HEAP_FLAG_ALLOW_ALL_BUFFERS_AND_TEXTURES; - - ComPtr Heap; - if (auto Err = - HR::toError(Device->CreateHeap(&HeapDesc, IID_PPV_ARGS(&Heap)), - "Failed to create memory heap.")) - return Err; + ComPtr PSO; + if (auto Err = HR::toError( + Device->CreateComputePipelineState(&Desc, IID_PPV_ARGS(&PSO)), + "Failed to create PSO.")) + return Err; - const std::wstring WStr(Name.begin(), Name.end()); - Heap->SetName(WStr.c_str()); + return std::make_unique(Name, RootSig, std::move(Layout), + PSO, std::nullopt); +} - return std::make_unique(Heap, Name); - } -}; - -class DXQueue : public offloadtest::Queue { -public: - using Queue::submit; - - ComPtr Queue; - std::unique_ptr SubmitFence; - uint64_t FenceCounter = 0; - // Batches of command buffers submitted to the GPU that may still be - // in-flight. The ID3D12CommandAllocator owns the backing memory for - // recorded commands, so it must outlive GPU execution. Each batch - // records the fence value it signals so we can non-blockingly query - // progress and release completed batches. - struct InFlightBatch { - uint64_t FenceValue; - llvm::SmallVector> CBs; - }; - llvm::SmallVector InFlightBatches; - - DXQueue(ComPtr Queue, - std::unique_ptr SubmitFence) - : Queue(Queue), SubmitFence(std::move(SubmitFence)) {} - DXQueue(DXQueue &&) = default; - ~DXQueue() override {} - - static llvm::Expected - createGraphicsQueue(ComPtr Device) { - const D3D12_COMMAND_QUEUE_DESC Desc = {D3D12_COMMAND_LIST_TYPE_DIRECT, 0, - D3D12_COMMAND_QUEUE_FLAG_NONE, 0}; - ComPtr CmdQueue; - if (auto Err = HR::toError( - Device->CreateCommandQueue(&Desc, IID_PPV_ARGS(&CmdQueue)), - "Failed to create command queue.")) - return Err; - auto FenceOrErr = DXFence::create(Device.Get(), "QueueSubmitFence"); - if (!FenceOrErr) - return FenceOrErr.takeError(); - return DXQueue(CmdQueue, std::move(*FenceOrErr)); - } +llvm::Expected> +DXDevice::createTraditionalRasterPipeline( + llvm::StringRef Name, const BindingsDesc &BndDesc, + const TraditionalRasterPipelineCreateDesc &Desc) { + assert(Desc.RTFormats.size() <= 8); - llvm::Expected - submit(llvm::SmallVector> CBs) - override; + ComPtr RootSig; + DescriptorSetsLayout Layout; + if (auto Err = createRootSignature(Name, BndDesc, Desc.VS, + /*IsGraphics=*/true, RootSig, Layout)) + return Err; - llvm::Expected - updateTileMappings(offloadtest::Buffer &Resource, - llvm::ArrayRef Mappings) override { - return updateTileMappingsImpl(llvm::cast(Resource).Buffer.Get(), - Mappings); - } + std::vector DXInputLayout; + DXInputLayout.reserve(Desc.InputLayout.size()); + for (const InputLayoutDesc &Elem : Desc.InputLayout) { + assert(!Elem.InstanceStepRate && + "Instance step rate is currently not supported."); + + D3D12_INPUT_ELEMENT_DESC ElementDesc = {}; + ElementDesc.SemanticName = Elem.Name.c_str(); + ElementDesc.SemanticIndex = 0; + ElementDesc.Format = getDXGIFormat(Elem.Fmt); + ElementDesc.InputSlot = 0; + ElementDesc.AlignedByteOffset = Elem.OffsetInBytes; + ElementDesc.InputSlotClass = D3D12_INPUT_CLASSIFICATION_PER_VERTEX_DATA; + DXInputLayout.push_back(ElementDesc); + } + + D3D12_GRAPHICS_PIPELINE_STATE_DESC PSODesc = {}; + PSODesc.InputLayout = {DXInputLayout.data(), (UINT)DXInputLayout.size()}; + PSODesc.pRootSignature = RootSig.Get(); + PSODesc.VS = {Desc.VS.Shader->getBuffer().data(), + Desc.VS.Shader->getBuffer().size()}; + PSODesc.PS = {Desc.PS.Shader->getBuffer().data(), + Desc.PS.Shader->getBuffer().size()}; + if (PSODesc.VS.BytecodeLength == 0 || PSODesc.PS.BytecodeLength == 0) + return llvm::createStringError(std::errc::invalid_argument, + "Graphics pipeline requires both a vertex " + "shader and a pixel shader."); + if (Desc.HS) + PSODesc.HS = {Desc.HS->Shader->getBuffer().data(), + Desc.HS->Shader->getBuffer().size()}; + if (Desc.DS) + PSODesc.DS = {Desc.DS->Shader->getBuffer().data(), + Desc.DS->Shader->getBuffer().size()}; + if (Desc.GS) + PSODesc.GS = {Desc.GS->Shader->getBuffer().data(), + Desc.GS->Shader->getBuffer().size()}; + + PSODesc.RasterizerState = CD3DX12_RASTERIZER_DESC(D3D12_DEFAULT); + PSODesc.RasterizerState.CullMode = D3D12_CULL_MODE_NONE; + PSODesc.RasterizerState.FrontCounterClockwise = TRUE; + PSODesc.BlendState = CD3DX12_BLEND_DESC(D3D12_DEFAULT); + PSODesc.DepthStencilState.DepthEnable = true; + PSODesc.DepthStencilState.DepthWriteMask = D3D12_DEPTH_WRITE_MASK_ALL; + PSODesc.DepthStencilState.DepthFunc = D3D12_COMPARISON_FUNC_LESS; + PSODesc.DepthStencilState.StencilEnable = false; + PSODesc.SampleMask = UINT_MAX; + PSODesc.PrimitiveTopologyType = getDXPrimitiveTopologyType(Desc.Topology); + PSODesc.NumRenderTargets = static_cast(Desc.RTFormats.size()); + if (Desc.DSFormat) + PSODesc.DSVFormat = getDXGIFormat(*Desc.DSFormat); + for (size_t I = 0; I < Desc.RTFormats.size(); ++I) + PSODesc.RTVFormats[I] = getDXGIFormat(Desc.RTFormats[I]); + PSODesc.SampleDesc.Count = 1; - llvm::Expected - updateTileMappings(offloadtest::Texture &Resource, - llvm::ArrayRef Mappings) override { - return updateTileMappingsImpl( - llvm::cast(Resource).Resource.Get(), Mappings); - } + ComPtr PSO; + if (auto Err = HR::toError( + Device->CreateGraphicsPipelineState(&PSODesc, IID_PPV_ARGS(&PSO)), + "Failed to create graphics PSO.")) + return Err; - // Shared tile-mapping path for buffers and textures: they differ only in how - // a TileRegion maps onto the underlying ID3D12Resource, so both variants - // resolve to a native resource and funnel through here. - llvm::Expected - updateTileMappingsImpl(ID3D12Resource *Resource, - llvm::ArrayRef Mappings) { - for (const TileMapping &M : Mappings) { - const D3D12_TILED_RESOURCE_COORDINATE StartCoord = { - M.Region.TileOffsetX, M.Region.TileOffsetY, M.Region.TileOffsetZ, - M.Region.Subresource}; - - D3D12_TILE_REGION_SIZE RegionSize = {}; - RegionSize.Width = M.Region.NumTilesX; - RegionSize.Height = static_cast(M.Region.NumTilesY); - RegionSize.Depth = static_cast(M.Region.NumTilesZ); - // A box is only needed for multi-dimensional (texture) regions; a buffer - // is always a single linear run of tiles along X. - RegionSize.UseBox = M.Region.NumTilesY > 1 || M.Region.NumTilesZ > 1; - RegionSize.NumTiles = - M.Region.NumTilesX * M.Region.NumTilesY * M.Region.NumTilesZ; - - // A null heap unbinds the region (reads return zero); otherwise bind it - // to the heap starting at BackingTileOffset. - ID3D12Heap *Heap = nullptr; - D3D12_TILE_RANGE_FLAGS RangeFlag = D3D12_TILE_RANGE_FLAG_NULL; - UINT HeapRangeStartOffset = 0; - if (M.Backing) { - Heap = llvm::cast(M.Backing)->Heap.Get(); - RangeFlag = D3D12_TILE_RANGE_FLAG_NONE; - HeapRangeStartOffset = static_cast(M.BackingTileOffset); - } - const UINT RangeTileCount = RegionSize.NumTiles; + return std::make_unique( + Name, RootSig, std::move(Layout), PSO, + getDXPrimitiveTopology(Desc.Topology, Desc.PatchControlPoints)); +} - Queue->UpdateTileMappings(Resource, 1, &StartCoord, &RegionSize, Heap, 1, - &RangeFlag, &HeapRangeStartOffset, - &RangeTileCount, D3D12_TILE_MAPPING_FLAG_NONE); - } +llvm::Expected> +DXDevice::createMeshShaderRasterPipeline( + llvm::StringRef Name, const BindingsDesc &BindingsDesc, + const MeshShaderRasterPipelineCreateDesc &Desc) { + assert(Desc.RTFormats.size() <= 8); - // UpdateTileMappings runs on the queue timeline (not a command list), so - // signal the submit fence afterwards exactly like submit() does. - const uint64_t CurrentCounter = ++FenceCounter; - if (auto Err = - HR::toError(Queue->Signal(SubmitFence->Fence.Get(), CurrentCounter), - "Failed to signal after tile mapping update.")) - return Err; + ComPtr RootSig; + DescriptorSetsLayout Layout; + if (auto Err = createRootSignature(Name, BindingsDesc, Desc.MS, + /*IsGraphics=*/true, RootSig, Layout)) + return Err; - return offloadtest::SubmitResult{SubmitFence.get(), CurrentCounter}; - } -}; - -class DXDevice; // forward decl — defined below in this same anon ns - -class DXCommandBuffer : public offloadtest::CommandBuffer { -public: - ComPtr Allocator; - ComPtr CmdList; - /// Back-pointer to the owning device. Used by encoders that need access to - /// device-level resources (e.g. allocating AS scratch buffers). - DXDevice *Dev = nullptr; - /// Whether a UAV barrier is pending from a prior compute command. - bool PendingUAVBarrier = false; - llvm::SmallVector PendingTransitions; - /// Buffers that must outlive command-buffer submission (e.g. AS scratch - /// and TLAS instance buffers used during builds). - llvm::SmallVector> KeepAliveOwned; - - static llvm::Expected> - create(ComPtr Device) { - auto CB = std::unique_ptr(new DXCommandBuffer()); - if (auto Err = HR::toError( - Device->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT, - IID_PPV_ARGS(&CB->Allocator)), - "Failed to create command allocator.")) - return Err; - if (auto Err = HR::toError( - Device->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT, - CB->Allocator.Get(), nullptr, - IID_PPV_ARGS(&CB->CmdList)), - "Failed to create command list.")) - return Err; - return CB; - } + const D3D12_SHADER_BYTECODE MSBytecode = {Desc.MS.Shader->getBuffer().data(), + Desc.MS.Shader->getBuffer().size()}; + if (MSBytecode.BytecodeLength == 0) + return llvm::createStringError( + std::errc::invalid_argument, + "Mesh shader pipeline requires a mesh shader."); + + // The amplification (task) shader is optional. + D3D12_SHADER_BYTECODE ASBytecode = {}; + if (Desc.AS) { + assert((*Desc.AS).Shader->getBufferSize() > 0 && + "The passed task/amplification shader was empty."); + ASBytecode = {(*Desc.AS).Shader->getBuffer().data(), + (*Desc.AS).Shader->getBuffer().size()}; + } + + // The pixel shader is optional + D3D12_SHADER_BYTECODE PSBytecode = {}; + if (Desc.PS) { + assert((*Desc.PS).Shader->getBufferSize() > 0 && + "The passed pixel shader was empty."); + PSBytecode = {(*Desc.PS).Shader->getBuffer().data(), + (*Desc.PS).Shader->getBuffer().size()}; + } + + D3D12_RT_FORMAT_ARRAY RTArray = {}; + RTArray.NumRenderTargets = static_cast(Desc.RTFormats.size()); + for (size_t I = 0; I < Desc.RTFormats.size(); ++I) + RTArray.RTFormats[I] = getDXGIFormat(Desc.RTFormats[I]); + + CD3DX12_DEPTH_STENCIL_DESC1 DepthStencil(D3D12_DEFAULT); + DepthStencil.DepthEnable = true; + DepthStencil.DepthWriteMask = D3D12_DEPTH_WRITE_MASK_ALL; + DepthStencil.DepthFunc = D3D12_COMPARISON_FUNC_LESS; + DepthStencil.StencilEnable = false; + + DXGI_SAMPLE_DESC SampleDesc = {}; + SampleDesc.Count = 1; + + CD3DX12_PIPELINE_MESH_STATE_STREAM Stream; + Stream.pRootSignature = RootSig.Get(); + Stream.AS = ASBytecode; + Stream.MS = MSBytecode; + Stream.PS = PSBytecode; + Stream.RasterizerState = CD3DX12_RASTERIZER_DESC(D3D12_DEFAULT); + Stream.BlendState = CD3DX12_BLEND_DESC(D3D12_DEFAULT); + Stream.DepthStencilState = DepthStencil; + Stream.SampleMask = UINT_MAX; + Stream.PrimitiveTopologyType = getDXPrimitiveTopologyType(Desc.Topology); + Stream.RTVFormats = RTArray; + if (Desc.DSFormat) + Stream.DSVFormat = getDXGIFormat(*Desc.DSFormat); + Stream.SampleDesc = SampleDesc; + + const D3D12_PIPELINE_STATE_STREAM_DESC StreamDesc = {sizeof(Stream), &Stream}; + + ComPtr PSO; + if (auto Err = HR::toError( + Device->CreatePipelineState(&StreamDesc, IID_PPV_ARGS(&PSO)), + "Failed to create mesh shader PSO.")) + return Err; + + return std::make_unique(Name, RootSig, std::move(Layout), + PSO, std::nullopt); +} - ~DXCommandBuffer() override = default; +llvm::Expected> +DXDevice::createPipelineRT(llvm::StringRef Name, const BindingsDesc &BndDesc, + const RayTracingPipelineCreateDesc &Desc) { + if (!Desc.Library) + return llvm::createStringError(std::errc::invalid_argument, + "RayTracingPipelineCreateDesc.Library is " + "null — backend needs a DXIL blob."); + + // Global root signature: try the library's embedded RTS0 part first; + // fall back to building one from BindingsDesc. + ShaderContainer LibContainer = {}; + LibContainer.Shader = Desc.Library; + ComPtr RootSig; + DescriptorSetsLayout Layout; + if (auto Err = createRootSignature(Name, BndDesc, LibContainer, + /*IsGraphics=*/false, RootSig, Layout)) + return Err; - static bool classof(const CommandBuffer *CB) { - return CB->getKind() == GPUAPI::DirectX; + CD3DX12_STATE_OBJECT_DESC SODesc(D3D12_STATE_OBJECT_TYPE_RAYTRACING_PIPELINE); + + // DXIL library subobject — add every Shader's entry point as an export. + // Wide-string storage must outlive SODesc since the subobject only stores + // pointers into it. + auto *Lib = SODesc.CreateSubobject(); + const llvm::StringRef LibBytes = Desc.Library->getBuffer(); + const D3D12_SHADER_BYTECODE Bytecode = {LibBytes.data(), LibBytes.size()}; + Lib->SetDXILLibrary(&Bytecode); + llvm::SmallVector WideNames; + WideNames.reserve(Desc.Shaders.size() + Desc.HitGroups.size()); + for (const auto &Sh : Desc.Shaders) { + WideNames.push_back(widen(Sh.EntryPoint)); + Lib->DefineExport(WideNames.back().c_str()); + } + + // One hit-group subobject per HitGroup entry. + for (const auto &HG : Desc.HitGroups) { + auto *HGObj = SODesc.CreateSubobject(); + HGObj->SetHitGroupType(getDXHitGroupType(HG.Type)); + WideNames.push_back(widen(HG.Name)); + HGObj->SetHitGroupExport(WideNames.back().c_str()); + WideNames.push_back(widen(HG.ClosestHit)); + HGObj->SetClosestHitShaderImport(WideNames.back().c_str()); + if (HG.AnyHit) { + WideNames.push_back(widen(*HG.AnyHit)); + HGObj->SetAnyHitShaderImport(WideNames.back().c_str()); + } + if (HG.Intersection) { + WideNames.push_back(widen(*HG.Intersection)); + HGObj->SetIntersectionShaderImport(WideNames.back().c_str()); + } } - void addPendingUAVBarrier() { PendingUAVBarrier = true; } - void addResourceTransition(ID3D12Resource *Resource, - D3D12_RESOURCE_STATES StateBefore, - D3D12_RESOURCE_STATES StateAfter) { + // Pipeline-wide shader config (max payload + max attribute bytes). + auto *ShaderCfg = + SODesc.CreateSubobject(); + ShaderCfg->Config(Desc.Config.MaxPayloadSizeInBytes, + Desc.Config.MaxAttributeSizeInBytes); - for (auto TransIt = PendingTransitions.begin(), - End = PendingTransitions.end(); - TransIt != End; ++TransIt) { - if (TransIt->Transition.pResource != Resource) - continue; + // Pipeline-wide config (max recursion depth). + auto *PipelineCfg = + SODesc.CreateSubobject(); + PipelineCfg->Config(Desc.Config.MaxTraceRecursionDepth); - assert(StateBefore == TransIt->Transition.StateAfter); + // Global root signature. + auto *GlobalRS = + SODesc.CreateSubobject(); + GlobalRS->SetRootSignature(RootSig.Get()); - if (StateAfter == TransIt->Transition.StateBefore) { - // We actually need the resource to be in the original state! - // Remove the transition from the list of Pending Transitions. - PendingTransitions.erase(TransIt); + ComPtr StateObject; + if (auto Err = HR::toError( + Device->CreateStateObject(SODesc, IID_PPV_ARGS(&StateObject)), + "Failed to create raytracing state object.")) + return Err; - // NOTE: This could cause an execution barrier issue since legacy - // barriers don't allow us to express execution barriers, cache - // flushes, and layout transitions separately. - // As a workaround we add a UAV barrier so an execution barrier and - // cache flush is inserted. - addPendingUAVBarrier(); - } else { - TransIt->Transition.StateAfter = StateAfter; - } + ComPtr Properties; + if (auto Err = HR::toError( + StateObject.As(&Properties), + "Failed to query ID3D12StateObjectProperties from state object.")) + return Err; - return; + auto State = std::make_unique( + Name, RootSig, Layout, StateObject, Properties); + // Cache identifiers up-front. The driver-owned blobs are alive for + // Properties' lifetime, which lives on the PSO. + // + // GetShaderIdentifier only returns non-null for entries that are + // directly bindable from an SBT record: raygen / miss / callable + // shaders and hit groups. Closest-hit / any-hit / intersection are + // bound *through* a hit-group subobject and aren't separately + // addressable, so skip them. + for (const auto &Sh : Desc.Shaders) { + switch (Sh.Stage) { + case Stages::RayGeneration: + case Stages::Miss: + case Stages::Callable: + break; + default: + continue; } - - PendingTransitions.push_back(CD3DX12_RESOURCE_BARRIER::Transition( - Resource, StateBefore, StateAfter)); + const std::wstring W = widen(Sh.EntryPoint); + const void *Id = Properties->GetShaderIdentifier(W.c_str()); + if (!Id) + return llvm::createStringError( + "GetShaderIdentifier returned null for shader '%s'", + Sh.EntryPoint.c_str()); + State->ShaderIdentifiers[Sh.EntryPoint] = Id; + } + for (const auto &HG : Desc.HitGroups) { + const std::wstring W = widen(HG.Name); + const void *Id = Properties->GetShaderIdentifier(W.c_str()); + if (!Id) + return llvm::createStringError( + "GetShaderIdentifier returned null for hit group '%s'", + HG.Name.c_str()); + State->ShaderIdentifiers[HG.Name] = Id; } + return State; +} - void flushBarrier() { +llvm::Expected> +DXDevice::createShaderBindingTable(const PipelineState &PSO, + const ShaderBindingTableDesc &Desc) { + if (!llvm::isa(&PSO)) + return llvm::createStringError( + std::errc::invalid_argument, + "createShaderBindingTable requires a RayTracing PipelineState"); + const auto &DXRTPSO = llvm::cast(PSO); - if (PendingUAVBarrier) { - PendingTransitions.push_back(CD3DX12_RESOURCE_BARRIER::UAV(nullptr)); - PendingUAVBarrier = false; - } + constexpr uint32_t IdSize = D3D12_SHADER_IDENTIFIER_SIZE_IN_BYTES; + const SBTLayout Layout = + computeSBTLayout(IdSize, D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT, + D3D12_RAYTRACING_SHADER_TABLE_BYTE_ALIGNMENT, Desc); + const uint32_t TotalSize = Layout.TotalSize; + const llvm::ArrayRef RGEntries(&Desc.RayGen, 1); + + // Upload heap so the CPU can write the SBT directly. The state-object + // identifiers don't need to live in default heap; using upload keeps + // PR3 simple. A staging copy to default heap is a follow-up. + const D3D12_HEAP_PROPERTIES HeapProps = + CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_UPLOAD); + const D3D12_RESOURCE_DESC BufDesc = CD3DX12_RESOURCE_DESC::Buffer(TotalSize); + ComPtr Buffer; + if (auto Err = HR::toError(Device->CreateCommittedResource( + &HeapProps, D3D12_HEAP_FLAG_NONE, &BufDesc, + D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, + IID_PPV_ARGS(&Buffer)), + "Failed to create SBT buffer.")) + return Err; - if (!PendingTransitions.empty()) { - CmdList->ResourceBarrier(PendingTransitions.size(), - PendingTransitions.data()); - PendingTransitions.clear(); + void *Mapped = nullptr; + const D3D12_RANGE ReadRange{0, 0}; + if (auto Err = HR::toError(Buffer->Map(0, &ReadRange, &Mapped), + "Failed to map SBT buffer.")) + return Err; + std::memset(Mapped, 0, TotalSize); + auto *MappedBytes = static_cast(Mapped); + + auto WriteEntries = [&](uint8_t *Region, llvm::ArrayRef Entries, + uint32_t Stride) -> llvm::Error { + for (size_t I = 0; I < Entries.size(); ++I) { + const auto &E = Entries[I]; + auto It = DXRTPSO.ShaderIdentifiers.find(E.ShaderName); + if (It == DXRTPSO.ShaderIdentifiers.end()) + return llvm::createStringError( + std::errc::invalid_argument, + "SBT references unknown shader/hit-group name: '%s'", + E.ShaderName.c_str()); + uint8_t *Dst = Region + I * Stride; + std::memcpy(Dst, It->second, IdSize); + if (!E.LocalRootData.empty()) + std::memcpy(Dst + IdSize, E.LocalRootData.data(), + E.LocalRootData.size()); } - } + return llvm::Error::success(); + }; + + auto WriteRegion = [&](const SBTRegionLayout &R, + llvm::ArrayRef Entries) -> llvm::Error { + return WriteEntries(MappedBytes + R.Offset, Entries, R.Stride); + }; + auto UnmapAndReturn = [&](llvm::Error Err) { + Buffer->Unmap(0, nullptr); + return Err; + }; + if (auto Err = WriteRegion(Layout.RayGen, RGEntries)) + return UnmapAndReturn(std::move(Err)); + if (auto Err = WriteRegion(Layout.Miss, Desc.Miss)) + return UnmapAndReturn(std::move(Err)); + if (auto Err = WriteRegion(Layout.HitGroup, Desc.HitGroup)) + return UnmapAndReturn(std::move(Err)); + if (auto Err = WriteRegion(Layout.Callable, Desc.Callable)) + return UnmapAndReturn(std::move(Err)); + Buffer->Unmap(0, nullptr); + + // D3D12_GPU_VIRTUAL_ADDRESS_RANGE / …_AND_STRIDE expect a zero address + // for empty regions, matching the helper's Size == 0 sentinel. + const D3D12_GPU_VIRTUAL_ADDRESS Base = Buffer->GetGPUVirtualAddress(); + auto MakeRange = [&](const SBTRegionLayout &R) { + return D3D12_GPU_VIRTUAL_ADDRESS_RANGE{R.Size ? Base + R.Offset : 0, + R.Size}; + }; + auto MakeRangeAndStride = [&](const SBTRegionLayout &R) { + return D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE{ + R.Size ? Base + R.Offset : 0, R.Size, R.Stride}; + }; + return std::make_unique( + Buffer, MakeRange(Layout.RayGen), MakeRangeAndStride(Layout.Miss), + MakeRangeAndStride(Layout.HitGroup), MakeRangeAndStride(Layout.Callable)); +} + +llvm::Expected> +DXDevice::createFence(llvm::StringRef Name) { + return DXFence::create(Device.Get(), Name); +} - llvm::Expected> - createComputeEncoder() override; +llvm::Expected> +DXDevice::createMemoryHeap(std::string Name, size_t SizeInBytes) { + return DXMemoryHeap::create(Device.Get(), Name, SizeInBytes); +} - llvm::Expected> - createRenderEncoder(const offloadtest::RenderPassBeginDesc &Desc) override; +llvm::Expected> +DXDevice::createBuffer(std::string Name, const BufferCreateDesc &Desc, + size_t SizeInBytes) { + const D3D12_HEAP_TYPE HeapType = getDXHeapType(Desc.Location); + // This flag is only allowed on GpuOnly memory. + const D3D12_RESOURCE_FLAGS Flags = + Desc.Location == MemoryLocation::GpuOnly + ? D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS + : D3D12_RESOURCE_FLAG_NONE; + + // Modify the size if needed + UINT64 CounterOffsetInBytes = 0; + UINT64 BufferSizeInBytes = SizeInBytes; + if (Desc.Usage == BufferUsage::ConstantBuffer) { + if (Desc.HasCounter) + return llvm::createStringError( + "Constant Buffers are not allowed to have a counter."); -private: - DXCommandBuffer() : CommandBuffer(GPUAPI::DirectX) {} -}; + BufferSizeInBytes = getCBVSize(BufferSizeInBytes); + } else if (Desc.HasCounter) { + if (Desc.AccessType == BufferShaderAccessType::Raw) + return llvm::createStringError( + "Raw Resources are not allowed to have a counter."); -struct DescriptorAllocator { - ComPtr Heap; - std::atomic NextIndex{0}; - uint32_t DescIncSize; - uint32_t Capacity; - - static llvm::Expected - create(ID3D12DeviceX *Device, D3D12_DESCRIPTOR_HEAP_TYPE Type, - uint32_t Capacity) { - ComPtr Heap; - const D3D12_DESCRIPTOR_HEAP_DESC HeapDesc = { - Type, Capacity, D3D12_DESCRIPTOR_HEAP_FLAG_NONE, 0}; - if (auto Err = HR::toError( - Device->CreateDescriptorHeap(&HeapDesc, IID_PPV_ARGS(&Heap)), - "Failed to create descriptor heap for DescriptorAllocator.")) - return Err; - const uint32_t DescIncSize = Device->GetDescriptorHandleIncrementSize(Type); - return DescriptorAllocator(Heap, DescIncSize, Capacity); + CounterOffsetInBytes = + llvm::alignTo(BufferSizeInBytes, D3D12_UAV_COUNTER_PLACEMENT_ALIGNMENT); + BufferSizeInBytes = CounterOffsetInBytes + sizeof(uint32_t); } - llvm::Expected allocate() { - // TODO(manon): Use a better allocator that can also free descriptors. - const uint32_t Index = NextIndex.fetch_add(1, std::memory_order_relaxed); - if (Index >= Capacity) - return llvm::createStringError(std::errc::not_enough_memory, - "Descriptor heap allocator exhausted."); - D3D12_CPU_DESCRIPTOR_HANDLE Handle = - Heap->GetCPUDescriptorHandleForHeapStart(); - Handle.ptr += Index * DescIncSize; - return Handle; - } + const D3D12_HEAP_PROPERTIES HeapProps = CD3DX12_HEAP_PROPERTIES(HeapType); + const D3D12_RESOURCE_DESC BufferDesc = + CD3DX12_RESOURCE_DESC::Buffer(BufferSizeInBytes, Flags); - DescriptorAllocator(DescriptorAllocator &&Other) - : Heap(std::move(Other.Heap)), - NextIndex(Other.NextIndex.load(std::memory_order_relaxed)), - DescIncSize(Other.DescIncSize), Capacity(Other.Capacity) {} - DescriptorAllocator &operator=(DescriptorAllocator &&) = delete; - DescriptorAllocator(const DescriptorAllocator &) = delete; - DescriptorAllocator &operator=(const DescriptorAllocator &) = delete; - - DescriptorAllocator(ComPtr Heap, uint32_t DescIncSize, - uint32_t Capacity) - : Heap(Heap), DescIncSize(DescIncSize), Capacity(Capacity) {} -}; - -class DXComputeEncoder : public offloadtest::ComputeEncoder { - DXCommandBuffer &CB; - - void addUAVBarrier() { - CB.addPendingUAVBarrier(); - CB.flushBarrier(); + D3D12_RESOURCE_STATES PreferredState = D3D12_RESOURCE_STATE_COMMON; + ComPtr BufferObject; + if (Desc.Backing == MemoryBacking::Sparse) { + if (auto Err = HR::toError(Device->CreateReservedResource( + &BufferDesc, D3D12_RESOURCE_STATE_COMMON, + nullptr, IID_PPV_ARGS(&BufferObject)), + "Failed to create reserved buffer.")) + return Err; + } else { + D3D12_RESOURCE_STATES InitialState = D3D12_RESOURCE_STATE_COMMON; + if (HeapType == D3D12_HEAP_TYPE_UPLOAD) + InitialState = D3D12_RESOURCE_STATE_GENERIC_READ; + else if (HeapType == D3D12_HEAP_TYPE_READBACK) + // As per the readback heap docs + // > Resources in this heap must be created with + // > D3D12_RESOURCE_STATE_COPY_DEST, and cannot be changed away from + // this. + InitialState = D3D12_RESOURCE_STATE_COPY_DEST; + PreferredState = InitialState; + if (auto Err = + HR::toError(Device->CreateCommittedResource( + &HeapProps, D3D12_HEAP_FLAG_NONE, &BufferDesc, + InitialState, nullptr, IID_PPV_ARGS(&BufferObject)), + "Failed to create buffer.")) + return Err; } -public: - DXComputeEncoder(DXCommandBuffer &CB) - : ComputeEncoder(GPUAPI::DirectX), CB(CB) {} + const std::wstring WStr(Name.begin(), Name.end()); + BufferObject->SetName(WStr.c_str()); + + D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle = {}; + { + auto SRVHandleOrErr = CSUAllocator.allocate(); + if (!SRVHandleOrErr) + return SRVHandleOrErr.takeError(); + SRVHandle = *SRVHandleOrErr; - ~DXComputeEncoder() override { endEncoding(); } + D3D12_SHADER_RESOURCE_VIEW_DESC SRVDesc = {}; + SRVDesc.ViewDimension = D3D12_SRV_DIMENSION_BUFFER; + SRVDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING; + switch (Desc.AccessType) { + case BufferShaderAccessType::Raw: + SRVDesc.Format = DXGI_FORMAT_R32_TYPELESS; + SRVDesc.Buffer.NumElements = static_cast(SizeInBytes / 4); + SRVDesc.Buffer.Flags = D3D12_BUFFER_SRV_FLAG_RAW; + break; + case BufferShaderAccessType::Typed: + SRVDesc.Format = getDXGIFormat(Desc.AccessTypeParams.Fmt); + SRVDesc.Buffer.NumElements = static_cast( + SizeInBytes / getFormatSizeInBytes(Desc.AccessTypeParams.Fmt)); + break; + case BufferShaderAccessType::Structured: + assert(Desc.AccessTypeParams.StructureStride > 0 && + "Structured buffers must have a Structure Stride."); + SRVDesc.Format = DXGI_FORMAT_UNKNOWN; + SRVDesc.Buffer.NumElements = static_cast( + SizeInBytes / Desc.AccessTypeParams.StructureStride); + SRVDesc.Buffer.StructureByteStride = + Desc.AccessTypeParams.StructureStride; + break; + } - static bool classof(const CommandEncoder *E) { - return E->getAPI() == GPUAPI::DirectX; + Device->CreateShaderResourceView(BufferObject.Get(), &SRVDesc, SRVHandle); } - // D3D12 debug labels require WinPixEventRuntime for the proper event - // encoding. Without it, BeginEvent/EndEvent/SetMarker with metadata type 0 - // crash the D3D12 debug layer, so leave these as no-ops for now. - void pushDebugGroup(llvm::StringRef Label) override {} - void popDebugGroup() override {} - void insertDebugSignpost(llvm::StringRef Label) override {} - - llvm::Error dispatch(const offloadtest::PipelineState &PSO, - uint32_t GroupCountX, uint32_t GroupCountY, - uint32_t GroupCountZ) override { - const auto &DXPSO = llvm::cast(PSO); - addUAVBarrier(); - insertDebugSignpost(llvm::formatv("Dispatch [{0},{1},{2}]", GroupCountX, - GroupCountY, GroupCountZ) - .str()); - CB.CmdList->SetPipelineState(DXPSO.PSO.Get()); - CB.CmdList->Dispatch(GroupCountX, GroupCountY, GroupCountZ); - return llvm::Error::success(); + D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle = {}; + if ((Flags & D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS) != 0) { + auto UAVHandleOrErr = CSUAllocator.allocate(); + if (!UAVHandleOrErr) + return UAVHandleOrErr.takeError(); + UAVHandle = *UAVHandleOrErr; + + D3D12_UNORDERED_ACCESS_VIEW_DESC UAVDesc = {}; + UAVDesc.ViewDimension = D3D12_UAV_DIMENSION_BUFFER; + switch (Desc.AccessType) { + case BufferShaderAccessType::Raw: + UAVDesc.Format = DXGI_FORMAT_R32_TYPELESS; + UAVDesc.Buffer.NumElements = static_cast(SizeInBytes / 4); + UAVDesc.Buffer.Flags = D3D12_BUFFER_UAV_FLAG_RAW; + break; + case BufferShaderAccessType::Typed: + UAVDesc.Format = getDXGIFormat(Desc.AccessTypeParams.Fmt); + UAVDesc.Buffer.NumElements = static_cast( + SizeInBytes / getFormatSizeInBytes(Desc.AccessTypeParams.Fmt)); + break; + case BufferShaderAccessType::Structured: + assert(Desc.AccessTypeParams.StructureStride > 0 && + "Structured buffers must have a Structure Stride."); + UAVDesc.Format = DXGI_FORMAT_UNKNOWN; + UAVDesc.Buffer.NumElements = static_cast( + SizeInBytes / Desc.AccessTypeParams.StructureStride); + UAVDesc.Buffer.StructureByteStride = + Desc.AccessTypeParams.StructureStride; + break; + } + + ID3D12Resource *CounterObject = nullptr; + if (Desc.HasCounter) { + UAVDesc.Buffer.CounterOffsetInBytes = CounterOffsetInBytes; + CounterObject = BufferObject.Get(); + } + + Device->CreateUnorderedAccessView(BufferObject.Get(), CounterObject, + &UAVDesc, UAVHandle); } - llvm::Error copyBufferToBuffer(offloadtest::Buffer &Src, size_t SrcOffset, - offloadtest::Buffer &Dst, size_t DstOffset, - size_t Size) override { - auto &DXSrc = static_cast(Src); - auto &DXDst = static_cast(Dst); - - // NOTE: Edge case in case of all the following being the case - // - multiple calls of copyBufferToBuffer with the same Dst Buffer - // - The Dst Buffer having a PreferredState of - // D3D12_RESOURCE_STATE_COPY_DEST - // - Each Src Buffer having a PreferredState of - // D3D12_RESOURCE_STATE_COPY_SOURCE - // In that case no barrier would be emitted - // and a race condition would occur. There are ways to solve this with - // legacy barriers, but switching to enhanced barriers is a better solution - // to this problem. - - if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) - CB.addResourceTransition(DXSrc.Buffer.Get(), DXSrc.PreferredState, - D3D12_RESOURCE_STATE_COPY_SOURCE); - if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) - CB.addResourceTransition(DXDst.Buffer.Get(), DXDst.PreferredState, - D3D12_RESOURCE_STATE_COPY_DEST); - CB.flushBarrier(); - - insertDebugSignpost(llvm::formatv("CopyBuffer {0}B", Size).str()); - CB.CmdList->CopyBufferRegion(DXDst.Buffer.Get(), DstOffset, - DXSrc.Buffer.Get(), SrcOffset, Size); - - if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) - CB.addResourceTransition(DXSrc.Buffer.Get(), - D3D12_RESOURCE_STATE_COPY_SOURCE, - DXSrc.PreferredState); - if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) - CB.addResourceTransition(DXDst.Buffer.Get(), - D3D12_RESOURCE_STATE_COPY_DEST, - DXDst.PreferredState); + D3D12_CPU_DESCRIPTOR_HANDLE CBVHandle = {}; + if (Desc.Usage == BufferUsage::ConstantBuffer) { + auto CBVHandleOrErr = CSUAllocator.allocate(); + if (!CBVHandleOrErr) + return CBVHandleOrErr.takeError(); + CBVHandle = *CBVHandleOrErr; - return llvm::Error::success(); + D3D12_CONSTANT_BUFFER_VIEW_DESC CBVDesc = {}; + CBVDesc.BufferLocation = BufferObject->GetGPUVirtualAddress(); + CBVDesc.SizeInBytes = BufferSizeInBytes; + + Device->CreateConstantBufferView(&CBVDesc, CBVHandle); } - llvm::Error copyBufferToTexture(Buffer &Src, Texture &Dst) override { - auto &DXSrc = llvm::cast(Src); - auto &DXDst = llvm::cast(Dst); - - if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) - CB.addResourceTransition(DXSrc.Buffer.Get(), DXSrc.PreferredState, - D3D12_RESOURCE_STATE_COPY_SOURCE); - - if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) - CB.addResourceTransition(DXDst.Resource.Get(), DXDst.PreferredState, - D3D12_RESOURCE_STATE_COPY_DEST); - CB.flushBarrier(); - - const uint32_t ElementSize = getFormatSizeInBytes(DXDst.Desc.Fmt); - const D3D12_PLACED_SUBRESOURCE_FOOTPRINT Footprint{ - 0, - CD3DX12_SUBRESOURCE_FOOTPRINT( - getDXGIFormat(DXDst.Desc.Fmt), DXDst.Desc.Width, DXDst.Desc.Height, - 1, getAlignedTexturePitch(DXDst.Desc.Width, ElementSize))}; - const CD3DX12_TEXTURE_COPY_LOCATION DstLoc(DXDst.Resource.Get(), 0); - const CD3DX12_TEXTURE_COPY_LOCATION SrcLoc(DXSrc.Buffer.Get(), Footprint); - CB.CmdList->CopyTextureRegion(&DstLoc, 0, 0, 0, &SrcLoc, nullptr); - - if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) - CB.addResourceTransition(DXSrc.Buffer.Get(), - D3D12_RESOURCE_STATE_COPY_SOURCE, - DXSrc.PreferredState); - - if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) - CB.addResourceTransition(DXDst.Resource.Get(), - D3D12_RESOURCE_STATE_COPY_DEST, - DXDst.PreferredState); + return std::make_unique(BufferObject, Name, Desc, SizeInBytes, + CounterOffsetInBytes, PreferredState, + SRVHandle, UAVHandle, CBVHandle); +} - return llvm::Error::success(); +llvm::Expected> +DXDevice::createTexture(std::string Name, const TextureCreateDesc &Desc) { + if (auto Err = validateTextureCreateDesc(Desc)) + return Err; + + const D3D12_HEAP_PROPERTIES HeapProps = + CD3DX12_HEAP_PROPERTIES(getDXHeapType(Desc.Location)); + + D3D12_RESOURCE_DESC TexDesc = {}; + TexDesc.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D; + TexDesc.Width = Desc.Width; + TexDesc.Height = Desc.Height; + TexDesc.DepthOrArraySize = 1; + TexDesc.MipLevels = static_cast(Desc.MipLevels); + TexDesc.Format = getDXGIFormat(Desc.Fmt); + TexDesc.SampleDesc.Count = 1; + if (Desc.Location == MemoryLocation::GpuOnly) { + if (Desc.Backing == MemoryBacking::Sparse) + TexDesc.Layout = D3D12_TEXTURE_LAYOUT_64KB_UNDEFINED_SWIZZLE; + else + TexDesc.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN; + } else { + TexDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; + } + TexDesc.Flags = getDXResourceFlags(Desc.Usage); + + const D3D12_CLEAR_VALUE *ClearValuePtr = nullptr; + D3D12_CLEAR_VALUE ClearValue = {}; + if (Desc.OptimizedClearValue) { + ClearValue.Format = TexDesc.Format; + std::visit( + [&ClearValue](auto &&V) { + using T = std::decay_t; + if constexpr (std::is_same_v) { + ClearValue.Color[0] = V.R; + ClearValue.Color[1] = V.G; + ClearValue.Color[2] = V.B; + ClearValue.Color[3] = V.A; + } else { + ClearValue.DepthStencil.Depth = V.Depth; + ClearValue.DepthStencil.Stencil = V.Stencil; + } + }, + *Desc.OptimizedClearValue); + ClearValuePtr = &ClearValue; + } + + D3D12_RESOURCE_STATES InitialState = + D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE | + D3D12_RESOURCE_STATE_NON_PIXEL_SHADER_RESOURCE; + if ((Desc.Usage & TextureUsage::Storage)) + InitialState = D3D12_RESOURCE_STATE_UNORDERED_ACCESS; + + ComPtr TextureObject; + if (Desc.Backing == MemoryBacking::Sparse) { + if (auto Err = HR::toError(Device->CreateReservedResource( + &TexDesc, InitialState, ClearValuePtr, + IID_PPV_ARGS(&TextureObject)), + "Failed to create reserved texture.")) + return Err; + } else { + if (auto Err = HR::toError(Device->CreateCommittedResource( + &HeapProps, D3D12_HEAP_FLAG_NONE, &TexDesc, + InitialState, ClearValuePtr, + IID_PPV_ARGS(&TextureObject)), + "Failed to create texture.")) + return Err; } - llvm::Error copyCounterToBuffer(offloadtest::Buffer &Src, - offloadtest::Buffer &Dst) override { - auto &DXSrc = llvm::cast(Src); - auto &DXDst = llvm::cast(Dst); + D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle = {}; + { + auto SRVHandleOrErr = CSUAllocator.allocate(); + if (!SRVHandleOrErr) + return SRVHandleOrErr.takeError(); + SRVHandle = *SRVHandleOrErr; - if (!DXSrc.Desc.HasCounter) - return llvm::createStringError( - "Counter resource passed does not hvae a counter."); - - if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) - CB.addResourceTransition(DXSrc.Buffer.Get(), DXSrc.PreferredState, - D3D12_RESOURCE_STATE_COPY_SOURCE); - if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) - CB.addResourceTransition(DXDst.Buffer.Get(), DXDst.PreferredState, - D3D12_RESOURCE_STATE_COPY_DEST); - CB.flushBarrier(); - - insertDebugSignpost("copyCounterToBuffer 4B"); - CB.CmdList->CopyBufferRegion(DXDst.Buffer.Get(), 0, DXSrc.Buffer.Get(), - DXSrc.CounterOffsetInBytes, sizeof(uint32_t)); - - if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) - CB.addResourceTransition(DXSrc.Buffer.Get(), - D3D12_RESOURCE_STATE_COPY_SOURCE, - DXSrc.PreferredState); - if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) - CB.addResourceTransition(DXDst.Buffer.Get(), - D3D12_RESOURCE_STATE_COPY_DEST, - DXDst.PreferredState); + D3D12_SHADER_RESOURCE_VIEW_DESC SRVDesc = {}; + SRVDesc.ViewDimension = + D3D12_SRV_DIMENSION_TEXTURE2D; // assume this is correct for now. + SRVDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING; + SRVDesc.Format = getDXGIFormatSRV(Desc.Fmt); + SRVDesc.Texture2D.MostDetailedMip = 0; + SRVDesc.Texture2D.MipLevels = Desc.MipLevels; + SRVDesc.Texture2D.PlaneSlice = 0; + SRVDesc.Texture2D.ResourceMinLODClamp = 0.0f; - return llvm::Error::success(); + Device->CreateShaderResourceView(TextureObject.Get(), &SRVDesc, SRVHandle); } - llvm::Error copyTextureToBuffer(Texture &Src, Buffer &Dst) override { - auto &DXSrc = llvm::cast(Src); - auto &DXDst = llvm::cast(Dst); + D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle = {}; + if ((TexDesc.Flags & D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS) != 0) { + auto UAVHandleOrErr = CSUAllocator.allocate(); + if (!UAVHandleOrErr) + return UAVHandleOrErr.takeError(); + UAVHandle = *UAVHandleOrErr; + + D3D12_UNORDERED_ACCESS_VIEW_DESC UAVDesc = {}; + UAVDesc.ViewDimension = D3D12_UAV_DIMENSION_TEXTURE2D; + UAVDesc.Texture2D.MipSlice = 0; + UAVDesc.Texture2D.PlaneSlice = 0; + + Device->CreateUnorderedAccessView(TextureObject.Get(), nullptr, &UAVDesc, + UAVHandle); + } + const D3D12_RESOURCE_STATES PreferredState = InitialState; + auto Tex = std::make_unique(TextureObject, Name, Desc, + PreferredState, SRVHandle, UAVHandle); + + const bool IsRT = (Desc.Usage & TextureUsage::RenderTarget) != 0; + const bool IsDS = (Desc.Usage & TextureUsage::DepthStencil) != 0; + if (IsRT) { + auto HandleOrErr = RTVAllocator.allocate(); + if (!HandleOrErr) + return HandleOrErr.takeError(); + Tex->RTVHandle = *HandleOrErr; + Device->CreateRenderTargetView(TextureObject.Get(), nullptr, + Tex->RTVHandle); + } + if (IsDS) { + auto HandleOrErr = DSVAllocator.allocate(); + if (!HandleOrErr) + return HandleOrErr.takeError(); + Tex->DSVHandle = *HandleOrErr; + Device->CreateDepthStencilView(TextureObject.Get(), nullptr, + Tex->DSVHandle); + } + + return Tex; +} - if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) - CB.addResourceTransition(DXSrc.Resource.Get(), DXSrc.PreferredState, - D3D12_RESOURCE_STATE_COPY_SOURCE); +llvm::Expected> +DXDevice::createSampler(std::string Name, const SamplerCreateDesc &Desc) { + + auto HandleOrErr = SamplerAllocator.allocate(); + if (!HandleOrErr) + return HandleOrErr.takeError(); + const D3D12_CPU_DESCRIPTOR_HANDLE Handle = *HandleOrErr; + + const D3D12_TEXTURE_ADDRESS_MODE AddressMode = + getDXTextureAddressMode(Desc.Address); + + bool IsComparison = false; + D3D12_COMPARISON_FUNC ComparisonFunc = D3D12_COMPARISON_FUNC_NONE; + if (Desc.Kind == SamplerKind::SamplerComparison) { + IsComparison = true; + ComparisonFunc = getDXComparisonFunc(Desc.ComparisonOp); + } + + const D3D12_SAMPLER_DESC SamplerDesc = { + getDXFilterMode(Desc.MinFilter, Desc.MagFilter, IsComparison), + AddressMode, // U + AddressMode, // V + AddressMode, // W + Desc.MipLODBias, + 0, // MaxAnisotropy + ComparisonFunc, + {0.0f, 0.0f, 0.0f, 0.0f}, // BorderColor + Desc.MinLOD, + Desc.MaxLOD, + }; - if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) - CB.addResourceTransition(DXDst.Buffer.Get(), DXDst.PreferredState, - D3D12_RESOURCE_STATE_COPY_DEST); + Device->CreateSampler(&SamplerDesc, Handle); - CB.flushBarrier(); + return std::make_unique(Name, Desc, Handle); +} - const uint32_t ElementSize = getFormatSizeInBytes(DXSrc.Desc.Fmt); - const D3D12_PLACED_SUBRESOURCE_FOOTPRINT Footprint{ - 0, - CD3DX12_SUBRESOURCE_FOOTPRINT( - getDXGIFormat(DXSrc.Desc.Fmt), DXSrc.Desc.Width, DXSrc.Desc.Height, - 1, getAlignedTexturePitch(DXSrc.Desc.Width, ElementSize))}; - const CD3DX12_TEXTURE_COPY_LOCATION DstLoc(DXDst.Buffer.Get(), Footprint); - const CD3DX12_TEXTURE_COPY_LOCATION SrcLoc(DXSrc.Resource.Get(), 0); - CB.CmdList->CopyTextureRegion(&DstLoc, 0, 0, 0, &SrcLoc, nullptr); +uint32_t DXDevice::getTextureUploadRowStrideInBytes( + const TextureCreateDesc &Desc) const { + return getAlignedTexturePitch(Desc.Width, getFormatSizeInBytes(Desc.Fmt)); +} - if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) - CB.addResourceTransition(DXSrc.Resource.Get(), - D3D12_RESOURCE_STATE_COPY_SOURCE, - DXSrc.PreferredState); +TextureUploadLayout +DXDevice::getTextureUploadLayout(const TextureCreateDesc &Desc) const { + // Only the fields GetCopyableFootprints consults are needed here; layout, + // flags, and clear value do not affect the copyable footprint. + D3D12_RESOURCE_DESC TexDesc = {}; + TexDesc.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D; + TexDesc.Width = Desc.Width; + TexDesc.Height = Desc.Height; + TexDesc.DepthOrArraySize = 1; + TexDesc.MipLevels = static_cast(Desc.MipLevels); + TexDesc.Format = getDXGIFormat(Desc.Fmt); + TexDesc.SampleDesc.Count = 1; + + const uint32_t NumSubresources = Desc.MipLevels; + llvm::SmallVector Footprints( + NumSubresources); + llvm::SmallVector NumRows(NumSubresources); + llvm::SmallVector RowSizes(NumSubresources); + UINT64 TotalBytes = 0; + Device->GetCopyableFootprints(&TexDesc, 0, NumSubresources, 0, + Footprints.data(), NumRows.data(), + RowSizes.data(), &TotalBytes); + + TextureUploadLayout Layout; + Layout.TotalSizeInBytes = TotalBytes; + Layout.Subresources.reserve(NumSubresources); + for (uint32_t I = 0; I < NumSubresources; ++I) { + SubresourceFootprint Sub; + Sub.Offset = Footprints[I].Offset; + Sub.RowPitchInBytes = Footprints[I].Footprint.RowPitch; + Sub.RowSizeInBytes = static_cast(RowSizes[I]); + Sub.NumRows = NumRows[I]; + Layout.Subresources.push_back(Sub); + } + return Layout; +} - if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) - CB.addResourceTransition(DXDst.Buffer.Get(), - D3D12_RESOURCE_STATE_COPY_DEST, - DXDst.PreferredState); +llvm::Expected> +DXDevice::create(ComPtr Adapter, const DeviceConfig &Config) { + ComPtr Device; + if (auto Err = + HR::toError(D3D12CreateDevice(Adapter.Get(), D3D_FEATURE_LEVEL_11_0, + IID_PPV_ARGS(&Device)), + "Failed to create D3D device")) + return Err; - return llvm::Error::success(); + static std::once_flag SignalHandlerRegistered; + std::call_once(SignalHandlerRegistered, [] { + llvm::sys::AddSignalHandler(dumpD3DInfoQueues, nullptr); + }); + { + const std::lock_guard Lock(SignalHandlerMutex); + SignalHandlerDevices.push_back(Device.Get()); + } + + assert( + Adapter->IsPropertySupported(DXCoreAdapterProperty::DriverDescription)); + size_t BufferSize; + Adapter->GetPropertySize(DXCoreAdapterProperty::DriverDescription, + &BufferSize); + std::vector DescVec(BufferSize); + Adapter->GetProperty(DXCoreAdapterProperty::DriverDescription, BufferSize, + (void *)DescVec.data()); + + std::string DriverVer; + if (Adapter->IsPropertySupported(DXCoreAdapterProperty::DriverVersion)) { + uint64_t Packed = 0; + if (SUCCEEDED(Adapter->GetProperty(DXCoreAdapterProperty::DriverVersion, + sizeof(Packed), &Packed))) { + const uint16_t Major = static_cast((Packed >> 48) & 0xFFFF); + const uint16_t Minor = static_cast((Packed >> 32) & 0xFFFF); + const uint16_t Build = static_cast((Packed >> 16) & 0xFFFF); + const uint16_t Revision = static_cast(Packed & 0xFFFF); + DriverVer = (llvm::Twine(Major) + "." + llvm::Twine(Minor) + "." + + llvm::Twine(Build) + "." + llvm::Twine(Revision)) + .str(); + } } - // Defined out-of-line below — needs DXDevice's full type for access to the - // ID3D12Device5 entry point and helper allocators. - llvm::Error batchBuildAS(llvm::ArrayRef Items) override; - - // Defined out-of-line below — needs DXDevice's full type for access to - // Device5 and the DXRayTracingPipelineState definition. - llvm::Error dispatchRays(const PipelineState &PSO, - const ShaderBindingTable &SBT, uint32_t Width, - uint32_t Height, uint32_t Depth) override; + if (Config.EnableDebugLayer || Config.EnableValidationLayer) + if (auto Err = configureInfoQueue(Device.Get())) + return Err; - void endEncodingImpl() override { popDebugGroup(); } -}; + auto GraphicsQueueOrErr = DXQueue::createGraphicsQueue(Device); + if (!GraphicsQueueOrErr) + return GraphicsQueueOrErr.takeError(); + + auto RTVHeapOrErr = DescriptorAllocator::create( + Device.Get(), D3D12_DESCRIPTOR_HEAP_TYPE_RTV, 256); + if (!RTVHeapOrErr) + return RTVHeapOrErr.takeError(); + + auto DSVHeapOrErr = DescriptorAllocator::create( + Device.Get(), D3D12_DESCRIPTOR_HEAP_TYPE_DSV, 256); + if (!DSVHeapOrErr) + return DSVHeapOrErr.takeError(); + + auto CSUHeapOrErr = DescriptorAllocator::create( + Device.Get(), D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV, 4096); + if (!CSUHeapOrErr) + return CSUHeapOrErr.takeError(); + + auto SamplerHeapOrErr = DescriptorAllocator::create( + Device.Get(), D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER, 256); + if (!SamplerHeapOrErr) + return SamplerHeapOrErr.takeError(); + + return std::make_unique( + Adapter, Device, std::move(*GraphicsQueueOrErr), std::move(*RTVHeapOrErr), + std::move(*DSVHeapOrErr), std::move(*CSUHeapOrErr), + std::move(*SamplerHeapOrErr), std::string(DescVec.data()), + std::move(DriverVer)); +} -llvm::Expected> -DXCommandBuffer::createComputeEncoder() { - auto Enc = std::make_unique(*this); - Enc->pushDebugGroup("ComputeEncoder"); - return Enc; +const Capabilities &DXDevice::getCapabilities() { + if (Caps.empty()) + queryCapabilities(); + return Caps; } -class DXRenderPass final : public offloadtest::RenderPass { -public: - offloadtest::RenderPassDesc Desc; - - explicit DXRenderPass(offloadtest::RenderPassDesc Desc) - : RenderPass(GPUAPI::DirectX), Desc(std::move(Desc)) {} - - static bool classof(const offloadtest::RenderPass *RP) { - return RP->getAPI() == GPUAPI::DirectX; - } -}; - -class DXRenderEncoder : public offloadtest::RenderEncoder { - DXCommandBuffer &CB; - offloadtest::RenderPassBeginDesc Desc; - - // Encoder contract: viewport and scissor must both be set before draw(). - bool ViewportSet = false; - bool ScissorSet = false; - - llvm::Error bindCommonDrawState(const offloadtest::PipelineState &PSO) { - if (!ViewportSet) - return llvm::createStringError(std::errc::invalid_argument, - "Viewport must be set before drawing."); - if (!ScissorSet) - return llvm::createStringError(std::errc::invalid_argument, - "Scissor must be set before drawing."); - - const auto &DXPSO = llvm::cast(PSO); - CB.CmdList->SetGraphicsRootSignature(DXPSO.RootSig.Get()); - CB.CmdList->SetPipelineState(DXPSO.PSO.Get()); - // Mesh-shader pipelines bypass the input assembler and carry no IA - // topology; only bind one when the pipeline actually has one. - if (DXPSO.Topology) - CB.CmdList->IASetPrimitiveTopology(*DXPSO.Topology); - return llvm::Error::success(); - } - -public: - DXRenderEncoder(DXCommandBuffer &CB, - const offloadtest::RenderPassBeginDesc &Desc) - : RenderEncoder(GPUAPI::DirectX), CB(CB), Desc(Desc) {} - DXRenderEncoder(const DXRenderEncoder &CB) = delete; - DXRenderEncoder(DXRenderEncoder &&CB) = delete; - DXRenderEncoder &operator=(DXRenderEncoder &CB) = delete; - DXRenderEncoder &operator=(const DXRenderEncoder &&CB) = delete; - - ~DXRenderEncoder() override { endEncoding(); } - - static bool classof(const CommandEncoder *E) { - return E->getAPI() == GPUAPI::DirectX; - } - - // See DXComputeEncoder for why these are no-ops. - void pushDebugGroup(llvm::StringRef Label) override {} - void popDebugGroup() override {} - void insertDebugSignpost(llvm::StringRef Label) override {} - - void setViewport(const offloadtest::Viewport &VP) override { - D3D12_VIEWPORT DXVP = {}; - DXVP.TopLeftX = VP.X; - DXVP.TopLeftY = VP.Y; - DXVP.Width = VP.Width; - DXVP.Height = VP.Height; - DXVP.MinDepth = VP.MinDepth; - DXVP.MaxDepth = VP.MaxDepth; - CB.CmdList->RSSetViewports(1, &DXVP); - ViewportSet = true; - } - - void setScissor(const offloadtest::ScissorRect &Rect) override { - const D3D12_RECT DXRect = {Rect.X, Rect.Y, - static_cast(Rect.X + Rect.Width), - static_cast(Rect.Y + Rect.Height)}; - CB.CmdList->RSSetScissorRects(1, &DXRect); - ScissorSet = true; - } - - void setVertexBuffer(uint32_t Slot, offloadtest::Buffer *VB, size_t Offset, - uint32_t Stride) override { - assert(Slot < D3D12_IA_VERTEX_INPUT_RESOURCE_SLOT_COUNT && - "Vertex buffer slot exceeds D3D12 IA input resource slot count"); - assert(Slot == 0 && "Pipeline input layout only describes slot 0"); - if (VB) { - auto &DXVB = llvm::cast(*VB); - D3D12_VERTEX_BUFFER_VIEW VBView = {}; - VBView.BufferLocation = DXVB.Buffer->GetGPUVirtualAddress() + Offset; - VBView.SizeInBytes = static_cast(DXVB.getSizeInBytes() - Offset); - VBView.StrideInBytes = Stride; - CB.CmdList->IASetVertexBuffers(Slot, 1, &VBView); - } else { - CB.CmdList->IASetVertexBuffers(Slot, 1, nullptr); - } - } - - llvm::Error drawInstanced(const offloadtest::PipelineState &PSO, - uint32_t VertexCount, uint32_t InstanceCount, - uint32_t FirstVertex, - uint32_t FirstInstance) override { - if (auto Err = bindCommonDrawState(PSO)) - return Err; - CB.CmdList->DrawInstanced(VertexCount, InstanceCount, FirstVertex, - FirstInstance); - return llvm::Error::success(); - } - - llvm::Error dispatchMesh(const offloadtest::PipelineState &PSO, - uint32_t GroupCountX, uint32_t GroupCountY, - uint32_t GroupCountZ) override { - if (auto Err = bindCommonDrawState(PSO)) - return Err; - CB.CmdList->DispatchMesh(GroupCountX, GroupCountY, GroupCountZ); - return llvm::Error::success(); - } - - void endEncodingImpl() override { - // State transitions - for (offloadtest::Texture *Tex : Desc.ColorAttachments) { - auto &DXTex = llvm::cast(*Tex); - if (DXTex.PreferredState != D3D12_RESOURCE_STATE_RENDER_TARGET) - CB.addResourceTransition(DXTex.Resource.Get(), - D3D12_RESOURCE_STATE_RENDER_TARGET, - DXTex.PreferredState); - } - if (Desc.DepthStencil) { - auto &DXTex = llvm::cast(*Desc.DepthStencil); - if (DXTex.PreferredState != D3D12_RESOURCE_STATE_DEPTH_WRITE) - CB.addResourceTransition(DXTex.Resource.Get(), - D3D12_RESOURCE_STATE_DEPTH_WRITE, - DXTex.PreferredState); - } - - popDebugGroup(); - } -}; - -llvm::Expected> -DXCommandBuffer::createRenderEncoder( - const offloadtest::RenderPassBeginDesc &Desc) { - // The pass carries format / load / store policy; the begin desc supplies - // the actual textures. Walk both in lockstep. - if (!Desc.Pass) - return llvm::createStringError( - std::errc::invalid_argument, - "RenderPassBeginDesc is missing its RenderPass."); - auto &DXPass = llvm::cast(*Desc.Pass); - const offloadtest::RenderPassDesc &PassDesc = DXPass.Desc; - - if (Desc.ColorAttachments.size() != PassDesc.ColorAttachments.size()) - return llvm::createStringError( - std::errc::invalid_argument, - "RenderPassBeginDesc color attachment count does not match its " - "RenderPass."); - if (PassDesc.DepthStencil.has_value() != (Desc.DepthStencil != nullptr)) - return llvm::createStringError(std::errc::invalid_argument, - "RenderPassBeginDesc depth-stencil " - "presence does not match its RenderPass."); - - if (auto Err = findAndValidateRenderPassTextureSize(Desc, nullptr, nullptr)) - return Err; - - // Validate attachments and gather the RTV / DSV CPU handles. RT and DSV - // descriptors are owned by the textures themselves; this just collects - // them for OMSetRenderTargets. - llvm::SmallVector RTTextures; - llvm::SmallVector RTVHandles; - RTTextures.reserve(Desc.ColorAttachments.size()); - RTVHandles.reserve(Desc.ColorAttachments.size()); - for (offloadtest::Texture *Tex : Desc.ColorAttachments) { - if (!Tex) - return llvm::createStringError( - std::errc::invalid_argument, - "RenderPassBeginDesc has a null color attachment texture."); - auto &DXTex = llvm::cast(*Tex); - if (DXTex.RTVHandle.ptr == 0) - return llvm::createStringError( - std::errc::invalid_argument, - "Color attachment texture was not created with RenderTarget usage."); - RTTextures.push_back(&DXTex); - RTVHandles.push_back(DXTex.RTVHandle); - } - - DXTexture *DSTexture = nullptr; - D3D12_CPU_DESCRIPTOR_HANDLE DSVHandle = {}; - if (Desc.DepthStencil) { - auto &DXDS = llvm::cast(*Desc.DepthStencil); - if (DXDS.DSVHandle.ptr == 0) - return llvm::createStringError( - std::errc::invalid_argument, - "Depth-stencil texture was not created with DepthStencil usage."); - DSTexture = &DXDS; - DSVHandle = DXDS.DSVHandle; - } - - // State transitions - for (offloadtest::Texture *Tex : Desc.ColorAttachments) { - auto &DXTex = llvm::cast(*Tex); - if (DXTex.PreferredState != D3D12_RESOURCE_STATE_RENDER_TARGET) - this->addResourceTransition(DXTex.Resource.Get(), DXTex.PreferredState, - D3D12_RESOURCE_STATE_RENDER_TARGET); - } - if (Desc.DepthStencil) { - auto &DXTex = llvm::cast(*Desc.DepthStencil); - if (DXTex.PreferredState != D3D12_RESOURCE_STATE_DEPTH_WRITE) - this->addResourceTransition(DXTex.Resource.Get(), DXTex.PreferredState, - D3D12_RESOURCE_STATE_DEPTH_WRITE); - } - - this->flushBarrier(); - - CmdList->OMSetRenderTargets(static_cast(RTVHandles.size()), - RTVHandles.data(), - /*RTsSingleHandleToDescriptorRange=*/false, - Desc.DepthStencil ? &DSVHandle : nullptr); - - for (size_t I = 0; I < PassDesc.ColorAttachments.size(); ++I) { - if (PassDesc.ColorAttachments[I].Load != offloadtest::LoadAction::Clear) - continue; - if (!RTTextures[I]->Desc.OptimizedClearValue) - return llvm::createStringError( - std::errc::invalid_argument, - "LoadAction::Clear requires the render target to have been " - "created with an OptimizedClearValue."); - const auto *CV = - std::get_if(&*RTTextures[I]->Desc.OptimizedClearValue); - assert(CV && "RenderTarget OptimizedClearValue must be a ClearColor"); - const float ClearArr[4] = {CV->R, CV->G, CV->B, CV->A}; - CmdList->ClearRenderTargetView(RTVHandles[I], ClearArr, 0, nullptr); - } - if (PassDesc.DepthStencil) { - D3D12_CLEAR_FLAGS Flags = static_cast(0); - if (PassDesc.DepthStencil->DepthLoad == offloadtest::LoadAction::Clear) - Flags |= D3D12_CLEAR_FLAG_DEPTH; - if (PassDesc.DepthStencil->StencilLoad == offloadtest::LoadAction::Clear) - Flags |= D3D12_CLEAR_FLAG_STENCIL; - if (Flags != 0) { - if (!DSTexture->Desc.OptimizedClearValue) - return llvm::createStringError( - std::errc::invalid_argument, - "LoadAction::Clear requires the depth-stencil texture to have " - "been created with an OptimizedClearValue."); - const auto *CV = - std::get_if(&*DSTexture->Desc.OptimizedClearValue); - assert(CV && - "DepthStencil OptimizedClearValue must be a ClearDepthStencil"); - CmdList->ClearDepthStencilView(DSVHandle, Flags, CV->Depth, CV->Stencil, - 0, nullptr); - } - } - - auto Enc = std::make_unique(*this, Desc); - Enc->pushDebugGroup("RenderEncoder"); - return Enc; -} - -class DXDevice : public offloadtest::Device { -public: - ComPtr Adapter; - ComPtr Device; - DXQueue GraphicsQueue; - Capabilities Caps; - DescriptorAllocator RTVAllocator; - DescriptorAllocator DSVAllocator; - DescriptorAllocator CSUAllocator; - - struct ResourceSet { - std::unique_ptr BackingMemory; - std::unique_ptr Buffer; - std::unique_ptr Texture; - std::unique_ptr Readback; - std::unique_ptr CounterReadback; - - // AS-only; mutually exclusive with the buffer/texture fields above. - AccelerationStructure *AS = nullptr; - - ResourceSet(std::unique_ptr Buffer, - std::unique_ptr BackingMemory, - std::unique_ptr Readback, - std::unique_ptr CounterReadback) - : BackingMemory(std::move(BackingMemory)), Buffer(std::move(Buffer)), - Readback(std::move(Readback)), - CounterReadback(std::move(CounterReadback)) {} - ResourceSet(std::unique_ptr Texture, - std::unique_ptr BackingMemory, - std::unique_ptr Readback) - : BackingMemory(std::move(BackingMemory)), Texture(std::move(Texture)), - Readback(std::move(Readback)) {} - explicit ResourceSet(AccelerationStructure *AS) : AS(AS) {} - - ResourceSet(const ResourceSet &) = delete; - ResourceSet &operator=(const ResourceSet &) = delete; - - ResourceSet(ResourceSet &&A) - : BackingMemory(std::move(A.BackingMemory)), - Buffer(std::move(A.Buffer)), Texture(std::move(A.Texture)), - Readback(std::move(A.Readback)), - CounterReadback(std::move(A.CounterReadback)), AS(A.AS) {} - ResourceSet &operator=(ResourceSet &&A) { - BackingMemory = std::move(A.BackingMemory); - Buffer = std::move(A.Buffer); - Texture = std::move(A.Texture); - Readback = std::move(A.Readback); - CounterReadback = std::move(A.CounterReadback); - AS = A.AS; - return *this; - } - }; - - // ResourceBundle will contain one ResourceSet for a singular resource - // or multiple ResourceSets for resource array. - using ResourceBundle = llvm::SmallVector; - using ResourcePair = std::pair; - - struct DescriptorTable { - llvm::SmallVector Resources; - }; - - struct InvocationState { - ComPtr DescHeap; - std::unique_ptr CB; - std::unique_ptr Pipeline; - // Lifetime-tied to the pipeline; only set for RT pipelines. - std::unique_ptr SBT; - - // Resources for graphics pipelines. - std::unique_ptr RenderPass; - std::unique_ptr RenderTarget; - std::unique_ptr RTReadback; - std::unique_ptr DepthStencil; - std::unique_ptr VB; - - llvm::SmallVector> KeepAliveBuffers; - - llvm::SmallVector DescTables; - llvm::SmallVector RootResources; - - // Parallel-indexed to `P.AccelStructs.BLAS`. - llvm::SmallVector> - BLASes; - // Keyed by `TLASDesc::Name`. - llvm::StringMap> TLASes; - // Vertex/index buffers consumed during AS builds; must outlive submission. - llvm::SmallVector> ASInputBuffers; - }; - - DXDevice(ComPtr A, ComPtr D, DXQueue Q, - DescriptorAllocator RTVAllocator, DescriptorAllocator DSVAllocator, - DescriptorAllocator CSUAllocator, std::string Desc, - std::string DriverVer) - : Adapter(A), Device(D), GraphicsQueue(std::move(Q)), - RTVAllocator(std::move(RTVAllocator)), - DSVAllocator(std::move(DSVAllocator)), - CSUAllocator(std::move(CSUAllocator)) { - Description = std::move(Desc); - DriverVersion = std::move(DriverVer); - DriverName = "DirectX"; - - DXCoreHardwareID HardwareID; - if (SUCCEEDED(Adapter->GetProperty(DXCoreAdapterProperty::HardwareID, - &HardwareID))) { - // 0x8086 is the Vendor ID for Intel - if (HardwareID.vendorID == 0x8086) { - FamilyPrefix = static_cast(HardwareID.deviceID) & 0xFF00; - const IntelGpuEra Era = - getIntelGpuEra(static_cast(HardwareID.deviceID)); - if (Era == IntelGpuEra::Gen7_to_10) - GPUGeneration = "Intel Gen7-10"; - else if (Era == IntelGpuEra::Gen11_to_14_and_Xe) - GPUGeneration = "Intel Gen11-14/Xe"; - else - GPUGeneration = "Intel Unknown"; - } else { - // We don't have a need yet to identify other GPU vendors. - GPUGeneration = "Unknown"; - } - } - } - DXDevice(const DXDevice &) = delete; - DXDevice &operator=(const DXDevice &) = delete; - - ~DXDevice() override { - const std::lock_guard Lock(SignalHandlerMutex); - llvm::erase(SignalHandlerDevices, Device.Get()); - } - - llvm::StringRef getAPIName() const override { return "DirectX"; } - GPUAPI getAPI() const override { return GPUAPI::DirectX; } - - static bool classof(const offloadtest::Device *D) { - return D->getAPI() == GPUAPI::DirectX; - } - - Queue &getGraphicsQueue() override { return GraphicsQueue; } - - llvm::Error - createRootSignatureFromShader(llvm::StringRef, const ShaderContainer &Shader, - ComPtr &OutRootSignature) { - // Try pulling a root signature from the DXIL first - auto ExContainer = - llvm::object::DXContainer::create(Shader.Shader->getMemBufferRef()); - // If this fails we really have a problem... - if (!ExContainer) - return ExContainer.takeError(); - - bool HasRootSigPart = false; - for (const auto &Part : *ExContainer) - if (memcmp(Part.Part.Name, "RTS0", 4) == 0) - HasRootSigPart = true; - - if (HasRootSigPart) { - const llvm::StringRef Binary = Shader.Shader->getBuffer(); - if (auto Err = HR::toError( - Device->CreateRootSignature(0, Binary.data(), Binary.size(), - IID_PPV_ARGS(&OutRootSignature)), - "Failed to create root signature.")) - return Err; - } - - return llvm::Error::success(); - } - - llvm::Error createRootSignatureFromBindingsDesc( - llvm::StringRef, const BindingsDesc &BndDesc, bool IsGraphics, - ComPtr &OutRootSignature) { - uint32_t DescriptorCount = 0; - for (auto &D : BndDesc.DescriptorSetDescs) - DescriptorCount += D.ResourceBindings.size(); - - std::vector RootParams; - const std::unique_ptr Ranges( - new D3D12_DESCRIPTOR_RANGE[DescriptorCount]); - uint32_t RangeIdx = 0; - for (const auto &Set : BndDesc.DescriptorSetDescs) { - uint32_t DescriptorIdx = 0; - const uint32_t StartRangeIdx = RangeIdx; - for (const auto &Binding : Set.ResourceBindings) { - switch (getDescriptorKind(Binding.Kind)) { - case DescriptorKind::SRV: - Ranges.get()[RangeIdx].RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_SRV; - break; - case DescriptorKind::UAV: - Ranges.get()[RangeIdx].RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_UAV; - break; - case DescriptorKind::CBV: - Ranges.get()[RangeIdx].RangeType = D3D12_DESCRIPTOR_RANGE_TYPE_CBV; - break; - case DescriptorKind::SAMPLER: - llvm_unreachable("Not implemented yet."); // Requires a separate heap - } - Ranges.get()[RangeIdx].NumDescriptors = Binding.DescriptorCount; - Ranges.get()[RangeIdx].BaseShaderRegister = Binding.DXBinding.Register; - Ranges.get()[RangeIdx].RegisterSpace = Binding.DXBinding.Space; - Ranges.get()[RangeIdx].OffsetInDescriptorsFromTableStart = - DescriptorIdx; - RangeIdx++; - DescriptorIdx += Binding.DescriptorCount; - } - RootParams.push_back(D3D12_ROOT_PARAMETER{ - D3D12_ROOT_PARAMETER_TYPE_DESCRIPTOR_TABLE, - {D3D12_ROOT_DESCRIPTOR_TABLE{ - static_cast(Set.ResourceBindings.size()), - &Ranges.get()[StartRangeIdx]}}, - D3D12_SHADER_VISIBILITY_ALL}); - } - - CD3DX12_ROOT_SIGNATURE_DESC Desc; - Desc.Init(static_cast(RootParams.size()), RootParams.data(), 0, - nullptr, - IsGraphics - ? D3D12_ROOT_SIGNATURE_FLAG_ALLOW_INPUT_ASSEMBLER_INPUT_LAYOUT - : D3D12_ROOT_SIGNATURE_FLAG_NONE); - - ComPtr Signature; - ComPtr Error; - if (auto Err = HR::toError( - D3D12SerializeRootSignature(&Desc, D3D_ROOT_SIGNATURE_VERSION_1, - &Signature, &Error), - "Failed to seialize root signature.")) { - const std::string Msg = - std::string(reinterpret_cast(Error->GetBufferPointer()), - Error->GetBufferSize() / sizeof(char)); - return joinErrors( - std::move(Err), - llvm::createStringError(std::errc::protocol_error, Msg.c_str())); - } - - if (auto Err = HR::toError( - Device->CreateRootSignature(0, Signature->GetBufferPointer(), - Signature->GetBufferSize(), - IID_PPV_ARGS(&OutRootSignature)), - "Failed to create root signature.")) - return Err; - - return llvm::Error::success(); - } - - llvm::Error - createRootSignature(llvm::StringRef Name, const BindingsDesc &BndDesc, - const ShaderContainer &Shader, bool IsGraphics, - ComPtr &OutRootSignature) { - assert(OutRootSignature.Get() == nullptr); - - if (auto Err = - createRootSignatureFromShader(Name, Shader, OutRootSignature)) - return Err; - - if (OutRootSignature.Get() != nullptr) - return llvm::Error::success(); - - return createRootSignatureFromBindingsDesc(Name, BndDesc, IsGraphics, - OutRootSignature); - } - - llvm::Expected> - createPipelineCs(llvm::StringRef Name, const BindingsDesc &BndDesc, - ShaderContainer CS) override { - ComPtr RootSig; - if (auto Err = createRootSignature(Name, BndDesc, CS, - /*IsGraphics=*/false, RootSig)) - return Err; - - auto DXIL = CS.Shader->getBuffer(); - const D3D12_COMPUTE_PIPELINE_STATE_DESC Desc = { - RootSig.Get(), - {DXIL.data(), DXIL.size()}, - 0, - { - nullptr, - 0, - }, - D3D12_PIPELINE_STATE_FLAG_NONE}; - - ComPtr PSO; - if (auto Err = HR::toError( - Device->CreateComputePipelineState(&Desc, IID_PPV_ARGS(&PSO)), - "Failed to create PSO.")) - return Err; - - return std::make_unique(Name, RootSig, PSO, std::nullopt); - } - - llvm::Expected> - createTraditionalRasterPipeline( - llvm::StringRef Name, const BindingsDesc &BndDesc, - const TraditionalRasterPipelineCreateDesc &Desc) override { - assert(Desc.RTFormats.size() <= 8); - - ComPtr RootSig; - if (auto Err = createRootSignature(Name, BndDesc, Desc.VS, - /*IsGraphics=*/true, RootSig)) - return Err; - - std::vector DXInputLayout; - DXInputLayout.reserve(Desc.InputLayout.size()); - for (const InputLayoutDesc &Elem : Desc.InputLayout) { - assert(!Elem.InstanceStepRate && - "Instance step rate is currently not supported."); - - D3D12_INPUT_ELEMENT_DESC ElementDesc = {}; - ElementDesc.SemanticName = Elem.Name.c_str(); - ElementDesc.SemanticIndex = 0; - ElementDesc.Format = getDXGIFormat(Elem.Fmt); - ElementDesc.InputSlot = 0; - ElementDesc.AlignedByteOffset = Elem.OffsetInBytes; - ElementDesc.InputSlotClass = D3D12_INPUT_CLASSIFICATION_PER_VERTEX_DATA; - DXInputLayout.push_back(ElementDesc); - } - - D3D12_GRAPHICS_PIPELINE_STATE_DESC PSODesc = {}; - PSODesc.InputLayout = {DXInputLayout.data(), (UINT)DXInputLayout.size()}; - PSODesc.pRootSignature = RootSig.Get(); - PSODesc.VS = {Desc.VS.Shader->getBuffer().data(), - Desc.VS.Shader->getBuffer().size()}; - PSODesc.PS = {Desc.PS.Shader->getBuffer().data(), - Desc.PS.Shader->getBuffer().size()}; - if (PSODesc.VS.BytecodeLength == 0 || PSODesc.PS.BytecodeLength == 0) - return llvm::createStringError(std::errc::invalid_argument, - "Graphics pipeline requires both a vertex " - "shader and a pixel shader."); - if (Desc.HS) - PSODesc.HS = {Desc.HS->Shader->getBuffer().data(), - Desc.HS->Shader->getBuffer().size()}; - if (Desc.DS) - PSODesc.DS = {Desc.DS->Shader->getBuffer().data(), - Desc.DS->Shader->getBuffer().size()}; - if (Desc.GS) - PSODesc.GS = {Desc.GS->Shader->getBuffer().data(), - Desc.GS->Shader->getBuffer().size()}; - - PSODesc.RasterizerState = CD3DX12_RASTERIZER_DESC(D3D12_DEFAULT); - PSODesc.RasterizerState.CullMode = D3D12_CULL_MODE_NONE; - PSODesc.RasterizerState.FrontCounterClockwise = TRUE; - PSODesc.BlendState = CD3DX12_BLEND_DESC(D3D12_DEFAULT); - PSODesc.DepthStencilState.DepthEnable = true; - PSODesc.DepthStencilState.DepthWriteMask = D3D12_DEPTH_WRITE_MASK_ALL; - PSODesc.DepthStencilState.DepthFunc = D3D12_COMPARISON_FUNC_LESS; - PSODesc.DepthStencilState.StencilEnable = false; - PSODesc.SampleMask = UINT_MAX; - PSODesc.PrimitiveTopologyType = getDXPrimitiveTopologyType(Desc.Topology); - PSODesc.NumRenderTargets = static_cast(Desc.RTFormats.size()); - if (Desc.DSFormat) - PSODesc.DSVFormat = getDXGIFormat(*Desc.DSFormat); - for (size_t I = 0; I < Desc.RTFormats.size(); ++I) - PSODesc.RTVFormats[I] = getDXGIFormat(Desc.RTFormats[I]); - PSODesc.SampleDesc.Count = 1; - - ComPtr PSO; - if (auto Err = HR::toError( - Device->CreateGraphicsPipelineState(&PSODesc, IID_PPV_ARGS(&PSO)), - "Failed to create graphics PSO.")) - return Err; - - return std::make_unique( - Name, RootSig, PSO, - getDXPrimitiveTopology(Desc.Topology, Desc.PatchControlPoints)); - } - - llvm::Expected> createMeshShaderRasterPipeline( - llvm::StringRef Name, const BindingsDesc &BindingsDesc, - const MeshShaderRasterPipelineCreateDesc &Desc) override { - assert(Desc.RTFormats.size() <= 8); - - ComPtr RootSig; - if (auto Err = createRootSignature(Name, BindingsDesc, Desc.MS, - /*IsGraphics=*/true, RootSig)) - return Err; - - const D3D12_SHADER_BYTECODE MSBytecode = { - Desc.MS.Shader->getBuffer().data(), Desc.MS.Shader->getBuffer().size()}; - if (MSBytecode.BytecodeLength == 0) - return llvm::createStringError( - std::errc::invalid_argument, - "Mesh shader pipeline requires a mesh shader."); - - // The amplification (task) shader is optional. - D3D12_SHADER_BYTECODE ASBytecode = {}; - if (Desc.AS) { - assert((*Desc.AS).Shader->getBufferSize() > 0 && - "The passed task/amplification shader was empty."); - ASBytecode = {(*Desc.AS).Shader->getBuffer().data(), - (*Desc.AS).Shader->getBuffer().size()}; - } - - // The pixel shader is optional - D3D12_SHADER_BYTECODE PSBytecode = {}; - if (Desc.PS) { - assert((*Desc.PS).Shader->getBufferSize() > 0 && - "The passed pixel shader was empty."); - PSBytecode = {(*Desc.PS).Shader->getBuffer().data(), - (*Desc.PS).Shader->getBuffer().size()}; - } - - D3D12_RT_FORMAT_ARRAY RTArray = {}; - RTArray.NumRenderTargets = static_cast(Desc.RTFormats.size()); - for (size_t I = 0; I < Desc.RTFormats.size(); ++I) - RTArray.RTFormats[I] = getDXGIFormat(Desc.RTFormats[I]); - - CD3DX12_DEPTH_STENCIL_DESC1 DepthStencil(D3D12_DEFAULT); - DepthStencil.DepthEnable = true; - DepthStencil.DepthWriteMask = D3D12_DEPTH_WRITE_MASK_ALL; - DepthStencil.DepthFunc = D3D12_COMPARISON_FUNC_LESS; - DepthStencil.StencilEnable = false; - - DXGI_SAMPLE_DESC SampleDesc = {}; - SampleDesc.Count = 1; - - CD3DX12_PIPELINE_MESH_STATE_STREAM Stream; - Stream.pRootSignature = RootSig.Get(); - Stream.AS = ASBytecode; - Stream.MS = MSBytecode; - Stream.PS = PSBytecode; - Stream.RasterizerState = CD3DX12_RASTERIZER_DESC(D3D12_DEFAULT); - Stream.BlendState = CD3DX12_BLEND_DESC(D3D12_DEFAULT); - Stream.DepthStencilState = DepthStencil; - Stream.SampleMask = UINT_MAX; - Stream.PrimitiveTopologyType = getDXPrimitiveTopologyType(Desc.Topology); - Stream.RTVFormats = RTArray; - if (Desc.DSFormat) - Stream.DSVFormat = getDXGIFormat(*Desc.DSFormat); - Stream.SampleDesc = SampleDesc; - - const D3D12_PIPELINE_STATE_STREAM_DESC StreamDesc = {sizeof(Stream), - &Stream}; - - ComPtr PSO; - if (auto Err = HR::toError( - Device->CreatePipelineState(&StreamDesc, IID_PPV_ARGS(&PSO)), - "Failed to create mesh shader PSO.")) - return Err; - - return std::make_unique(Name, RootSig, PSO, std::nullopt); - } - - static std::wstring widen(llvm::StringRef S) { - // Entry-point names and hit-group names are ASCII; a straight 1:1 widen - // is sufficient. - return std::wstring(S.begin(), S.end()); - } - - static D3D12_HIT_GROUP_TYPE getDXHitGroupType(HitGroupType T) { - switch (T) { - case HitGroupType::Triangles: - return D3D12_HIT_GROUP_TYPE_TRIANGLES; - case HitGroupType::Procedural: - return D3D12_HIT_GROUP_TYPE_PROCEDURAL_PRIMITIVE; - } - llvm_unreachable("All HitGroupType cases handled"); - } - - llvm::Expected> - createPipelineRT(llvm::StringRef Name, const BindingsDesc &BndDesc, - const RayTracingPipelineCreateDesc &Desc) override { - if (!Desc.Library) - return llvm::createStringError(std::errc::invalid_argument, - "RayTracingPipelineCreateDesc.Library is " - "null — backend needs a DXIL blob."); - - // Global root signature: try the library's embedded RTS0 part first; - // fall back to building one from BindingsDesc. - ShaderContainer LibContainer = {}; - LibContainer.Shader = Desc.Library; - ComPtr RootSig; - if (auto Err = createRootSignature(Name, BndDesc, LibContainer, - /*IsGraphics=*/false, RootSig)) - return Err; - - CD3DX12_STATE_OBJECT_DESC SODesc( - D3D12_STATE_OBJECT_TYPE_RAYTRACING_PIPELINE); - - // DXIL library subobject — add every Shader's entry point as an export. - // Wide-string storage must outlive SODesc since the subobject only stores - // pointers into it. - auto *Lib = SODesc.CreateSubobject(); - const llvm::StringRef LibBytes = Desc.Library->getBuffer(); - const D3D12_SHADER_BYTECODE Bytecode = {LibBytes.data(), LibBytes.size()}; - Lib->SetDXILLibrary(&Bytecode); - llvm::SmallVector WideNames; - WideNames.reserve(Desc.Shaders.size() + Desc.HitGroups.size()); - for (const auto &Sh : Desc.Shaders) { - WideNames.push_back(widen(Sh.EntryPoint)); - Lib->DefineExport(WideNames.back().c_str()); - } - - // One hit-group subobject per HitGroup entry. - for (const auto &HG : Desc.HitGroups) { - auto *HGObj = SODesc.CreateSubobject(); - HGObj->SetHitGroupType(getDXHitGroupType(HG.Type)); - WideNames.push_back(widen(HG.Name)); - HGObj->SetHitGroupExport(WideNames.back().c_str()); - WideNames.push_back(widen(HG.ClosestHit)); - HGObj->SetClosestHitShaderImport(WideNames.back().c_str()); - if (HG.AnyHit) { - WideNames.push_back(widen(*HG.AnyHit)); - HGObj->SetAnyHitShaderImport(WideNames.back().c_str()); - } - if (HG.Intersection) { - WideNames.push_back(widen(*HG.Intersection)); - HGObj->SetIntersectionShaderImport(WideNames.back().c_str()); - } - } - - // Pipeline-wide shader config (max payload + max attribute bytes). - auto *ShaderCfg = - SODesc.CreateSubobject(); - ShaderCfg->Config(Desc.Config.MaxPayloadSizeInBytes, - Desc.Config.MaxAttributeSizeInBytes); - - // Pipeline-wide config (max recursion depth). - auto *PipelineCfg = - SODesc.CreateSubobject(); - PipelineCfg->Config(Desc.Config.MaxTraceRecursionDepth); - - // Global root signature. - auto *GlobalRS = - SODesc.CreateSubobject(); - GlobalRS->SetRootSignature(RootSig.Get()); - - ComPtr StateObject; - if (auto Err = HR::toError( - Device->CreateStateObject(SODesc, IID_PPV_ARGS(&StateObject)), - "Failed to create raytracing state object.")) - return Err; - - ComPtr Properties; - if (auto Err = HR::toError( - StateObject.As(&Properties), - "Failed to query ID3D12StateObjectProperties from state object.")) - return Err; - - auto State = std::make_unique( - Name, RootSig, StateObject, Properties); - // Cache identifiers up-front. The driver-owned blobs are alive for - // Properties' lifetime, which lives on the PSO. - // - // GetShaderIdentifier only returns non-null for entries that are - // directly bindable from an SBT record: raygen / miss / callable - // shaders and hit groups. Closest-hit / any-hit / intersection are - // bound *through* a hit-group subobject and aren't separately - // addressable, so skip them. - for (const auto &Sh : Desc.Shaders) { - switch (Sh.Stage) { - case Stages::RayGeneration: - case Stages::Miss: - case Stages::Callable: - break; - default: - continue; - } - const std::wstring W = widen(Sh.EntryPoint); - const void *Id = Properties->GetShaderIdentifier(W.c_str()); - if (!Id) - return llvm::createStringError( - "GetShaderIdentifier returned null for shader '%s'", - Sh.EntryPoint.c_str()); - State->ShaderIdentifiers[Sh.EntryPoint] = Id; - } - for (const auto &HG : Desc.HitGroups) { - const std::wstring W = widen(HG.Name); - const void *Id = Properties->GetShaderIdentifier(W.c_str()); - if (!Id) - return llvm::createStringError( - "GetShaderIdentifier returned null for hit group '%s'", - HG.Name.c_str()); - State->ShaderIdentifiers[HG.Name] = Id; - } - return State; - } - - llvm::Expected> - createShaderBindingTable(const PipelineState &PSO, - const ShaderBindingTableDesc &Desc) override { - if (!llvm::isa(&PSO)) - return llvm::createStringError( - std::errc::invalid_argument, - "createShaderBindingTable requires a RayTracing PipelineState"); - const auto &DXRTPSO = llvm::cast(PSO); - - constexpr uint32_t IdSize = D3D12_SHADER_IDENTIFIER_SIZE_IN_BYTES; - const SBTLayout Layout = - computeSBTLayout(IdSize, D3D12_RAYTRACING_SHADER_RECORD_BYTE_ALIGNMENT, - D3D12_RAYTRACING_SHADER_TABLE_BYTE_ALIGNMENT, Desc); - const uint32_t TotalSize = Layout.TotalSize; - const llvm::ArrayRef RGEntries(&Desc.RayGen, 1); - - // Upload heap so the CPU can write the SBT directly. The state-object - // identifiers don't need to live in default heap; using upload keeps - // PR3 simple. A staging copy to default heap is a follow-up. - const D3D12_HEAP_PROPERTIES HeapProps = - CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_UPLOAD); - const D3D12_RESOURCE_DESC BufDesc = - CD3DX12_RESOURCE_DESC::Buffer(TotalSize); - ComPtr Buffer; - if (auto Err = HR::toError(Device->CreateCommittedResource( - &HeapProps, D3D12_HEAP_FLAG_NONE, &BufDesc, - D3D12_RESOURCE_STATE_GENERIC_READ, nullptr, - IID_PPV_ARGS(&Buffer)), - "Failed to create SBT buffer.")) - return Err; - - void *Mapped = nullptr; - const D3D12_RANGE ReadRange{0, 0}; - if (auto Err = HR::toError(Buffer->Map(0, &ReadRange, &Mapped), - "Failed to map SBT buffer.")) - return Err; - std::memset(Mapped, 0, TotalSize); - auto *MappedBytes = static_cast(Mapped); - - auto WriteEntries = [&](uint8_t *Region, llvm::ArrayRef Entries, - uint32_t Stride) -> llvm::Error { - for (size_t I = 0; I < Entries.size(); ++I) { - const auto &E = Entries[I]; - auto It = DXRTPSO.ShaderIdentifiers.find(E.ShaderName); - if (It == DXRTPSO.ShaderIdentifiers.end()) - return llvm::createStringError( - std::errc::invalid_argument, - "SBT references unknown shader/hit-group name: '%s'", - E.ShaderName.c_str()); - uint8_t *Dst = Region + I * Stride; - std::memcpy(Dst, It->second, IdSize); - if (!E.LocalRootData.empty()) - std::memcpy(Dst + IdSize, E.LocalRootData.data(), - E.LocalRootData.size()); - } - return llvm::Error::success(); - }; - - auto WriteRegion = [&](const SBTRegionLayout &R, - llvm::ArrayRef Entries) -> llvm::Error { - return WriteEntries(MappedBytes + R.Offset, Entries, R.Stride); - }; - auto UnmapAndReturn = [&](llvm::Error Err) { - Buffer->Unmap(0, nullptr); - return Err; - }; - if (auto Err = WriteRegion(Layout.RayGen, RGEntries)) - return UnmapAndReturn(std::move(Err)); - if (auto Err = WriteRegion(Layout.Miss, Desc.Miss)) - return UnmapAndReturn(std::move(Err)); - if (auto Err = WriteRegion(Layout.HitGroup, Desc.HitGroup)) - return UnmapAndReturn(std::move(Err)); - if (auto Err = WriteRegion(Layout.Callable, Desc.Callable)) - return UnmapAndReturn(std::move(Err)); - Buffer->Unmap(0, nullptr); - - // D3D12_GPU_VIRTUAL_ADDRESS_RANGE / …_AND_STRIDE expect a zero address - // for empty regions, matching the helper's Size == 0 sentinel. - const D3D12_GPU_VIRTUAL_ADDRESS Base = Buffer->GetGPUVirtualAddress(); - auto MakeRange = [&](const SBTRegionLayout &R) { - return D3D12_GPU_VIRTUAL_ADDRESS_RANGE{R.Size ? Base + R.Offset : 0, - R.Size}; - }; - auto MakeRangeAndStride = [&](const SBTRegionLayout &R) { - return D3D12_GPU_VIRTUAL_ADDRESS_RANGE_AND_STRIDE{ - R.Size ? Base + R.Offset : 0, R.Size, R.Stride}; - }; - return std::make_unique( - Buffer, MakeRange(Layout.RayGen), MakeRangeAndStride(Layout.Miss), - MakeRangeAndStride(Layout.HitGroup), - MakeRangeAndStride(Layout.Callable)); - } - - llvm::Expected> - createFence(llvm::StringRef Name) override { - return DXFence::create(Device.Get(), Name); - } - - llvm::Expected> - createMemoryHeap(std::string Name, size_t SizeInBytes) override { - return DXMemoryHeap::create(Device.Get(), Name, SizeInBytes); - } - - llvm::Expected> - createBuffer(std::string Name, const BufferCreateDesc &Desc, - size_t SizeInBytes) override { - const D3D12_HEAP_TYPE HeapType = getDXHeapType(Desc.Location); - // This flag is only allowed on GpuOnly memory. - const D3D12_RESOURCE_FLAGS Flags = - Desc.Location == MemoryLocation::GpuOnly - ? D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS - : D3D12_RESOURCE_FLAG_NONE; - - // Modify the size if needed - UINT64 CounterOffsetInBytes = 0; - UINT64 BufferSizeInBytes = SizeInBytes; - if (Desc.Usage == BufferUsage::ConstantBuffer) { - if (Desc.HasCounter) - return llvm::createStringError( - "Constant Buffers are not allowed to have a counter."); - - BufferSizeInBytes = getCBVSize(BufferSizeInBytes); - } else if (Desc.HasCounter) { - if (Desc.AccessType == BufferShaderAccessType::Raw) - return llvm::createStringError( - "Raw Resources are not allowed to have a counter."); - - CounterOffsetInBytes = llvm::alignTo( - BufferSizeInBytes, D3D12_UAV_COUNTER_PLACEMENT_ALIGNMENT); - BufferSizeInBytes = CounterOffsetInBytes + sizeof(uint32_t); - } - - const D3D12_HEAP_PROPERTIES HeapProps = CD3DX12_HEAP_PROPERTIES(HeapType); - const D3D12_RESOURCE_DESC BufferDesc = - CD3DX12_RESOURCE_DESC::Buffer(BufferSizeInBytes, Flags); - - D3D12_RESOURCE_STATES PreferredState = D3D12_RESOURCE_STATE_COMMON; - ComPtr BufferObject; - if (Desc.Backing == MemoryBacking::Sparse) { - if (auto Err = HR::toError(Device->CreateReservedResource( - &BufferDesc, D3D12_RESOURCE_STATE_COMMON, - nullptr, IID_PPV_ARGS(&BufferObject)), - "Failed to create reserved buffer.")) - return Err; - } else { - D3D12_RESOURCE_STATES InitialState = D3D12_RESOURCE_STATE_COMMON; - if (HeapType == D3D12_HEAP_TYPE_UPLOAD) - InitialState = D3D12_RESOURCE_STATE_GENERIC_READ; - else if (HeapType == D3D12_HEAP_TYPE_READBACK) - // As per the readback heap docs - // > Resources in this heap must be created with - // > D3D12_RESOURCE_STATE_COPY_DEST, and cannot be changed away from - // this. - InitialState = D3D12_RESOURCE_STATE_COPY_DEST; - PreferredState = InitialState; - if (auto Err = HR::toError(Device->CreateCommittedResource( - &HeapProps, D3D12_HEAP_FLAG_NONE, - &BufferDesc, InitialState, nullptr, - IID_PPV_ARGS(&BufferObject)), - "Failed to create buffer.")) - return Err; - } - - const std::wstring WStr(Name.begin(), Name.end()); - BufferObject->SetName(WStr.c_str()); - - D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle = {}; - { - auto SRVHandleOrErr = CSUAllocator.allocate(); - if (!SRVHandleOrErr) - return SRVHandleOrErr.takeError(); - SRVHandle = *SRVHandleOrErr; - - D3D12_SHADER_RESOURCE_VIEW_DESC SRVDesc = {}; - SRVDesc.ViewDimension = D3D12_SRV_DIMENSION_BUFFER; - SRVDesc.Shader4ComponentMapping = - D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING; - switch (Desc.AccessType) { - case BufferShaderAccessType::Raw: - SRVDesc.Format = DXGI_FORMAT_R32_TYPELESS; - SRVDesc.Buffer.NumElements = static_cast(SizeInBytes / 4); - SRVDesc.Buffer.Flags = D3D12_BUFFER_SRV_FLAG_RAW; - break; - case BufferShaderAccessType::Typed: - SRVDesc.Format = getDXGIFormat(Desc.AccessTypeParams.Fmt); - SRVDesc.Buffer.NumElements = static_cast( - SizeInBytes / getFormatSizeInBytes(Desc.AccessTypeParams.Fmt)); - break; - case BufferShaderAccessType::Structured: - assert(Desc.AccessTypeParams.StructureStride > 0 && - "Structured buffers must have a Structure Stride."); - SRVDesc.Format = DXGI_FORMAT_UNKNOWN; - SRVDesc.Buffer.NumElements = static_cast( - SizeInBytes / Desc.AccessTypeParams.StructureStride); - SRVDesc.Buffer.StructureByteStride = - Desc.AccessTypeParams.StructureStride; - break; - } - - Device->CreateShaderResourceView(BufferObject.Get(), &SRVDesc, SRVHandle); - } - - D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle = {}; - if ((Flags & D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS) != 0) { - auto UAVHandleOrErr = CSUAllocator.allocate(); - if (!UAVHandleOrErr) - return UAVHandleOrErr.takeError(); - UAVHandle = *UAVHandleOrErr; - - D3D12_UNORDERED_ACCESS_VIEW_DESC UAVDesc = {}; - UAVDesc.ViewDimension = D3D12_UAV_DIMENSION_BUFFER; - switch (Desc.AccessType) { - case BufferShaderAccessType::Raw: - UAVDesc.Format = DXGI_FORMAT_R32_TYPELESS; - UAVDesc.Buffer.NumElements = static_cast(SizeInBytes / 4); - UAVDesc.Buffer.Flags = D3D12_BUFFER_UAV_FLAG_RAW; - break; - case BufferShaderAccessType::Typed: - UAVDesc.Format = getDXGIFormat(Desc.AccessTypeParams.Fmt); - UAVDesc.Buffer.NumElements = static_cast( - SizeInBytes / getFormatSizeInBytes(Desc.AccessTypeParams.Fmt)); - break; - case BufferShaderAccessType::Structured: - assert(Desc.AccessTypeParams.StructureStride > 0 && - "Structured buffers must have a Structure Stride."); - UAVDesc.Format = DXGI_FORMAT_UNKNOWN; - UAVDesc.Buffer.NumElements = static_cast( - SizeInBytes / Desc.AccessTypeParams.StructureStride); - UAVDesc.Buffer.StructureByteStride = - Desc.AccessTypeParams.StructureStride; - break; - } - - ID3D12Resource *CounterObject = nullptr; - if (Desc.HasCounter) { - UAVDesc.Buffer.CounterOffsetInBytes = CounterOffsetInBytes; - CounterObject = BufferObject.Get(); - } - - Device->CreateUnorderedAccessView(BufferObject.Get(), CounterObject, - &UAVDesc, UAVHandle); - } - - D3D12_CPU_DESCRIPTOR_HANDLE CBVHandle = {}; - if (Desc.Usage == BufferUsage::ConstantBuffer) { - auto CBVHandleOrErr = CSUAllocator.allocate(); - if (!CBVHandleOrErr) - return CBVHandleOrErr.takeError(); - CBVHandle = *CBVHandleOrErr; - - D3D12_CONSTANT_BUFFER_VIEW_DESC CBVDesc = {}; - CBVDesc.BufferLocation = BufferObject->GetGPUVirtualAddress(); - CBVDesc.SizeInBytes = BufferSizeInBytes; - - Device->CreateConstantBufferView(&CBVDesc, CBVHandle); - } - - return std::make_unique(BufferObject, Name, Desc, SizeInBytes, - CounterOffsetInBytes, PreferredState, - SRVHandle, UAVHandle, CBVHandle); - } - - llvm::Expected> - createTexture(std::string Name, const TextureCreateDesc &Desc) override { - if (auto Err = validateTextureCreateDesc(Desc)) - return Err; - - const D3D12_HEAP_PROPERTIES HeapProps = - CD3DX12_HEAP_PROPERTIES(getDXHeapType(Desc.Location)); - - D3D12_RESOURCE_DESC TexDesc = {}; - TexDesc.Dimension = D3D12_RESOURCE_DIMENSION_TEXTURE2D; - TexDesc.Width = Desc.Width; - TexDesc.Height = Desc.Height; - TexDesc.DepthOrArraySize = 1; - TexDesc.MipLevels = static_cast(Desc.MipLevels); - TexDesc.Format = getDXGIFormat(Desc.Fmt); - TexDesc.SampleDesc.Count = 1; - if (Desc.Location == MemoryLocation::GpuOnly) { - if (Desc.Backing == MemoryBacking::Sparse) - TexDesc.Layout = D3D12_TEXTURE_LAYOUT_64KB_UNDEFINED_SWIZZLE; - else - TexDesc.Layout = D3D12_TEXTURE_LAYOUT_UNKNOWN; - } else { - TexDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR; - } - TexDesc.Flags = getDXResourceFlags(Desc.Usage); - - const D3D12_CLEAR_VALUE *ClearValuePtr = nullptr; - D3D12_CLEAR_VALUE ClearValue = {}; - if (Desc.OptimizedClearValue) { - ClearValue.Format = TexDesc.Format; - std::visit( - [&ClearValue](auto &&V) { - using T = std::decay_t; - if constexpr (std::is_same_v) { - ClearValue.Color[0] = V.R; - ClearValue.Color[1] = V.G; - ClearValue.Color[2] = V.B; - ClearValue.Color[3] = V.A; - } else { - ClearValue.DepthStencil.Depth = V.Depth; - ClearValue.DepthStencil.Stencil = V.Stencil; - } - }, - *Desc.OptimizedClearValue); - ClearValuePtr = &ClearValue; - } - - D3D12_RESOURCE_STATES InitialState = - D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE | - D3D12_RESOURCE_STATE_NON_PIXEL_SHADER_RESOURCE; - if ((Desc.Usage & TextureUsage::Storage)) - InitialState = D3D12_RESOURCE_STATE_UNORDERED_ACCESS; - - ComPtr TextureObject; - if (Desc.Backing == MemoryBacking::Sparse) { - if (auto Err = HR::toError(Device->CreateReservedResource( - &TexDesc, InitialState, ClearValuePtr, - IID_PPV_ARGS(&TextureObject)), - "Failed to create reserved texture.")) - return Err; - } else { - if (auto Err = HR::toError(Device->CreateCommittedResource( - &HeapProps, D3D12_HEAP_FLAG_NONE, &TexDesc, - InitialState, ClearValuePtr, - IID_PPV_ARGS(&TextureObject)), - "Failed to create texture.")) - return Err; - } - - D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle = {}; - { - auto SRVHandleOrErr = CSUAllocator.allocate(); - if (!SRVHandleOrErr) - return SRVHandleOrErr.takeError(); - SRVHandle = *SRVHandleOrErr; - - D3D12_SHADER_RESOURCE_VIEW_DESC SRVDesc = {}; - SRVDesc.ViewDimension = - D3D12_SRV_DIMENSION_TEXTURE2D; // assume this is correct for now. - SRVDesc.Shader4ComponentMapping = - D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING; - SRVDesc.Format = getDXGIFormatSRV(Desc.Fmt); - SRVDesc.Texture2D.MostDetailedMip = 0; - SRVDesc.Texture2D.MipLevels = Desc.MipLevels; - SRVDesc.Texture2D.PlaneSlice = 0; - SRVDesc.Texture2D.ResourceMinLODClamp = 0.0f; - - Device->CreateShaderResourceView(TextureObject.Get(), &SRVDesc, - SRVHandle); - } - - D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle = {}; - if ((TexDesc.Flags & D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS) != 0) { - auto UAVHandleOrErr = CSUAllocator.allocate(); - if (!UAVHandleOrErr) - return UAVHandleOrErr.takeError(); - UAVHandle = *UAVHandleOrErr; - - D3D12_UNORDERED_ACCESS_VIEW_DESC UAVDesc = {}; - UAVDesc.ViewDimension = D3D12_UAV_DIMENSION_TEXTURE2D; - UAVDesc.Texture2D.MipSlice = 0; - UAVDesc.Texture2D.PlaneSlice = 0; - - Device->CreateUnorderedAccessView(TextureObject.Get(), nullptr, &UAVDesc, - UAVHandle); - } - const D3D12_RESOURCE_STATES PreferredState = InitialState; - auto Tex = std::make_unique( - TextureObject, Name, Desc, PreferredState, SRVHandle, UAVHandle); - - const bool IsRT = (Desc.Usage & TextureUsage::RenderTarget) != 0; - const bool IsDS = (Desc.Usage & TextureUsage::DepthStencil) != 0; - if (IsRT) { - auto HandleOrErr = RTVAllocator.allocate(); - if (!HandleOrErr) - return HandleOrErr.takeError(); - Tex->RTVHandle = *HandleOrErr; - Device->CreateRenderTargetView(TextureObject.Get(), nullptr, - Tex->RTVHandle); - } - if (IsDS) { - auto HandleOrErr = DSVAllocator.allocate(); - if (!HandleOrErr) - return HandleOrErr.takeError(); - Tex->DSVHandle = *HandleOrErr; - Device->CreateDepthStencilView(TextureObject.Get(), nullptr, - Tex->DSVHandle); - } - - return Tex; - } - - uint32_t getTextureUploadRowStrideInBytes( - const TextureCreateDesc &Desc) const override { - return getAlignedTexturePitch(Desc.Width, getFormatSizeInBytes(Desc.Fmt)); - } - - static llvm::Expected> - create(ComPtr Adapter, const DeviceConfig &Config) { - ComPtr Device; - if (auto Err = - HR::toError(D3D12CreateDevice(Adapter.Get(), D3D_FEATURE_LEVEL_11_0, - IID_PPV_ARGS(&Device)), - "Failed to create D3D device")) - return Err; - - static std::once_flag SignalHandlerRegistered; - std::call_once(SignalHandlerRegistered, [] { - llvm::sys::AddSignalHandler(dumpD3DInfoQueues, nullptr); - }); - { - const std::lock_guard Lock(SignalHandlerMutex); - SignalHandlerDevices.push_back(Device.Get()); - } - - assert( - Adapter->IsPropertySupported(DXCoreAdapterProperty::DriverDescription)); - size_t BufferSize; - Adapter->GetPropertySize(DXCoreAdapterProperty::DriverDescription, - &BufferSize); - std::vector DescVec(BufferSize); - Adapter->GetProperty(DXCoreAdapterProperty::DriverDescription, BufferSize, - (void *)DescVec.data()); - - std::string DriverVer; - if (Adapter->IsPropertySupported(DXCoreAdapterProperty::DriverVersion)) { - uint64_t Packed = 0; - if (SUCCEEDED(Adapter->GetProperty(DXCoreAdapterProperty::DriverVersion, - sizeof(Packed), &Packed))) { - const uint16_t Major = static_cast((Packed >> 48) & 0xFFFF); - const uint16_t Minor = static_cast((Packed >> 32) & 0xFFFF); - const uint16_t Build = static_cast((Packed >> 16) & 0xFFFF); - const uint16_t Revision = static_cast(Packed & 0xFFFF); - DriverVer = (llvm::Twine(Major) + "." + llvm::Twine(Minor) + "." + - llvm::Twine(Build) + "." + llvm::Twine(Revision)) - .str(); - } - } - - if (Config.EnableDebugLayer || Config.EnableValidationLayer) - if (auto Err = configureInfoQueue(Device.Get())) - return Err; - - auto GraphicsQueueOrErr = DXQueue::createGraphicsQueue(Device); - if (!GraphicsQueueOrErr) - return GraphicsQueueOrErr.takeError(); - - auto RTVHeapOrErr = DescriptorAllocator::create( - Device.Get(), D3D12_DESCRIPTOR_HEAP_TYPE_RTV, 256); - if (!RTVHeapOrErr) - return RTVHeapOrErr.takeError(); - - auto DSVHeapOrErr = DescriptorAllocator::create( - Device.Get(), D3D12_DESCRIPTOR_HEAP_TYPE_DSV, 256); - if (!DSVHeapOrErr) - return DSVHeapOrErr.takeError(); - - auto CSUHeapOrErr = DescriptorAllocator::create( - Device.Get(), D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV, 4096); - if (!CSUHeapOrErr) - return CSUHeapOrErr.takeError(); - - return std::make_unique( - Adapter, Device, std::move(*GraphicsQueueOrErr), - std::move(*RTVHeapOrErr), std::move(*DSVHeapOrErr), - std::move(*CSUHeapOrErr), std::string(DescVec.data()), - std::move(DriverVer)); - } - - const Capabilities &getCapabilities() override { - if (Caps.empty()) - queryCapabilities(); - return Caps; - } - - void queryCapabilities() { - CD3DX12FeatureSupport Features; - Features.Init(Device.Get()); +void DXDevice::queryCapabilities() { + CD3DX12FeatureSupport Features; + Features.Init(Device.Get()); #define D3D_FEATURE_BOOL(Name) \ Caps.insert( \ @@ -2288,1303 +1414,194 @@ class DXDevice : public offloadtest::Device { makeCapability(#Name, static_cast(Features.Name())))); #include "DXFeatures.def" - } - - static llvm::Error configureInfoQueue(ID3D12DeviceX *Device) { -#ifdef _WIN32 - ComPtr InfoQueue; - if (auto Err = HR::toError(Device->QueryInterface(InfoQueue.GetAddressOf()), - "Error initializing info queue")) - return Err; - InfoQueue->SetBreakOnSeverity(D3D12_MESSAGE_SEVERITY_CORRUPTION, TRUE); - InfoQueue->SetBreakOnSeverity(D3D12_MESSAGE_SEVERITY_ERROR, TRUE); - InfoQueue->SetBreakOnSeverity(D3D12_MESSAGE_SEVERITY_WARNING, TRUE); -#endif - return llvm::Error::success(); - } - - llvm::Error createDescriptorHeap(Pipeline &P, InvocationState &State) { - if (P.getDescriptorCount() == 0) - return llvm::Error::success(); - const D3D12_DESCRIPTOR_HEAP_DESC HeapDesc = { - D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV, - P.getDescriptorCountWithFlattenedArrays(), - D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE, 0}; - if (auto Err = HR::toError(Device->CreateDescriptorHeap( - &HeapDesc, IID_PPV_ARGS(&State.DescHeap)), - "Failed to create descriptor heap.")) - return Err; - return llvm::Error::success(); - } - - llvm::Expected> - createCommandBuffer() override { - auto CBOrErr = DXCommandBuffer::create(Device); - if (!CBOrErr) - return CBOrErr.takeError(); - (*CBOrErr)->Dev = this; - return std::unique_ptr(std::move(*CBOrErr)); - } - - llvm::Expected> - createRenderPass(const offloadtest::RenderPassDesc &Desc) override { - return std::make_unique(Desc); - } - - llvm::Expected - getBLASBuildSizes(llvm::ArrayRef Triangles) override { - if (auto Err = validateBLASGeometry(Triangles)) - return Err; - - llvm::SmallVector GeomDescs; - GeomDescs.reserve(Triangles.size()); - for (const auto &T : Triangles) { - D3D12_RAYTRACING_GEOMETRY_DESC GD = {}; - GD.Type = D3D12_RAYTRACING_GEOMETRY_TYPE_TRIANGLES; - if (T.Opaque) - GD.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_OPAQUE; - - auto &Tri = GD.Triangles; - // GPU addresses are not needed for the prebuild size query; they will - // be populated at build time. - Tri.VertexBuffer.StrideInBytes = T.VertexStride; - Tri.VertexCount = T.VertexCount; - Tri.VertexFormat = getDXGIFormat(T.VertexFormat); - - if (T.IndexBuffer) { - Tri.IndexCount = T.IndexCount; - Tri.IndexFormat = getDXGIIndexFormat(T.IdxFormat); - } - - GeomDescs.push_back(GD); - } - return queryBLASPrebuildSize(GeomDescs); - } - - llvm::Expected - getBLASBuildSizes(llvm::ArrayRef AABBs) override { - if (auto Err = validateBLASGeometry(AABBs)) - return Err; - - llvm::SmallVector GeomDescs; - GeomDescs.reserve(AABBs.size()); - for (const auto &A : AABBs) { - D3D12_RAYTRACING_GEOMETRY_DESC GD = {}; - GD.Type = D3D12_RAYTRACING_GEOMETRY_TYPE_PROCEDURAL_PRIMITIVE_AABBS; - if (A.Opaque) - GD.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_OPAQUE; - - GD.AABBs.AABBs.StrideInBytes = A.AABBStride; - GD.AABBs.AABBCount = A.AABBCount; - - GeomDescs.push_back(GD); - } - return queryBLASPrebuildSize(GeomDescs); - } - - AccelerationStructureSizes queryBLASPrebuildSize( - llvm::ArrayRef GeomDescs) { - D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS Inputs = {}; - Inputs.Type = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL; - Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY; - Inputs.NumDescs = GeomDescs.size(); - Inputs.pGeometryDescs = GeomDescs.data(); - - D3D12_RAYTRACING_ACCELERATION_STRUCTURE_PREBUILD_INFO Info = {}; - Device->GetRaytracingAccelerationStructurePrebuildInfo(&Inputs, &Info); - - return {Info.ResultDataMaxSizeInBytes, Info.ScratchDataSizeInBytes, - Info.UpdateScratchDataSizeInBytes}; - } - - llvm::Expected> - allocateAS(const AccelerationStructureSizes &Sizes, const char *Kind) { - const uint64_t AlignedSize = - llvm::alignTo(Sizes.ResultDataMaxSizeInBytes, - D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BYTE_ALIGNMENT); - const D3D12_HEAP_PROPERTIES HeapProps = - CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_DEFAULT); - const D3D12_RESOURCE_DESC BufferDesc = CD3DX12_RESOURCE_DESC::Buffer( - AlignedSize, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS); - - ComPtr ASBuffer; - if (auto Err = HR::toError( - Device->CreateCommittedResource( - &HeapProps, D3D12_HEAP_FLAG_NONE, &BufferDesc, - D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE, nullptr, - IID_PPV_ARGS(&ASBuffer)), - "Failed to create " + llvm::Twine(Kind) + " resource.")) - return Err; - - return std::make_unique(ASBuffer, Sizes); - } - - llvm::Expected - getTLASBuildSizes(uint32_t InstanceCount) override { - D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS Inputs = {}; - Inputs.Type = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL; - Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY; - Inputs.NumDescs = InstanceCount; - - D3D12_RAYTRACING_ACCELERATION_STRUCTURE_PREBUILD_INFO Info = {}; - Device->GetRaytracingAccelerationStructurePrebuildInfo(&Inputs, &Info); - - return AccelerationStructureSizes{Info.ResultDataMaxSizeInBytes, - Info.ScratchDataSizeInBytes, - Info.UpdateScratchDataSizeInBytes}; - } - - llvm::Expected> - createBLAS(const AccelerationStructureSizes &Sizes) override { - return allocateAS(Sizes, "BLAS"); - } - - llvm::Expected> - createTLAS(const AccelerationStructureSizes &Sizes) override { - auto TLASOrErr = allocateAS(Sizes, "TLAS"); - if (!TLASOrErr) - return TLASOrErr.takeError(); - auto TLAS = std::move(*TLASOrErr); - DXAccelerationStructure &TLASDX = - llvm::cast(*TLAS.get()); - - auto SRVHandleOrErr = CSUAllocator.allocate(); - if (!SRVHandleOrErr) - return SRVHandleOrErr.takeError(); - TLASDX.SRVHandle = *SRVHandleOrErr; - - D3D12_SHADER_RESOURCE_VIEW_DESC SRVDesc = {}; - SRVDesc.Format = DXGI_FORMAT_UNKNOWN; - SRVDesc.ViewDimension = - D3D12_SRV_DIMENSION_RAYTRACING_ACCELERATION_STRUCTURE; - SRVDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING; - SRVDesc.RaytracingAccelerationStructure.Location = - TLASDX.getGPUVirtualAddress(); - - // AS SRVs are created with a null resource; the AS lives in the - // buffer referenced by Location. - Device->CreateShaderResourceView(nullptr, &SRVDesc, TLASDX.SRVHandle); - - return TLAS; - } - - static size_t getCBVSize(size_t Sz) { - return (Sz + 255u) & 0xFFFFFFFFFFFFFF00; - } - - llvm::Expected> createAS(Resource &R) { - assert(R.TLASPtr && "AS resource must be resolved to a TLAS"); - assert(R.getArraySize() == 1 && "AS arrays not yet supported"); - auto SizesOrErr = - getTLASBuildSizes(static_cast(R.TLASPtr->Instances.size())); - if (!SizesOrErr) - return SizesOrErr.takeError(); - return createTLAS(*SizesOrErr); - } - - llvm::Error createBuffers(Pipeline &P, InvocationState &IS) { - auto EncOrErr = IS.CB->createComputeEncoder(); - if (!EncOrErr) - return EncOrErr.takeError(); - auto Enc = std::move(*EncOrErr); - - auto CreateBuffer = - [&Enc, &IS, - this](Resource &R, - llvm::SmallVectorImpl &Resources) -> llvm::Error { - ResourceBundle ResBundle; - if (R.isBuffer()) { - BufferCreateDesc CreateDesc = {}; - CreateDesc.Location = MemoryLocation::GpuOnly; - CreateDesc.Backing = - R.IsReserved ? MemoryBacking::Sparse : MemoryBacking::Automatic; - CreateDesc.Usage = bufferUsageFromResourceKind(R.Kind); - CreateDesc.AccessType = bufferShaderAccessTypeFromResourceKind( - R, CreateDesc.AccessTypeParams); - CreateDesc.HasCounter = R.HasCounter; - - for (auto &Data : R.BufferPtr->Data) { - std::unique_ptr UploadBuffer; - std::unique_ptr BackingMemoryHeap; - - std::unique_ptr Buffer; - if (R.IsReserved) { - auto BufferOrErr = createSparseBufferWithData( - *this, GraphicsQueue, "Sparse Buffer", CreateDesc, R.size(), - R.TilesMapped, Data.get(), R.size(), *Enc.get(), UploadBuffer, - BackingMemoryHeap); - if (!BufferOrErr) - return BufferOrErr.takeError(); - - Buffer = std::move(*BufferOrErr); - } else { - auto BufferOrErr = - createBufferWithData(*this, "Buffer", CreateDesc, Data.get(), - R.size(), Enc.get(), &UploadBuffer); - if (!BufferOrErr) - return BufferOrErr.takeError(); - - Buffer = std::move(*BufferOrErr); - } - - std::unique_ptr ReadbackBuffer; - std::unique_ptr CounterReadbackBuffer; - if (getDescriptorKind(R.Kind) == DescriptorKind::UAV) { - const BufferCreateDesc ReadbackDesc = - BufferCreateDesc::readbackBuffer(); - auto ReadbackOrErr = createBuffer("Readback", ReadbackDesc, - Buffer->getSizeInBytes()); - if (!ReadbackOrErr) - return ReadbackOrErr.takeError(); - ReadbackBuffer = std::move(*ReadbackOrErr); - - if (R.HasCounter) { - auto CounterReadbackOrErr = - createBuffer("Readback", ReadbackDesc, sizeof(uint32_t)); - if (!CounterReadbackOrErr) - return CounterReadbackOrErr.takeError(); - CounterReadbackBuffer = std::move(*CounterReadbackOrErr); - } - } - - IS.KeepAliveBuffers.push_back(std::move(UploadBuffer)); - ResourceSet RSet(std::move(Buffer), std::move(BackingMemoryHeap), - std::move(ReadbackBuffer), - std::move(CounterReadbackBuffer)); - ResBundle.push_back(std::move(RSet)); - } - } else if (R.isTexture()) { - if (R.BufferPtr->OutputProps.MipLevels != 1) - return llvm::createStringError(std::errc::not_supported, - "Multiple mip levels are not yet " - "supported for DirectX textures."); - - auto FormatOrErr = toFormat(R.BufferPtr->Format, R.BufferPtr->Channels); - if (!FormatOrErr) - return FormatOrErr.takeError(); - - LLVM_ENABLE_BITMASK_ENUMS_IN_NAMESPACE(); - - TextureCreateDesc CreateDesc = {}; - CreateDesc.Location = MemoryLocation::GpuOnly; - CreateDesc.Backing = - R.IsReserved ? MemoryBacking::Sparse : MemoryBacking::Automatic; - CreateDesc.Usage = TextureUsage::Sampled; - if (R.Kind == ResourceKind::RWTexture2D) - CreateDesc.Usage |= TextureUsage::Storage; - CreateDesc.Fmt = *FormatOrErr; - CreateDesc.Width = R.BufferPtr->OutputProps.Width; - CreateDesc.Height = R.BufferPtr->OutputProps.Height; - CreateDesc.MipLevels = 1; - - for (auto &Data : R.BufferPtr->Data) { - std::unique_ptr UploadBuffer; - std::unique_ptr BackingMemoryHeap; - - std::unique_ptr Texture; - if (R.IsReserved) { - auto TextureOrErr = createSparseTextureWithData( - *this, GraphicsQueue, "Sparse Texture", CreateDesc, Data.get(), - R.size(), *Enc.get(), UploadBuffer, BackingMemoryHeap); - if (!TextureOrErr) - return TextureOrErr.takeError(); - - Texture = std::move(*TextureOrErr); - } else { - auto TextureOrErr = - createTextureWithData(*this, "Texture", CreateDesc, Data.get(), - R.size(), Enc.get(), &UploadBuffer); - if (!TextureOrErr) - return TextureOrErr.takeError(); - - Texture = std::move(*TextureOrErr); - } - - std::unique_ptr ReadbackBuffer; - if (getDescriptorKind(R.Kind) == DescriptorKind::UAV) { - const BufferCreateDesc ReadbackDesc = - BufferCreateDesc::readbackBuffer(); - auto ReadbackOrErr = - createBuffer("Readback", ReadbackDesc, - Texture->calculateLinearSizeInBytes(*this)); - if (!ReadbackOrErr) - return ReadbackOrErr.takeError(); - ReadbackBuffer = std::move(*ReadbackOrErr); - } - - IS.KeepAliveBuffers.push_back(std::move(UploadBuffer)); - ResourceSet RSet(std::move(Texture), std::move(BackingMemoryHeap), - std::move(ReadbackBuffer)); - ResBundle.push_back(std::move(RSet)); - } - } else if (R.isAccelerationStructure()) { - auto ASOrErr = createAS(R); - if (!ASOrErr) - return ASOrErr.takeError(); - ResBundle.emplace_back(ASOrErr->get()); - auto Inserted = - IS.TLASes.try_emplace(R.TLASPtr->Name, std::move(*ASOrErr)); - assert(Inserted.second && "TLAS bound to multiple resources NYI"); - (void)Inserted; - } else { - return llvm::createStringError( - std::errc::not_supported, - "Samplers are not yet implemented for DirectX."); - } - - Resources.push_back(std::make_pair(&R, std::move(ResBundle))); - return llvm::Error::success(); - }; - - for (auto &D : P.Sets) { - IS.DescTables.emplace_back(DescriptorTable()); - DescriptorTable &Table = IS.DescTables.back(); - for (auto &R : D.Resources) - if (auto Err = CreateBuffer(R, Table.Resources)) - return Err; - } - - Enc->endEncoding(); - - // Bind descriptors in descriptor tables. - if (IS.DescHeap) { - uint32_t HeapIndex = 0; - const D3D12_CPU_DESCRIPTOR_HANDLE HeapStart = - IS.DescHeap->GetCPUDescriptorHandleForHeapStart(); - const uint32_t DescHandleIncSize = - Device->GetDescriptorHandleIncrementSize( - D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); - for (auto &T : IS.DescTables) { - for (auto &R : T.Resources) { - for (const auto &Set : R.second) { - D3D12_CPU_DESCRIPTOR_HANDLE DescriptorHandle = {}; - if (Set.Buffer != nullptr) { - const DXBuffer &BufferDX = - llvm::cast(*Set.Buffer.get()); - switch (getDescriptorKind(R.first->Kind)) { - case DescriptorKind::SRV: - assert(BufferDX.SRVHandle.ptr != 0 && - "Missing SRV Descriptor. Is BufferUsage correct?"); - DescriptorHandle = BufferDX.SRVHandle; - break; - case DescriptorKind::UAV: - assert(BufferDX.UAVHandle.ptr != 0 && - "Missing UAV Descriptor. Is BufferUsage correct?"); - DescriptorHandle = BufferDX.UAVHandle; - break; - case DescriptorKind::CBV: - assert(BufferDX.CBVHandle.ptr != 0 && - "Missing CBV Descriptor. Is BufferUsage correct?"); - DescriptorHandle = BufferDX.CBVHandle; - break; - default: - llvm_unreachable("Invalid DescriptorKind for a Buffer."); - break; - } - } else if (Set.Texture != nullptr) { - const DXTexture &TextureDX = - llvm::cast(*Set.Texture.get()); - switch (getDescriptorKind(R.first->Kind)) { - case DescriptorKind::SRV: - assert(TextureDX.SRVHandle.ptr != 0 && - "Missing SRV Descriptor. Is TextureUsage correct?"); - DescriptorHandle = TextureDX.SRVHandle; - break; - case DescriptorKind::UAV: - assert(TextureDX.UAVHandle.ptr != 0 && - "Missing UAV Descriptor. Is TextureUsage correct?"); - DescriptorHandle = TextureDX.UAVHandle; - break; - default: - llvm_unreachable("Invalid DescriptorKind for a Texture."); - break; - } - } else if (Set.AS != nullptr) { - const DXAccelerationStructure &AccelerationStructureDX = - llvm::cast(*Set.AS); - DescriptorHandle = AccelerationStructureDX.SRVHandle; - } else { - llvm_unreachable("Resource was a texture nor buffer. Samplers " - "are unsupported"); - } - - assert(DescriptorHandle.ptr != 0 && - "Somehow got a null descriptor :("); - Device->CopyDescriptorsSimple( - 1, {HeapStart.ptr + HeapIndex * DescHandleIncSize}, - DescriptorHandle, D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); - HeapIndex += 1; - } - } - } - } - - // Setup root descriptors - for (auto &R : P.Settings.DX.RootParams) { - if (R.Kind != dx::RootParamKind::RootDescriptor) - continue; - auto &Resource = std::get(R.Data); - if (!Resource.IsReserved && Resource.TilesMapped.has_value()) { - return llvm::createStringError( - std::errc::invalid_argument, - "Error: Cannot define tiles mapped without declaring resource as " - "reserved."); - } - if (auto Err = CreateBuffer(Resource, IS.RootResources)) - return Err; - } - - if (P.isTraditionalRaster() && P.Bindings.VertexBufferPtr) { - const CPUBuffer *VBuffer = P.Bindings.VertexBufferPtr; - - BufferCreateDesc BufDesc = {}; - BufDesc.Location = MemoryLocation::CpuToGpu; - BufDesc.Usage = BufferUsage::VertexBuffer; - auto BufOrErr = createBufferWithData(*this, "VertexBuffer", BufDesc, - VBuffer->Data[0].get(), - VBuffer->size(), nullptr, nullptr); - if (!BufOrErr) - return BufOrErr.takeError(); - IS.VB = std::move(*BufOrErr); - llvm::outs() << "Vertex buffer created.\n"; - } - - return llvm::Error::success(); - } - - static llvm::Error - copyBackResource(offloadtest::ComputeEncoder &ReadbackEncoder, - ResourcePair &R) { - if (R.first->isTexture()) { - for (const ResourceSet &RS : R.second) { - if (RS.Readback == nullptr) - continue; - - if (auto Err = - ReadbackEncoder.copyTextureToBuffer(*RS.Texture, *RS.Readback)) - return Err; - } - } else if (R.first->isBuffer()) { - for (const ResourceSet &RS : R.second) { - if (RS.Readback == nullptr) - continue; - - if (auto Err = ReadbackEncoder.copyBufferToBuffer( - *RS.Buffer, 0, *RS.Readback, 0, RS.Buffer->getSizeInBytes())) - return Err; - - if (!RS.Buffer->getDesc().HasCounter) - continue; - - if (auto Err = ReadbackEncoder.copyCounterToBuffer(*RS.Buffer, - *RS.CounterReadback)) - return Err; - } - } - - return llvm::Error::success(); - } - - llvm::Error createComputeCommands(Pipeline &P, InvocationState &IS) { - CD3DX12_GPU_DESCRIPTOR_HANDLE Handle; - if (IS.DescHeap) { - ID3D12DescriptorHeap *const Heaps[] = {IS.DescHeap.Get()}; - IS.CB->CmdList->SetDescriptorHeaps(1, Heaps); - Handle = IS.DescHeap->GetGPUDescriptorHandleForHeapStart(); - } - const DXPipelineState &DXPipeline = - llvm::cast(*IS.Pipeline.get()); - IS.CB->CmdList->SetComputeRootSignature(DXPipeline.RootSig.Get()); - - const uint32_t Inc = Device->GetDescriptorHandleIncrementSize( - D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); - - if (P.Settings.DX.RootParams.size() > 0) { - uint32_t ConstantOffset = 0u; - uint32_t RootParamIndex = 0u; - uint32_t DescriptorTableIndex = 0u; - auto *RootDescIt = IS.RootResources.begin(); - for (const auto &Param : P.Settings.DX.RootParams) { - switch (Param.Kind) { - case dx::RootParamKind::Constant: { - auto &Constant = std::get(Param.Data); - if (Constant.BufferPtr->ArraySize != 1) - return llvm::createStringError( - std::errc::value_too_large, - "Root constant cannot refer to resource arrays."); - const uint32_t NumValues = - Constant.BufferPtr->size() / sizeof(uint32_t); - IS.CB->CmdList->SetComputeRoot32BitConstants( - RootParamIndex++, NumValues, - Constant.BufferPtr->Data.back().get(), ConstantOffset); - ConstantOffset += NumValues; - break; - } - case dx::RootParamKind::DescriptorTable: - IS.CB->CmdList->SetComputeRootDescriptorTable(RootParamIndex++, - Handle); - Handle.Offset(P.Sets[DescriptorTableIndex++].Resources.size(), Inc); - break; - case dx::RootParamKind::RootDescriptor: - assert(RootDescIt != IS.RootResources.end()); - if (RootDescIt->first->getArraySize() != 1) - return llvm::createStringError( - std::errc::value_too_large, - "Root descriptor cannot refer to resource arrays."); - - const DXBuffer *BufferDX = llvm::cast_if_present( - RootDescIt->second.back().Buffer.get()); - if (!BufferDX) { - return llvm::createStringError( - std::errc::value_too_large, - "Root descriptor can only refer to buffers."); - } - - const D3D12_GPU_VIRTUAL_ADDRESS VirtualAddress = - BufferDX->Buffer->GetGPUVirtualAddress(); - switch (getDescriptorKind(RootDescIt->first->Kind)) { - case DescriptorKind::SRV: - IS.CB->CmdList->SetComputeRootShaderResourceView(RootParamIndex++, - VirtualAddress); - break; - case DescriptorKind::UAV: - IS.CB->CmdList->SetComputeRootUnorderedAccessView(RootParamIndex++, - VirtualAddress); - break; - case DescriptorKind::CBV: - IS.CB->CmdList->SetComputeRootConstantBufferView(RootParamIndex++, - VirtualAddress); - break; - case DescriptorKind::SAMPLER: - llvm_unreachable("Not implemented yet."); - } - ++RootDescIt; - break; - } - } - } else { - // If no explicit root parameters are provided, fall back to using the - // descriptor set layout. This is to make it easier to write tests that - // don't need complicated root signatures. - for (uint32_t Idx = 0u; Idx < P.Sets.size(); ++Idx) { - IS.CB->CmdList->SetComputeRootDescriptorTable(Idx, Handle); - Handle.Offset(P.Sets[Idx].Resources.size(), Inc); - } - } - - { - auto EncoderOrErr = IS.CB->createComputeEncoder(); - if (!EncoderOrErr) - return EncoderOrErr.takeError(); - auto &Encoder = *EncoderOrErr.get(); - if (P.isRayTracing()) { - if (auto Err = Encoder.dispatchRays( - *IS.Pipeline, *IS.SBT, - P.DispatchParameters.DispatchGroupCount[0], - P.DispatchParameters.DispatchGroupCount[1], - P.DispatchParameters.DispatchGroupCount[2])) - return Err; - } else if (auto Err = Encoder.dispatch( - *IS.Pipeline.get(), - P.DispatchParameters.DispatchGroupCount[0], - P.DispatchParameters.DispatchGroupCount[1], - P.DispatchParameters.DispatchGroupCount[2])) { - return Err; - } - Encoder.endEncoding(); - } - - auto EncoderOrErr = IS.CB->createComputeEncoder(); - if (!EncoderOrErr) - return EncoderOrErr.takeError(); - auto ReadbackEncoder = std::move(*EncoderOrErr); - - for (auto &Table : IS.DescTables) - for (auto &R : Table.Resources) - if (auto Err = DXDevice::copyBackResource(*ReadbackEncoder, R)) - return Err; - - for (auto &R : IS.RootResources) - if (auto Err = DXDevice::copyBackResource(*ReadbackEncoder, R)) - return Err; - - ReadbackEncoder->endEncoding(); - - return llvm::Error::success(); - } - - llvm::Error readBack(Pipeline &P, InvocationState &IS) { - auto MemCpyBack = [this](ResourcePair &R) -> llvm::Error { - if (!R.first->isReadWrite()) - return llvm::Error::success(); - - auto *RSIt = R.second.begin(); - auto *DataIt = R.first->BufferPtr->Data.begin(); - for (; RSIt != R.second.end() && DataIt != R.first->BufferPtr->Data.end(); - ++RSIt, ++DataIt) { - offloadtest::Buffer &Readback = *RSIt->Readback; - auto DataPtrOrErr = Readback.map(); - if (!DataPtrOrErr) - return DataPtrOrErr.takeError(); - const void *DataPtr = *DataPtrOrErr; - - if (R.first->isTexture()) { - const TextureCreateDesc &Desc = RSIt->Texture->getDesc(); - const uint32_t SrcStrideInBytes = - getTextureUploadRowStrideInBytes(Desc); - const uint32_t DstStrideInBytes = - Desc.Width * getFormatSizeInBytes(Desc.Fmt); - assert(DstStrideInBytes <= SrcStrideInBytes && - "Destination should not have padding and thus should be <= " - "than SrcStride where we do expect potential padding."); - uint8_t *Dst = (uint8_t *)DataIt->get(); - const uint8_t *Src = (const uint8_t *)DataPtr; - - for (uint32_t Y = 0; Y < Desc.Height; ++Y) { - memcpy(Dst, Src, DstStrideInBytes); - Dst += DstStrideInBytes; - Src += SrcStrideInBytes; - } - } else { - memcpy(DataIt->get(), DataPtr, R.first->size()); - } - - Readback.unmap(); - - if (R.first->HasCounter) { - offloadtest::Buffer &CounterReadback = *RSIt->CounterReadback; - auto CounterPtrOrErr = CounterReadback.map(); - if (!CounterPtrOrErr) - return CounterPtrOrErr.takeError(); - const uint32_t *CounterPtr = (const uint32_t *)*CounterPtrOrErr; - R.first->BufferPtr->Counters.push_back(*CounterPtr); - CounterReadback.unmap(); - } - } - - return llvm::Error::success(); - }; - - for (auto &Table : IS.DescTables) - for (auto &R : Table.Resources) - if (auto Err = MemCpyBack(R)) - return Err; - - for (auto &R : IS.RootResources) - if (auto Err = MemCpyBack(R)) - return Err; - - // If there is no render target, return early. - if (!IS.RTReadback) - return llvm::Error::success(); - - auto DataPtrOrErr = IS.RTReadback->map(); - if (!DataPtrOrErr) - return DataPtrOrErr.takeError(); - const void *Mapped = *DataPtrOrErr; - - const uint32_t SrcStrideInBytes = - getTextureUploadRowStrideInBytes(IS.RenderTarget->getDesc()); - - P.Bindings.RTargetBufferPtr->copyFromTexture(Mapped, SrcStrideInBytes); - IS.RTReadback->unmap(); - return llvm::Error::success(); - } - - llvm::Error createRenderTarget(Pipeline &P, InvocationState &IS) { - if (!P.Bindings.RTargetBufferPtr) - return llvm::createStringError( - std::errc::invalid_argument, - "No render target bound for graphics pipeline."); - const CPUBuffer &OutBuf = *P.Bindings.RTargetBufferPtr; - - auto TexOrErr = offloadtest::createRenderTargetFromCPUBuffer(*this, OutBuf); - if (!TexOrErr) - return TexOrErr.takeError(); - - IS.RenderTarget = std::move(*TexOrErr); - - // Create readback buffer sized for the pixel data with row pitch padded - // up to D3D12_TEXTURE_DATA_PITCH_ALIGNMENT, which is what D3D12 requires - // for the placed footprint used by CopyTextureRegion. The compaction - // back to a tight layout happens in readBack() via GetCopyableFootprints. - BufferCreateDesc BufDesc = {}; - BufDesc.Location = MemoryLocation::GpuToCpu; - BufDesc.Usage = BufferUsage::Storage; - auto BufOrErr = createBuffer("RTReadback", BufDesc, - getAlignedTextureBufferSize(OutBuf)); - if (!BufOrErr) - return BufOrErr.takeError(); - IS.RTReadback = std::move(*BufOrErr); +} - return llvm::Error::success(); - } +llvm::Expected> +DXDevice::createCommandBuffer() { + auto CBOrErr = DXCommandBuffer::create(Device); + if (!CBOrErr) + return CBOrErr.takeError(); + (*CBOrErr)->Dev = this; + return std::unique_ptr(std::move(*CBOrErr)); +} - llvm::Error createDepthStencil(Pipeline &P, InvocationState &IS) { - auto TexOrErr = offloadtest::createDefaultDepthStencilTarget( - *this, P.Bindings.RTargetBufferPtr->OutputProps.Width, - P.Bindings.RTargetBufferPtr->OutputProps.Height); - if (!TexOrErr) - return TexOrErr.takeError(); - IS.DepthStencil = std::move(*TexOrErr); - return llvm::Error::success(); - } +llvm::Expected> +DXDevice::createRenderPass(const offloadtest::RenderPassDesc &Desc) { + return std::make_unique(Desc); +} - llvm::Error createGraphicsCommands(Pipeline &P, InvocationState &IS) { - const DXPipelineState &DXPipeline = - llvm::cast(*IS.Pipeline.get()); - IS.CB->CmdList->SetGraphicsRootSignature(DXPipeline.RootSig.Get()); - if (IS.DescHeap) { - ID3D12DescriptorHeap *const Heaps[] = {IS.DescHeap.Get()}; - IS.CB->CmdList->SetDescriptorHeaps(1, Heaps); - IS.CB->CmdList->SetGraphicsRootDescriptorTable( - 0, IS.DescHeap->GetGPUDescriptorHandleForHeapStart()); - } +llvm::Expected +DXDevice::getBLASBuildSizes(llvm::ArrayRef Triangles) { + if (auto Err = validateBLASGeometry(Triangles)) + return Err; - RenderPassBeginDesc BeginDesc = {}; - BeginDesc.Pass = IS.RenderPass.get(); - BeginDesc.ColorAttachments.push_back(IS.RenderTarget.get()); - BeginDesc.DepthStencil = IS.DepthStencil.get(); - - auto EncOrErr = IS.CB->createRenderEncoder(BeginDesc); - if (!EncOrErr) - return EncOrErr.takeError(); - auto &Encoder = *EncOrErr.get(); - - Viewport VP; - VP.Width = - static_cast(P.Bindings.RTargetBufferPtr->OutputProps.Width); - VP.Height = - static_cast(P.Bindings.RTargetBufferPtr->OutputProps.Height); - Encoder.setViewport(VP); - - ScissorRect Scissor; - Scissor.Width = static_cast(VP.Width); - Scissor.Height = static_cast(VP.Height); - Encoder.setScissor(Scissor); - - if (P.isTraditionalRaster()) { - if (IS.VB) - Encoder.setVertexBuffer(0, IS.VB.get(), 0, - P.Bindings.getVertexStride()); - - if (auto Err = - Encoder.drawInstanced(*IS.Pipeline.get(), P.getVertexCount(), - /*InstanceCount=*/1)) - return Err; - } else { - if (auto Err = Encoder.dispatchMesh( - *IS.Pipeline.get(), P.DispatchParameters.DispatchGroupCount[0], - P.DispatchParameters.DispatchGroupCount[1], - P.DispatchParameters.DispatchGroupCount[2])) - return Err; + llvm::SmallVector GeomDescs; + GeomDescs.reserve(Triangles.size()); + for (const auto &T : Triangles) { + D3D12_RAYTRACING_GEOMETRY_DESC GD = {}; + GD.Type = D3D12_RAYTRACING_GEOMETRY_TYPE_TRIANGLES; + if (T.Opaque) + GD.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_OPAQUE; + + auto &Tri = GD.Triangles; + // GPU addresses are not needed for the prebuild size query; they will + // be populated at build time. + Tri.VertexBuffer.StrideInBytes = T.VertexStride; + Tri.VertexCount = T.VertexCount; + Tri.VertexFormat = getDXGIFormat(T.VertexFormat); + + if (T.IndexBuffer) { + Tri.IndexCount = T.IndexCount; + Tri.IndexFormat = getDXGIIndexFormat(T.IdxFormat); } - Encoder.endEncoding(); - - auto EncoderOrErr = IS.CB->createComputeEncoder(); - if (!EncoderOrErr) - return EncoderOrErr.takeError(); - auto ReadbackEncoder = std::move(*EncoderOrErr); - - if (auto Err = ReadbackEncoder->copyTextureToBuffer(*IS.RenderTarget, - *IS.RTReadback)) - return Err; - - for (auto &Table : IS.DescTables) - for (auto &R : Table.Resources) - if (auto Err = DXDevice::copyBackResource(*ReadbackEncoder, R)) - return Err; - - for (auto &R : IS.RootResources) - if (auto Err = DXDevice::copyBackResource(*ReadbackEncoder, R)) - return Err; - - return llvm::Error::success(); + GeomDescs.push_back(GD); } + return queryBLASPrebuildSize(GeomDescs); +} - llvm::Error executeProgram(Pipeline &P) override { - InvocationState State; - llvm::outs() << "Configuring execution on device: " << Description << "\n"; - if (auto Err = createDescriptorHeap(P, State)) - return Err; - llvm::outs() << "Descriptor heap created.\n"; - - auto CBOrErr = DXCommandBuffer::create(Device); - if (!CBOrErr) - return CBOrErr.takeError(); - State.CB = std::move(*CBOrErr); - State.CB->Dev = this; - llvm::outs() << "Command buffer created.\n"; - - if (auto Err = createBuffers(P, State)) - return Err; - llvm::outs() << "Buffers created.\n"; - - if (!P.AccelStructs.BLAS.empty() || !P.AccelStructs.TLAS.empty()) { - auto EncOrErr = State.CB->createComputeEncoder(); - if (!EncOrErr) - return EncOrErr.takeError(); - if (auto Err = offloadtest::buildPipelineAccelerationStructures( - *this, **EncOrErr, P, State.BLASes, State.TLASes, - State.ASInputBuffers)) - return Err; - (*EncOrErr)->endEncoding(); - } - - BindingsDesc BndDesc = {}; - for (auto &S : P.Sets) { - DescriptorSetLayoutDesc Layout; - for (auto &R : S.Resources) { - ResourceBindingDesc ResourceBinding = {}; - ResourceBinding.Kind = R.Kind; - ResourceBinding.DXBinding.Register = R.DXBinding.Register; - ResourceBinding.DXBinding.Space = R.DXBinding.Space; - ResourceBinding.VKBinding = R.VKBinding; - ResourceBinding.DescriptorCount = R.getArraySize(); - - Layout.ResourceBindings.push_back(ResourceBinding); - } - - BndDesc.DescriptorSetDescs.push_back(Layout); - } - - if (P.isRaster()) { - // Create render target and depth/stencil - if (auto Err = createRenderTarget(P, State)) - return Err; - llvm::outs() << "Render target created.\n"; - // TODO: Always created for graphics pipelines. Consider making this - // conditional on the pipeline definition. - if (auto Err = createDepthStencil(P, State)) - return Err; - llvm::outs() << "Depth stencil created.\n"; - } - - if (P.isCompute()) { - // This is an arbitrary distinction that we could alter in the future. - if (P.Shaders.size() != 1 || P.Shaders[0].Stage != Stages::Compute) - return llvm::createStringError( - std::errc::invalid_argument, - "Compute pipeline must have exactly one compute shader."); - - ShaderContainer CS = {}; - CS.EntryPoint = P.Shaders[0].Entry; - CS.Shader = P.Shaders[0].Shader.get(); - - auto PipelineStateOrErr = - createPipelineCs("Compute Pipeline State", BndDesc, CS); - if (!PipelineStateOrErr) - return PipelineStateOrErr.takeError(); - State.Pipeline = std::move(*PipelineStateOrErr); - llvm::outs() << "Compute Pipeline created.\n"; - if (auto Err = createComputeCommands(P, State)) - return Err; - llvm::outs() << "Compute command list created.\n"; - - } else if (P.isRaster()) { - - // Begin a render pass: bind RT/DSV and clear depth-stencil. Color - // load action is Load — the existing inline code didn't clear color. - ColorAttachmentFormatDesc ColorAttachment = {}; - ColorAttachment.Fmt = State.RenderTarget->getDesc().Fmt; - ColorAttachment.Load = LoadAction::Load; - ColorAttachment.Store = StoreAction::Store; - - DepthStencilAttachmentFormatDesc DSAttachment = {}; - DSAttachment.Fmt = State.DepthStencil->getDesc().Fmt; - DSAttachment.DepthLoad = LoadAction::Clear; - DSAttachment.DepthStore = StoreAction::Store; - DSAttachment.StencilLoad = LoadAction::DontCare; - DSAttachment.StencilStore = StoreAction::DontCare; - - RenderPassDesc PassDesc; - PassDesc.ColorAttachments.push_back(ColorAttachment); - PassDesc.DepthStencil = DSAttachment; - - auto RenderPassOrErr = createRenderPass(PassDesc); - if (!RenderPassOrErr) - return RenderPassOrErr.takeError(); - State.RenderPass = std::move(*RenderPassOrErr); - llvm::outs() << "Render pass created.\n"; - - if (P.isTraditionalRaster()) { - ShaderContainer VS = {}; - ShaderContainer PS = {}; - for (auto &Shader : P.Shaders) { - if (Shader.Stage == Stages::Vertex) { - VS.EntryPoint = Shader.Entry; - VS.Shader = Shader.Shader.get(); - } else if (Shader.Stage == Stages::Pixel) { - PS.EntryPoint = Shader.Entry; - PS.Shader = Shader.Shader.get(); - } - } - - TraditionalRasterPipelineCreateDesc PipelineDesc = {}; - PipelineDesc.Topology = P.Bindings.Topology; - PipelineDesc.PatchControlPoints = P.Bindings.PatchControlPoints; - PipelineDesc.DSFormat = Format::D32FloatS8Uint; - for (auto &Shader : P.Shaders) { - ShaderContainer SC = {}; - SC.EntryPoint = Shader.Entry; - SC.Shader = Shader.Shader.get(); - PipelineDesc.setShader(Shader.Stage, std::move(SC)); - } - - // Create the input layout based on the vertex attributes. - for (auto &Attr : P.Bindings.VertexAttributes) { - auto FormatOrErr = toFormat(Attr.Format, Attr.Channels); - if (!FormatOrErr) - return FormatOrErr.takeError(); - - InputLayoutDesc Layout = {}; - Layout.Name = Attr.Name; - Layout.Fmt = *FormatOrErr; - Layout.OffsetInBytes = Attr.Offset; - PipelineDesc.InputLayout.push_back(Layout); - } - - auto FormatOrErr = toFormat(P.Bindings.RTargetBufferPtr->Format, - P.Bindings.RTargetBufferPtr->Channels); - if (!FormatOrErr) - return FormatOrErr.takeError(); - PipelineDesc.RTFormats.push_back(*FormatOrErr); - - auto PipelineStateOrErr = createTraditionalRasterPipeline( - "Graphics Pipeline State", BndDesc, PipelineDesc); - if (!PipelineStateOrErr) - return PipelineStateOrErr.takeError(); - State.Pipeline = std::move(*PipelineStateOrErr); - llvm::outs() << "Traditional Raster Pipeline created.\n"; - - } else if (P.isMeshShaderRaster()) { - MeshShaderRasterPipelineCreateDesc PipelineDesc = {}; - PipelineDesc.Topology = P.Bindings.Topology; - PipelineDesc.DSFormat = Format::D32FloatS8Uint; - for (auto &Shader : P.Shaders) { - ShaderContainer SC = {}; - SC.EntryPoint = Shader.Entry; - SC.Shader = Shader.Shader.get(); - PipelineDesc.setShader(Shader.Stage, std::move(SC)); - } - - auto FormatOrErr = toFormat(P.Bindings.RTargetBufferPtr->Format, - P.Bindings.RTargetBufferPtr->Channels); - if (!FormatOrErr) - return FormatOrErr.takeError(); - PipelineDesc.RTFormats.push_back(*FormatOrErr); - - auto PipelineStateOrErr = createMeshShaderRasterPipeline( - "Mesh Shader Pipeline State", BndDesc, PipelineDesc); - - if (!PipelineStateOrErr) - return PipelineStateOrErr.takeError(); - State.Pipeline = std::move(*PipelineStateOrErr); - llvm::outs() << "Mesh Shader Pipeline created.\n"; - } +llvm::Expected +DXDevice::getBLASBuildSizes(llvm::ArrayRef AABBs) { + if (auto Err = validateBLASGeometry(AABBs)) + return Err; - if (auto Err = createGraphicsCommands(P, State)) - return Err; - llvm::outs() << "Graphics command list created complete.\n"; - } else if (P.isRayTracing()) { - if (P.Shaders.empty() || !P.SBT || !P.RTConfig) - return llvm::createStringError( - std::errc::invalid_argument, - "RayTracing pipeline requires Shaders, " - "ShaderBindingTable, and RayTracingPipelineConfig."); - - RayTracingPipelineCreateDesc RTDesc{}; - RTDesc.Library = P.Shaders.front().Shader.get(); - RTDesc.HitGroups = P.HitGroups; - RTDesc.Config = *P.RTConfig; - RTDesc.Shaders.reserve(P.Shaders.size()); - for (const auto &Sh : P.Shaders) - RTDesc.Shaders.push_back({Sh.Stage, Sh.Entry}); - - auto PSOOrErr = - createPipelineRT("RayTracing Pipeline State", BndDesc, RTDesc); - if (!PSOOrErr) - return PSOOrErr.takeError(); - State.Pipeline = std::move(*PSOOrErr); - llvm::outs() << "RayTracing Pipeline created.\n"; - - auto SBTOrErr = createShaderBindingTable(*State.Pipeline, *P.SBT); - if (!SBTOrErr) - return SBTOrErr.takeError(); - State.SBT = std::move(*SBTOrErr); - llvm::outs() << "Shader Binding Table created.\n"; - - if (auto Err = createComputeCommands(P, State)) - return Err; - llvm::outs() << "RayTracing command list created.\n"; - } else { - return llvm::createStringError("Pipeline was neither Compute nor Raster"); - } + llvm::SmallVector GeomDescs; + GeomDescs.reserve(AABBs.size()); + for (const auto &A : AABBs) { + D3D12_RAYTRACING_GEOMETRY_DESC GD = {}; + GD.Type = D3D12_RAYTRACING_GEOMETRY_TYPE_PROCEDURAL_PRIMITIVE_AABBS; + if (A.Opaque) + GD.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_OPAQUE; - auto SubmitResult = GraphicsQueue.submit(std::move(State.CB)); - if (!SubmitResult) - return SubmitResult.takeError(); - llvm::outs() << "Compute commands executed.\n"; - if (auto Err = SubmitResult->waitForCompletion()) - return Err; - if (auto Err = readBack(P, State)) - return Err; - llvm::outs() << "Read data back.\n"; + GD.AABBs.AABBs.StrideInBytes = A.AABBStride; + GD.AABBs.AABBCount = A.AABBCount; - return llvm::Error::success(); + GeomDescs.push_back(GD); } -}; - -TileShape DXTexture::querySparseTileShape(const Device &Dev) const { - const DXDevice &DeviceDX = llvm::cast(Dev); - - D3D12_TILE_SHAPE Shape = {}; - UINT NumTiles = 0; - UINT NumSubresourceTilings = 0; - DeviceDX.Device->GetResourceTiling(Resource.Get(), &NumTiles, nullptr, &Shape, - &NumSubresourceTilings, 0, nullptr); - return TileShape{Shape.WidthInTexels, Shape.HeightInTexels, - Shape.DepthInTexels}; + return queryBLASPrebuildSize(GeomDescs); } -llvm::Error DXComputeEncoder::batchBuildAS(llvm::ArrayRef Items) { - if (Items.empty()) - return llvm::Error::success(); - if (!CB.Dev || !CB.Dev->Device) - return llvm::createStringError( - std::errc::not_supported, - "Ray tracing not supported on this command buffer's device."); - DXDevice *Dev = CB.Dev; - - // BuildRaytracingAccelerationStructure() lives on ID3D12GraphicsCommandList4. - ComPtr CmdList4; - if (auto Err = HR::toError(CB.CmdList.As(&CmdList4), - "Failed to query ID3D12GraphicsCommandList4.")) - return Err; - - // Flush a pending barrier before reading, like dispatch(): a TLAS build must - // observe BLASes built in the previous batch. - CB.flushBarrier(); - - // Per the ComputeEncoder::batchBuildAS() contract, the caller guarantees no - // inter-item memory dependencies within a batch (BLAS and TLAS go in - // separate batches, so a TLAS never sees BLASes from the same call). Each - // item also gets its own scratch resource, so there's no aliasing between - // the builds — no intra-loop UAV barrier is needed. - for (const auto &Item : Items) { - D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_DESC Desc = {}; - llvm::SmallVector GeomDescs; - uint64_t ScratchSize = 0; - - if (const auto *BLAS = llvm::dyn_cast(Item)) { - auto *DXAS = llvm::cast(BLAS->AS); - Desc.DestAccelerationStructureData = DXAS->getGPUVirtualAddress(); - Desc.Inputs.Type = - D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL; - Desc.Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY; - - if (const auto *Tris = - std::get_if>( - &BLAS->Geometry)) { - GeomDescs.reserve(Tris->size()); - for (const auto &T : *Tris) { - D3D12_RAYTRACING_GEOMETRY_DESC GD = {}; - GD.Type = D3D12_RAYTRACING_GEOMETRY_TYPE_TRIANGLES; - if (T.Opaque) - GD.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_OPAQUE; - auto *VB = llvm::cast(T.VertexBuffer); - GD.Triangles.VertexBuffer.StartAddress = - VB->Buffer->GetGPUVirtualAddress() + T.VertexBufferOffset; - GD.Triangles.VertexBuffer.StrideInBytes = T.VertexStride; - GD.Triangles.VertexCount = T.VertexCount; - GD.Triangles.VertexFormat = getDXGIFormat(T.VertexFormat); - if (T.IndexBuffer) { - auto *IB = llvm::cast(T.IndexBuffer); - GD.Triangles.IndexBuffer = - IB->Buffer->GetGPUVirtualAddress() + T.IndexBufferOffset; - GD.Triangles.IndexCount = T.IndexCount; - GD.Triangles.IndexFormat = getDXGIIndexFormat(T.IdxFormat); - } - GeomDescs.push_back(GD); - } - } else { - const auto &AABBs = - std::get>(BLAS->Geometry); - GeomDescs.reserve(AABBs.size()); - for (const auto &A : AABBs) { - D3D12_RAYTRACING_GEOMETRY_DESC GD = {}; - GD.Type = D3D12_RAYTRACING_GEOMETRY_TYPE_PROCEDURAL_PRIMITIVE_AABBS; - if (A.Opaque) - GD.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_OPAQUE; - auto *AB = llvm::cast(A.AABBBuffer); - GD.AABBs.AABBs.StartAddress = - AB->Buffer->GetGPUVirtualAddress() + A.AABBBufferOffset; - GD.AABBs.AABBs.StrideInBytes = A.AABBStride; - GD.AABBs.AABBCount = A.AABBCount; - GeomDescs.push_back(GD); - } - } - Desc.Inputs.NumDescs = static_cast(GeomDescs.size()); - Desc.Inputs.pGeometryDescs = GeomDescs.data(); - ScratchSize = BLAS->AS->getSizes().ScratchDataSizeInBytes; - } else { - const auto *TLAS = llvm::cast(Item); - auto *DXAS = llvm::cast(TLAS->AS); - Desc.DestAccelerationStructureData = DXAS->getGPUVirtualAddress(); - Desc.Inputs.Type = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL; - Desc.Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY; - Desc.Inputs.NumDescs = static_cast(TLAS->Instances.size()); - - // D3D12_RAYTRACING_INSTANCE_DESC has the same byte layout as - // VkAccelerationStructureInstanceKHR. Serialize and upload via the - // shared abstract-API helper using an upload-heap (CpuToGpu) buffer. - llvm::SmallVector Native; - Native.reserve(TLAS->Instances.size()); - for (const auto &Inst : TLAS->Instances) { - D3D12_RAYTRACING_INSTANCE_DESC NI = {}; - static_assert(sizeof(NI.Transform) == sizeof(Inst.Transform), - "Transform layout mismatch"); - memcpy(&NI.Transform, Inst.Transform, sizeof(Inst.Transform)); - // D3D12_RAYTRACING_INSTANCE_DESC packs InstanceID into a 24-bit - // bitfield; truncate explicitly so the value matches the VK path - // (vkInstanceCustomIndex is likewise 24-bit) instead of relying on - // silent narrowing. - NI.InstanceID = Inst.InstanceID & 0xFFFFFFu; - NI.InstanceMask = Inst.InstanceMask; - NI.InstanceContributionToHitGroupIndex = - Inst.InstanceContributionToHitGroupIndex & 0xFFFFFFu; - // Bits in AccelerationStructureInstanceFlags match - // D3D12_RAYTRACING_INSTANCE_FLAGS by design. - NI.Flags = static_cast(Inst.Flags); - auto *BLASPtr = llvm::cast(Inst.BLAS); - NI.AccelerationStructure = BLASPtr->getGPUVirtualAddress(); - Native.push_back(NI); - } - const size_t Bytes = - Native.size() * sizeof(D3D12_RAYTRACING_INSTANCE_DESC); - - const BufferCreateDesc UploadDesc = BufferCreateDesc::uploadBuffer(); - auto InstBufOrErr = offloadtest::createBufferWithData( - *Dev, "TLAS-Instances", UploadDesc, Native.data(), Bytes, nullptr, - nullptr); - if (!InstBufOrErr) - return InstBufOrErr.takeError(); - auto *DXInstBuf = llvm::cast(InstBufOrErr->get()); - Desc.Inputs.InstanceDescs = DXInstBuf->Buffer->GetGPUVirtualAddress(); - - CB.KeepAliveOwned.push_back(std::move(*InstBufOrErr)); - ScratchSize = TLAS->AS->getSizes().ScratchDataSizeInBytes; - } +AccelerationStructureSizes DXDevice::queryBLASPrebuildSize( + llvm::ArrayRef GeomDescs) { + D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS Inputs = {}; + Inputs.Type = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL; + Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY; + Inputs.NumDescs = GeomDescs.size(); + Inputs.pGeometryDescs = GeomDescs.data(); - const BufferCreateDesc ScratchDesc = BufferCreateDesc::scratchBuffer(); - auto ScratchOrErr = - Dev->createBuffer("AS-Scratch", ScratchDesc, ScratchSize); - if (!ScratchOrErr) - return ScratchOrErr.takeError(); - auto *DXScratchBuf = llvm::cast(ScratchOrErr->get()); - Desc.ScratchAccelerationStructureData = - DXScratchBuf->Buffer->GetGPUVirtualAddress(); - CB.KeepAliveOwned.push_back(std::move(*ScratchOrErr)); - - insertDebugSignpost("BuildRaytracingAccelerationStructure"); - CmdList4->BuildRaytracingAccelerationStructure(&Desc, 0, nullptr); - } + D3D12_RAYTRACING_ACCELERATION_STRUCTURE_PREBUILD_INFO Info = {}; + Device->GetRaytracingAccelerationStructurePrebuildInfo(&Inputs, &Info); - // Signal that this batch's AS writes need a barrier before the next reader. - CB.addPendingUAVBarrier(); - return llvm::Error::success(); + return {Info.ResultDataMaxSizeInBytes, Info.ScratchDataSizeInBytes, + Info.UpdateScratchDataSizeInBytes}; } -llvm::Error DXComputeEncoder::dispatchRays(const PipelineState &PSO, - const ShaderBindingTable &SBT, - uint32_t Width, uint32_t Height, - uint32_t Depth) { - if (!llvm::isa(&PSO)) - return llvm::createStringError( - std::errc::invalid_argument, - "dispatchRays requires a RayTracing PipelineState."); - if (!llvm::isa(&SBT)) - return llvm::createStringError( - std::errc::invalid_argument, - "dispatchRays requires a DirectX ShaderBindingTable."); - const auto &DXRTPSO = llvm::cast(PSO); - const auto &DXSBT = llvm::cast(SBT); - - // SetPipelineState1 and DispatchRays live on ID3D12GraphicsCommandList4. - // The AS-build path (line ~3000 above) follows the same query pattern. - ComPtr CmdList4; - if (auto Err = - HR::toError(CB.CmdList.As(&CmdList4), - "ID3D12GraphicsCommandList4 query failed; raytracing " - "is unsupported on this command list.")) +llvm::Expected> +DXDevice::allocateAS(const AccelerationStructureSizes &Sizes, + const char *Kind) { + const uint64_t AlignedSize = + llvm::alignTo(Sizes.ResultDataMaxSizeInBytes, + D3D12_RAYTRACING_ACCELERATION_STRUCTURE_BYTE_ALIGNMENT); + const D3D12_HEAP_PROPERTIES HeapProps = + CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_DEFAULT); + const D3D12_RESOURCE_DESC BufferDesc = CD3DX12_RESOURCE_DESC::Buffer( + AlignedSize, D3D12_RESOURCE_FLAG_ALLOW_UNORDERED_ACCESS); + + ComPtr ASBuffer; + if (auto Err = HR::toError( + Device->CreateCommittedResource( + &HeapProps, D3D12_HEAP_FLAG_NONE, &BufferDesc, + D3D12_RESOURCE_STATE_RAYTRACING_ACCELERATION_STRUCTURE, nullptr, + IID_PPV_ARGS(&ASBuffer)), + "Failed to create " + llvm::Twine(Kind) + " resource.")) return Err; - addUAVBarrier(); - insertDebugSignpost( - llvm::formatv("DispatchRays [{0},{1},{2}]", Width, Height, Depth).str()); - - // Global root signature is shared with the compute bind point; bind it on - // the underlying command list before SetPipelineState1. - CB.CmdList->SetComputeRootSignature(DXRTPSO.RootSig.Get()); - CmdList4->SetPipelineState1(DXRTPSO.StateObject.Get()); - - D3D12_DISPATCH_RAYS_DESC RaysDesc{}; - RaysDesc.RayGenerationShaderRecord = DXSBT.RayGenRange; - RaysDesc.MissShaderTable = DXSBT.MissRange; - RaysDesc.HitGroupTable = DXSBT.HitGroupRange; - RaysDesc.CallableShaderTable = DXSBT.CallableRange; - RaysDesc.Width = Width; - RaysDesc.Height = Height; - RaysDesc.Depth = Depth; - CmdList4->DispatchRays(&RaysDesc); - return llvm::Error::success(); + return std::make_unique(ASBuffer, Sizes); } -} // namespace -llvm::Expected DXQueue::submit( - llvm::SmallVector> CBs) { - // Non-blocking: query how far the GPU has progressed and release - // command buffers from completed submissions. - { - const uint64_t Completed = SubmitFence->getFenceValue(); - llvm::erase_if(InFlightBatches, [Completed](const InFlightBatch &B) { - return B.FenceValue <= Completed; - }); - } +llvm::Expected +DXDevice::getTLASBuildSizes(uint32_t InstanceCount) { + D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_INPUTS Inputs = {}; + Inputs.Type = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL; + Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY; + Inputs.NumDescs = InstanceCount; - llvm::SmallVector CmdLists; - CmdLists.reserve(CBs.size()); + D3D12_RAYTRACING_ACCELERATION_STRUCTURE_PREBUILD_INFO Info = {}; + Device->GetRaytracingAccelerationStructurePrebuildInfo(&Inputs, &Info); - // GPU-side wait so that back-to-back submits don't overlap on the GPU. - // Skip on first submit since Wait(fence, 0) triggers a D3D12 validation - // warning. - if (FenceCounter > 0) - if (auto Err = - HR::toError(Queue->Wait(SubmitFence->Fence.Get(), FenceCounter), - "Failed to wait on previous submit.")) - return Err; + return AccelerationStructureSizes{Info.ResultDataMaxSizeInBytes, + Info.ScratchDataSizeInBytes, + Info.UpdateScratchDataSizeInBytes}; +} - for (auto &CB : CBs) { - auto &DCB = *llvm::cast(CB.get()); - if (auto Err = - HR::toError(DCB.CmdList->Close(), "Failed to close command list.")) - return Err; - CmdLists.push_back(DCB.CmdList.Get()); - } +llvm::Expected> +DXDevice::createBLAS(const AccelerationStructureSizes &Sizes) { + return allocateAS(Sizes, "BLAS"); +} - Queue->ExecuteCommandLists(CmdLists.size(), CmdLists.data()); +llvm::Expected> +DXDevice::createTLAS(const AccelerationStructureSizes &Sizes, + uint32_t /*InstanceCount*/) { + auto TLASOrErr = allocateAS(Sizes, "TLAS"); + if (!TLASOrErr) + return TLASOrErr.takeError(); + auto TLAS = std::move(*TLASOrErr); + DXAccelerationStructure &TLASDX = + llvm::cast(*TLAS.get()); + + auto SRVHandleOrErr = CSUAllocator.allocate(); + if (!SRVHandleOrErr) + return SRVHandleOrErr.takeError(); + TLASDX.SRVHandle = *SRVHandleOrErr; + + D3D12_SHADER_RESOURCE_VIEW_DESC SRVDesc = {}; + SRVDesc.Format = DXGI_FORMAT_UNKNOWN; + SRVDesc.ViewDimension = D3D12_SRV_DIMENSION_RAYTRACING_ACCELERATION_STRUCTURE; + SRVDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING; + SRVDesc.RaytracingAccelerationStructure.Location = + TLASDX.getGPUVirtualAddress(); + + // AS SRVs are created with a null resource; the AS lives in the + // buffer referenced by Location. + Device->CreateShaderResourceView(nullptr, &SRVDesc, TLASDX.SRVHandle); + + return TLAS; +} - const uint64_t CurrentCounter = ++FenceCounter; - if (auto Err = - HR::toError(Queue->Signal(SubmitFence->Fence.Get(), CurrentCounter), - "Failed to add signal.")) - return Err; +llvm::Expected> +DXDevice::createDescriptorPool() { + return DXDescriptorPool::create(Device); +} - // Keep submitted command buffers alive until the GPU is done with them. - InFlightBatches.push_back({CurrentCounter, std::move(CBs)}); +llvm::Expected> +DXDevice::createDescriptorSetsBuilder(DescriptorPool &Pool, + const PipelineState &Pipeline) { + DXDescriptorPool &PoolDX = llvm::cast(Pool); + const DXPipelineState &PipelineDX = llvm::cast(Pipeline); + + llvm::SmallVector Sets; + for (const auto &Counts : PipelineDX.Layout.Sets) { + DXDescriptorSet Set = {}; + if (Counts.DescriptorCount > 0) + PoolDX.allocateDescriptors(Counts.DescriptorCount, Set.CSUHandle, + Set.CSUHandleGPU); + if (Counts.SamplerCount > 0) + PoolDX.allocateSamplers(Counts.SamplerCount, Set.SamplerHandle, + Set.SamplerHandleGPU); + Sets.push_back(Set); + } + + const uint32_t CSUIncSize = Device->GetDescriptorHandleIncrementSize( + D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV); + const uint32_t SamplerIncSize = Device->GetDescriptorHandleIncrementSize( + D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER); + + return std::make_unique(Device, std::move(Sets), + CSUIncSize, SamplerIncSize); +} - return offloadtest::SubmitResult{SubmitFence.get(), CurrentCounter}; +llvm::Error DXDevice::executeProgram(Pipeline &P) { + return executeUnitTest(*this, P); } llvm::Error offloadtest::initializeDX12Devices( diff --git a/lib/API/DX/Encoder.cpp b/lib/API/DX/Encoder.cpp new file mode 100644 index 000000000..66bcae601 --- /dev/null +++ b/lib/API/DX/Encoder.cpp @@ -0,0 +1,579 @@ +//===- DX/Encoder.cpp - DirectX Encoder API -------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#include + +#include "API/DX/AccelerationStructure.h" +#include "API/DX/Buffer.h" +#include "API/DX/CommandBuffer.h" +#include "API/DX/Descriptors.h" +#include "API/DX/Device.h" +#include "API/DX/Encoder.h" +#include "API/DX/PipelineState.h" +#include "API/DX/SBT.h" +#include "API/DX/Texture.h" +#include "Support/WinError.h" + +#include "DXResources.h" + +#include "llvm/Support/FormatVariadic.h" + +using namespace offloadtest; + +void DXComputeEncoder::addUAVBarrier() { + CB.addPendingUAVBarrier(); + CB.flushBarrier(); +} + +DXComputeEncoder::~DXComputeEncoder() { endEncoding(); } + +void DXComputeEncoder::bindDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) { + const auto &DXPSO = llvm::cast(PSO); + const DXDescriptorSets &DSetsDX = llvm::cast(DSets); + assert(DSetsDX.Sets.size() == DXPSO.Layout.Sets.size() && + "Descriptor layout must match DescriptorSets"); + + CB.CmdList->SetComputeRootSignature(DXPSO.RootSig.Get()); + + for (size_t I = 0, J = 0; I < DSetsDX.Sets.size(); ++I) { + const auto &Set = DSetsDX.Sets[I]; + const auto &SetLayout = DXPSO.Layout.Sets[I]; + + if (SetLayout.DescriptorCount > 0) { + CB.CmdList->SetComputeRootDescriptorTable(J, Set.CSUHandleGPU); + assert(DXPSO.Layout.RSigLayout[J].ParameterType == + RootParameterType::DescriptorTable && + "Descriptor layout doesn't match root signature."); + J += 1; + } + + if (SetLayout.SamplerCount > 0) { + CB.CmdList->SetComputeRootDescriptorTable(J, Set.SamplerHandleGPU); + assert(DXPSO.Layout.RSigLayout[J].ParameterType == + RootParameterType::SamplerTable && + "Descriptor layout doesn't match root signature."); + J += 1; + } + } +} + +void DXComputeEncoder::bindComputeDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) { + return bindDescriptorSets(PSO, DSets); +} + +void DXComputeEncoder::bindRayTracingDescriptorSets( + const PipelineState &PSO, const DescriptorSets &DSets) { + return bindDescriptorSets(PSO, DSets); +} + +llvm::Error DXComputeEncoder::dispatch(const offloadtest::PipelineState &PSO, + uint32_t GroupCountX, + uint32_t GroupCountY, + uint32_t GroupCountZ) { + const auto &DXPSO = llvm::cast(PSO); + addUAVBarrier(); + insertDebugSignpost(llvm::formatv("Dispatch [{0},{1},{2}]", GroupCountX, + GroupCountY, GroupCountZ) + .str()); + CB.CmdList->SetPipelineState(DXPSO.PSO.Get()); + CB.CmdList->Dispatch(GroupCountX, GroupCountY, GroupCountZ); + return llvm::Error::success(); +} + +llvm::Error DXComputeEncoder::copyBufferToBuffer(offloadtest::Buffer &Src, + size_t SrcOffset, + offloadtest::Buffer &Dst, + size_t DstOffset, + size_t Size) { + auto &DXSrc = static_cast(Src); + auto &DXDst = static_cast(Dst); + + // NOTE: Edge case in case of all the following being the case + // - multiple calls of copyBufferToBuffer with the same Dst Buffer + // - The Dst Buffer having a PreferredState of + // D3D12_RESOURCE_STATE_COPY_DEST + // - Each Src Buffer having a PreferredState of + // D3D12_RESOURCE_STATE_COPY_SOURCE + // In that case no barrier would be emitted + // and a race condition would occur. There are ways to solve this with + // legacy barriers, but switching to enhanced barriers is a better solution + // to this problem. + + if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) + CB.addResourceTransition(DXSrc.Buffer.Get(), DXSrc.PreferredState, + D3D12_RESOURCE_STATE_COPY_SOURCE); + if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) + CB.addResourceTransition(DXDst.Buffer.Get(), DXDst.PreferredState, + D3D12_RESOURCE_STATE_COPY_DEST); + CB.flushBarrier(); + + insertDebugSignpost(llvm::formatv("CopyBuffer {0}B", Size).str()); + CB.CmdList->CopyBufferRegion(DXDst.Buffer.Get(), DstOffset, + DXSrc.Buffer.Get(), SrcOffset, Size); + + if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) + CB.addResourceTransition(DXSrc.Buffer.Get(), + D3D12_RESOURCE_STATE_COPY_SOURCE, + DXSrc.PreferredState); + if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) + CB.addResourceTransition(DXDst.Buffer.Get(), D3D12_RESOURCE_STATE_COPY_DEST, + DXDst.PreferredState); + + return llvm::Error::success(); +} + +llvm::Error DXComputeEncoder::copyBufferToTexture(Buffer &Src, Texture &Dst) { + auto &DXSrc = llvm::cast(Src); + auto &DXDst = llvm::cast(Dst); + + if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) + CB.addResourceTransition(DXSrc.Buffer.Get(), DXSrc.PreferredState, + D3D12_RESOURCE_STATE_COPY_SOURCE); + + if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) + CB.addResourceTransition(DXDst.Resource.Get(), DXDst.PreferredState, + D3D12_RESOURCE_STATE_COPY_DEST); + CB.flushBarrier(); + + const D3D12_RESOURCE_DESC TexDesc = DXDst.Resource->GetDesc(); + const uint32_t NumSubresources = TexDesc.MipLevels; + llvm::SmallVector Layouts( + NumSubresources); + ComPtr Device; + DXDst.Resource->GetDevice(IID_PPV_ARGS(&Device)); + Device->GetCopyableFootprints(&TexDesc, 0, NumSubresources, 0, Layouts.data(), + nullptr, nullptr, nullptr); + for (uint32_t Sub = 0; Sub < NumSubresources; ++Sub) { + const CD3DX12_TEXTURE_COPY_LOCATION DstLoc(DXDst.Resource.Get(), Sub); + const CD3DX12_TEXTURE_COPY_LOCATION SrcLoc(DXSrc.Buffer.Get(), + Layouts[Sub]); + CB.CmdList->CopyTextureRegion(&DstLoc, 0, 0, 0, &SrcLoc, nullptr); + } + + if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) + CB.addResourceTransition(DXSrc.Buffer.Get(), + D3D12_RESOURCE_STATE_COPY_SOURCE, + DXSrc.PreferredState); + + if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) + CB.addResourceTransition(DXDst.Resource.Get(), + D3D12_RESOURCE_STATE_COPY_DEST, + DXDst.PreferredState); + + return llvm::Error::success(); +} + +llvm::Error DXComputeEncoder::copyCounterToBuffer(offloadtest::Buffer &Src, + offloadtest::Buffer &Dst) { + auto &DXSrc = llvm::cast(Src); + auto &DXDst = llvm::cast(Dst); + + if (!DXSrc.Desc.HasCounter) + return llvm::createStringError( + "Counter resource passed does not hvae a counter."); + + if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) + CB.addResourceTransition(DXSrc.Buffer.Get(), DXSrc.PreferredState, + D3D12_RESOURCE_STATE_COPY_SOURCE); + if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) + CB.addResourceTransition(DXDst.Buffer.Get(), DXDst.PreferredState, + D3D12_RESOURCE_STATE_COPY_DEST); + CB.flushBarrier(); + + insertDebugSignpost("copyCounterToBuffer 4B"); + CB.CmdList->CopyBufferRegion(DXDst.Buffer.Get(), 0, DXSrc.Buffer.Get(), + DXSrc.CounterOffsetInBytes, sizeof(uint32_t)); + + if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) + CB.addResourceTransition(DXSrc.Buffer.Get(), + D3D12_RESOURCE_STATE_COPY_SOURCE, + DXSrc.PreferredState); + if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) + CB.addResourceTransition(DXDst.Buffer.Get(), D3D12_RESOURCE_STATE_COPY_DEST, + DXDst.PreferredState); + + return llvm::Error::success(); +} + +llvm::Error DXComputeEncoder::copyTextureToBuffer(Texture &Src, Buffer &Dst) { + auto &DXSrc = llvm::cast(Src); + auto &DXDst = llvm::cast(Dst); + + if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) + CB.addResourceTransition(DXSrc.Resource.Get(), DXSrc.PreferredState, + D3D12_RESOURCE_STATE_COPY_SOURCE); + + if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) + CB.addResourceTransition(DXDst.Buffer.Get(), DXDst.PreferredState, + D3D12_RESOURCE_STATE_COPY_DEST); + + CB.flushBarrier(); + + const uint32_t ElementSize = getFormatSizeInBytes(DXSrc.Desc.Fmt); + const D3D12_PLACED_SUBRESOURCE_FOOTPRINT Footprint{ + 0, CD3DX12_SUBRESOURCE_FOOTPRINT( + getDXGIFormat(DXSrc.Desc.Fmt), DXSrc.Desc.Width, DXSrc.Desc.Height, + 1, getAlignedTexturePitch(DXSrc.Desc.Width, ElementSize))}; + const CD3DX12_TEXTURE_COPY_LOCATION DstLoc(DXDst.Buffer.Get(), Footprint); + const CD3DX12_TEXTURE_COPY_LOCATION SrcLoc(DXSrc.Resource.Get(), 0); + CB.CmdList->CopyTextureRegion(&DstLoc, 0, 0, 0, &SrcLoc, nullptr); + + if (DXSrc.PreferredState != D3D12_RESOURCE_STATE_COPY_SOURCE) + CB.addResourceTransition(DXSrc.Resource.Get(), + D3D12_RESOURCE_STATE_COPY_SOURCE, + DXSrc.PreferredState); + + if (DXDst.PreferredState != D3D12_RESOURCE_STATE_COPY_DEST) + CB.addResourceTransition(DXDst.Buffer.Get(), D3D12_RESOURCE_STATE_COPY_DEST, + DXDst.PreferredState); + + return llvm::Error::success(); +} + +llvm::Error DXComputeEncoder::batchBuildAS(llvm::ArrayRef Items) { + if (Items.empty()) + return llvm::Error::success(); + if (!CB.Dev || !CB.Dev->Device) + return llvm::createStringError( + std::errc::not_supported, + "Ray tracing not supported on this command buffer's device."); + DXDevice *Dev = CB.Dev; + + // BuildRaytracingAccelerationStructure() lives on ID3D12GraphicsCommandList4. + ComPtr CmdList4; + if (auto Err = HR::toError(CB.CmdList.As(&CmdList4), + "Failed to query ID3D12GraphicsCommandList4.")) + return Err; + + // Flush a pending barrier before reading, like dispatch(): a TLAS build must + // observe BLASes built in the previous batch. + CB.flushBarrier(); + + // Per the ComputeEncoder::batchBuildAS() contract, the caller guarantees no + // inter-item memory dependencies within a batch (BLAS and TLAS go in + // separate batches, so a TLAS never sees BLASes from the same call). Each + // item also gets its own scratch resource, so there's no aliasing between + // the builds — no intra-loop UAV barrier is needed. + for (const auto &Item : Items) { + D3D12_BUILD_RAYTRACING_ACCELERATION_STRUCTURE_DESC Desc = {}; + llvm::SmallVector GeomDescs; + uint64_t ScratchSize = 0; + + if (const auto *BLAS = llvm::dyn_cast(Item)) { + auto *DXAS = llvm::cast(BLAS->AS); + Desc.DestAccelerationStructureData = DXAS->getGPUVirtualAddress(); + Desc.Inputs.Type = + D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL; + Desc.Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY; + + if (const auto *Tris = + std::get_if>( + &BLAS->Geometry)) { + GeomDescs.reserve(Tris->size()); + for (const auto &T : *Tris) { + D3D12_RAYTRACING_GEOMETRY_DESC GD = {}; + GD.Type = D3D12_RAYTRACING_GEOMETRY_TYPE_TRIANGLES; + if (T.Opaque) + GD.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_OPAQUE; + auto *VB = llvm::cast(T.VertexBuffer); + GD.Triangles.VertexBuffer.StartAddress = + VB->Buffer->GetGPUVirtualAddress() + T.VertexBufferOffset; + GD.Triangles.VertexBuffer.StrideInBytes = T.VertexStride; + GD.Triangles.VertexCount = T.VertexCount; + GD.Triangles.VertexFormat = getDXGIFormat(T.VertexFormat); + if (T.IndexBuffer) { + auto *IB = llvm::cast(T.IndexBuffer); + GD.Triangles.IndexBuffer = + IB->Buffer->GetGPUVirtualAddress() + T.IndexBufferOffset; + GD.Triangles.IndexCount = T.IndexCount; + GD.Triangles.IndexFormat = getDXGIIndexFormat(T.IdxFormat); + } + GeomDescs.push_back(GD); + } + } else { + const auto &AABBs = + std::get>(BLAS->Geometry); + GeomDescs.reserve(AABBs.size()); + for (const auto &A : AABBs) { + D3D12_RAYTRACING_GEOMETRY_DESC GD = {}; + GD.Type = D3D12_RAYTRACING_GEOMETRY_TYPE_PROCEDURAL_PRIMITIVE_AABBS; + if (A.Opaque) + GD.Flags = D3D12_RAYTRACING_GEOMETRY_FLAG_OPAQUE; + auto *AB = llvm::cast(A.AABBBuffer); + GD.AABBs.AABBs.StartAddress = + AB->Buffer->GetGPUVirtualAddress() + A.AABBBufferOffset; + GD.AABBs.AABBs.StrideInBytes = A.AABBStride; + GD.AABBs.AABBCount = A.AABBCount; + GeomDescs.push_back(GD); + } + } + Desc.Inputs.NumDescs = static_cast(GeomDescs.size()); + Desc.Inputs.pGeometryDescs = GeomDescs.data(); + ScratchSize = BLAS->AS->getSizes().ScratchDataSizeInBytes; + } else { + const auto *TLAS = llvm::cast(Item); + auto *DXAS = llvm::cast(TLAS->AS); + Desc.DestAccelerationStructureData = DXAS->getGPUVirtualAddress(); + Desc.Inputs.Type = D3D12_RAYTRACING_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL; + Desc.Inputs.DescsLayout = D3D12_ELEMENTS_LAYOUT_ARRAY; + Desc.Inputs.NumDescs = static_cast(TLAS->Instances.size()); + + // D3D12_RAYTRACING_INSTANCE_DESC has the same byte layout as + // VkAccelerationStructureInstanceKHR. Serialize and upload via the + // shared abstract-API helper using an upload-heap (CpuToGpu) buffer. + llvm::SmallVector Native; + Native.reserve(TLAS->Instances.size()); + for (const auto &Inst : TLAS->Instances) { + D3D12_RAYTRACING_INSTANCE_DESC NI = {}; + static_assert(sizeof(NI.Transform) == sizeof(Inst.Transform), + "Transform layout mismatch"); + memcpy(&NI.Transform, Inst.Transform, sizeof(Inst.Transform)); + // D3D12_RAYTRACING_INSTANCE_DESC packs InstanceID into a 24-bit + // bitfield; truncate explicitly so the value matches the VK path + // (vkInstanceCustomIndex is likewise 24-bit) instead of relying on + // silent narrowing. + NI.InstanceID = Inst.InstanceID & 0xFFFFFFu; + NI.InstanceMask = Inst.InstanceMask; + NI.InstanceContributionToHitGroupIndex = + Inst.InstanceContributionToHitGroupIndex & 0xFFFFFFu; + // Bits in AccelerationStructureInstanceFlags match + // D3D12_RAYTRACING_INSTANCE_FLAGS by design. + NI.Flags = static_cast(Inst.Flags); + auto *BLASPtr = llvm::cast(Inst.BLAS); + NI.AccelerationStructure = BLASPtr->getGPUVirtualAddress(); + Native.push_back(NI); + } + const size_t Bytes = + Native.size() * sizeof(D3D12_RAYTRACING_INSTANCE_DESC); + + const BufferCreateDesc UploadDesc = BufferCreateDesc::uploadBuffer(); + auto InstBufOrErr = offloadtest::createBufferWithData( + *Dev, "TLAS-Instances", UploadDesc, Native.data(), Bytes, nullptr, + nullptr); + if (!InstBufOrErr) + return InstBufOrErr.takeError(); + auto *DXInstBuf = llvm::cast(InstBufOrErr->get()); + Desc.Inputs.InstanceDescs = DXInstBuf->Buffer->GetGPUVirtualAddress(); + + CB.KeepAliveOwned.push_back(std::move(*InstBufOrErr)); + ScratchSize = TLAS->AS->getSizes().ScratchDataSizeInBytes; + } + + const BufferCreateDesc ScratchDesc = BufferCreateDesc::scratchBuffer(); + auto ScratchOrErr = + Dev->createBuffer("AS-Scratch", ScratchDesc, ScratchSize); + if (!ScratchOrErr) + return ScratchOrErr.takeError(); + auto *DXScratchBuf = llvm::cast(ScratchOrErr->get()); + Desc.ScratchAccelerationStructureData = + DXScratchBuf->Buffer->GetGPUVirtualAddress(); + CB.KeepAliveOwned.push_back(std::move(*ScratchOrErr)); + + insertDebugSignpost("BuildRaytracingAccelerationStructure"); + CmdList4->BuildRaytracingAccelerationStructure(&Desc, 0, nullptr); + } + + // Signal that this batch's AS writes need a barrier before the next reader. + CB.addPendingUAVBarrier(); + return llvm::Error::success(); +} + +llvm::Error DXComputeEncoder::dispatchRays(const PipelineState &PSO, + const ShaderBindingTable &SBT, + uint32_t Width, uint32_t Height, + uint32_t Depth) { + if (!llvm::isa(&PSO)) + return llvm::createStringError( + std::errc::invalid_argument, + "dispatchRays requires a RayTracing PipelineState."); + if (!llvm::isa(&SBT)) + return llvm::createStringError( + std::errc::invalid_argument, + "dispatchRays requires a DirectX ShaderBindingTable."); + const auto &DXRTPSO = llvm::cast(PSO); + const auto &DXSBT = llvm::cast(SBT); + + // SetPipelineState1 and DispatchRays live on ID3D12GraphicsCommandList4. + // The AS-build path (line ~3000 above) follows the same query pattern. + ComPtr CmdList4; + if (auto Err = + HR::toError(CB.CmdList.As(&CmdList4), + "ID3D12GraphicsCommandList4 query failed; raytracing " + "is unsupported on this command list.")) + return Err; + + addUAVBarrier(); + insertDebugSignpost( + llvm::formatv("DispatchRays [{0},{1},{2}]", Width, Height, Depth).str()); + + // Global root signature is shared with the compute bind point; bind it on + // the underlying command list before SetPipelineState1. + CB.CmdList->SetComputeRootSignature(DXRTPSO.RootSig.Get()); + CmdList4->SetPipelineState1(DXRTPSO.StateObject.Get()); + + D3D12_DISPATCH_RAYS_DESC RaysDesc{}; + RaysDesc.RayGenerationShaderRecord = DXSBT.RayGenRange; + RaysDesc.MissShaderTable = DXSBT.MissRange; + RaysDesc.HitGroupTable = DXSBT.HitGroupRange; + RaysDesc.CallableShaderTable = DXSBT.CallableRange; + RaysDesc.Width = Width; + RaysDesc.Height = Height; + RaysDesc.Depth = Depth; + CmdList4->DispatchRays(&RaysDesc); + return llvm::Error::success(); +} + +void DXComputeEncoder::endEncodingImpl() { popDebugGroup(); } + +llvm::Error +DXRenderEncoder::bindCommonDrawState(const offloadtest::PipelineState &PSO) { + if (!ViewportSet) + return llvm::createStringError(std::errc::invalid_argument, + "Viewport must be set before drawing."); + if (!ScissorSet) + return llvm::createStringError(std::errc::invalid_argument, + "Scissor must be set before drawing."); + + const auto &DXPSO = llvm::cast(PSO); + CB.CmdList->SetGraphicsRootSignature(DXPSO.RootSig.Get()); + CB.CmdList->SetPipelineState(DXPSO.PSO.Get()); + // Mesh-shader pipelines bypass the input assembler and carry no IA + // topology; only bind one when the pipeline actually has one. + if (DXPSO.Topology) + CB.CmdList->IASetPrimitiveTopology(*DXPSO.Topology); + return llvm::Error::success(); +} + +DXRenderEncoder::DXRenderEncoder(DXCommandBuffer &CB, + const offloadtest::RenderPassBeginDesc &Desc) + : RenderEncoder(GPUAPI::DirectX), CB(CB), Desc(Desc) {} + +DXRenderEncoder::~DXRenderEncoder() { endEncoding(); } + +bool DXRenderEncoder::classof(const CommandEncoder *E) { + return E->getAPI() == GPUAPI::DirectX; +} + +// See DXComputeEncoder for why these are no-ops. +void DXRenderEncoder::pushDebugGroup(llvm::StringRef Label) {} +void DXRenderEncoder::popDebugGroup() {} +void DXRenderEncoder::insertDebugSignpost(llvm::StringRef Label) {} + +void DXRenderEncoder::setViewport(const offloadtest::Viewport &VP) { + D3D12_VIEWPORT DXVP = {}; + DXVP.TopLeftX = VP.X; + DXVP.TopLeftY = VP.Y; + DXVP.Width = VP.Width; + DXVP.Height = VP.Height; + DXVP.MinDepth = VP.MinDepth; + DXVP.MaxDepth = VP.MaxDepth; + CB.CmdList->RSSetViewports(1, &DXVP); + ViewportSet = true; +} + +void DXRenderEncoder::setScissor(const offloadtest::ScissorRect &Rect) { + const D3D12_RECT DXRect = {Rect.X, Rect.Y, + static_cast(Rect.X + Rect.Width), + static_cast(Rect.Y + Rect.Height)}; + CB.CmdList->RSSetScissorRects(1, &DXRect); + ScissorSet = true; +} + +void DXRenderEncoder::setVertexBuffer(uint32_t Slot, offloadtest::Buffer *VB, + size_t Offset, uint32_t Stride) { + assert(Slot < D3D12_IA_VERTEX_INPUT_RESOURCE_SLOT_COUNT && + "Vertex buffer slot exceeds D3D12 IA input resource slot count"); + assert(Slot == 0 && "Pipeline input layout only describes slot 0"); + if (VB) { + auto &DXVB = llvm::cast(*VB); + D3D12_VERTEX_BUFFER_VIEW VBView = {}; + VBView.BufferLocation = DXVB.Buffer->GetGPUVirtualAddress() + Offset; + VBView.SizeInBytes = static_cast(DXVB.getSizeInBytes() - Offset); + VBView.StrideInBytes = Stride; + CB.CmdList->IASetVertexBuffers(Slot, 1, &VBView); + } else { + CB.CmdList->IASetVertexBuffers(Slot, 1, nullptr); + } +} + +void DXRenderEncoder::bindDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) { + const auto &DXPSO = llvm::cast(PSO); + const DXDescriptorSets &DSetsDX = llvm::cast(DSets); + assert(DSetsDX.Sets.size() == DXPSO.Layout.Sets.size() && + "Descriptor layout must match DescriptorSets"); + + CB.CmdList->SetGraphicsRootSignature(DXPSO.RootSig.Get()); + + for (size_t I = 0, J = 0; I < DSetsDX.Sets.size(); ++I) { + const auto &Set = DSetsDX.Sets[I]; + const auto &SetLayout = DXPSO.Layout.Sets[I]; + + if (SetLayout.DescriptorCount > 0) { + CB.CmdList->SetGraphicsRootDescriptorTable(J, Set.CSUHandleGPU); + assert(DXPSO.Layout.RSigLayout[J].ParameterType == + RootParameterType::DescriptorTable && + "Descriptor layout doesn't match root signature."); + J += 1; + } + + if (SetLayout.SamplerCount > 0) { + CB.CmdList->SetGraphicsRootDescriptorTable(J, Set.SamplerHandleGPU); + assert(DXPSO.Layout.RSigLayout[J].ParameterType == + RootParameterType::SamplerTable && + "Descriptor layout doesn't match root signature."); + J += 1; + } + } +} + +llvm::Error +DXRenderEncoder::drawInstanced(const offloadtest::PipelineState &PSO, + uint32_t VertexCount, uint32_t InstanceCount, + uint32_t FirstVertex, uint32_t FirstInstance) { + if (auto Err = bindCommonDrawState(PSO)) + return Err; + CB.CmdList->DrawInstanced(VertexCount, InstanceCount, FirstVertex, + FirstInstance); + return llvm::Error::success(); +} + +llvm::Error DXRenderEncoder::dispatchMesh(const offloadtest::PipelineState &PSO, + uint32_t GroupCountX, + uint32_t GroupCountY, + uint32_t GroupCountZ) { + if (auto Err = bindCommonDrawState(PSO)) + return Err; + CB.CmdList->DispatchMesh(GroupCountX, GroupCountY, GroupCountZ); + return llvm::Error::success(); +} + +void DXRenderEncoder::endEncodingImpl() { + // State transitions + for (offloadtest::Texture *Tex : Desc.ColorAttachments) { + auto &DXTex = llvm::cast(*Tex); + if (DXTex.PreferredState != D3D12_RESOURCE_STATE_RENDER_TARGET) + CB.addResourceTransition(DXTex.Resource.Get(), + D3D12_RESOURCE_STATE_RENDER_TARGET, + DXTex.PreferredState); + } + if (Desc.DepthStencil) { + auto &DXTex = llvm::cast(*Desc.DepthStencil); + if (DXTex.PreferredState != D3D12_RESOURCE_STATE_DEPTH_WRITE) + CB.addResourceTransition(DXTex.Resource.Get(), + D3D12_RESOURCE_STATE_DEPTH_WRITE, + DXTex.PreferredState); + } + + popDebugGroup(); +} diff --git a/lib/API/DX/MemoryHeap.cpp b/lib/API/DX/MemoryHeap.cpp new file mode 100644 index 000000000..760544585 --- /dev/null +++ b/lib/API/DX/MemoryHeap.cpp @@ -0,0 +1,47 @@ +//===- DX/MemoryHeap.cpp - DirectX MemoryHeap API -------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#include + +#include "API/DX/MemoryHeap.h" +#include "Support/WinError.h" + +using namespace offloadtest; + +DXMemoryHeap::DXMemoryHeap(ComPtr Heap, llvm::StringRef Name) + : offloadtest::MemoryHeap(GPUAPI::DirectX), Name(Name), Heap(Heap) {} + +bool DXMemoryHeap::classof(const offloadtest::MemoryHeap *H) { + return H->getAPI() == GPUAPI::DirectX; +} + +llvm::Expected> +DXMemoryHeap::create(ID3D12DeviceX *Device, llvm::StringRef Name, + size_t SizeInBytes) { + D3D12_HEAP_DESC HeapDesc = {}; + HeapDesc.Properties = CD3DX12_HEAP_PROPERTIES(D3D12_HEAP_TYPE_DEFAULT); + HeapDesc.Alignment = D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT; + // Heap sizes must be a multiple of the placement alignment (64KB), which + // is also the tiled-resource tile size. + HeapDesc.SizeInBytes = + llvm::alignTo(SizeInBytes, D3D12_DEFAULT_RESOURCE_PLACEMENT_ALIGNMENT); + HeapDesc.Flags = D3D12_HEAP_FLAG_ALLOW_ALL_BUFFERS_AND_TEXTURES; + + ComPtr Heap; + if (auto Err = HR::toError(Device->CreateHeap(&HeapDesc, IID_PPV_ARGS(&Heap)), + "Failed to create memory heap.")) + return Err; + + const std::wstring WStr(Name.begin(), Name.end()); + Heap->SetName(WStr.c_str()); + + return std::make_unique(Heap, Name); +} diff --git a/lib/API/DX/PipelineState.cpp b/lib/API/DX/PipelineState.cpp new file mode 100644 index 000000000..334d40b63 --- /dev/null +++ b/lib/API/DX/PipelineState.cpp @@ -0,0 +1,43 @@ +//===- DX/PipelineState.cpp - DirectX PipelineState API -------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#include "API/DX/PipelineState.h" + +using namespace offloadtest; + +DXPipelineState::DXPipelineState(llvm::StringRef Name, + ComPtr RootSig, + DescriptorSetsLayout Layout, + ComPtr PSO, + std::optional Topology, + bool IsRT) + : offloadtest::PipelineState(GPUAPI::DirectX), Name(Name), RootSig(RootSig), + Layout(std::move(Layout)), PSO(PSO), Topology(Topology), + IsRayTracing(IsRT) {} + +bool DXPipelineState::classof(const offloadtest::PipelineState *B) { + return B->getAPI() == GPUAPI::DirectX; +} + +DXRayTracingPipelineState::DXRayTracingPipelineState( + llvm::StringRef Name, ComPtr RootSig, + DescriptorSetsLayout Layout, ComPtr SO, + ComPtr Props) + : DXPipelineState(Name, RootSig, std::move(Layout), /*PSO=*/nullptr, + std::nullopt, + /*IsRT=*/true), + StateObject(SO), Properties(Props) {} + +bool DXRayTracingPipelineState::classof(const offloadtest::PipelineState *B) { + if (B->getAPI() != GPUAPI::DirectX) + return false; + return static_cast(B)->IsRayTracing; +} diff --git a/lib/API/DX/Queue.cpp b/lib/API/DX/Queue.cpp new file mode 100644 index 000000000..9c85fb734 --- /dev/null +++ b/lib/API/DX/Queue.cpp @@ -0,0 +1,201 @@ +//===- DX/Queue.cpp - DirectX Queue API -----------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#include "API/DX/Queue.h" +#include "API/DX/Buffer.h" +#include "API/DX/CommandBuffer.h" +#include "API/DX/MemoryHeap.h" +#include "API/DX/Texture.h" + +#include "Support/WinError.h" + +using namespace offloadtest; + +llvm::Expected> DXFence::create(ID3D12DeviceX *Device, + llvm::StringRef Name) { + ComPtr Fence; + if (auto Err = HR::toError( + Device->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&Fence)), + "Failed to create Fence.")) + return Err; + +#ifdef _WIN32 + HANDLE Event = CreateEventA(nullptr, false, false, nullptr); + if (!Event) +#else // WSL + int Event = eventfd(0, EFD_CLOEXEC | EFD_NONBLOCK); + if (Event == -1) +#endif + return llvm::createStringError(std::errc::device_or_resource_busy, + "Failed to create event."); + + return std::make_unique(Fence, Event, Name); +} + +DXFence::~DXFence() { +#ifdef _WIN32 + CloseHandle(Event); +#else // WSL + close(Event); +#endif +} + +uint64_t DXFence::getFenceValue() { return Fence->GetCompletedValue(); } + +llvm::Error DXFence::waitForCompletion(uint64_t SignalValue) { + if (Fence->GetCompletedValue() >= SignalValue) + return llvm::Error::success(); + +#ifdef _WIN32 + if (auto Err = HR::toError(Fence->SetEventOnCompletion(SignalValue, Event), + "Failed to register end event.")) + return Err; + WaitForSingleObject(Event, INFINITE); +#else // WSL + if (auto Err = HR::toError(Fence->SetEventOnCompletion( + SignalValue, reinterpret_cast(Event)), + "Failed to register end event.")) + return Err; + pollfd PollEvent; + PollEvent.fd = Event; + PollEvent.events = POLLIN; + PollEvent.revents = 0; + if (poll(&PollEvent, 1, -1) == -1) + return llvm::createStringError( + std::error_code(errno, std::system_category()), strerror(errno)); +#endif + return llvm::Error::success(); +} + +DXQueue::~DXQueue() {} + +llvm::Expected +DXQueue::createGraphicsQueue(ComPtr Device) { + const D3D12_COMMAND_QUEUE_DESC Desc = {D3D12_COMMAND_LIST_TYPE_DIRECT, 0, + D3D12_COMMAND_QUEUE_FLAG_NONE, 0}; + ComPtr CmdQueue; + if (auto Err = HR::toError( + Device->CreateCommandQueue(&Desc, IID_PPV_ARGS(&CmdQueue)), + "Failed to create command queue.")) + return Err; + auto FenceOrErr = DXFence::create(Device.Get(), "QueueSubmitFence"); + if (!FenceOrErr) + return FenceOrErr.takeError(); + return DXQueue(CmdQueue, std::move(*FenceOrErr)); +} + +llvm::Expected DXQueue::submit( + llvm::SmallVector> CBs) { + // Non-blocking: query how far the GPU has progressed and release + // command buffers from completed submissions. + { + const uint64_t Completed = SubmitFence->getFenceValue(); + llvm::erase_if(InFlightBatches, [Completed](const InFlightBatch &B) { + return B.FenceValue <= Completed; + }); + } + + llvm::SmallVector CmdLists; + CmdLists.reserve(CBs.size()); + + // GPU-side wait so that back-to-back submits don't overlap on the GPU. + // Skip on first submit since Wait(fence, 0) triggers a D3D12 validation + // warning. + if (FenceCounter > 0) + if (auto Err = + HR::toError(Queue->Wait(SubmitFence->Fence.Get(), FenceCounter), + "Failed to wait on previous submit.")) + return Err; + + for (auto &CB : CBs) { + auto &DCB = *llvm::cast(CB.get()); + if (auto Err = + HR::toError(DCB.CmdList->Close(), "Failed to close command list.")) + return Err; + CmdLists.push_back(DCB.CmdList.Get()); + } + + Queue->ExecuteCommandLists(CmdLists.size(), CmdLists.data()); + + const uint64_t CurrentCounter = ++FenceCounter; + if (auto Err = + HR::toError(Queue->Signal(SubmitFence->Fence.Get(), CurrentCounter), + "Failed to add signal.")) + return Err; + + // Keep submitted command buffers alive until the GPU is done with them. + InFlightBatches.push_back({CurrentCounter, std::move(CBs)}); + + return offloadtest::SubmitResult{SubmitFence.get(), CurrentCounter}; +} + +llvm::Expected +DXQueue::updateTileMappings(offloadtest::Buffer &Resource, + llvm::ArrayRef Mappings) { + return updateTileMappingsImpl(llvm::cast(Resource).Buffer.Get(), + Mappings); +} + +llvm::Expected +DXQueue::updateTileMappings(offloadtest::Texture &Resource, + llvm::ArrayRef Mappings) { + return updateTileMappingsImpl(llvm::cast(Resource).Resource.Get(), + Mappings); +} + +// Shared tile-mapping path for buffers and textures: they differ only in how +// a TileRegion maps onto the underlying ID3D12Resource, so both variants +// resolve to a native resource and funnel through here. +llvm::Expected +DXQueue::updateTileMappingsImpl(ID3D12Resource *Resource, + llvm::ArrayRef Mappings) { + for (const TileMapping &M : Mappings) { + const D3D12_TILED_RESOURCE_COORDINATE StartCoord = { + M.Region.TileOffsetX, M.Region.TileOffsetY, M.Region.TileOffsetZ, + M.Region.Subresource}; + + D3D12_TILE_REGION_SIZE RegionSize = {}; + RegionSize.Width = M.Region.NumTilesX; + RegionSize.Height = static_cast(M.Region.NumTilesY); + RegionSize.Depth = static_cast(M.Region.NumTilesZ); + // A box is only needed for multi-dimensional (texture) regions; a buffer + // is always a single linear run of tiles along X. + RegionSize.UseBox = M.Region.NumTilesY > 1 || M.Region.NumTilesZ > 1; + RegionSize.NumTiles = + M.Region.NumTilesX * M.Region.NumTilesY * M.Region.NumTilesZ; + + // A null heap unbinds the region (reads return zero); otherwise bind it + // to the heap starting at BackingTileOffset. + ID3D12Heap *Heap = nullptr; + D3D12_TILE_RANGE_FLAGS RangeFlag = D3D12_TILE_RANGE_FLAG_NULL; + UINT HeapRangeStartOffset = 0; + if (M.Backing) { + Heap = llvm::cast(M.Backing)->Heap.Get(); + RangeFlag = D3D12_TILE_RANGE_FLAG_NONE; + HeapRangeStartOffset = static_cast(M.BackingTileOffset); + } + const UINT RangeTileCount = RegionSize.NumTiles; + + Queue->UpdateTileMappings(Resource, 1, &StartCoord, &RegionSize, Heap, 1, + &RangeFlag, &HeapRangeStartOffset, + &RangeTileCount, D3D12_TILE_MAPPING_FLAG_NONE); + } + + // UpdateTileMappings runs on the queue timeline (not a command list), so + // signal the submit fence afterwards exactly like submit() does. + const uint64_t CurrentCounter = ++FenceCounter; + if (auto Err = + HR::toError(Queue->Signal(SubmitFence->Fence.Get(), CurrentCounter), + "Failed to signal after tile mapping update.")) + return Err; + + return offloadtest::SubmitResult{SubmitFence.get(), CurrentCounter}; +} diff --git a/lib/API/DX/Sampler.cpp b/lib/API/DX/Sampler.cpp new file mode 100644 index 000000000..34133af19 --- /dev/null +++ b/lib/API/DX/Sampler.cpp @@ -0,0 +1,25 @@ +//===- DX/Sampler.cpp - DirectX Sampler API -------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#include "API/DX/Sampler.h" + +using namespace offloadtest; + +DXSampler::DXSampler(llvm::StringRef Name, SamplerCreateDesc Desc, + D3D12_CPU_DESCRIPTOR_HANDLE Handle) + : offloadtest::Sampler(GPUAPI::DirectX), Handle(Handle), Name(Name), + Desc(Desc) {} + +const SamplerCreateDesc &DXSampler::getDesc() const { return Desc; } + +bool DXSampler::classof(const offloadtest::Sampler *S) { + return S->getAPI() == GPUAPI::DirectX; +} diff --git a/lib/API/DX/Texture.cpp b/lib/API/DX/Texture.cpp new file mode 100644 index 000000000..44e5db06c --- /dev/null +++ b/lib/API/DX/Texture.cpp @@ -0,0 +1,43 @@ +//===- DX/Texture.cpp - DirectX Texture API -------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// +//===----------------------------------------------------------------------===// + +#include "API/DX/Texture.h" +#include "API/DX/Device.h" +#include "Support/WinError.h" + +using namespace offloadtest; + +DXTexture::DXTexture(ComPtr Resource, llvm::StringRef Name, + TextureCreateDesc Desc, + D3D12_RESOURCE_STATES PreferredState, + D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle, + D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle) + : offloadtest::Texture(GPUAPI::DirectX), Resource(Resource), + PreferredState(PreferredState), SRVHandle(SRVHandle), + UAVHandle(UAVHandle), Name(Name), Desc(Desc) {} + +TileShape DXTexture::querySparseTileShape(const Device &Dev) const { + const DXDevice &DeviceDX = llvm::cast(Dev); + + D3D12_TILE_SHAPE Shape = {}; + UINT NumTiles = 0; + UINT NumSubresourceTilings = 0; + DeviceDX.Device->GetResourceTiling(Resource.Get(), &NumTiles, nullptr, &Shape, + &NumSubresourceTilings, 0, nullptr); + return TileShape{Shape.WidthInTexels, Shape.HeightInTexels, + Shape.DepthInTexels}; +} + +const TextureCreateDesc &DXTexture::getDesc() const { return Desc; } + +bool DXTexture::classof(const offloadtest::Texture *T) { + return T->getAPI() == GPUAPI::DirectX; +} diff --git a/lib/API/Device.cpp b/lib/API/Device.cpp index 33260a6cf..70f5f71cd 100644 --- a/lib/API/Device.cpp +++ b/lib/API/Device.cpp @@ -36,6 +36,8 @@ Queue::~Queue() {} Texture::~Texture() {} +Sampler::~Sampler() {} + MemoryHeap::~MemoryHeap() {} RenderPass::~RenderPass() {} @@ -528,32 +530,34 @@ offloadtest::createTextureWithData( Device &Dev, std::string Name, const TextureCreateDesc &Desc, const void *Data, size_t SizeInBytes, ComputeEncoder *Encoder, std::unique_ptr *OutUploadBuffer) { - - const uint64_t PackedRowStrideInBytes = - Desc.Width * getFormatSizeInBytes(Desc.Fmt); - if (SizeInBytes < PackedRowStrideInBytes * Desc.Height) + if (Encoder == nullptr) return llvm::createStringError( - "Data upload is not enough for texture size."); + "An encoder is required to upload texture data."); + if (OutUploadBuffer == nullptr) + return llvm::createStringError( + "An upload buffer is required to create a texture with data."); auto TextureOrErr = Dev.createTexture(Name, Desc); if (!TextureOrErr) return TextureOrErr.takeError(); auto Texture = std::move(*TextureOrErr); - if (OutUploadBuffer == nullptr) - return llvm::createStringError("An upload buffer is required to create a " - "GpuOnly texture with data."); + const TextureUploadLayout Layout = Dev.getTextureUploadLayout(Desc); - const uint64_t TexRowStrideInBytes = - Dev.getTextureUploadRowStrideInBytes(Desc); - const uint64_t UploadBufferSizeInBytes = - (Desc.Height - 1) * TexRowStrideInBytes + PackedRowStrideInBytes; + // The source data is tightly packed across mips, so its required size is the + // sum of each subresource's tight row size times its row count, independent + // of any backend row/offset padding in the upload buffer. + uint64_t PackedSizeInBytes = 0; + for (const SubresourceFootprint &Sub : Layout.Subresources) + PackedSizeInBytes += uint64_t(Sub.RowSizeInBytes) * Sub.NumRows; + if (SizeInBytes < PackedSizeInBytes) + return llvm::createStringError( + "Data upload is not enough for texture size."); - // Create Upload buffer const BufferCreateDesc UploadDesc = BufferCreateDesc::uploadBuffer(); const std::string UploadBufferName = Name + " (Upload Buffer)"; auto UploadBufferOrErr = - Dev.createBuffer(UploadBufferName, UploadDesc, UploadBufferSizeInBytes); + Dev.createBuffer(UploadBufferName, UploadDesc, Layout.TotalSizeInBytes); if (!UploadBufferOrErr) return UploadBufferOrErr.takeError(); *OutUploadBuffer = std::move(*UploadBufferOrErr); @@ -561,18 +565,19 @@ offloadtest::createTextureWithData( auto MappedPtrOrErr = (*OutUploadBuffer)->map(); if (!MappedPtrOrErr) return MappedPtrOrErr.takeError(); - - uint8_t *DstPtr = (uint8_t *)*MappedPtrOrErr; - const uint8_t *SrcPtr = (const uint8_t *)Data; - - for (uint32_t Y = 0; Y < Desc.Height; ++Y) { - memcpy(DstPtr, SrcPtr, PackedRowStrideInBytes); - DstPtr += TexRowStrideInBytes; - SrcPtr += PackedRowStrideInBytes; + auto *const DstBase = static_cast(*MappedPtrOrErr); + const auto *SrcPtr = static_cast(Data); + + for (const SubresourceFootprint &Sub : Layout.Subresources) { + uint8_t *DstPtr = DstBase + Sub.Offset; + for (uint32_t Row = 0; Row < Sub.NumRows; ++Row) { + memcpy(DstPtr, SrcPtr, Sub.RowSizeInBytes); + DstPtr += Sub.RowPitchInBytes; + SrcPtr += Sub.RowSizeInBytes; + } } (*OutUploadBuffer)->unmap(); - // Copy Buffer to Texture if (auto Err = Encoder->copyBufferToTexture(**OutUploadBuffer, *Texture)) return Err; diff --git a/lib/API/MTL/MTLDescriptorHeap.cpp b/lib/API/MTL/MTLDescriptorHeap.cpp index bbf3a39c5..5bc0c4119 100644 --- a/lib/API/MTL/MTLDescriptorHeap.cpp +++ b/lib/API/MTL/MTLDescriptorHeap.cpp @@ -8,6 +8,7 @@ #include "MTLDescriptorHeap.h" #include "MetalIRConverter.h" +#include "ResidencyTracker.h" using namespace offloadtest; @@ -28,9 +29,9 @@ MTLGPUDescriptorHandle::addOffset(int32_t OffsetInDescriptors) { return *this; } -llvm::Expected> -MTLDescriptorHeap::create(MTL::Device *Device, - const MTLDescriptorHeapDesc &Desc) { +llvm::Expected> MTLDescriptorHeap::create( + MTL::Device *Device, const MTLDescriptorHeapDesc &Desc, + std::shared_ptr ResidencyTracker) { if (!Device) return llvm::createStringError(std::errc::invalid_argument, "Invalid MTL::Device pointer."); @@ -45,12 +46,18 @@ MTLDescriptorHeap::create(MTL::Device *Device, if (!Buf) return llvm::createStringError(std::errc::not_enough_memory, "Failed to create MTLDescriptorHeap."); - return std::make_unique(Desc, Buf); + ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->addAllocation(Buf); }); + return std::make_unique(Desc, Buf, + std::move(ResidencyTracker)); } MTLDescriptorHeap::~MTLDescriptorHeap() { - if (Buffer) + if (Buffer) { + ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->removeAllocation(Buffer); }); Buffer->release(); + } } MTLGPUDescriptorHandle @@ -65,7 +72,6 @@ MTLDescriptorHeap::getEntryHandle(uint32_t Index) const { } void MTLDescriptorHeap::bind(MTL::RenderCommandEncoder *Encoder) { - Encoder->useResource(Buffer, MTL::ResourceUsageRead); // Dynamic resource indexing const NS::UInteger BindPoint = getDescriptorHeapBindPoint(Desc.Type); Encoder->setVertexBuffer(Buffer, 0, BindPoint); @@ -73,7 +79,6 @@ void MTLDescriptorHeap::bind(MTL::RenderCommandEncoder *Encoder) { } void MTLDescriptorHeap::bind(MTL::ComputeCommandEncoder *Encoder) { - Encoder->useResource(Buffer, MTL::ResourceUsageRead); // Dynamic resource indexing const NS::UInteger BindPoint = getDescriptorHeapBindPoint(Desc.Type); Encoder->setBuffer(Buffer, 0, BindPoint); diff --git a/lib/API/MTL/MTLDescriptorHeap.h b/lib/API/MTL/MTLDescriptorHeap.h index ca04dddc8..f7c2d9a8f 100644 --- a/lib/API/MTL/MTLDescriptorHeap.h +++ b/lib/API/MTL/MTLDescriptorHeap.h @@ -22,6 +22,8 @@ class ComputeCommandEncoder; struct IRDescriptorTableEntry; namespace offloadtest { +struct MetalResidencyTracker; + struct MTLGPUDescriptorHandle { MTLGPUDescriptorHandle &addOffset(int32_t OffsetInDescriptors); @@ -43,15 +45,19 @@ struct MTLDescriptorHeapDesc { // by the argument buffer for shader resource binding with the explicit root // signature layout. class MTLDescriptorHeap { +public: MTLDescriptorHeapDesc Desc; MTL::Buffer *Buffer; + std::shared_ptr ResidencyTracker; -public: static llvm::Expected> - create(MTL::Device *Device, const MTLDescriptorHeapDesc &Desc); + create(MTL::Device *Device, const MTLDescriptorHeapDesc &Desc, + std::shared_ptr ResidencyTracker); - MTLDescriptorHeap(const MTLDescriptorHeapDesc &Desc, MTL::Buffer *Buffer) - : Desc(Desc), Buffer(Buffer) {} + MTLDescriptorHeap(const MTLDescriptorHeapDesc &Desc, MTL::Buffer *Buffer, + std::shared_ptr ResidencyTracker) + : Desc(Desc), Buffer(Buffer), + ResidencyTracker(std::move(ResidencyTracker)) {} ~MTLDescriptorHeap(); MTLGPUDescriptorHandle getGPUDescriptorHandleForHeapStart() const; diff --git a/lib/API/MTL/MTLDevice.cpp b/lib/API/MTL/MTLDevice.cpp index bf3554bb2..b7813874b 100644 --- a/lib/API/MTL/MTLDevice.cpp +++ b/lib/API/MTL/MTLDevice.cpp @@ -24,9 +24,11 @@ #include "API/Device.h" #include "API/Encoder.h" #include "API/FormatConversion.h" +#include "API/Util.h" #include "MTLDescriptorHeap.h" #include "MTLResources.h" #include "MTLTopLevelArgumentBuffer.h" +#include "ResidencyTracker.h" #include "Support/Pipeline.h" #include "llvm/ADT/ScopeExit.h" @@ -36,13 +38,17 @@ #include "llvm/Support/JSON.h" #include "llvm/Support/raw_ostream.h" -#include "../Util.h" +#include "../Support/OffloadMigration.h" #include #include using namespace offloadtest; +static constexpr MTL::RenderStages AllRenderStages = + MTL::RenderStageVertex | MTL::RenderStageFragment | MTL::RenderStageTile | + MTL::RenderStageObject | MTL::RenderStageMesh; + static llvm::Error toError(NS::Error *Err) { if (!Err) return llvm::Error::success(); @@ -102,34 +108,6 @@ static llvm::Error toError(const IRError *Err, llvm::StringRef Context) { return llvm::createStringError(EC, ErrMsg); } -#define MTLFormats(FMT) \ - if (Channels == 1) \ - return MTL::PixelFormatR##FMT; \ - if (Channels == 2) \ - return MTL::PixelFormatRG##FMT; \ - if (Channels == 4) \ - return MTL::PixelFormatRGBA##FMT; - -static MTL::PixelFormat getMTLFormat(DataFormat Format, int Channels) { - switch (Format) { - case DataFormat::Int32: - MTLFormats(32Sint) break; - case DataFormat::Float32: - MTLFormats(32Float) break; - case DataFormat::UInt64: - case DataFormat::Int64: - if (Channels == 1) - return MTL::PixelFormatRG32Uint; - if (Channels == 2) - return MTL::PixelFormatRGBA32Uint; - llvm_unreachable("Unsupported channel count for 64-bit format"); - - default: - llvm_unreachable("Unsupported Resource format specified"); - } - return MTL::PixelFormatInvalid; -} - static IRShaderStage getShaderStage(Stages Stage) { switch (Stage) { case Stages::Compute: @@ -193,6 +171,60 @@ struct MetalIR { IRShaderReflectionPtr Reflection; }; +struct MetalDescriptorSet { + IRDescriptorTableEntry *CSUHandle; + MTLGPUDescriptorHandle CSUHandleGPU; +}; + +class MetalDescriptorPool : public DescriptorPool { +public: + std::unique_ptr CSUHeap; + std::atomic CSUAllocator = 0; + + MetalDescriptorPool(std::unique_ptr CSUHeap) + : DescriptorPool(GPUAPI::Metal), CSUHeap(std::move(CSUHeap)) {} + + static llvm::Expected> + create(MTL::Device *Dev, + const std::shared_ptr &ResidencyTracker) { + const uint32_t DescriptorCount = 4096; + + const MTLDescriptorHeapDesc HeapDesc = {MTLDescriptorHeapType::CBV_SRV_UAV, + DescriptorCount}; + + auto DescHeapOrErr = + MTLDescriptorHeap::create(Dev, HeapDesc, ResidencyTracker); + if (!DescHeapOrErr) + return DescHeapOrErr.takeError(); + + return std::make_unique(std::move(*DescHeapOrErr)); + } + + void allocateDescriptors(uint32_t Count, IRDescriptorTableEntry *&CPU, + MTLGPUDescriptorHandle &GPU) { + const uint32_t Offset = CSUAllocator.fetch_add(Count); + CPU = CSUHeap->getEntryHandle(Offset); + GPU = CSUHeap->getGPUDescriptorHandleForHeapStart().addOffset(Offset); + } + + void reset() override { CSUAllocator.store(0); } + + static bool classof(const DescriptorPool *P) { + return P->getAPI() == GPUAPI::Metal; + } +}; + +class MetalDescriptorSets : public DescriptorSets { +public: + llvm::SmallVector Sets; + MetalDescriptorSets(llvm::SmallVector Sets) + : DescriptorSets(GPUAPI::Metal), Sets(std::move(Sets)) {} + + static bool classof(const DescriptorSets *S) { + return S->getAPI() == GPUAPI::Metal; + } +}; + class MTLFence : public offloadtest::Fence { public: MTLFence(MTL::SharedEvent *Event, llvm::StringRef Name) @@ -231,6 +263,7 @@ class MTLQueue : public offloadtest::Queue { MTL::CommandQueue *Queue; std::unique_ptr SubmitFence; uint64_t FenceCounter = 0; + std::shared_ptr ResidencyTracker; // Batches of command buffers submitted to the GPU that may still be // in-flight. Each batch records the fence value it signals so we can @@ -241,8 +274,12 @@ class MTLQueue : public offloadtest::Queue { }; llvm::SmallVector InFlightBatches; - MTLQueue(MTL::CommandQueue *Queue, std::unique_ptr SubmitFence) - : Queue(Queue), SubmitFence(std::move(SubmitFence)) {} + MTLQueue(MTL::CommandQueue *Queue, std::unique_ptr SubmitFence, + std::shared_ptr ResidencyTracker) + : Queue(Queue), SubmitFence(std::move(SubmitFence)), + ResidencyTracker(std::move(ResidencyTracker)) { + Queue->addResidencySet(this->ResidencyTracker->ResidencySet); + } ~MTLQueue() override { if (Queue) Queue->release(); @@ -269,11 +306,31 @@ class MTLQueue : public offloadtest::Queue { } }; +struct RootSignatureLayout { + uint32_t IsSampler : 1; + uint32_t Count : 31; + + RootSignatureLayout(bool IsSampler, uint32_t Count) + : IsSampler(IsSampler), Count(Count) {} + RootSignatureLayout() = delete; +}; + +struct DescriptorCountPair { + uint32_t DescriptorCount; + uint32_t SamplerCount; +}; + +struct DescriptorSetsLayout { + llvm::SmallVector RSigLayout; + llvm::SmallVector Sets; +}; + class MTLPipelineState : public offloadtest::PipelineState { public: std::string Name; IRRootSignaturePtr RootSig; std::unique_ptr ArgBuffer; + DescriptorSetsLayout Layout; MTL::ComputePipelineState *ComputePipeline = nullptr; MTL::RenderPipelineState *RenderPipeline = nullptr; @@ -298,14 +355,17 @@ class MTLPipelineState : public offloadtest::PipelineState { MTLPipelineState(llvm::StringRef Name, IRRootSignaturePtr RootSig, std::unique_ptr ArgBuffer, + DescriptorSetsLayout Layout, MTL::ComputePipelineState *ComputePipeline, MTL::Size ThreadsPerGroup) : offloadtest::PipelineState(GPUAPI::Metal), Name(Name), RootSig(std::move(RootSig)), ArgBuffer(std::move(ArgBuffer)), - ComputePipeline(ComputePipeline), ThreadsPerGroup(ThreadsPerGroup) {} + Layout(std::move(Layout)), ComputePipeline(ComputePipeline), + ThreadsPerGroup(ThreadsPerGroup) {} MTLPipelineState(llvm::StringRef Name, IRRootSignaturePtr RootSig, std::unique_ptr ArgBuffer, + DescriptorSetsLayout Layout, MTL::RenderPipelineState *RenderPipeline, MTL::DepthStencilState *DepthStencilState, MTL::CullMode CullMode, @@ -313,8 +373,8 @@ class MTLPipelineState : public offloadtest::PipelineState { MTL::Size ObjectThreadsPerThreadgroup = {1, 1, 1}) : offloadtest::PipelineState(GPUAPI::Metal), Name(Name), RootSig(std::move(RootSig)), ArgBuffer(std::move(ArgBuffer)), - RenderPipeline(RenderPipeline), DepthStencilState(DepthStencilState), - CullMode(CullMode), + Layout(std::move(Layout)), RenderPipeline(RenderPipeline), + DepthStencilState(DepthStencilState), CullMode(CullMode), MeshThreadsPerThreadgroup(MeshThreadsPerThreadgroup), ObjectThreadsPerThreadgroup(ObjectThreadsPerThreadgroup) {} @@ -336,10 +396,10 @@ class MTLPipelineState : public offloadtest::PipelineState { // the rest of the layout (Name, root signature, argument buffer). MTLPipelineState(llvm::StringRef Name, IRRootSignaturePtr RootSig, std::unique_ptr ArgBuffer, - bool IsRT) + DescriptorSetsLayout Layout, bool IsRT) : offloadtest::PipelineState(GPUAPI::Metal), Name(Name), RootSig(std::move(RootSig)), ArgBuffer(std::move(ArgBuffer)), - IsRayTracing(IsRT) {} + Layout(std::move(Layout)), IsRayTracing(IsRT) {} }; /// Ray tracing pipeline state. Layered on top of MTLPipelineState so the @@ -362,6 +422,7 @@ class MTLRayTracingPipelineState : public MTLPipelineState { // ResourceID-based callable tables wired into IRDispatchRaysArgument. MTL::VisibleFunctionTable *VFT = nullptr; MTL::IntersectionFunctionTable *IFT = nullptr; + std::shared_ptr ResidencyTracker; // Per shader entry / hit-group: pre-built IRShaderIdentifier the SBT // builder memcpys into each record. Keys are EntryPoint strings for @@ -373,12 +434,24 @@ class MTLRayTracingPipelineState : public MTLPipelineState { llvm::SmallVector Libraries; llvm::SmallVector Functions; - MTLRayTracingPipelineState(llvm::StringRef Name, IRRootSignaturePtr RootSig, - std::unique_ptr ArgBuf) + MTLRayTracingPipelineState( + llvm::StringRef Name, IRRootSignaturePtr RootSig, + std::unique_ptr ArgBuf, + DescriptorSetsLayout Layout, + std::shared_ptr ResidencyTracker) : MTLPipelineState(Name, std::move(RootSig), std::move(ArgBuf), - /*IsRT=*/true) {} + std::move(Layout), + /*IsRT=*/true), + ResidencyTracker(std::move(ResidencyTracker)) {} ~MTLRayTracingPipelineState() override { + ResidencyTracker->withLock([&](MTL::ResidencySet *RS) { + if (VFT) + RS->removeAllocation(VFT); + if (IFT) + RS->removeAllocation(IFT); + }); + if (VFT) VFT->release(); if (IFT) @@ -409,18 +482,27 @@ class MTLShaderBindingTable : public offloadtest::ShaderBindingTable { IRVirtualAddressRangeAndStride MissRegion{}; IRVirtualAddressRangeAndStride HitGroupRegion{}; IRVirtualAddressRangeAndStride CallableRegion{}; + std::shared_ptr ResidencyTracker; MTLShaderBindingTable(MTL::Buffer *Buf, IRVirtualAddressRange RG, IRVirtualAddressRangeAndStride MS, IRVirtualAddressRangeAndStride HG, - IRVirtualAddressRangeAndStride CL) + IRVirtualAddressRangeAndStride CL, + std::shared_ptr ResidencyTracker) : offloadtest::ShaderBindingTable(GPUAPI::Metal), Buffer(Buf), RayGenRegion(RG), MissRegion(MS), HitGroupRegion(HG), - CallableRegion(CL) {} + CallableRegion(CL), ResidencyTracker(std::move(ResidencyTracker)) { + + this->ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->addAllocation(Buffer); }); + } ~MTLShaderBindingTable() override { - if (Buffer) + if (Buffer) { + ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->removeAllocation(Buffer); }); Buffer->release(); + } } static bool classof(const offloadtest::ShaderBindingTable *S) { @@ -430,15 +512,22 @@ class MTLShaderBindingTable : public offloadtest::ShaderBindingTable { class MTLBuffer : public offloadtest::Buffer { public: - MTL::Buffer *Buf; + MTL::Resource + *Resource; // MTL::Texture* for typed buffer, otherwise MTL::Buffer* std::string Name; BufferCreateDesc Desc; size_t SizeInBytes; + std::shared_ptr ResidencyTracker; - MTLBuffer(MTL::Buffer *Buf, llvm::StringRef Name, BufferCreateDesc Desc, - size_t SizeInBytes) - : offloadtest::Buffer(GPUAPI::Metal), Buf(Buf), Name(Name), Desc(Desc), - SizeInBytes(SizeInBytes) {} + MTLBuffer(MTL::Resource *Resource, llvm::StringRef Name, + BufferCreateDesc Desc, size_t SizeInBytes, + std::shared_ptr ResidencyTracker) + : offloadtest::Buffer(GPUAPI::Metal), Resource(Resource), Name(Name), + Desc(Desc), SizeInBytes(SizeInBytes), + ResidencyTracker(std::move(ResidencyTracker)) { + this->ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->addAllocation(Resource); }); + } MTLBuffer(const MTLBuffer &) = delete; MTLBuffer(MTLBuffer &&) = delete; MTLBuffer &operator=(const MTLBuffer &) = delete; @@ -448,11 +537,28 @@ class MTLBuffer : public offloadtest::Buffer { size_t querySparseTileSizeInBytes(const Device &Dev) const override; + // Only valid for non-typed buffers + MTL::Buffer *getBufferPtr() const { + assert(Desc.AccessType != BufferShaderAccessType::Typed); + return static_cast(Resource); + } + + // Only valid for typed buffers + MTL::Texture *getTexturePtr() const { + assert(Desc.AccessType == BufferShaderAccessType::Typed); + return static_cast(Resource); + } + llvm::Expected map() override { if (Desc.Location == MemoryLocation::GpuOnly) return llvm::createStringError(std::errc::invalid_argument, "Cannot map a GpuOnly buffer."); - return Buf->contents(); + if (Desc.AccessType == BufferShaderAccessType::Typed) + return llvm::createStringError( + std::errc::not_supported, + "Metal does not support mapping typed buffers."); + + return getBufferPtr()->contents(); } void unmap() override { @@ -460,12 +566,15 @@ class MTLBuffer : public offloadtest::Buffer { // propagate CPU-side writes to the GPU. Shared storage (GpuToCpu) is // coherent and needs no action. if (Desc.Location == MemoryLocation::CpuToGpu) - Buf->didModifyRange(NS::Range::Make(0, SizeInBytes)); + getBufferPtr()->didModifyRange(NS::Range::Make(0, SizeInBytes)); } ~MTLBuffer() override { - if (Buf) - Buf->release(); + if (Resource) { + ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->removeAllocation(Resource); }); + Resource->release(); + } } const BufferCreateDesc &getDesc() const override { return Desc; } @@ -480,13 +589,22 @@ class MTLTexture : public offloadtest::Texture { MTL::Texture *Tex; std::string Name; TextureCreateDesc Desc; + std::shared_ptr ResidencyTracker; - MTLTexture(MTL::Texture *Tex, llvm::StringRef Name, TextureCreateDesc Desc) - : offloadtest::Texture(GPUAPI::Metal), Tex(Tex), Name(Name), Desc(Desc) {} + MTLTexture(MTL::Texture *Tex, llvm::StringRef Name, TextureCreateDesc Desc, + std::shared_ptr ResidencyTracker) + : offloadtest::Texture(GPUAPI::Metal), Tex(Tex), Name(Name), Desc(Desc), + ResidencyTracker(std::move(ResidencyTracker)) { + this->ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->addAllocation(Tex); }); + } ~MTLTexture() override { - if (Tex) + if (Tex) { + ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->removeAllocation(Tex); }); Tex->release(); + } } TileShape querySparseTileShape(const Device &Dev) const override; @@ -498,6 +616,17 @@ class MTLTexture : public offloadtest::Texture { } }; +class MTLSampler : public offloadtest::Sampler { +public: + SamplerCreateDesc Desc; + + const SamplerCreateDesc &getDesc() const override { return Desc; } + + static bool classof(const offloadtest::Sampler *S) { + return S->getAPI() == GPUAPI::Metal; + } +}; + /// Metal has no standalone render-pass object: render pass info lives on /// MTLRenderPassDescriptor and is consumed when a render command encoder /// is created. We therefore just stash the descriptor for the encoder to @@ -519,6 +648,7 @@ class MTLDevice; // forward decl — defined below in this same anon ns class MTLCommandBuffer : public offloadtest::CommandBuffer { public: MTL::CommandBuffer *CmdBuffer = nullptr; + MTL::Fence *Fence = nullptr; /// Back-pointer to the owning device; used by encoders that need to /// allocate scratch / instance buffers for AS builds. MTLDevice *Dev = nullptr; @@ -526,6 +656,8 @@ class MTLCommandBuffer : public offloadtest::CommandBuffer { /// and TLAS instance buffers used during builds). llvm::SmallVector> KeepAliveOwned; + MTL::Buffer *BoundCSUHeapBuffer = nullptr; + static llvm::Expected> create(MTL::CommandQueue *Queue) { auto CB = std::unique_ptr(new MTLCommandBuffer()); @@ -536,34 +668,71 @@ class MTLCommandBuffer : public offloadtest::CommandBuffer { return CB; } - ~MTLCommandBuffer() override = default; + ~MTLCommandBuffer() override { + if (Fence) + Fence->release(); + } + + MTL::Fence *ensureFence(); static bool classof(const CommandBuffer *CB) { return CB->getKind() == GPUAPI::Metal; } + void bindPool(const DescriptorPool &Pool) override { + const MetalDescriptorPool &PoolMTL = llvm::cast(Pool); + BoundCSUHeapBuffer = PoolMTL.CSUHeap->Buffer; + } + llvm::Expected> createComputeEncoder() override; llvm::Expected> createRenderEncoder(const offloadtest::RenderPassBeginDesc &Desc) override; -private: MTLCommandBuffer() : CommandBuffer(GPUAPI::Metal) {} }; class MetalAccelerationStructure : public offloadtest::AccelerationStructure { public: MTL::AccelerationStructure *AccelStruct; - - MetalAccelerationStructure(MTL::AccelerationStructure *AccelStruct, - const AccelerationStructureSizes &Sizes) + std::shared_ptr ResidencyTracker; + std::unique_ptr HeaderBuffer; // TLAS Only + std::unique_ptr ContribBuffer; // TLAS Only + + // BLAS + MetalAccelerationStructure( + MTL::AccelerationStructure *AccelStruct, + const AccelerationStructureSizes &Sizes, + std::shared_ptr ResidencyTracker) : offloadtest::AccelerationStructure(GPUAPI::Metal, Sizes), - AccelStruct(AccelStruct) {} + AccelStruct(AccelStruct), + ResidencyTracker(std::move(ResidencyTracker)) { + this->ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->addAllocation(AccelStruct); }); + } + + // TLAS + MetalAccelerationStructure( + MTL::AccelerationStructure *AccelStruct, + const AccelerationStructureSizes &Sizes, + std::shared_ptr ResidencyTracker, + std::unique_ptr HeaderBuffer, + std::unique_ptr ContribBuffer) + : offloadtest::AccelerationStructure(GPUAPI::Metal, Sizes), + AccelStruct(AccelStruct), ResidencyTracker(std::move(ResidencyTracker)), + HeaderBuffer(std::move(HeaderBuffer)), + ContribBuffer(std::move(ContribBuffer)) { + this->ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->addAllocation(AccelStruct); }); + } ~MetalAccelerationStructure() override { - if (AccelStruct) + if (AccelStruct) { + ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->removeAllocation(AccelStruct); }); AccelStruct->release(); + } } static bool classof(const offloadtest::AccelerationStructure *AS) { @@ -571,6 +740,120 @@ class MetalAccelerationStructure : public offloadtest::AccelerationStructure { } }; +class MetalDescriptorSetsBuilder : public DescriptorSetsBuilder { +public: + llvm::SmallVector Sets; + llvm::SmallVector SetStates; + + MetalDescriptorSetsBuilder(llvm::SmallVector Sets) + : DescriptorSetsBuilder(GPUAPI::Metal), Sets(std::move(Sets)) { + SetStates.reserve(this->Sets.size()); + for (const auto &Set : this->Sets) { + SetStates.push_back(Set.CSUHandle); + } + } + + DescriptorSetsBuilder &bindBuffers(uint32_t SetIndex, + llvm::ArrayRef B) { + IRDescriptorTableEntry *EntryPtr = SetStates[SetIndex]; + SetStates[SetIndex] += B.size(); + + for (size_t I = 0, N = B.size(); I < N; ++I) { + const MTLBuffer &BufferMTL = llvm::cast(*B[I]); + if (BufferMTL.Desc.AccessType != BufferShaderAccessType::Typed) { + IRBufferView View = {}; + View.buffer = BufferMTL.getBufferPtr(); + View.bufferSize = BufferMTL.SizeInBytes; + IRDescriptorTableSetBufferView(EntryPtr + I, &View); + } else { + IRDescriptorTableSetTexture(EntryPtr + I, BufferMTL.getTexturePtr(), 0, + 0); + } + } + + return *this; + } + + DescriptorSetsBuilder &bindTextures(uint32_t SetIndex, + llvm::ArrayRef T, + llvm::ArrayRef S) { + assert((S.empty() || S.size() == T.size()) && + "Sampler list must either be empty or match " + "texture list when binding descriptors."); + + IRDescriptorTableEntry *EntryPtr = SetStates[SetIndex]; + SetStates[SetIndex] += T.size(); + + for (size_t I = 0, N = T.size(); I < N; ++I) { + assert((S.empty() || S[I] == nullptr) && + "Metal does not support combined image samplers."); + + const MTLTexture &TextureMTL = llvm::cast(*T[I]); + IRDescriptorTableSetTexture(EntryPtr + I, TextureMTL.Tex, 0, 0); + } + + return *this; + } + + DescriptorSetsBuilder &constant(uint32_t SetIndex, + llvm::ArrayRef B, + VKBind) override { + return bindBuffers(SetIndex, B); + } + + DescriptorSetsBuilder & + read(uint32_t SetIndex, llvm::ArrayRef B, VKBind) override { + return bindBuffers(SetIndex, B); + } + DescriptorSetsBuilder &read(uint32_t SetIndex, + llvm::ArrayRef T, + llvm::ArrayRef S, + VKBind) override { + return bindTextures(SetIndex, T, S); + } + DescriptorSetsBuilder &read(uint32_t SetIndex, + llvm::ArrayRef A, + VKBind) override { + IRDescriptorTableEntry *EntryPtr = SetStates[SetIndex]; + SetStates[SetIndex] += A.size(); + + for (size_t I = 0, N = A.size(); I < N; ++I) { + const MetalAccelerationStructure &AccelStructMTL = + llvm::cast(*A[I]); + const MTLBuffer &HeaderBufferMTL = + llvm::cast(*AccelStructMTL.HeaderBuffer.get()); + IRDescriptorTableSetAccelerationStructure( + EntryPtr + I, HeaderBufferMTL.getBufferPtr()->gpuAddress()); + } + + return *this; + } + + DescriptorSetsBuilder & + write(uint32_t SetIndex, llvm::ArrayRef B, VKBind) override { + return bindBuffers(SetIndex, B); + } + DescriptorSetsBuilder & + write(uint32_t SetIndex, llvm::ArrayRef T, VKBind) override { + return bindTextures(SetIndex, T, {}); + } + + DescriptorSetsBuilder &sampler(uint32_t SetIndex, + llvm::ArrayRef S, + VKBind) override { + assert(false && "Samplers are not implemented on Metal."); + return *this; + } + + std::unique_ptr build() override { + return std::make_unique(std::move(Sets)); + } + + static bool classof(const DescriptorSetsBuilder *B) { + return B->getAPI() == GPUAPI::Metal; + } +}; + llvm::Expected MTLQueue::submit( llvm::SmallVector> CBs) { // Non-blocking: query how far the GPU has progressed and release @@ -586,6 +869,16 @@ llvm::Expected MTLQueue::submit( // so no explicit wait on prior work is needed here. const uint64_t SignalValue = ++FenceCounter; + // Update residency. + // From a correctness perspective this is the least complex way + // From a performance perspective this is not ideal, because requestResidency + // may take some time to execute. Preferably this would be done earlier in the + // frame. + ResidencyTracker->withLock([&](MTL::ResidencySet *RS) { + RS->commit(); + RS->requestResidency(); + }); + for (size_t I = 0; I < CBs.size(); ++I) { auto &MCB = llvm::cast(*CBs[I].get()); // Signal the submit fence when the last command buffer completes. @@ -612,12 +905,16 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { /// Accumulated barrier scope from commands recorded since the last barrier. MTL::BarrierScope PendingScope = MTL::BarrierScope(0); - /// Record that a command touched the given resource types. The accumulated - /// scope is flushed as a memoryBarrier before the next command. + // Commands that touch resources via descriptors require manual memory + // bariers. Use addBarrierScope on resource access, call flushBarrier before + // executing a command that accesses resources via descriptors. + // Note: This is only valid on ComputeEnc. BlitEnc and ASEnc track hazards + // automatically. void addBarrierScope(MTL::BarrierScope Scope) { PendingScope |= Scope; } - void flushBarrier() { - if (ComputeEnc && PendingScope != MTL::BarrierScope(0)) { + assert(ComputeEnc && + "flushBarrier must be called after ensureComputeEncoder"); + if (PendingScope != MTL::BarrierScope(0)) { ComputeEnc->memoryBarrier(PendingScope); PendingScope = MTL::BarrierScope(0); } @@ -634,6 +931,8 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { if (!ComputeEnc) return llvm::createStringError(std::errc::device_or_resource_busy, "Failed to create Metal compute encoder."); + if (CB->Fence) + ComputeEnc->waitForFence(CB->Fence); ComputeEnc->pushDebugGroup( NS::String::string("ComputeEncoder", NS::UTF8StringEncoding)); return llvm::Error::success(); @@ -648,14 +947,14 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { if (!BlitEnc) return llvm::createStringError(std::errc::device_or_resource_busy, "Failed to create Metal blit encoder."); + if (CB->Fence) + BlitEnc->waitForFence(CB->Fence); return llvm::Error::success(); } public: - MTLComputeEncoder(MTLCommandBuffer *CB, MTL::CommandBuffer *CmdBuffer, - MTL::ComputeCommandEncoder *Encoder) - : ComputeEncoder(GPUAPI::Metal), CB(CB), CmdBuffer(CmdBuffer), - ComputeEnc(Encoder) {} + MTLComputeEncoder(MTLCommandBuffer *CB, MTL::CommandBuffer *CmdBuffer) + : ComputeEncoder(GPUAPI::Metal), CB(CB), CmdBuffer(CmdBuffer) {} ~MTLComputeEncoder() override { endEncoding(); } @@ -663,12 +962,18 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { return E->getAPI() == GPUAPI::Metal; } - MTL::ComputeCommandEncoder *getNative() const { return ComputeEnc; } + llvm::Expected getNative() { + if (auto Err = ensureComputeEncoder()) + return Err; + return ComputeEnc; + } MTL::CommandEncoder *getActiveEncoder() const { if (ComputeEnc) return ComputeEnc; - return BlitEnc; + if (BlitEnc) + return BlitEnc; + return ASEnc; } void pushDebugGroup(llvm::StringRef Label) override { @@ -688,6 +993,41 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { NS::String::string(Label.data(), NS::UTF8StringEncoding)); } + void bindDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) { + const auto &MTLPSO = llvm::cast(PSO); + const MetalDescriptorSets &DSetsMTL = + llvm::cast(DSets); + + if (auto Err = ensureComputeEncoder()) + llvm::report_fatal_error(std::move(Err)); + + ComputeEnc->setBuffer(CB->BoundCSUHeapBuffer, 0, + kIRDescriptorHeapBindPoint); + + // NOTE(manon): This is a problem when we dispatch the same pipeline + // multiple times. This buffer should not be living in the PSO. + for (uint32_t I = 0, N = DSetsMTL.Sets.size(); I < N; ++I) { + assert(!MTLPSO.Layout.RSigLayout[I].IsSampler && + "Descriptor layout doesn't match root signature."); + MTLPSO.ArgBuffer->setRootDescriptorTable(I, + DSetsMTL.Sets[I].CSUHandleGPU); + } + + ComputeEnc->setBuffer(MTLPSO.ArgBuffer->Buffer, 0, + kIRArgumentBufferBindPoint); + } + + void bindComputeDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) override { + return bindDescriptorSets(PSO, DSets); + } + + void bindRayTracingDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) override { + return bindDescriptorSets(PSO, DSets); + } + llvm::Error dispatch(const offloadtest::PipelineState &PSO, uint32_t GroupCountX, uint32_t GroupCountY, uint32_t GroupCountZ) override { @@ -720,8 +1060,44 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { auto &MTLSrc = static_cast(Src); auto &MTLDst = static_cast(Dst); insertDebugSignpost(llvm::formatv("CopyBuffer {0}B", Size).str()); - BlitEnc->copyFromBuffer(MTLSrc.Buf, SrcOffset, MTLDst.Buf, DstOffset, Size); - addBarrierScope(MTL::BarrierScopeBuffers); + + if (MTLSrc.Desc.AccessType == BufferShaderAccessType::Typed) { + if (MTLDst.Desc.AccessType == BufferShaderAccessType::Typed) { + const uint32_t SrcElementSize = + getFormatSizeInBytes(MTLSrc.Desc.AccessTypeParams.Fmt); + const uint32_t DstElementSize = + getFormatSizeInBytes(MTLDst.Desc.AccessTypeParams.Fmt); + BlitEnc->copyFromTexture( + MTLSrc.getTexturePtr(), 0 /*sourceSlice (unused)*/, + 0 /*sourceLevel (unused)*/, + MTL::Origin(SrcOffset / SrcElementSize, 0, 0), + MTL::Size(Size / SrcElementSize, 1, 1), MTLDst.getTexturePtr(), + 0 /*destinationSlice (unused)*/, 0 /*destinationLevel (unused)*/, + MTL::Origin(DstOffset / DstElementSize, 0, 0)); + } else { + const uint32_t ElementSize = + getFormatSizeInBytes(MTLSrc.Desc.AccessTypeParams.Fmt); + BlitEnc->copyFromTexture( + MTLSrc.getTexturePtr(), 0 /*sourceSlice (unused)*/, + 0 /*sourceLevel (unused)*/, + MTL::Origin(SrcOffset / ElementSize, 0, 0), + MTL::Size(Size / ElementSize, 1, 1), MTLDst.getBufferPtr(), + DstOffset, 0 /*destinationBytesPerRow (unused)*/, + 0 /*destinationBytesPerImage (unused)*/); + } + } else if (MTLDst.Desc.AccessType == BufferShaderAccessType::Typed) { + const uint32_t ElementSize = + getFormatSizeInBytes(MTLDst.Desc.AccessTypeParams.Fmt); + BlitEnc->copyFromBuffer( + MTLSrc.getBufferPtr(), SrcOffset, 0 /*sourceBytesPerRow (unused)*/, + 0 /*sourceBytesPerImage (unused)*/, + MTL::Size(Size / ElementSize, 1, 1), MTLDst.getTexturePtr(), + 0 /*destinationSlice (unused)*/, 0 /*destinationLevel (unused)*/, + MTL::Origin(DstOffset / ElementSize, 0, 0)); + } else { + BlitEnc->copyFromBuffer(MTLSrc.getBufferPtr(), SrcOffset, + MTLDst.getBufferPtr(), DstOffset, Size); + } return llvm::Error::success(); } @@ -731,24 +1107,30 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { return Err; auto &MTLSrc = static_cast(Src); auto &MTLDst = static_cast(Dst); + assert(MTLSrc.Desc.AccessType != BufferShaderAccessType::Typed && + "TODO(manon): Support typed buffer copies."); - // The upload buffer is laid out with a tightly packed row stride matching - // getTextureUploadRowStrideInBytes(), so the source bytes-per-row is the - // texture width times the element size. + // The upload buffer holds tightly packed texel data for every mip level + // (see createTextureWithData): each mip's rows are contiguous with no + // padding, and the mips follow one another. Copy one mip at a time, with + // the source bytes-per-row being that mip's width times the element size. const size_t ElemSize = getFormatSizeInBytes(MTLDst.Desc.Fmt); - const size_t RowBytes = MTLDst.Desc.Width * ElemSize; - const size_t ImageBytes = RowBytes * MTLDst.Desc.Height; - const MTL::Size CopySize(MTLDst.Desc.Width, MTLDst.Desc.Height, 1); insertDebugSignpost(llvm::formatv("copyBufferToTexture {0} -> {1}", MTLSrc.Name, MTLDst.Name) .str()); - BlitEnc->copyFromBuffer(MTLSrc.Buf, /*sourceOffset=*/0, RowBytes, - ImageBytes, CopySize, MTLDst.Tex, - /*destinationSlice=*/0, /*destinationLevel=*/0, - MTL::Origin(0, 0, 0)); - addBarrierScope(MTL::BarrierScopeTextures); - + size_t CurrentOffset = 0; + for (uint32_t I = 0; I < MTLDst.Desc.MipLevels; ++I) { + const uint32_t MipWidth = std::max(1u, MTLDst.Desc.Width >> I); + const uint32_t MipHeight = std::max(1u, MTLDst.Desc.Height >> I); + const size_t RowBytes = MipWidth * ElemSize; + const size_t ImageBytes = RowBytes * MipHeight; + BlitEnc->copyFromBuffer( + MTLSrc.getBufferPtr(), CurrentOffset, RowBytes, ImageBytes, + MTL::Size(MipWidth, MipHeight, 1), MTLDst.Tex, + /*destinationSlice=*/0, /*destinationLevel=*/I, MTL::Origin(0, 0, 0)); + CurrentOffset += ImageBytes; + } return llvm::Error::success(); } @@ -765,6 +1147,8 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { return Err; auto &MTLSrc = static_cast(Src); auto &MTLDst = static_cast(Dst); + assert(MTLDst.Desc.AccessType != BufferShaderAccessType::Typed && + "TODO(manon): Support typed buffer copies."); // The readback buffer is linear with a tightly packed row stride, so the // destination bytes-per-row is the texture width times the element size. @@ -777,9 +1161,9 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { MTLSrc.Name, MTLDst.Name) .str()); BlitEnc->copyFromTexture(MTLSrc.Tex, /*sourceSlice=*/0, /*sourceLevel=*/0, - MTL::Origin(0, 0, 0), CopySize, MTLDst.Buf, + MTL::Origin(0, 0, 0), CopySize, + MTLDst.getBufferPtr(), /*destinationOffset=*/0, RowBytes, ImageBytes); - addBarrierScope(MTL::BarrierScopeBuffers); return llvm::Error::success(); } @@ -793,35 +1177,9 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { // and the SBT buffer) must already be set on the active compute encoder by // the caller — this method only binds the pipeline state and issues the // dispatch. - llvm::Error dispatchRays(const PipelineState &PSO, const ShaderBindingTable &, - uint32_t Width, uint32_t Height, - uint32_t Depth) override { - if (!llvm::isa(&PSO)) - return llvm::createStringError( - std::errc::invalid_argument, - "dispatchRays requires a RayTracing PipelineState."); - const auto &RTPSO = llvm::cast(PSO); - if (!RTPSO.ComputePipeline) - return llvm::createStringError( - std::errc::invalid_argument, - "RayTracing PipelineState has no compute pipeline state."); - if (auto Err = ensureComputeEncoder()) - return Err; - flushBarrier(); - insertDebugSignpost( - llvm::formatv("DispatchRays [{0},{1},{2}]", Width, Height, Depth) - .str()); - ComputeEnc->setComputePipelineState(RTPSO.ComputePipeline); - - // DispatchRays(W, H, D) launches W*H*D rays; tid in the irconverter raygen - // kernel is the per-ray index. Pass grid as raw (W, H, D) and let Metal - // ceil-divide by ThreadsPerGroup to compute threadgroup count. - const MTL::Size GridSize(Width, Height, Depth); - ComputeEnc->dispatchThreads(GridSize, RTPSO.ThreadsPerGroup); - addBarrierScope(MTL::BarrierScopeBuffers | MTL::BarrierScopeTextures); - return llvm::Error::success(); - } - + llvm::Error dispatchRays(const PipelineState &PSO, + const ShaderBindingTable &SBT, uint32_t Width, + uint32_t Height, uint32_t Depth) override; /// Lazily transition into an AccelerationStructureCommandEncoder; mirrors /// the existing compute↔blit lazy switch. llvm::Error ensureASEncoder() { @@ -833,21 +1191,28 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { return llvm::createStringError( std::errc::device_or_resource_busy, "Failed to create Metal acceleration-structure encoder."); + if (CB->Fence) + ASEnc->waitForFence(CB->Fence); return llvm::Error::success(); } void endEncodingImpl() override { if (ComputeEnc) { - flushBarrier(); + // Clear PendingScope, no actual barrier needed. + PendingScope = MTL::BarrierScope(0); + ComputeEnc->popDebugGroup(); + ComputeEnc->updateFence(CB->ensureFence()); ComputeEnc->endEncoding(); ComputeEnc = nullptr; } if (BlitEnc) { + BlitEnc->updateFence(CB->ensureFence()); BlitEnc->endEncoding(); BlitEnc = nullptr; } if (ASEnc) { + ASEnc->updateFence(CB->ensureFence()); ASEnc->endEncoding(); ASEnc = nullptr; } @@ -856,15 +1221,7 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder { llvm::Expected> MTLCommandBuffer::createComputeEncoder() { - MTL::ComputeCommandEncoder *NativeEncoder = - CmdBuffer->computeCommandEncoder(); - if (!NativeEncoder) - return llvm::createStringError( - std::errc::device_or_resource_busy, - "Failed to create Metal compute command encoder."); - NativeEncoder->pushDebugGroup( - NS::String::string("ComputeEncoder", NS::UTF8StringEncoding)); - return std::make_unique(this, CmdBuffer, NativeEncoder); + return std::make_unique(this, CmdBuffer); } static MTL::LoadAction getMTLLoadAction(offloadtest::LoadAction Action) { @@ -891,6 +1248,7 @@ static MTL::StoreAction getMTLStoreAction(offloadtest::StoreAction Action) { class MTLRenderEncoder : public offloadtest::RenderEncoder { MTL::RenderCommandEncoder *RenderEnc = nullptr; + MTLCommandBuffer *CB = nullptr; // Encoder contract: viewport and scissor must both be set before // drawInstanced(). @@ -898,8 +1256,8 @@ class MTLRenderEncoder : public offloadtest::RenderEncoder { bool ScissorSet = false; public: - MTLRenderEncoder(MTL::RenderCommandEncoder *Enc) - : RenderEncoder(GPUAPI::Metal), RenderEnc(Enc) {} + MTLRenderEncoder(MTL::RenderCommandEncoder *Enc, MTLCommandBuffer *CB) + : RenderEncoder(GPUAPI::Metal), RenderEnc(Enc), CB(CB) {} MTLRenderEncoder(const MTLRenderEncoder &CB) = delete; MTLRenderEncoder(MTLRenderEncoder &&CB) = delete; MTLRenderEncoder &operator=(MTLRenderEncoder &CB) = delete; @@ -965,12 +1323,38 @@ class MTLRenderEncoder : public offloadtest::RenderEncoder { assert(Slot == 0 && "Pipeline vertex descriptor only describes slot 0"); if (VB) { auto &MTLVB = llvm::cast(*VB); - RenderEnc->setVertexBuffer(MTLVB.Buf, Offset, BufIdx); + RenderEnc->setVertexBuffer(MTLVB.getBufferPtr(), Offset, BufIdx); } else { RenderEnc->setVertexBuffer(nullptr, 0, BufIdx); } } + void bindDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) override { + const auto &MTLPSO = llvm::cast(PSO); + const MetalDescriptorSets &DSetsMTL = + llvm::cast(DSets); + + RenderEnc->setVertexBuffer(CB->BoundCSUHeapBuffer, 0, + kIRDescriptorHeapBindPoint); + RenderEnc->setFragmentBuffer(CB->BoundCSUHeapBuffer, 0, + kIRDescriptorHeapBindPoint); + + // NOTE(manon): This is a problem when we dispatch the same pipeline + // multiple times. This buffer should not be living in the PSO. + for (uint32_t I = 0, N = DSetsMTL.Sets.size(); I < N; ++I) { + assert(!MTLPSO.Layout.RSigLayout[I].IsSampler && + "Descriptor layout doesn't match root signature."); + MTLPSO.ArgBuffer->setRootDescriptorTable(I, + DSetsMTL.Sets[I].CSUHandleGPU); + } + + RenderEnc->setVertexBuffer(MTLPSO.ArgBuffer->Buffer, 0, + kIRArgumentBufferBindPoint); + RenderEnc->setFragmentBuffer(MTLPSO.ArgBuffer->Buffer, 0, + kIRArgumentBufferBindPoint); + } + llvm::Error drawInstanced(const offloadtest::PipelineState &PSO, uint32_t VertexCount, uint32_t InstanceCount, uint32_t FirstVertex, @@ -1036,6 +1420,9 @@ class MTLRenderEncoder : public offloadtest::RenderEncoder { void endEncodingImpl() override { assert(RenderEnc); + // Perform worst-case sync, something that could be optimized in the future. + // Probably minimal perf impact. + RenderEnc->updateFence(CB->ensureFence(), AllRenderStages); RenderEnc->popDebugGroup(); RenderEnc->endEncoding(); RenderEnc = nullptr; @@ -1142,67 +1529,29 @@ MTLCommandBuffer::createRenderEncoder( return llvm::createStringError( std::errc::device_or_resource_busy, "Failed to create Metal render command encoder."); + + // Perform worst-case sync, something that could be optimized in the future. + // Probably minimal perf impact. + if (Fence) + NativeEncoder->waitForFence(Fence, AllRenderStages); + NativeEncoder->pushDebugGroup( NS::String::string("RenderEncoder", NS::UTF8StringEncoding)); - return std::make_unique(NativeEncoder); + return std::make_unique(NativeEncoder, this); } class MTLDevice : public offloadtest::Device { public: Capabilities Caps; MTL::Device *Device; + std::shared_ptr ResidencyTracker; MTLQueue GraphicsQueue; - struct ResourceSet { - MTLPtr Resource; - // AS-only; mutually exclusive with Resource above. - MetalAccelerationStructure *AS = nullptr; - explicit ResourceSet(MTL::Resource *Resource) : Resource(Resource) {} - explicit ResourceSet(MetalAccelerationStructure *AS) : AS(AS) {} - }; - - // ResourceBundle will contain one ResourceSet for a singular resource - // or multiple ResourceSets for resource array. - using ResourceBundle = llvm::SmallVector; - using ResourcePair = std::pair; - - struct DescriptorTable { - llvm::SmallVector Resources; - }; - - struct InvocationState { - InvocationState() { Pool = NS::AutoreleasePool::alloc()->init(); } - ~InvocationState() { Pool->release(); } - - NS::AutoreleasePool *Pool = nullptr; - std::unique_ptr DescHeap; - std::unique_ptr VB; - std::unique_ptr RenderTarget; - std::unique_ptr FrameBufferReadback; - std::unique_ptr DepthStencil; - std::unique_ptr CB; - std::unique_ptr Pipeline; - std::unique_ptr SBT; - std::unique_ptr RenderPass; - - llvm::SmallVector DescTables; - // TODO: Support RootResources? - - // Parallel-indexed to `P.AccelStructs.BLAS`. - llvm::SmallVector> - BLASes; - // Keyed by `TLASDesc::Name`. - llvm::StringMap> TLASes; - // Vertex/index buffers consumed during AS builds; must outlive submission. - llvm::SmallVector> ASInputBuffers; - // Per-AS header + contributions buffers; resident at dispatch. - llvm::SmallVector> ASDescriptorBuffers; - }; - - llvm::Error createRootSignature( - const BindingsDesc &BindingsDesc, bool IsGraphics, - IRRootSignaturePtr &OutRootSig, - std::unique_ptr &OutArgBuffer) { + llvm::Error + createRootSignature(const BindingsDesc &BindingsDesc, bool IsGraphics, + IRRootSignaturePtr &OutRootSig, + std::unique_ptr &OutArgBuffer, + DescriptorSetsLayout &Layout) { uint32_t DescriptorCount = 0; for (auto &D : BindingsDesc.DescriptorSetDescs) DescriptorCount += D.ResourceBindings.size(); @@ -1254,6 +1603,13 @@ class MTLDevice : public offloadtest::Device { static_cast(Set.ResourceBindings.size()); Param.DescriptorTable.pDescriptorRanges = &Ranges.get()[StartRangeIdx]; Param.ShaderVisibility = IRShaderVisibilityAll; + + Layout.RSigLayout.push_back(RootSignatureLayout(false, DescriptorIdx)); + + Layout.Sets.push_back({ + DescriptorIdx, // DescriptorCount + 0 // SamplerCount + }); } // NOTE: Attempting to create a RS with version 1.0 seems to fail @@ -1277,8 +1633,8 @@ class MTLDevice : public offloadtest::Device { OutRootSig = std::move(RootSig); - auto ArgBufferOrErr = - MTLTopLevelArgumentBuffer::create(Device, OutRootSig.get()); + auto ArgBufferOrErr = MTLTopLevelArgumentBuffer::create( + Device, OutRootSig.get(), ResidencyTracker); if (!ArgBufferOrErr) return ArgBufferOrErr.takeError(); @@ -1286,26 +1642,6 @@ class MTLDevice : public offloadtest::Device { return llvm::Error::success(); } - llvm::Error createDescriptorHeap(Pipeline &P, InvocationState &State) { - if (P.getDescriptorCount() == 0) { - llvm::outs() - << "No descriptors found, skipping descriptor heap creation.\n"; - return llvm::Error::success(); - } - const uint32_t DescriptorCount = P.getDescriptorCountWithFlattenedArrays(); - const MTLDescriptorHeapDesc HeapDesc = {MTLDescriptorHeapType::CBV_SRV_UAV, - DescriptorCount}; - - auto DescHeapOrErr = MTLDescriptorHeap::create(Device, HeapDesc); - if (!DescHeapOrErr) - return DescHeapOrErr.takeError(); - - State.DescHeap = std::move(*DescHeapOrErr); - llvm::outs() << "Descriptor heap created with " << DescriptorCount - << " descriptors.\n"; - return llvm::Error::success(); - } - llvm::Expected convertToMetalIR(Stages Stage, bool IsGraphics, IRRootSignature *RootSig, const ShaderContainer &SC) { @@ -1417,516 +1753,68 @@ class MTLDevice : public offloadtest::Device { return MetalIR{std::move(MetalLib), std::move(Reflection)}; } - // Creates a Metal resource (buffer or texture) for the given Resource at the - // specified array index. - llvm::Expected - createResource(Resource &R, size_t ResourceArrayIndex = 0) { - const offloadtest::CPUBuffer &B = *R.BufferPtr; - - if (R.isRaw()) { - MTL::Buffer *Buf = - Device->newBuffer(B.Data[ResourceArrayIndex].get(), R.size(), - MTL::ResourceStorageModeManaged); - Buf->didModifyRange(NS::Range::Make(0, Buf->length())); - return Buf; - } - const uint64_t Width = - R.isTexture() ? B.OutputProps.Width : R.size() / R.getElementSize(); - const uint64_t Height = R.isTexture() ? B.OutputProps.Height : 1; - MTL::TextureUsage UsageFlags = MTL::ResourceUsageRead; - if (R.isReadWrite()) - UsageFlags |= MTL::ResourceUsageWrite; - MTL::TextureDescriptor *Desc = nullptr; - const MTL::PixelFormat Format = getMTLFormat(B.Format, B.Channels); - switch (R.Kind) { - case ResourceKind::Buffer: - case ResourceKind::RWBuffer: - Desc = MTL::TextureDescriptor::textureBufferDescriptor( - Format, Width, MTL::ResourceStorageModeManaged, UsageFlags); - break; - case ResourceKind::Texture2D: - case ResourceKind::RWTexture2D: - Desc = MTL::TextureDescriptor::texture2DDescriptor(Format, Width, Height, - false); - break; - case ResourceKind::Sampler: - llvm_unreachable("Not implemented yet."); - case ResourceKind::SampledTexture2D: - llvm_unreachable("SampledTextures aren't supported in Metal."); - case ResourceKind::StructuredBuffer: - case ResourceKind::RWStructuredBuffer: - case ResourceKind::ByteAddressBuffer: - case ResourceKind::RWByteAddressBuffer: - case ResourceKind::ConstantBuffer: - llvm_unreachable("Raw is checked above"); - case ResourceKind::AccelerationStructure: - llvm_unreachable("Acceleration structures use a separate path!"); - } - - MTL::Texture *NewTex = Device->newTexture(Desc); - NewTex->replaceRegion(MTL::Region(0, 0, Width, Height), 0, - B.Data[ResourceArrayIndex].get(), - Width * R.getElementSize()); - return NewTex; - } - - llvm::Expected createSRV(Resource &R) { - ResourceBundle Bundle; - - for (size_t RegOffset = 0; RegOffset < R.BufferPtr->Data.size(); - ++RegOffset) { - llvm::outs() << "Creating SRV: { Size = " << R.size() << ", Register = t" - << R.DXBinding.Register + RegOffset - << ", Space = " << R.DXBinding.Space; - llvm::outs() << " }\n"; + llvm::Error createComputeCommands(Pipeline &P, SharedInvocationState &IS, + DescriptorPool &Pool) { + IS.CB->bindPool(Pool); - auto ResourceOrErr = createResource(R, RegOffset); - if (!ResourceOrErr) - return ResourceOrErr.takeError(); - - Bundle.emplace_back(ResourceOrErr.get()); - } - return Bundle; - } + auto DescSetsOrErr = + buildDescriptorSets(*this, Pool, *IS.Pipeline, IS.DescTables); + if (!DescSetsOrErr) + return DescSetsOrErr.takeError(); + auto DescSets = std::move(*DescSetsOrErr); - // TODO: counter buffer via IRRuntimeCreateAppendBufferView? - llvm::Expected createUAV(Resource &R) { - ResourceBundle Bundle; - - for (size_t RegOffset = 0; RegOffset < R.BufferPtr->Data.size(); - ++RegOffset) { - llvm::outs() << "Creating UAV: { Size = " << R.size() << ", Register = u" - << R.DXBinding.Register + RegOffset - << ", Space = " << R.DXBinding.Space - << ", HasCounter = " << R.HasCounter; - llvm::outs() << " }\n"; - - auto ResourceOrErr = createResource(R, RegOffset); - if (!ResourceOrErr) - return ResourceOrErr.takeError(); - - Bundle.emplace_back(ResourceOrErr.get()); - } - return Bundle; - } - - llvm::Expected createCBV(Resource &R) { - ResourceBundle Bundle; - - for (size_t RegOffset = 0; RegOffset < R.BufferPtr->Data.size(); - ++RegOffset) { - llvm::outs() << "Creating CBV: { Size = " << R.size() << ", Register = b" - << R.DXBinding.Register + RegOffset - << ", Space = " << R.DXBinding.Space << " }\n"; - - auto ResourceOrErr = createResource(R, RegOffset); - if (!ResourceOrErr) - return ResourceOrErr.takeError(); - - Bundle.emplace_back(ResourceOrErr.get()); - } - return Bundle; - } - - void createDescriptor(Resource &R, MTL::Resource *Resource, - IRDescriptorTableEntry *Entry) { - if (R.isRaw()) { - IRBufferView View = {}; - View.buffer = static_cast(Resource); - View.bufferSize = R.size(); - IRDescriptorTableSetBufferView(Entry, &View); - } else { - MTL::Texture *Tex = static_cast(Resource); - IRDescriptorTableSetTexture(Entry, Tex, 0, 0); - } - } - - // returns the next available HeapIdx - uint32_t bindSRV(Resource &R, InvocationState &IS, uint32_t HeapIdx, - const ResourceBundle &ResBundle) { - const uint32_t EltSize = R.getElementSize(); - const uint32_t NumElts = R.size() / EltSize; - - for (const ResourceSet &RS : ResBundle) { - llvm::outs() << "SRV: HeapIdx = " << HeapIdx << " EltSize = " << EltSize - << " NumElts = " << NumElts << "\n"; - createDescriptor(R, RS.Resource.get(), - IS.DescHeap->getEntryHandle(HeapIdx)); - HeapIdx++; - } - return HeapIdx; - } - - // returns the next available HeapIdx - uint32_t bindUAV(Resource &R, InvocationState &IS, uint32_t HeapIdx, - const ResourceBundle &ResBundle) { - const uint32_t EltSize = R.getElementSize(); - const uint32_t NumElts = R.size() / EltSize; - for (const ResourceSet &RS : ResBundle) { - llvm::outs() << "UAV: HeapIdx = " << HeapIdx << " EltSize = " << EltSize - << " NumElts = " << NumElts << "\n"; - createDescriptor(R, RS.Resource.get(), - IS.DescHeap->getEntryHandle(HeapIdx)); - HeapIdx++; - } - return HeapIdx; - } - - // returns the next available HeapIdx - uint32_t bindCBV(Resource &R, InvocationState &IS, uint32_t HeapIdx, - const ResourceBundle &ResBundle) { - for (const ResourceSet &RS : ResBundle) { - llvm::outs() << "CBV: HeapIdx = " << HeapIdx << " Size = " << R.size() - << "\n"; - createDescriptor(R, RS.Resource.get(), - IS.DescHeap->getEntryHandle(HeapIdx)); - HeapIdx++; - } - return HeapIdx; - } - - llvm::Expected> createAS(Resource &R) { - assert(R.TLASPtr && "AS resource must be resolved to a TLAS"); - assert(R.getArraySize() == 1 && "AS arrays not yet supported"); - auto SizesOrErr = - getTLASBuildSizes(static_cast(R.TLASPtr->Instances.size())); - if (!SizesOrErr) - return SizesOrErr.takeError(); - return createTLAS(*SizesOrErr); - } - - llvm::Error createBuffers(Pipeline &P, InvocationState &IS) { - auto CreateBuffer = - [&IS, - this](Resource &R, - llvm::SmallVectorImpl &Resources) -> llvm::Error { - if (R.isAccelerationStructure()) { - auto ASOrErr = createAS(R); - if (!ASOrErr) - return ASOrErr.takeError(); - ResourceBundle Bundle; - Bundle.emplace_back( - llvm::cast(ASOrErr->get())); - auto Inserted = - IS.TLASes.try_emplace(R.TLASPtr->Name, std::move(*ASOrErr)); - assert(Inserted.second && "TLAS bound to multiple resources NYI"); - (void)Inserted; - Resources.emplace_back(&R, std::move(Bundle)); - return llvm::Error::success(); - } - switch (getDescriptorKind(R.Kind)) { - case DescriptorKind::SRV: { - auto ExRes = createSRV(R); - if (!ExRes) - return ExRes.takeError(); - Resources.emplace_back(&R, std::move(*ExRes)); - break; - } - case DescriptorKind::UAV: { - auto ExRes = createUAV(R); - if (!ExRes) - return ExRes.takeError(); - Resources.emplace_back(&R, std::move(*ExRes)); - break; - } - case DescriptorKind::CBV: { - auto ExRes = createCBV(R); - if (!ExRes) - return ExRes.takeError(); - Resources.emplace_back(&R, std::move(*ExRes)); - break; - } - case DescriptorKind::SAMPLER: - return llvm::createStringError( - std::errc::not_supported, - "Samplers are not yet implemented for Metal."); - } - return llvm::Error::success(); - }; - - for (auto &D : P.Sets) { - IS.DescTables.emplace_back(DescriptorTable()); - DescriptorTable &Table = IS.DescTables.back(); - for (auto &R : D.Resources) - if (auto Err = CreateBuffer(R, Table.Resources)) - return Err; - } - - // Bind descriptors in descriptor tables. - uint32_t HeapIndex = 0; - for (auto &T : IS.DescTables) { - for (auto &R : T.Resources) { - if (MetalAccelerationStructure *MTLAS = R.second[0].AS) { - // The Metal shader converter binds the AS indirectly through an - // `IRRaytracingAccelerationStructureGPUHeader` buffer carrying the - // AS's `gpuResourceID` and a pointer to an instance-contributions - // array (one `uint32` per instance, equivalent to D3D12's - // `InstanceContributionToHitGroupIndex`). - const uint32_t InstCount = - static_cast(R.first->TLASPtr->Instances.size()); - llvm::SmallVector Contributions; - Contributions.reserve(InstCount); - for (const auto &Inst : R.first->TLASPtr->Instances) - Contributions.push_back(Inst.InstanceContributionToHitGroupIndex & - 0xFFFFFFu); - const BufferCreateDesc Desc = BufferCreateDesc::uploadBuffer(); - auto ContribBufOrErr = createBufferWithData( - *IS.CB->Dev, "AS-Contributions", Desc, Contributions.data(), - InstCount * sizeof(uint32_t), nullptr, nullptr); - if (!ContribBufOrErr) - return ContribBufOrErr.takeError(); - auto *MTLContrib = llvm::cast(ContribBufOrErr->get()); - auto HeaderBufOrErr = IS.CB->Dev->createBuffer( - "AS-Header", Desc, - sizeof(IRRaytracingAccelerationStructureGPUHeader)); - if (!HeaderBufOrErr) - return HeaderBufOrErr.takeError(); - auto *MTLHeader = llvm::cast(HeaderBufOrErr->get()); - IRRaytracingSetAccelerationStructure( - static_cast(MTLHeader->Buf->contents()), - MTLAS->AccelStruct->gpuResourceID(), - static_cast(MTLContrib->Buf->contents()), - MTLContrib->Buf->gpuAddress(), Contributions.data(), InstCount); - - IRDescriptorTableSetAccelerationStructure( - IS.DescHeap->getEntryHandle(HeapIndex), - MTLHeader->Buf->gpuAddress()); - - // The shader dereferences the contributions buffer through the - // header, so both must be resident at dispatch. - IS.ASDescriptorBuffers.push_back(std::move(*HeaderBufOrErr)); - IS.ASDescriptorBuffers.push_back(std::move(*ContribBufOrErr)); - HeapIndex += R.first->getArraySize(); - continue; - } - switch (getDescriptorKind(R.first->Kind)) { - case DescriptorKind::SRV: - HeapIndex = bindSRV(*(R.first), IS, HeapIndex, R.second); - break; - case DescriptorKind::UAV: - HeapIndex = bindUAV(*(R.first), IS, HeapIndex, R.second); - break; - case DescriptorKind::CBV: - HeapIndex = bindCBV(*(R.first), IS, HeapIndex, R.second); - break; - case DescriptorKind::SAMPLER: - llvm_unreachable("Not implemented yet."); - } - } - } - - if (P.isTraditionalRaster() && P.Bindings.VertexBufferPtr) { - const CPUBuffer *VBuffer = P.Bindings.VertexBufferPtr; - - BufferCreateDesc BufDesc = {}; - BufDesc.Location = MemoryLocation::CpuToGpu; - BufDesc.Usage = BufferUsage::VertexBuffer; - auto BufOrErr = createBufferWithData(*this, "VertexBuffer", BufDesc, - VBuffer->Data[0].get(), - VBuffer->size(), nullptr, nullptr); - if (!BufOrErr) - return BufOrErr.takeError(); - IS.VB = std::move(*BufOrErr); - llvm::outs() << "Vertex buffer created.\n"; - } - return llvm::Error::success(); - } - - llvm::Error createComputeCommands(Pipeline &P, InvocationState &IS) { auto EncoderOrErr = IS.CB->createComputeEncoder(); if (!EncoderOrErr) return EncoderOrErr.takeError(); - auto &Encoder = llvm::cast(*EncoderOrErr.get()); - MTL::ComputeCommandEncoder *NativeEncoder = Encoder.getNative(); - - const auto &PS = llvm::cast(IS.Pipeline.get()); - MTLGPUDescriptorHandle Handle = {}; - if (IS.DescHeap) { - IS.DescHeap->bind(NativeEncoder); - Handle = IS.DescHeap->getGPUDescriptorHandleForHeapStart(); - } + auto Encoder = std::move(*EncoderOrErr); - for (uint32_t Idx = 0u; Idx < P.Sets.size(); ++Idx) { - PS->ArgBuffer->setRootDescriptorTable(Idx, Handle); - Handle.addOffset(P.Sets[Idx].Resources.size()); - } + Encoder->bindComputeDescriptorSets(*IS.Pipeline, *DescSets); - PS->ArgBuffer->bind(NativeEncoder); - for (const auto &Table : IS.DescTables) - for (const auto &ResPair : Table.Resources) - for (const auto &ResSet : ResPair.second) - NativeEncoder->useResource(ResSet.Resource.get(), - MTL::ResourceUsageRead | - MTL::ResourceUsageWrite); - auto MarkASResident = - [&](const std::unique_ptr &AS) { - auto *MTLAS = llvm::cast(AS.get()); - NativeEncoder->useResource(MTLAS->AccelStruct, - MTL::ResourceUsageRead); - }; - for (auto &AS : IS.BLASes) - MarkASResident(AS); - for (auto &Entry : IS.TLASes) - MarkASResident(Entry.second); - for (auto &B : IS.ASDescriptorBuffers) - NativeEncoder->useResource(llvm::cast(B.get())->Buf, - MTL::ResourceUsageRead); - - if (auto Err = Encoder.dispatch(*IS.Pipeline.get(), - P.DispatchParameters.DispatchGroupCount[0], - P.DispatchParameters.DispatchGroupCount[1], - P.DispatchParameters.DispatchGroupCount[2])) + if (auto Err = Encoder->dispatch( + *IS.Pipeline.get(), P.DispatchParameters.DispatchGroupCount[0], + P.DispatchParameters.DispatchGroupCount[1], + P.DispatchParameters.DispatchGroupCount[2])) return Err; - Encoder.endEncoding(); + Encoder->endEncoding(); return llvm::Error::success(); } - llvm::Error createRayTracingCommands(Pipeline &P, InvocationState &IS) { + llvm::Error createRayTracingCommands(Pipeline &P, SharedInvocationState &IS, + DescriptorPool &Pool) { + IS.CB->bindPool(Pool); + + auto DescSetsOrErr = + buildDescriptorSets(*this, Pool, *IS.Pipeline, IS.DescTables); + if (!DescSetsOrErr) + return DescSetsOrErr.takeError(); + auto DescSets = std::move(*DescSetsOrErr); + auto EncoderOrErr = IS.CB->createComputeEncoder(); if (!EncoderOrErr) return EncoderOrErr.takeError(); - auto &Encoder = llvm::cast(*EncoderOrErr.get()); - MTL::ComputeCommandEncoder *NativeEncoder = Encoder.getNative(); - - const auto &RTPSO = - llvm::cast(*IS.Pipeline.get()); - const auto &SBT = llvm::cast(*IS.SBT.get()); - - // Bind the global descriptor heap + top-level argument buffer the same - // way the compute path does; the raygen kernel and any visible-function - // callees consume them at the same slots (kIRDescriptorHeapBindPoint and - // kIRArgumentBufferBindPoint). - MTLGPUDescriptorHandle Handle = {}; - if (IS.DescHeap) { - IS.DescHeap->bind(NativeEncoder); - Handle = IS.DescHeap->getGPUDescriptorHandleForHeapStart(); - } - for (uint32_t Idx = 0u; Idx < P.Sets.size(); ++Idx) { - RTPSO.ArgBuffer->setRootDescriptorTable(Idx, Handle); - Handle.addOffset(P.Sets[Idx].Resources.size()); - } - RTPSO.ArgBuffer->bind(NativeEncoder); - - // Populate the per-dispatch IRDispatchRaysArgument: SBT region addresses - // (RayGen / Miss / HitGroup / Callable), GPU pointers to the global - // root-signature argument buffer + descriptor heaps, plus resource IDs - // for the visible / intersection function tables. The raygen kernel - // reads this struct from the buffer bound at kIRRayDispatchArgumentsBind- - // Point and any visible-function callees inherit it through the same - // pointer. - IRDispatchRaysArgument Args{}; - Args.DispatchRaysDesc.RayGenerationShaderRecord = SBT.RayGenRegion; - Args.DispatchRaysDesc.MissShaderTable = SBT.MissRegion; - Args.DispatchRaysDesc.HitGroupTable = SBT.HitGroupRegion; - Args.DispatchRaysDesc.CallableShaderTable = SBT.CallableRegion; - Args.DispatchRaysDesc.Width = P.DispatchParameters.DispatchGroupCount[0]; - Args.DispatchRaysDesc.Height = P.DispatchParameters.DispatchGroupCount[1]; - Args.DispatchRaysDesc.Depth = P.DispatchParameters.DispatchGroupCount[2]; - Args.GRS = RTPSO.ArgBuffer->getGPUAddress(); - Args.ResDescHeap = - IS.DescHeap ? IS.DescHeap->getGPUDescriptorHandleForHeapStart().Ptr : 0; - Args.SmpDescHeap = 0; - Args.VisibleFunctionTable = - RTPSO.VFT ? RTPSO.VFT->gpuResourceID() : MTL::ResourceID{0}; - Args.IntersectionFunctionTable = - RTPSO.IFT ? RTPSO.IFT->gpuResourceID() : MTL::ResourceID{0}; - Args.IntersectionFunctionTables = 0; - - const BufferCreateDesc ArgsBufDesc = BufferCreateDesc::uploadBuffer(); - auto ArgsBufOrErr = offloadtest::createBufferWithData( - *IS.CB->Dev, "MTL Dispatch Rays Arguments", ArgsBufDesc, &Args, - sizeof(IRDispatchRaysArgument), nullptr, nullptr); - if (!ArgsBufOrErr) - return ArgsBufOrErr.takeError(); - - auto *MTLArgsBuf = llvm::cast(ArgsBufOrErr->get()); - IS.CB->KeepAliveOwned.push_back(std::move(*ArgsBufOrErr)); - - NativeEncoder->setBuffer(MTLArgsBuf->Buf, 0, - kIRRayDispatchArgumentsBindPoint); - NativeEncoder->useResource(MTLArgsBuf->Buf, MTL::ResourceUsageRead); - - // Mark every dispatch-side resource resident: descriptor-table bundles, - // acceleration structures + their irconverter header/contribution - // buffers (so RayQuery/TraceRay can read them), the SBT buffer (the - // raygen kernel dereferences SBT addresses), and the visible / - // intersection function tables. - for (const auto &Table : IS.DescTables) - for (const auto &ResPair : Table.Resources) - for (const auto &ResSet : ResPair.second) - NativeEncoder->useResource(ResSet.Resource.get(), - MTL::ResourceUsageRead | - MTL::ResourceUsageWrite); - auto MarkASResident = - [&](std::unique_ptr &AS) { - auto *MTLAS = llvm::cast(AS.get()); - NativeEncoder->useResource(MTLAS->AccelStruct, - MTL::ResourceUsageRead); - }; - for (auto &AS : IS.BLASes) - MarkASResident(AS); - for (auto &Entry : IS.TLASes) - MarkASResident(Entry.second); - for (auto &B : IS.ASDescriptorBuffers) - NativeEncoder->useResource(llvm::cast(B.get())->Buf, - MTL::ResourceUsageRead); - if (SBT.Buffer) - NativeEncoder->useResource(SBT.Buffer, MTL::ResourceUsageRead); - if (RTPSO.VFT) - NativeEncoder->useResource(RTPSO.VFT, MTL::ResourceUsageRead); - if (RTPSO.IFT) - NativeEncoder->useResource(RTPSO.IFT, MTL::ResourceUsageRead); + auto Encoder = std::move(*EncoderOrErr); + + Encoder->bindRayTracingDescriptorSets(*IS.Pipeline, *DescSets); if (auto Err = - Encoder.dispatchRays(*IS.Pipeline.get(), *IS.SBT.get(), - P.DispatchParameters.DispatchGroupCount[0], - P.DispatchParameters.DispatchGroupCount[1], - P.DispatchParameters.DispatchGroupCount[2])) + Encoder->dispatchRays(*IS.Pipeline.get(), *IS.SBT.get(), + P.DispatchParameters.DispatchGroupCount[0], + P.DispatchParameters.DispatchGroupCount[1], + P.DispatchParameters.DispatchGroupCount[2])) return Err; - Encoder.endEncoding(); - return llvm::Error::success(); - } - - llvm::Error createRenderTarget(Pipeline &P, InvocationState &IS) { - if (!P.Bindings.RTargetBufferPtr) - return llvm::createStringError( - std::errc::invalid_argument, - "No render target bound for graphics pipeline."); - const CPUBuffer &OutBuf = *P.Bindings.RTargetBufferPtr; - - auto TexOrErr = offloadtest::createRenderTargetFromCPUBuffer(*this, OutBuf); - if (!TexOrErr) - return TexOrErr.takeError(); - - IS.RenderTarget = std::move(*TexOrErr); - - // Create a readback buffer for copying render target data to the CPU. - const BufferCreateDesc BufDesc = BufferCreateDesc::readbackBuffer(); - auto BufOrErr = createBuffer("RTReadback", BufDesc, OutBuf.size()); - if (!BufOrErr) - return BufOrErr.takeError(); - IS.FrameBufferReadback = std::move(*BufOrErr); - + Encoder->endEncoding(); return llvm::Error::success(); } - llvm::Error createDepthStencil(Pipeline &P, InvocationState &IS) { - auto TexOrErr = offloadtest::createDefaultDepthStencilTarget( - *this, P.Bindings.RTargetBufferPtr->OutputProps.Width, - P.Bindings.RTargetBufferPtr->OutputProps.Height); - if (!TexOrErr) - return TexOrErr.takeError(); - IS.DepthStencil = std::move(*TexOrErr); - return llvm::Error::success(); - } + llvm::Error createGraphicsCommands(Pipeline &P, SharedInvocationState &IS, + DescriptorPool &Pool) { + IS.CB->bindPool(Pool); - llvm::Error createGraphicsCommands(Pipeline &P, InvocationState &IS) { - if (auto Err = createRenderTarget(P, IS)) - return Err; - // TODO: Always created for graphics pipelines. Consider making this - // conditional on the pipeline definition. - if (auto Err = createDepthStencil(P, IS)) - return Err; + auto DescSetsOrErr = + buildDescriptorSets(*this, Pool, *IS.Pipeline, IS.DescTables); + if (!DescSetsOrErr) + return DescSetsOrErr.takeError(); + auto DescSets = std::move(*DescSetsOrErr); const uint64_t Width = IS.RenderTarget->getDesc().Width; const uint64_t Height = IS.RenderTarget->getDesc().Height; @@ -1939,115 +1827,48 @@ class MTLDevice : public offloadtest::Device { auto EncOrErr = IS.CB->createRenderEncoder(BeginDesc); if (!EncOrErr) return EncOrErr.takeError(); - auto &Encoder = *EncOrErr.get(); + auto Encoder = std::move(*EncOrErr); - { - auto &MTLEncoder = llvm::cast(Encoder); - const auto &PS = llvm::cast(IS.Pipeline.get()); - auto *CmdEncoder = MTLEncoder.getNative(); - if (IS.DescHeap) { - IS.DescHeap->bind(CmdEncoder); - // NOTE: This code assumes 1 descriptor set (D3D12 backend also assumes - // this) - PS->ArgBuffer->setRootDescriptorTable( - 0, IS.DescHeap->getGPUDescriptorHandleForHeapStart()); - } - PS->ArgBuffer->bind(CmdEncoder); - for (const auto &Table : IS.DescTables) - for (const auto &ResPair : Table.Resources) - for (const auto &ResSet : ResPair.second) - CmdEncoder->useResource(ResSet.Resource.get(), - MTL::ResourceUsageRead | - MTL::ResourceUsageWrite); - } + Encoder->bindDescriptorSets(*IS.Pipeline, *DescSets); Viewport VP; VP.Width = static_cast(Width); VP.Height = static_cast(Height); - Encoder.setViewport(VP); + Encoder->setViewport(VP); ScissorRect Scissor; Scissor.Width = static_cast(Width); Scissor.Height = static_cast(Height); - Encoder.setScissor(Scissor); + Encoder->setScissor(Scissor); if (P.isTraditionalRaster()) { if (IS.VB) - Encoder.setVertexBuffer(0, IS.VB.get(), 0, - P.Bindings.getVertexStride()); + Encoder->setVertexBuffer(0, IS.VB.get(), 0, + P.Bindings.getVertexStride()); if (auto Err = - Encoder.drawInstanced(*IS.Pipeline.get(), P.getVertexCount(), - /*InstanceCount=*/1)) + Encoder->drawInstanced(*IS.Pipeline.get(), P.getVertexCount(), + /*InstanceCount=*/1)) return Err; } else { - if (auto Err = Encoder.dispatchMesh( + if (auto Err = Encoder->dispatchMesh( *IS.Pipeline.get(), P.DispatchParameters.DispatchGroupCount[0], P.DispatchParameters.DispatchGroupCount[1], P.DispatchParameters.DispatchGroupCount[2])) return Err; } - Encoder.endEncoding(); + Encoder->endEncoding(); - // Blit the render target into the readback buffer for CPU access. - auto &FBTex = llvm::cast(*IS.RenderTarget); - auto &FBReadback = llvm::cast(*IS.FrameBufferReadback); - MTL::BlitCommandEncoder *Blit = IS.CB->CmdBuffer->blitCommandEncoder(); - const size_t ElemSize = getFormatSizeInBytes(FBTex.Desc.Fmt); - const size_t RowBytes = Width * ElemSize; - Blit->copyFromTexture(FBTex.Tex, 0, 0, MTL::Origin(0, 0, 0), - MTL::Size(Width, Height, 1), FBReadback.Buf, 0, - RowBytes, 0); - Blit->endEncoding(); - - return llvm::Error::success(); - } - - llvm::Error copyBack(Pipeline &P, InvocationState &IS) { - auto MemCpyBack = [](ResourcePair &Pair) -> llvm::Error { - const Resource &R = *Pair.first; - if (!R.isReadWrite()) - return llvm::Error::success(); - - const CPUBuffer &B = *R.BufferPtr; - auto *RSIt = Pair.second.begin(); - auto *DataIt = B.Data.begin(); - for (; RSIt != Pair.second.end() && DataIt != B.Data.end(); - ++RSIt, ++DataIt) { - if (R.isRaw()) { - MTL::Buffer *Buf = static_cast(RSIt->Resource.get()); - memcpy(DataIt->get(), Buf->contents(), Buf->length()); - } else { - MTL::Texture *Tex = static_cast(RSIt->Resource.get()); - const uint64_t Width = R.isTexture() ? B.OutputProps.Width - : R.size() / R.getElementSize(); - const uint64_t Height = R.isTexture() ? B.OutputProps.Height : 1; - Tex->getBytes(DataIt->get(), Width * R.getElementSize(), - MTL::Region(0, 0, Width, Height), 0); - } - } - - return llvm::Error::success(); - }; - - for (auto &Table : IS.DescTables) - for (auto &R : Table.Resources) - if (auto Err = MemCpyBack(R)) - return Err; - - if (P.isRaster()) { - auto &FBReadback = llvm::cast(*IS.FrameBufferReadback); - auto *RT = P.Bindings.RTargetBufferPtr; - RT->copyFromTexture(FBReadback.Buf->contents(), RT->getImageRowBytes()); - } return llvm::Error::success(); } public: MTLDevice(MTL::Device *D, MTL::CommandQueue *Q, std::unique_ptr SubmitFence) - : Device(D), GraphicsQueue(Q, std::move(SubmitFence)) { + : Device(D), + ResidencyTracker(llvm::cantFail(MetalResidencyTracker::create(D))), + GraphicsQueue(Q, std::move(SubmitFence), ResidencyTracker) { Description = Device->name()->utf8String(); } const Capabilities &getCapabilities() override { @@ -2079,8 +1900,9 @@ class MTLDevice : public offloadtest::Device { IRRootSignaturePtr RootSig; std::unique_ptr ArgBuffer; - if (auto Err = - createRootSignature(BD, /*IsGraphics=*/false, RootSig, ArgBuffer)) + DescriptorSetsLayout Layout; + if (auto Err = createRootSignature(BD, /*IsGraphics=*/false, RootSig, + ArgBuffer, Layout)) return Err; // Configure the irconverter ray tracing pipeline. Raygen is compiled as a @@ -2105,7 +1927,8 @@ class MTLDevice : public offloadtest::Device { RTConfig.get(), IRIntersectionFunctionCompilationVisibleFunction); auto State = std::make_unique( - Name, std::move(RootSig), std::move(ArgBuffer)); + Name, std::move(RootSig), std::move(ArgBuffer), std::move(Layout), + ResidencyTracker); // Compile each entry point. Raygen lands in `RaygenFn` (becomes the // compute function); everything else gets linked in via LinkedFunctions @@ -2255,6 +2078,13 @@ class MTLDevice : public offloadtest::Device { State->VFT = VFT; } + ResidencyTracker->withLock([&](MTL::ResidencySet *RS) { + if (State->VFT) + RS->addAllocation(State->VFT); + if (State->IFT) + RS->addAllocation(State->IFT); + }); + State->ComputePipeline = PSO; State->ThreadsPerGroup = RaygenThreadsPerGroup; return State; @@ -2342,7 +2172,7 @@ class MTLDevice : public offloadtest::Device { return std::make_unique( Buffer, MakeRange(Layout.RayGen), MakeRangeAndStride(Layout.Miss), MakeRangeAndStride(Layout.HitGroup), - MakeRangeAndStride(Layout.Callable)); + MakeRangeAndStride(Layout.Callable), ResidencyTracker); } llvm::Expected> @@ -2370,12 +2200,30 @@ class MTLDevice : public offloadtest::Device { std::errc::not_supported, "Metal backend does not support sparse memory backing."); - MTL::Buffer *Buf = Device->newBuffer( - SizeInBytes, getMetalBufferResourceOptions(Desc.Location)); - if (!Buf) - return llvm::createStringError(std::errc::not_enough_memory, - "Failed to create Metal buffer."); - return std::make_unique(Buf, Name, Desc, SizeInBytes); + MTL::Resource *Res = nullptr; + if (Desc.AccessType == BufferShaderAccessType::Typed) { + MTL::TextureDescriptor *TDesc = + MTL::TextureDescriptor::textureBufferDescriptor( + getMetalPixelFormat(Desc.AccessTypeParams.Fmt), + SizeInBytes / getFormatSizeInBytes(Desc.AccessTypeParams.Fmt), + getMetalBufferResourceOptions(Desc.Location), + MTL::ResourceUsageRead | MTL::ResourceUsageWrite); + + Res = Device->newTexture(TDesc); + if (!Res) + return llvm::createStringError( + std::errc::not_enough_memory, + "Failed to create Metal typed buffer (texture)."); + } else { + Res = Device->newBuffer(SizeInBytes, + getMetalBufferResourceOptions(Desc.Location)); + if (!Res) + return llvm::createStringError(std::errc::not_enough_memory, + "Failed to create Metal buffer."); + } + + return std::make_unique(Res, Name, Desc, SizeInBytes, + ResidencyTracker); } llvm::Expected> @@ -2394,7 +2242,12 @@ class MTLDevice : public offloadtest::Device { if (!Tex) return llvm::createStringError(std::errc::not_enough_memory, "Failed to create Metal texture."); - return std::make_unique(Tex, Name, Desc); + return std::make_unique(Tex, Name, Desc, ResidencyTracker); + } + + llvm::Expected> + createSampler(std::string, const SamplerCreateDesc &) override { + return llvm::createStringError("createSampler is unimplemented on Metal."); } uint32_t getTextureUploadRowStrideInBytes( @@ -2402,6 +2255,12 @@ class MTLDevice : public offloadtest::Device { return Desc.Width * getFormatSizeInBytes(Desc.Fmt); } + TextureUploadLayout + getTextureUploadLayout(const TextureCreateDesc &Desc) const override { + // copyBufferToTexture consumes a tightly-packed staging buffer. + return computeTightTextureUploadLayout(Desc); + } + llvm::Expected> createCommandBuffer() override { auto CBOrErr = MTLCommandBuffer::create(GraphicsQueue.Queue); @@ -2421,8 +2280,9 @@ class MTLDevice : public offloadtest::Device { ShaderContainer CS) override { IRRootSignaturePtr RootSig; std::unique_ptr ArgBuffer; + DescriptorSetsLayout Layout; if (auto Err = createRootSignature(BindingsDesc, /*IsGraphics=*/false, - RootSig, ArgBuffer)) + RootSig, ArgBuffer, Layout)) return Err; auto MetalIR = convertToMetalIR(Stages::Compute, /*IsGraphics=*/false, @@ -2464,7 +2324,8 @@ class MTLDevice : public offloadtest::Device { IRShaderReflectionReleaseComputeInfo(&Info); return std::make_unique( - Name, std::move(RootSig), std::move(ArgBuffer), PSO, ThreadsPerGroup); + Name, std::move(RootSig), std::move(ArgBuffer), std::move(Layout), PSO, + ThreadsPerGroup); } llvm::Expected> @@ -2492,8 +2353,9 @@ class MTLDevice : public offloadtest::Device { IRRootSignaturePtr RootSig; std::unique_ptr ArgBuffer; + DescriptorSetsLayout Layout; if (auto Err = createRootSignature(BindingsDesc, /*IsGraphics=*/true, - RootSig, ArgBuffer)) + RootSig, ArgBuffer, Layout)) return Err; NS::Error *Error = nullptr; @@ -2651,9 +2513,9 @@ class MTLDevice : public offloadtest::Device { MTL::DepthStencilState *DSState = Device->newDepthStencilState(DSDesc); DSDesc->release(); - return std::make_unique(Name, std::move(RootSig), - std::move(ArgBuffer), PSO, - DSState, MTL::CullModeNone); + return std::make_unique( + Name, std::move(RootSig), std::move(ArgBuffer), std::move(Layout), PSO, + DSState, MTL::CullModeNone); } llvm::Expected> createMeshShaderRasterPipeline( @@ -2661,8 +2523,9 @@ class MTLDevice : public offloadtest::Device { const MeshShaderRasterPipelineCreateDesc &Desc) override { IRRootSignaturePtr RootSig; std::unique_ptr ArgBuffer; + DescriptorSetsLayout Layout; if (auto Err = createRootSignature(BindingsDesc, /*IsGraphics=*/true, - RootSig, ArgBuffer)) + RootSig, ArgBuffer, Layout)) return Err; NS::Error *Error = nullptr; @@ -2781,8 +2644,8 @@ class MTLDevice : public offloadtest::Device { } return std::make_unique( - Name, std::move(RootSig), std::move(ArgBuffer), PSO, DSState, - MTL::CullModeNone, MeshTGSize, ObjectTGSize); + Name, std::move(RootSig), std::move(ArgBuffer), std::move(Layout), PSO, + DSState, MTL::CullModeNone, MeshTGSize, ObjectTGSize); } llvm::Expected @@ -2801,7 +2664,7 @@ class MTLDevice : public offloadtest::Device { auto *TD = MTL::AccelerationStructureTriangleGeometryDescriptor::alloc()->init(); auto *VB = llvm::cast(T.VertexBuffer); - TD->setVertexBuffer(VB->Buf); + TD->setVertexBuffer(VB->getBufferPtr()); TD->setVertexBufferOffset(T.VertexBufferOffset); TD->setVertexStride(T.VertexStride); TD->setVertexFormat(getMetalPositionFormat(T.VertexFormat)); @@ -2809,7 +2672,7 @@ class MTLDevice : public offloadtest::Device { : T.VertexCount / 3); if (T.IndexBuffer) { auto *IB = llvm::cast(T.IndexBuffer); - TD->setIndexBuffer(IB->Buf); + TD->setIndexBuffer(IB->getBufferPtr()); TD->setIndexBufferOffset(T.IndexBufferOffset); TD->setIndexType(getMetalIndexType(T.IdxFormat)); } @@ -2840,7 +2703,7 @@ class MTLDevice : public offloadtest::Device { MTL::AccelerationStructureBoundingBoxGeometryDescriptor::alloc() ->init(); auto *BB = llvm::cast(A.AABBBuffer); - AD->setBoundingBoxBuffer(BB->Buf); + AD->setBoundingBoxBuffer(BB->getBufferPtr()); AD->setBoundingBoxBufferOffset(A.AABBBufferOffset); AD->setBoundingBoxStride(A.AABBStride); AD->setBoundingBoxCount(A.AABBCount); @@ -2872,22 +2735,6 @@ class MTLDevice : public offloadtest::Device { Sizes.refitScratchBufferSize}; } - llvm::Expected> - allocateAS(const AccelerationStructureSizes &Sizes, const char *Kind) { - if (!Device->supportsRaytracing()) - return llvm::createStringError( - std::errc::not_supported, - "Ray tracing is not supported on this device."); - - MTL::AccelerationStructure *AS = - Device->newAccelerationStructure(Sizes.ResultDataMaxSizeInBytes); - if (!AS) - return llvm::createStringError( - std::make_error_code(std::errc::not_enough_memory), - "Failed to create Metal " + llvm::Twine(Kind) + "."); - return std::make_unique(AS, Sizes); - } - llvm::Expected getTLASBuildSizes(uint32_t InstanceCount) override { if (!Device->supportsRaytracing()) @@ -2916,27 +2763,111 @@ class MTLDevice : public offloadtest::Device { llvm::Expected> createBLAS(const AccelerationStructureSizes &Sizes) override { - return allocateAS(Sizes, "BLAS"); + if (!Device->supportsRaytracing()) + return llvm::createStringError( + std::errc::not_supported, + "Ray tracing is not supported on this device."); + + MTL::AccelerationStructure *AS = + Device->newAccelerationStructure(Sizes.ResultDataMaxSizeInBytes); + if (!AS) + return llvm::createStringError( + std::make_error_code(std::errc::not_enough_memory), + "Failed to create Metal BLAS."); + + return std::make_unique(AS, Sizes, + ResidencyTracker); } llvm::Expected> - createTLAS(const AccelerationStructureSizes &Sizes) override { - return allocateAS(Sizes, "TLAS"); + createTLAS(const AccelerationStructureSizes &Sizes, + uint32_t InstanceCount) override { + if (!Device->supportsRaytracing()) + return llvm::createStringError( + std::errc::not_supported, + "Ray tracing is not supported on this device."); + + // TODO(manon): We would prefer these to live in GPUOnly memory in the + // future. + const BufferCreateDesc ContribBufferDesc = + BufferCreateDesc::gpuOnlyStorage(); + auto ContribBufferOrErr = + createBuffer("AS-Contributions", ContribBufferDesc, + InstanceCount * sizeof(uint32_t)); + if (!ContribBufferOrErr) + return ContribBufferOrErr.takeError(); + auto ContribBuffer = std::move(*ContribBufferOrErr); + + const MTLBuffer &ContribBufferMTL = + llvm::cast(*ContribBuffer.get()); + + MTL::AccelerationStructure *AS = + Device->newAccelerationStructure(Sizes.ResultDataMaxSizeInBytes); + if (!AS) + return llvm::createStringError( + std::make_error_code(std::errc::not_enough_memory), + "Failed to create Metal TLAS."); + + IRRaytracingAccelerationStructureGPUHeader Header = {}; + Header.accelerationStructureID = AS->gpuResourceID()._impl; + Header.addressOfInstanceContributions = + ContribBufferMTL.getBufferPtr()->gpuAddress(); + + const BufferCreateDesc HeaderBufferDesc = BufferCreateDesc::uploadBuffer(); + auto HeaderBufOrErr = + createBufferWithData(*this, "AS-Header", HeaderBufferDesc, &Header, + sizeof(Header), nullptr, nullptr); + if (!HeaderBufOrErr) { + AS->release(); + return HeaderBufOrErr.takeError(); + } + auto HeaderBuffer = std::move(*HeaderBufOrErr); + + return std::make_unique( + AS, Sizes, ResidencyTracker, std::move(HeaderBuffer), + std::move(ContribBuffer)); + } + + llvm::Expected> + createDescriptorPool() override { + return MetalDescriptorPool::create(Device, ResidencyTracker); + } + + llvm::Expected> + createDescriptorSetsBuilder(DescriptorPool &Pool, + const PipelineState &Pipeline) override { + MetalDescriptorPool &PoolMTL = llvm::cast(Pool); + const MTLPipelineState &PipelineMTL = + llvm::cast(Pipeline); + + llvm::SmallVector Sets; + for (const auto &Counts : PipelineMTL.Layout.Sets) { + MetalDescriptorSet Set = {}; + if (Counts.DescriptorCount > 0) + PoolMTL.allocateDescriptors(Counts.DescriptorCount, Set.CSUHandle, + Set.CSUHandleGPU); + Sets.push_back(Set); + } + return std::make_unique(std::move(Sets)); } llvm::Error executeProgram(Pipeline &P) override { - InvocationState IS; + auto DescriptorPoolOrErr = createDescriptorPool(); + if (!DescriptorPoolOrErr) + return DescriptorPoolOrErr.takeError(); + auto DescriptorPool = std::move(*DescriptorPoolOrErr); - auto CBOrErr = MTLCommandBuffer::create(GraphicsQueue.Queue); + SharedInvocationState IS; + + NS::AutoreleasePool *Pool = NS::AutoreleasePool::alloc()->init(); + auto PoolScope = llvm::scope_exit([&] { Pool->release(); }); + + auto CBOrErr = createCommandBuffer(); if (!CBOrErr) return CBOrErr.takeError(); IS.CB = std::move(*CBOrErr); - IS.CB->Dev = this; - - if (auto Err = createDescriptorHeap(P, IS)) - return Err; - if (auto Err = createBuffers(P, IS)) + if (auto Err = createResources(*this, P, IS)) return Err; if (!P.AccelStructs.BLAS.empty() || !P.AccelStructs.TLAS.empty()) { @@ -2981,7 +2912,7 @@ class MTLDevice : public offloadtest::Device { IS.Pipeline = std::move(*PipelineStateOrErr); llvm::outs() << "Compute Pipeline created.\n"; - if (auto Err = createComputeCommands(P, IS)) + if (auto Err = createComputeCommands(P, IS, *DescriptorPool)) return Err; } else if (P.isRaster()) { auto FormatOrErr = toFormat(P.Bindings.RTargetBufferPtr->Format, @@ -3062,7 +2993,7 @@ class MTLDevice : public offloadtest::Device { return RenderPassOrErr.takeError(); IS.RenderPass = std::move(*RenderPassOrErr); - if (auto Err = createGraphicsCommands(P, IS)) + if (auto Err = createGraphicsCommands(P, IS, *DescriptorPool)) return Err; } else if (P.isRayTracing()) { if (P.Shaders.empty() || !P.SBT || !P.RTConfig) @@ -3094,10 +3025,32 @@ class MTLDevice : public offloadtest::Device { IS.SBT = std::move(*SBTOrErr); llvm::outs() << "Shader Binding Table created.\n"; - if (auto Err = createRayTracingCommands(P, IS)) + if (auto Err = createRayTracingCommands(P, IS, *DescriptorPool)) + return Err; + } + + auto EncoderOrErr = IS.CB->createComputeEncoder(); + if (!EncoderOrErr) + return EncoderOrErr.takeError(); + auto ReadbackEncoder = std::move(*EncoderOrErr); + + if (IS.RenderTarget) { + if (auto Err = ReadbackEncoder->copyTextureToBuffer(*IS.RenderTarget, + *IS.RTReadback)) return Err; } + for (auto &Table : IS.DescTables) + for (auto &R : Table.Resources) + if (auto Err = copyBackResource(*ReadbackEncoder, R)) + return Err; + + for (auto &R : IS.RootResources) + if (auto Err = copyBackResource(*ReadbackEncoder, R)) + return Err; + + ReadbackEncoder->endEncoding(); + auto SubmitResult = GraphicsQueue.submit(std::move(IS.CB)); if (!SubmitResult) return SubmitResult.takeError(); @@ -3105,8 +3058,9 @@ class MTLDevice : public offloadtest::Device { if (auto Err = SubmitResult->waitForCompletion()) return Err; - if (auto Err = copyBack(P, IS)) + if (auto Err = readBack(*this, P, IS)) return Err; + llvm::outs() << "Read data back.\n"; return llvm::Error::success(); } @@ -3151,9 +3105,6 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef Items) { std::errc::not_supported, "Ray tracing is not supported on this Metal device."); - if (auto Err = ensureASEncoder()) - return Err; - for (const auto &Item : Items) { MetalAccelerationStructure *AS = nullptr; MTL::AccelerationStructureDescriptor *Desc = nullptr; @@ -3171,7 +3122,7 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef Items) { MTL::AccelerationStructureTriangleGeometryDescriptor::alloc() ->init(); auto *VB = llvm::cast(T.VertexBuffer); - TD->setVertexBuffer(VB->Buf); + TD->setVertexBuffer(VB->getBufferPtr()); TD->setVertexBufferOffset(T.VertexBufferOffset); TD->setVertexStride(T.VertexStride); TD->setVertexFormat(getMetalPositionFormat(T.VertexFormat)); @@ -3179,7 +3130,7 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef Items) { : T.VertexCount / 3); if (T.IndexBuffer) { auto *IB = llvm::cast(T.IndexBuffer); - TD->setIndexBuffer(IB->Buf); + TD->setIndexBuffer(IB->getBufferPtr()); TD->setIndexBufferOffset(T.IndexBufferOffset); TD->setIndexType(getMetalIndexType(T.IdxFormat)); } @@ -3195,7 +3146,7 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef Items) { MTL::AccelerationStructureBoundingBoxGeometryDescriptor::alloc() ->init(); auto *BB = llvm::cast(A.AABBBuffer); - AD->setBoundingBoxBuffer(BB->Buf); + AD->setBoundingBoxBuffer(BB->getBufferPtr()); AD->setBoundingBoxBufferOffset(A.AABBBufferOffset); AD->setBoundingBoxStride(A.AABBStride); AD->setBoundingBoxCount(A.AABBCount); @@ -3236,11 +3187,26 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef Items) { InstanceASIdx.push_back(Idx); } + const BufferCreateDesc UploadDesc = BufferCreateDesc::uploadBuffer(); + const uint32_t ContribBufferSize = AS->ContribBuffer->getSizeInBytes(); + auto ContribUploadBufferOrErr = CB->Dev->createBuffer( + "Contrib Upload Buffer", UploadDesc, ContribBufferSize); + if (!ContribUploadBufferOrErr) + return ContribUploadBufferOrErr.takeError(); + auto ContribUploadBuffer = std::move(*ContribUploadBufferOrErr); + + auto ContribPtrOrErr = ContribUploadBuffer->map(); + if (!ContribPtrOrErr) + return ContribPtrOrErr.takeError(); + uint32_t *ContribPtr = static_cast(*ContribPtrOrErr); + // Pack instance descriptors. Layout differs from VK/DX12: 32-byte // entries with an index instead of a GPU address. llvm::SmallVector Native; + llvm::SmallVector HitContributions; Native.reserve(TLAS->Instances.size()); + HitContributions.reserve(TLAS->Instances.size()); for (size_t I = 0; I < TLAS->Instances.size(); ++I) { const auto &Src = TLAS->Instances[I]; MTL::AccelerationStructureUserIDInstanceDescriptor D = {}; @@ -3258,22 +3224,31 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef Items) { D.accelerationStructureIndex = InstanceASIdx[I]; D.userID = Src.InstanceID; Native.push_back(D); + ContribPtr[I] = Src.InstanceContributionToHitGroupIndex & + 0xffffff; // cut-off to 24-bit to match dx12 and vulkan. } + ContribUploadBuffer->unmap(); + + if (auto Err = this->copyBufferToBuffer(*ContribUploadBuffer.get(), 0, + *AS->ContribBuffer.get(), 0, + ContribBufferSize)) + return Err; + + CB->KeepAliveOwned.push_back(std::move(ContribUploadBuffer)); + const size_t InstByteSize = Native.size() * sizeof(MTL::AccelerationStructureUserIDInstanceDescriptor); - const BufferCreateDesc UploadDesc = BufferCreateDesc::uploadBuffer(); auto InstBufOrErr = offloadtest::createBufferWithData( *CB->Dev, "TLAS-Instances", UploadDesc, Native.data(), InstByteSize, nullptr, nullptr); if (!InstBufOrErr) return InstBufOrErr.takeError(); - auto *MTLInstBuf = llvm::cast(InstBufOrErr->get()); - CB->KeepAliveOwned.push_back(std::move(*InstBufOrErr)); + const MTLBuffer &MTLInstBuf = llvm::cast(*InstBufOrErr->get()); auto *ID = MTL::InstanceAccelerationStructureDescriptor::alloc()->init(); - ID->setInstanceDescriptorBuffer(MTLInstBuf->Buf); + ID->setInstanceDescriptorBuffer(MTLInstBuf.getBufferPtr()); ID->setInstanceCount(TLAS->Instances.size()); ID->setInstanceDescriptorType( MTL::AccelerationStructureInstanceDescriptorTypeUserID); @@ -3283,8 +3258,13 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef Items) { ID->setInstancedAccelerationStructures(BLASArr); Desc = ID; ScratchSize = TLAS->AS->getSizes().ScratchDataSizeInBytes; + + CB->KeepAliveOwned.push_back(std::move(*InstBufOrErr)); } + if (auto Err = ensureASEncoder()) + return Err; + const BufferCreateDesc ScratchDesc = BufferCreateDesc::scratchBuffer(); auto ScratchOrErr = CB->Dev->createBuffer("AS-Scratch", ScratchDesc, ScratchSize); @@ -3296,13 +3276,85 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef Items) { CB->KeepAliveOwned.push_back(std::move(*ScratchOrErr)); insertDebugSignpost("BuildAccelerationStructure"); - ASEnc->buildAccelerationStructure(AS->AccelStruct, Desc, MTLScratch->Buf, - 0); + ASEnc->buildAccelerationStructure(AS->AccelStruct, Desc, + MTLScratch->getBufferPtr(), 0); Desc->release(); } return llvm::Error::success(); } + +MTL::Fence *MTLCommandBuffer::ensureFence() { + if (Fence) + return Fence; + + Fence = Dev->Device->newFence(); + if (!Fence) + llvm::report_fatal_error("Failed to create Metal fence."); + + return Fence; +} + +llvm::Error MTLComputeEncoder::dispatchRays(const PipelineState &PSO, + const ShaderBindingTable &SBT, + uint32_t Width, uint32_t Height, + uint32_t Depth) { + if (!llvm::isa(&PSO)) + return llvm::createStringError( + std::errc::invalid_argument, + "dispatchRays requires a RayTracing PipelineState."); + const auto &RTPSO = llvm::cast(PSO); + const auto &MTLSBT = llvm::cast(SBT); + if (!RTPSO.ComputePipeline) + return llvm::createStringError( + std::errc::invalid_argument, + "RayTracing PipelineState has no compute pipeline state."); + if (auto Err = ensureComputeEncoder()) + return Err; + flushBarrier(); + insertDebugSignpost( + llvm::formatv("DispatchRays [{0},{1},{2}]", Width, Height, Depth).str()); + + // Per-dispatch ray arguments, consumed at kIRRayDispatchArgumentsBindPoint. + IRDispatchRaysArgument Args{}; + Args.DispatchRaysDesc.RayGenerationShaderRecord = MTLSBT.RayGenRegion; + Args.DispatchRaysDesc.MissShaderTable = MTLSBT.MissRegion; + Args.DispatchRaysDesc.HitGroupTable = MTLSBT.HitGroupRegion; + Args.DispatchRaysDesc.CallableShaderTable = MTLSBT.CallableRegion; + Args.DispatchRaysDesc.Width = Width; + Args.DispatchRaysDesc.Height = Height; + Args.DispatchRaysDesc.Depth = Depth; + Args.GRS = RTPSO.ArgBuffer->getGPUAddress(); + Args.ResDescHeap = + CB->BoundCSUHeapBuffer ? CB->BoundCSUHeapBuffer->gpuAddress() : 0; + Args.SmpDescHeap = 0; + Args.VisibleFunctionTable = + RTPSO.VFT ? RTPSO.VFT->gpuResourceID() : MTL::ResourceID{0}; + Args.IntersectionFunctionTable = + RTPSO.IFT ? RTPSO.IFT->gpuResourceID() : MTL::ResourceID{0}; + Args.IntersectionFunctionTables = 0; + + const BufferCreateDesc ArgsBufDesc = BufferCreateDesc::uploadBuffer(); + auto ArgsBufOrErr = offloadtest::createBufferWithData( + *CB->Dev, "MTL Dispatch Rays Arguments", ArgsBufDesc, &Args, + sizeof(IRDispatchRaysArgument), nullptr, nullptr); + if (!ArgsBufOrErr) + return ArgsBufOrErr.takeError(); + auto *MTLArgsBuf = llvm::cast(ArgsBufOrErr->get()); + CB->KeepAliveOwned.push_back(std::move(*ArgsBufOrErr)); + + ComputeEnc->setBuffer(MTLArgsBuf->getBufferPtr(), 0, + kIRRayDispatchArgumentsBindPoint); + + ComputeEnc->setComputePipelineState(RTPSO.ComputePipeline); + + // dispatchThreads automatically takes care of bounds checking. + const MTL::Size GridSize(Width, Height, Depth); + ComputeEnc->dispatchThreads(GridSize, RTPSO.ThreadsPerGroup); + addBarrierScope(MTL::BarrierScopeBuffers | MTL::BarrierScopeTextures); + return llvm::Error::success(); +} + } // namespace llvm::Error offloadtest::initializeMetalDevices( diff --git a/lib/API/MTL/MTLResources.h b/lib/API/MTL/MTLResources.h index 02a1d3625..d54b8a558 100644 --- a/lib/API/MTL/MTLResources.h +++ b/lib/API/MTL/MTLResources.h @@ -85,12 +85,16 @@ inline MTL::PixelFormat getMetalPixelFormat(Format Format) { return MTL::PixelFormatRGBA32Uint; case Format::RGBA32Float: return MTL::PixelFormatRGBA32Float; - // Metal has no 64-bit-per-channel pixel formats. case Format::R64Uint: + return MTL::PixelFormatRG32Uint; // Metal has no R64, expects R32G32 case Format::R64Sint: + return MTL::PixelFormatRG32Sint; // Metal has no R64, expects R32G32 case Format::RG64Uint: + return MTL::PixelFormatRGBA32Uint; // Metal has no RG64, expects + // R32G32B32A32 case Format::RG64Sint: - llvm_unreachable("64-bit formats have no Metal pixel format equivalent"); + return MTL::PixelFormatRGBA32Sint; // Metal has no RG64, expects + // R32G32B32A32 case Format::D32Float: return MTL::PixelFormatDepth32Float; case Format::D32FloatS8Uint: diff --git a/lib/API/MTL/MTLTopLevelArgumentBuffer.cpp b/lib/API/MTL/MTLTopLevelArgumentBuffer.cpp index 063c55053..0ac200db5 100644 --- a/lib/API/MTL/MTLTopLevelArgumentBuffer.cpp +++ b/lib/API/MTL/MTLTopLevelArgumentBuffer.cpp @@ -7,6 +7,7 @@ //===----------------------------------------------------------------------===// #include "MTLTopLevelArgumentBuffer.h" +#include "ResidencyTracker.h" using namespace offloadtest; @@ -26,8 +27,9 @@ static llvm::StringRef getResourceTypeString(IRResourceType Type) { } llvm::Expected> -MTLTopLevelArgumentBuffer::create(MTL::Device *Device, - IRRootSignature *RootSig) { +MTLTopLevelArgumentBuffer::create( + MTL::Device *Device, IRRootSignature *RootSig, + std::shared_ptr ResidencyTracker) { if (!Device) return llvm::createStringError(std::errc::invalid_argument, "Invalid MTL::Device pointer."); @@ -40,7 +42,8 @@ MTLTopLevelArgumentBuffer::create(MTL::Device *Device, IRRootSignatureGetResourceCount(RootSig)); // Empty root signature is valid, bind methods will be no-ops if (ResourceLocs.empty()) { - return std::make_unique(ResourceLocs, nullptr); + return std::make_unique(ResourceLocs, nullptr, + nullptr); } IRRootSignatureGetResourceLocations(RootSig, ResourceLocs.data()); @@ -64,13 +67,19 @@ MTLTopLevelArgumentBuffer::create(MTL::Device *Device, return llvm::createStringError( std::errc::not_enough_memory, "Failed to create top-level argument buffer."); - return std::make_unique(std::move(ResourceLocs), - Buffer); + + ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->addAllocation(Buffer); }); + return std::make_unique( + std::move(ResourceLocs), Buffer, std::move(ResidencyTracker)); } MTLTopLevelArgumentBuffer::~MTLTopLevelArgumentBuffer() { - if (Buffer) + if (Buffer) { + ResidencyTracker->withLock( + [&](MTL::ResidencySet *RS) { RS->removeAllocation(Buffer); }); Buffer->release(); + } } bool MTLTopLevelArgumentBuffer::checkIndex(uint32_t Index) const { @@ -175,7 +184,6 @@ void MTLTopLevelArgumentBuffer::bind(MTL::RenderCommandEncoder *Encoder) const { if (!Buffer) return; - Encoder->useResource(Buffer, MTL::ResourceUsageRead); Encoder->setVertexBuffer(Buffer, 0, kIRArgumentBufferBindPoint); Encoder->setFragmentBuffer(Buffer, 0, kIRArgumentBufferBindPoint); } @@ -185,7 +193,6 @@ void MTLTopLevelArgumentBuffer::bind( if (!Buffer) return; - Encoder->useResource(Buffer, MTL::ResourceUsageRead); Encoder->setBuffer(Buffer, 0, kIRArgumentBufferBindPoint); } diff --git a/lib/API/MTL/MTLTopLevelArgumentBuffer.h b/lib/API/MTL/MTLTopLevelArgumentBuffer.h index 50c878603..ac07eb066 100644 --- a/lib/API/MTL/MTLTopLevelArgumentBuffer.h +++ b/lib/API/MTL/MTLTopLevelArgumentBuffer.h @@ -14,11 +14,15 @@ #include "llvm/ADT/SmallVector.h" namespace offloadtest { + +struct MetalResidencyTracker; // Manages a Metal buffer that serves as the top-level argument buffer for // shader resource binding with the explicit root signature layout. class MTLTopLevelArgumentBuffer { +public: llvm::SmallVector ResourceLocs; MTL::Buffer *Buffer = nullptr; + std::shared_ptr ResidencyTracker; bool checkIndex(uint32_t Index) const; bool checkResourceType(uint32_t Index, IRResourceType ExpectedType) const; @@ -27,7 +31,6 @@ class MTLTopLevelArgumentBuffer { template void setResource(uint32_t Index, T Resource) const; -public: /// @brief Creates a MTLTopLevelArgumentBuffer based on the given root /// signature. Empty root signature (zero resources) is allowed, bind methods /// will be no-op in that case. @@ -36,11 +39,14 @@ class MTLTopLevelArgumentBuffer { /// buffer. /// @return Created MTLTopLevelArgumentBuffer or error if creation failed. static llvm::Expected> - create(MTL::Device *Device, IRRootSignature *RootSig); + create(MTL::Device *Device, IRRootSignature *RootSig, + std::shared_ptr ResidencyTracker); - MTLTopLevelArgumentBuffer(llvm::SmallVector ResourceLocs, - MTL::Buffer *Buffer) - : ResourceLocs(std::move(ResourceLocs)), Buffer(Buffer) {} + MTLTopLevelArgumentBuffer( + llvm::SmallVector ResourceLocs, MTL::Buffer *Buffer, + std::shared_ptr ResidencyTracker) + : ResourceLocs(std::move(ResourceLocs)), Buffer(Buffer), + ResidencyTracker(std::move(ResidencyTracker)) {} ~MTLTopLevelArgumentBuffer(); // Binds 32-bit root constant(s) to the argument buffer. diff --git a/lib/API/MTL/ResidencyTracker.cpp b/lib/API/MTL/ResidencyTracker.cpp new file mode 100644 index 000000000..4edf132d8 --- /dev/null +++ b/lib/API/MTL/ResidencyTracker.cpp @@ -0,0 +1,42 @@ +//===- MTL/ResidencyTracker.cpp - Metal Residency Tracker ----------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "ResidencyTracker.h" + +#include "Foundation/Foundation.hpp" +#include "Metal/Metal.hpp" + +namespace offloadtest { + +llvm::Expected> +MetalResidencyTracker::create(MTL::Device *Device) { + std::shared_ptr Tracker = + std::make_shared(); + + MTL::ResidencySetDescriptor *Desc = + MTL::ResidencySetDescriptor::alloc()->init(); + Desc->setInitialCapacity(128); + + NS::Error *Err = nullptr; + + Tracker->ResidencySet = Device->newResidencySet(Desc, &Err); + + Desc->release(); + + if (!Tracker->ResidencySet) + return llvm::createStringError("Failed to create a metal residency set."); + + return Tracker; +} + +MetalResidencyTracker::~MetalResidencyTracker() { + if (ResidencySet) + ResidencySet->release(); +} + +} // namespace offloadtest diff --git a/lib/API/MTL/ResidencyTracker.h b/lib/API/MTL/ResidencyTracker.h new file mode 100644 index 000000000..29c8342ce --- /dev/null +++ b/lib/API/MTL/ResidencyTracker.h @@ -0,0 +1,45 @@ +//===- ResidencyTracker.h - Metal Residency Tracker -----------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_API_MTL_RESIDENCYTRACKER_H +#define OFFLOADTEST_API_MTL_RESIDENCYTRACKER_H + +#include "llvm/Support/Error.h" +#include +#include +#include + +// Forward declarations +namespace MTL { +class Device; +class ResidencySet; +} // namespace MTL + +namespace offloadtest { +struct MetalResidencyTracker { + std::mutex ResidencyMutex; + MTL::ResidencySet *ResidencySet = nullptr; + + template decltype(auto) withLock(ClosureT &&Closure) { + std::lock_guard Lock(ResidencyMutex); + return std::forward(Closure)(ResidencySet); + } + + static llvm::Expected> + create(MTL::Device *Device); + + MetalResidencyTracker() = default; + MetalResidencyTracker(const MetalResidencyTracker &) = delete; + MetalResidencyTracker(MetalResidencyTracker &&) = delete; + MetalResidencyTracker &operator=(const MetalResidencyTracker &) = delete; + MetalResidencyTracker &operator=(MetalResidencyTracker &&) = delete; + ~MetalResidencyTracker(); +}; +} // namespace offloadtest + +#endif // OFFLOADTEST_API_MTL_RESIDENCYTRACKER_H diff --git a/lib/API/Texture.cpp b/lib/API/Texture.cpp index c7d086b74..24f5387fd 100644 --- a/lib/API/Texture.cpp +++ b/lib/API/Texture.cpp @@ -1,6 +1,8 @@ #include "API/Texture.h" #include "API/Device.h" +#include + // Calculate the size in bytes of the texture data given a linear layout // Useful for calculating the size for an upload or readback buffer. size_t @@ -10,3 +12,24 @@ offloadtest::Texture::calculateLinearSizeInBytes(const Device &Dev) const { return (Desc.Height - 1) * Stride + Desc.Width * getFormatSizeInBytes(Desc.Fmt); } + +offloadtest::TextureUploadLayout +offloadtest::computeTightTextureUploadLayout(const TextureCreateDesc &Desc) { + const uint32_t ElementSize = getFormatSizeInBytes(Desc.Fmt); + TextureUploadLayout Layout; + Layout.Subresources.reserve(Desc.MipLevels); + uint64_t Offset = 0; + for (uint32_t I = 0; I < Desc.MipLevels; ++I) { + const uint32_t MipWidth = std::max(1u, Desc.Width >> I); + const uint32_t MipHeight = std::max(1u, Desc.Height >> I); + SubresourceFootprint Sub; + Sub.Offset = Offset; + Sub.RowSizeInBytes = MipWidth * ElementSize; + Sub.RowPitchInBytes = Sub.RowSizeInBytes; + Sub.NumRows = MipHeight; + Layout.Subresources.push_back(Sub); + Offset += uint64_t(Sub.RowSizeInBytes) * Sub.NumRows; + } + Layout.TotalSizeInBytes = Offset; + return Layout; +} diff --git a/lib/API/Util.cpp b/lib/API/Util.cpp index 3a8bf34d9..7126be0ef 100644 --- a/lib/API/Util.cpp +++ b/lib/API/Util.cpp @@ -1,5 +1,4 @@ -#include "Util.h" - +#include "API/Util.h" #include "API/Texture.h" llvm::Error offloadtest::findAndValidateRenderPassTextureSize( diff --git a/lib/API/VK/Device.cpp b/lib/API/VK/Device.cpp index 5e5f5c9f7..cf75aaa17 100644 --- a/lib/API/VK/Device.cpp +++ b/lib/API/VK/Device.cpp @@ -12,6 +12,7 @@ #include "API/Device.h" #include "API/Encoder.h" #include "API/FormatConversion.h" +#include "API/Util.h" #include "Support/Pipeline.h" #include "Support/VkError.h" #include "VKResources.h" @@ -21,7 +22,7 @@ #include "llvm/Support/FormatVariadic.h" #include "llvm/Support/MathExtras.h" -#include "../Util.h" +#include "../Support/OffloadMigration.h" #include #include @@ -32,43 +33,19 @@ using namespace offloadtest; -#define VKFormats(FMT, BITS) \ - if (Channels == 1) \ - return VK_FORMAT_R##BITS##_##FMT; \ - if (Channels == 2) \ - return VK_FORMAT_R##BITS##G##BITS##_##FMT; \ - if (Channels == 3) \ - return VK_FORMAT_R##BITS##G##BITS##B##BITS##_##FMT; \ - if (Channels == 4) \ - return VK_FORMAT_R##BITS##G##BITS##B##BITS##A##BITS##_##FMT; - -static VkFormat getVKFormat(DataFormat Format, int Channels) { - switch (Format) { - case DataFormat::Int16: - VKFormats(SINT, 16) break; - case DataFormat::UInt16: - VKFormats(UINT, 16) break; - case DataFormat::Int32: - VKFormats(SINT, 32) break; - case DataFormat::UInt32: - VKFormats(UINT, 32) break; - case DataFormat::Float32: - VKFormats(SFLOAT, 32) break; - case DataFormat::Int64: - VKFormats(SINT, 64) break; - case DataFormat::UInt64: - VKFormats(UINT, 64) break; - case DataFormat::Float64: - VKFormats(SFLOAT, 64) break; - case DataFormat::Depth32: - if (Channels != 1) - llvm_unreachable("Depth32 format only supports a single channel."); - return VK_FORMAT_D32_SFLOAT; - default: - llvm_unreachable("Unsupported Resource format specified"); - } - return VK_FORMAT_UNDEFINED; -} +struct DescriptorCounts { + // Counters to determine how many infos and counts to allocate in the builder. + // Must match for builder to function correctly. + uint32_t ImageInfoCount = 0; + uint32_t BufferInfoCount = 0; + uint32_t BufferViewCount = 0; + uint32_t ASInfoCount = 0; + uint32_t ASHandleCount = 0; + + // Size hint for how many write commands to allocate in the builder, doesn't + // need to be correct, but will prevent (re)allocations if guessed correctly + uint32_t DescriptorWriteHint = 0; +}; static VkDescriptorType getDescriptorType(const ResourceKind RK) { switch (RK) { @@ -150,61 +127,61 @@ static VkCompareOp getVKCompareOp(CompareFunction Func) { llvm_unreachable("All compare op cases handled"); } -static VkBufferUsageFlagBits getFlagBits(const ResourceKind RK) { - switch (RK) { - case ResourceKind::Buffer: - return VK_BUFFER_USAGE_UNIFORM_TEXEL_BUFFER_BIT; - case ResourceKind::RWBuffer: - return VK_BUFFER_USAGE_STORAGE_TEXEL_BUFFER_BIT; - case ResourceKind::ByteAddressBuffer: - case ResourceKind::RWByteAddressBuffer: - case ResourceKind::StructuredBuffer: - case ResourceKind::RWStructuredBuffer: - return VK_BUFFER_USAGE_STORAGE_BUFFER_BIT; - case ResourceKind::ConstantBuffer: - return VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT; - case ResourceKind::Texture2D: - case ResourceKind::RWTexture2D: - case ResourceKind::Sampler: - case ResourceKind::SampledTexture2D: - case ResourceKind::AccelerationStructure: - llvm_unreachable( - "Textures, samplers, and AS don't have buffer usage bits!"); - } - llvm_unreachable("All cases handled"); -} - -static VkImageViewType getImageViewType(const ResourceKind RK) { - switch (RK) { - case ResourceKind::Texture2D: - case ResourceKind::RWTexture2D: - case ResourceKind::SampledTexture2D: - return VK_IMAGE_VIEW_TYPE_2D; - case ResourceKind::Buffer: - case ResourceKind::RWBuffer: - case ResourceKind::ByteAddressBuffer: - case ResourceKind::RWByteAddressBuffer: - case ResourceKind::StructuredBuffer: - case ResourceKind::RWStructuredBuffer: - case ResourceKind::ConstantBuffer: - case ResourceKind::Sampler: - case ResourceKind::AccelerationStructure: - llvm_unreachable("Not an image view!"); - } - llvm_unreachable("All cases handled"); -} - -static VkImageType getVKImageType(const ResourceKind RK) { - switch (RK) { - case ResourceKind::Texture2D: - case ResourceKind::RWTexture2D: - case ResourceKind::SampledTexture2D: - return VK_IMAGE_TYPE_2D; - default: - llvm_unreachable("Unsupported image kind"); - } - llvm_unreachable("All cases handled"); -} +// static VkBufferUsageFlagBits getFlagBits(const ResourceKind RK) { +// switch (RK) { +// case ResourceKind::Buffer: +// return VK_BUFFER_USAGE_UNIFORM_TEXEL_BUFFER_BIT; +// case ResourceKind::RWBuffer: +// return VK_BUFFER_USAGE_STORAGE_TEXEL_BUFFER_BIT; +// case ResourceKind::ByteAddressBuffer: +// case ResourceKind::RWByteAddressBuffer: +// case ResourceKind::StructuredBuffer: +// case ResourceKind::RWStructuredBuffer: +// return VK_BUFFER_USAGE_STORAGE_BUFFER_BIT; +// case ResourceKind::ConstantBuffer: +// return VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT; +// case ResourceKind::Texture2D: +// case ResourceKind::RWTexture2D: +// case ResourceKind::Sampler: +// case ResourceKind::SampledTexture2D: +// case ResourceKind::AccelerationStructure: +// llvm_unreachable( +// "Textures, samplers, and AS don't have buffer usage bits!"); +// } +// llvm_unreachable("All cases handled"); +// } + +// static VkImageViewType getImageViewType(const ResourceKind RK) { +// switch (RK) { +// case ResourceKind::Texture2D: +// case ResourceKind::RWTexture2D: +// case ResourceKind::SampledTexture2D: +// return VK_IMAGE_VIEW_TYPE_2D; +// case ResourceKind::Buffer: +// case ResourceKind::RWBuffer: +// case ResourceKind::ByteAddressBuffer: +// case ResourceKind::RWByteAddressBuffer: +// case ResourceKind::StructuredBuffer: +// case ResourceKind::RWStructuredBuffer: +// case ResourceKind::ConstantBuffer: +// case ResourceKind::Sampler: +// case ResourceKind::AccelerationStructure: +// llvm_unreachable("Not an image view!"); +// } +// llvm_unreachable("All cases handled"); +// } + +// static VkImageType getVKImageType(const ResourceKind RK) { +// switch (RK) { +// case ResourceKind::Texture2D: +// case ResourceKind::RWTexture2D: +// case ResourceKind::SampledTexture2D: +// return VK_IMAGE_TYPE_2D; +// default: +// llvm_unreachable("Unsupported image kind"); +// } +// llvm_unreachable("All cases handled"); +// } static VkShaderStageFlagBits getShaderStageFlag(Stages Stage) { switch (Stage) { @@ -403,6 +380,72 @@ static bool isExtensionSupported( return false; } +class VulkanDescriptorPool : public DescriptorPool { +public: + VkDevice Dev; + VkDescriptorPool Pool; + + VulkanDescriptorPool(VkDevice Dev, VkDescriptorPool Pool) + : DescriptorPool(GPUAPI::Vulkan), Dev(Dev), Pool(Pool) {} + + ~VulkanDescriptorPool() override { + vkDestroyDescriptorPool(Dev, Pool, nullptr); + } + + static llvm::Expected> create(VkDevice Dev) { + constexpr VkDescriptorType DescriptorTypes[] = { + VK_DESCRIPTOR_TYPE_SAMPLER, + VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, + VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, + VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, + VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, + VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, + VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, + VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, + VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR}; + + llvm::SmallVector PoolSizes; + for (const VkDescriptorType Type : DescriptorTypes) { + VkDescriptorPoolSize PoolSize = {}; + PoolSize.type = Type; + PoolSize.descriptorCount = 1024; // Just allocate enough + PoolSizes.push_back(PoolSize); + } + + VkDescriptorPoolCreateInfo PoolCreateInfo = {}; + PoolCreateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO; + PoolCreateInfo.poolSizeCount = PoolSizes.size(); + PoolCreateInfo.pPoolSizes = PoolSizes.data(); + PoolCreateInfo.maxSets = 64; // Just allocate enough + + VkDescriptorPool Pool = VK_NULL_HANDLE; + if (auto Err = VK::toError( + vkCreateDescriptorPool(Dev, &PoolCreateInfo, nullptr, &Pool), + "Failed to create descriptor pool.")) + return Err; + + return std::make_unique(Dev, Pool); + } + + void reset() override { vkResetDescriptorPool(Dev, Pool, 0); } + + static bool classof(const DescriptorPool *P) { + return P->getAPI() == GPUAPI::Vulkan; + } +}; + +class VulkanDescriptorSets : public DescriptorSets { +public: + llvm::SmallVector Sets; + + VulkanDescriptorSets(llvm::SmallVector Sets) + : DescriptorSets(GPUAPI::Vulkan), Sets(std::move(Sets)) {} + + static bool classof(const DescriptorSets *S) { + return S->getAPI() == GPUAPI::Vulkan; + } +}; + struct VulkanInstance { VkInstance Instance; VkDebugUtilsMessengerEXT DebugMessenger; @@ -448,14 +491,16 @@ class VulkanBuffer : public offloadtest::Buffer { std::string Name; BufferCreateDesc Desc; size_t SizeInBytes; + VkBufferView View; VulkanBuffer(VkDevice Dev, VkBuffer Buffer, VkBuffer CounterBuffer, VkDeviceMemory Memory, VkDeviceAddress DeviceAddress, - llvm::StringRef Name, BufferCreateDesc Desc, size_t SizeInBytes) + llvm::StringRef Name, BufferCreateDesc Desc, size_t SizeInBytes, + VkBufferView View) : offloadtest::Buffer(GPUAPI::Vulkan), Dev(Dev), Buffer(Buffer), CounterBuffer(CounterBuffer), Memory(Memory), DeviceAddress(DeviceAddress), Name(Name), Desc(Desc), - SizeInBytes(SizeInBytes) {} + SizeInBytes(SizeInBytes), View(View) {} VulkanBuffer(const VulkanBuffer &) = delete; VulkanBuffer(VulkanBuffer &&) = delete; @@ -494,6 +539,8 @@ class VulkanBuffer : public offloadtest::Buffer { void unmap() override { vkUnmapMemory(Dev, Memory); } ~VulkanBuffer() override { + if (View != nullptr) + vkDestroyBufferView(Dev, View, nullptr); if (CounterBuffer != nullptr) vkDestroyBuffer(Dev, CounterBuffer, nullptr); vkDestroyBuffer(Dev, Buffer, nullptr); @@ -532,12 +579,12 @@ class VulkanTexture : public offloadtest::Texture { uint64_t SizeInBytes; VulkanTexture(VkDevice Dev, VkImage Image, VkDeviceMemory Memory, - llvm::StringRef Name, TextureCreateDesc Desc, + VkImageView View, llvm::StringRef Name, TextureCreateDesc Desc, VkImageLayout PreferredLayout, VkImageSubresourceRange FullRange, VkImageTiling Tiling, uint64_t SizeInBytes) : offloadtest::Texture(GPUAPI::Vulkan), Dev(Dev), Image(Image), - Memory(Memory), Name(Name), Desc(Desc), Tiling(Tiling), + Memory(Memory), View(View), Name(Name), Desc(Desc), Tiling(Tiling), PreferredLayout(PreferredLayout), FullRange(FullRange), SizeInBytes(SizeInBytes) {} @@ -570,6 +617,29 @@ class VulkanTexture : public offloadtest::Texture { } }; +class VulkanSampler : public offloadtest::Sampler { +public: + VkSampler Sampler; + VkDevice Device; + std::string Name; + SamplerCreateDesc Desc; + + VulkanSampler(std::string Name, const SamplerCreateDesc &Desc, + VkSampler Sampler, VkDevice Device) + : offloadtest::Sampler(GPUAPI::Vulkan), Sampler(Sampler), Device(Device), + Name(std::move(Name)), Desc(Desc) {} + ~VulkanSampler() override { + if (Sampler) + vkDestroySampler(Device, Sampler, nullptr); + } + + const SamplerCreateDesc &getDesc() const override { return Desc; } + + static bool classof(const offloadtest::Sampler *S) { + return S->getAPI() == GPUAPI::Vulkan; + } +}; + class VulkanAccelerationStructure : public offloadtest::AccelerationStructure { public: VkDevice Dev; @@ -658,6 +728,287 @@ class VulkanFence : public offloadtest::Fence { } }; +class VulkanDescriptorSetsBuilder : public DescriptorSetsBuilder { +public: + VkDevice Dev; + llvm::SmallVector Sets; + + llvm::SmallVector ImageInfos; + llvm::SmallVector BufferInfos; + llvm::SmallVector BufferViews; + llvm::SmallVector ASInfos; + llvm::SmallVector ASHandles; + + llvm::SmallVector WriteDescriptors; + + VulkanDescriptorSetsBuilder(VkDevice Dev, + llvm::SmallVector Sets, + const DescriptorCounts &DescCounts) + : DescriptorSetsBuilder(GPUAPI::Vulkan), Dev(Dev), Sets(std::move(Sets)) { + ImageInfos.reserve(DescCounts.ImageInfoCount); + BufferInfos.reserve(DescCounts.BufferInfoCount); + BufferViews.reserve(DescCounts.BufferViewCount); + ASInfos.reserve(DescCounts.ASInfoCount); + ASHandles.reserve(DescCounts.ASHandleCount); + } + + DescriptorSetsBuilder &bindBuffers(uint32_t SetIndex, + llvm::ArrayRef B, + VKBind Bnd, bool IsRead) { + const size_t BufferInfosCapacity = BufferInfos.capacity(); + const size_t BufferViewsCapacity = BufferViews.capacity(); + + const size_t Offset = BufferInfos.size(); + const size_t OffsetTyped = BufferViews.size(); + + uint32_t NumRawBuffers = 0; + uint32_t NumTypedBuffers = 0; + + for (const Buffer *Buf : B) { + const VulkanBuffer &BufferVk = llvm::cast(*Buf); + if (BufferVk.Desc.AccessType == BufferShaderAccessType::Typed) { + BufferViews.push_back(BufferVk.View); + NumTypedBuffers += 1; + } else { + const VkDescriptorBufferInfo BI = {BufferVk.Buffer, 0, VK_WHOLE_SIZE}; + BufferInfos.push_back(BI); + NumRawBuffers += 1; + } + } + + assert( + (NumRawBuffers == 0 || NumTypedBuffers == 0) && + "Cannot bind a mix of typed and raw buffers to a single bind point."); + + const size_t CountersOffset = BufferInfos.size(); + + uint32_t NumCounters = 0; + + for (const Buffer *Buf : B) { + const VulkanBuffer &BufferVk = llvm::cast(*Buf); + if (BufferVk.Desc.HasCounter) { + assert(BufferVk.Desc.AccessType != BufferShaderAccessType::Typed); + + const VkDescriptorBufferInfo BI = {BufferVk.CounterBuffer, 0, + VK_WHOLE_SIZE}; + BufferInfos.push_back(BI); + NumCounters += 1; + } + } + + if (NumCounters > 0) + assert((NumRawBuffers == NumCounters) && + "Cannot bind a mix of buffers with and without a counter to the " + "same bind point."); + + // Builder relies on no reallocations happening + assert(BufferInfos.size() <= BufferInfosCapacity && + "Too many BufferInfos inserted into DescriptorSets"); + assert(BufferViews.size() <= BufferViewsCapacity && + "Too many BufferViews inserted into DescriptorSets"); + + if (NumRawBuffers > 0) { + VkWriteDescriptorSet WDS = {}; + WDS.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; + WDS.dstSet = Sets[SetIndex]; + WDS.dstBinding = Bnd.Binding; + WDS.descriptorCount = NumRawBuffers; + WDS.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; + WDS.pBufferInfo = &BufferInfos[Offset]; + WriteDescriptors.push_back(WDS); + + if (NumCounters > 0) { + VkWriteDescriptorSet WDS = {}; + WDS.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; + WDS.dstSet = Sets[SetIndex]; + WDS.dstBinding = Bnd.CounterBinding; + WDS.descriptorCount = NumCounters; + WDS.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; + WDS.pBufferInfo = &BufferInfos[CountersOffset]; + WriteDescriptors.push_back(WDS); + } + } else if (NumTypedBuffers > 0) { + VkWriteDescriptorSet WDS = {}; + WDS.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; + WDS.dstSet = Sets[SetIndex]; + WDS.dstBinding = Bnd.Binding; + WDS.descriptorCount = NumTypedBuffers; + WDS.descriptorType = IsRead ? VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER + : VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER; + WDS.pTexelBufferView = &BufferViews[OffsetTyped]; + WriteDescriptors.push_back(WDS); + } + + return *this; + } + + DescriptorSetsBuilder &bindTextures(uint32_t SetIndex, + llvm::ArrayRef T, + llvm::ArrayRef S, + VKBind Bnd, bool IsRead) { + assert((S.empty() || S.size() == T.size()) && + "Sampler list must either be empty or match " + "texture list when binding descriptors."); + + const size_t ImageInfosCapacity = ImageInfos.capacity(); + const size_t Offset = ImageInfos.size(); + + for (size_t I = 0, N = T.size(); I < N; ++I) { + const VulkanTexture &TextureVk = llvm::cast(*T[I]); + VkSampler SamplerHandle = VK_NULL_HANDLE; + if (!S.empty() && S[I] != nullptr) { + const VulkanSampler &SamplerVk = llvm::cast(*S[I]); + SamplerHandle = SamplerVk.Sampler; + } + const VkDescriptorImageInfo ImageInfo = {SamplerHandle, TextureVk.View, + TextureVk.PreferredLayout}; + ImageInfos.push_back(ImageInfo); + } + + assert(ImageInfos.size() <= ImageInfosCapacity && + "Too many ImageInfos inserted into DescriptorSets"); + + VkWriteDescriptorSet WDS = {}; + WDS.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; + WDS.dstSet = Sets[SetIndex]; + WDS.dstBinding = Bnd.Binding; + WDS.descriptorCount = T.size(); + WDS.descriptorType = IsRead ? VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE + : VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; + WDS.pImageInfo = &ImageInfos[Offset]; + WriteDescriptors.push_back(WDS); + + return *this; + } + + DescriptorSetsBuilder &constant(uint32_t SetIndex, + llvm::ArrayRef B, + VKBind Bnd) override { + const size_t BufferInfosCapacity = BufferInfos.capacity(); + const size_t Offset = BufferInfos.size(); + for (const Buffer *Buf : B) { + const VulkanBuffer &BufferVk = llvm::cast(*Buf); + const VkDescriptorBufferInfo BI = {BufferVk.Buffer, 0, VK_WHOLE_SIZE}; + BufferInfos.push_back(BI); + } + + assert(BufferInfos.size() <= BufferInfosCapacity && + "Too many BufferInfos inserted into DescriptorSets"); + + VkWriteDescriptorSet WDS = {}; + WDS.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; + WDS.dstSet = Sets[SetIndex]; + WDS.dstBinding = Bnd.Binding; + WDS.descriptorCount = B.size(); + WDS.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER; + WDS.pBufferInfo = &BufferInfos[Offset]; + WriteDescriptors.push_back(WDS); + + return *this; + } + + DescriptorSetsBuilder &read(uint32_t SetIndex, + llvm::ArrayRef B, + VKBind Bnd) override { + return bindBuffers(SetIndex, B, Bnd, /*isRead=*/true); + } + DescriptorSetsBuilder &read(uint32_t SetIndex, + llvm::ArrayRef T, + llvm::ArrayRef S, + VKBind Bnd) override { + return bindTextures(SetIndex, T, S, Bnd, /*IsRead=*/true); + } + DescriptorSetsBuilder &read(uint32_t SetIndex, + llvm::ArrayRef A, + VKBind Bnd) override { + const size_t ASHandlesCapacity = ASHandles.capacity(); + const size_t ASInfosCapacity = ASInfos.capacity(); + const size_t HandleStart = ASHandles.size(); + + for (const AccelerationStructure *AS : A) { + const VulkanAccelerationStructure &AccelStructVulkan = + llvm::cast(*AS); + ASHandles.push_back(AccelStructVulkan.AccelStruct); + } + + VkWriteDescriptorSetAccelerationStructureKHR ASWrite = {}; + ASWrite.sType = + VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR; + ASWrite.accelerationStructureCount = A.size(); + ASWrite.pAccelerationStructures = &ASHandles[HandleStart]; + ASInfos.push_back(ASWrite); + + // Builder relies on no reallocations happening + assert( + ASHandles.size() <= ASHandlesCapacity && + "Too many Acceleration Structure Handles inserted into DescriptorSets"); + assert( + ASInfos.size() <= ASInfosCapacity && + "Too many Acceleration Structure Writes inserted into DescriptorSets"); + + VkWriteDescriptorSet WDS = {}; + WDS.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; + WDS.pNext = &ASInfos.back(); + WDS.dstSet = Sets[SetIndex]; + WDS.dstBinding = Bnd.Binding; + WDS.descriptorCount = A.size(); + WDS.descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR; + WriteDescriptors.push_back(WDS); + + return *this; + } + + DescriptorSetsBuilder &write(uint32_t SetIndex, + llvm::ArrayRef B, + VKBind Bnd) override { + return bindBuffers(SetIndex, B, Bnd, /*isRead=*/false); + } + DescriptorSetsBuilder &write(uint32_t SetIndex, + llvm::ArrayRef T, + VKBind Bnd) override { + return bindTextures(SetIndex, T, {}, Bnd, /*IsRead=*/false); + } + + DescriptorSetsBuilder &sampler(uint32_t SetIndex, + llvm::ArrayRef S, + VKBind Bnd) override { + const size_t ImageInfosCapacity = ImageInfos.capacity(); + const size_t Offset = ImageInfos.size(); + + for (const Sampler *Sampl : S) { + const VulkanSampler &SamplerVk = llvm::cast(*Sampl); + const VkDescriptorImageInfo ImageInfo = { + SamplerVk.Sampler, VK_NULL_HANDLE, + VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL}; + ImageInfos.push_back(ImageInfo); + } + + assert(ImageInfos.size() <= ImageInfosCapacity && + "Too many ImageInfos inserted into DescriptorSets"); + + VkWriteDescriptorSet WDS = {}; + WDS.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; + WDS.dstSet = Sets[SetIndex]; + WDS.dstBinding = Bnd.Binding; + WDS.descriptorCount = S.size(); + WDS.descriptorType = VK_DESCRIPTOR_TYPE_SAMPLER; + WDS.pImageInfo = &ImageInfos[Offset]; + WriteDescriptors.push_back(WDS); + + return *this; + } + + std::unique_ptr build() override { + vkUpdateDescriptorSets(Dev, WriteDescriptors.size(), + WriteDescriptors.data(), 0, nullptr); + return std::make_unique(std::move(Sets)); + } + + static bool classof(const DescriptorSetsBuilder *B) { + return B->getAPI() == GPUAPI::Vulkan; + } +}; + class VulkanQueue : public offloadtest::Queue { public: using Queue::submit; @@ -885,6 +1236,7 @@ class VulkanPipelineState : public offloadtest::PipelineState { VkPipeline Pipeline; VkPipelineLayout Layout; llvm::SmallVector SetLayouts; + DescriptorCounts DescCounts; // True for pipelines created via createPipelineRT — used by SBT / dispatch // code to safely downcast to VKRayTracingPipelineState. bool IsRayTracing = false; @@ -892,10 +1244,10 @@ class VulkanPipelineState : public offloadtest::PipelineState { VulkanPipelineState(llvm::StringRef Name, VkDevice Dev, VkPipeline Pipeline, VkPipelineLayout Layout, llvm::SmallVector SetLayouts, - bool IsRT = false) + const DescriptorCounts &DescCounts, bool IsRT = false) : offloadtest::PipelineState(GPUAPI::Vulkan), Name(Name.str()), Dev(Dev), Pipeline(Pipeline), Layout(Layout), SetLayouts(std::move(SetLayouts)), - IsRayTracing(IsRT) {} + DescCounts(DescCounts), IsRayTracing(IsRT) {} ~VulkanPipelineState() override { vkDestroyPipeline(Dev, Pipeline, nullptr); @@ -934,8 +1286,10 @@ class VKRayTracingPipelineState : public VulkanPipelineState { VKRayTracingPipelineState(llvm::StringRef Name, VkDevice Dev, VkPipeline Pipeline, VkPipelineLayout Layout, - llvm::SmallVector SetLayouts) + llvm::SmallVector SetLayouts, + const DescriptorCounts &DescCounts) : VulkanPipelineState(Name, Dev, Pipeline, Layout, std::move(SetLayouts), + DescCounts, /*IsRT=*/true) {} static bool classof(const offloadtest::PipelineState *B) { @@ -1000,6 +1354,30 @@ class VKComputeEncoder : public offloadtest::ComputeEncoder { CB.insertDebugSignpost(Label); } + void bindDescriptorSets(const PipelineState &PSO, const DescriptorSets &DSets, + VkPipelineBindPoint BindPoint) { + const VulkanPipelineState &VulkanPipeline = + llvm::cast(PSO); + const VulkanDescriptorSets &VulkanDSets = + llvm::cast(DSets); + + if (!VulkanDSets.Sets.empty()) + vkCmdBindDescriptorSets(CB.CmdBuffer, BindPoint, VulkanPipeline.Layout, 0, + VulkanDSets.Sets.size(), VulkanDSets.Sets.data(), + 0, 0); + } + + void bindComputeDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) override { + return bindDescriptorSets(PSO, DSets, VK_PIPELINE_BIND_POINT_COMPUTE); + } + + void bindRayTracingDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) override { + return bindDescriptorSets(PSO, DSets, + VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR); + } + llvm::Error dispatch(const offloadtest::PipelineState &PSO, uint32_t GroupCountX, uint32_t GroupCountY, uint32_t GroupCountZ) override { @@ -1050,11 +1428,32 @@ class VKComputeEncoder : public offloadtest::ComputeEncoder { VK_ACCESS_TRANSFER_WRITE_BIT); CB.flushBarrier(); + const VkImageAspectFlags AspectMask = isDepthFormat(VKDst.Desc.Fmt) + ? VK_IMAGE_ASPECT_DEPTH_BIT + : VK_IMAGE_ASPECT_COLOR_BIT; + const uint32_t ElementSize = getFormatSizeInBytes(VKDst.Desc.Fmt); + llvm::SmallVector Regions; + uint64_t CurrentOffset = 0; + for (uint32_t I = 0; I < VKDst.Desc.MipLevels; ++I) { + const uint32_t MipWidth = std::max(1u, VKDst.Desc.Width >> I); + const uint32_t MipHeight = std::max(1u, VKDst.Desc.Height >> I); + VkBufferImageCopy Region = {}; + Region.bufferOffset = CurrentOffset; + Region.imageSubresource.aspectMask = AspectMask; + Region.imageSubresource.mipLevel = I; + Region.imageSubresource.baseArrayLayer = 0; + Region.imageSubresource.layerCount = 1; + Region.imageExtent = {MipWidth, MipHeight, 1}; + Regions.push_back(Region); + CurrentOffset += uint64_t(MipWidth) * MipHeight * ElementSize; + } + insertDebugSignpost( - llvm::formatv("copyTextureToBuffer {0} -> {1}", VKSrc.Name, VKDst.Name) + llvm::formatv("copyBufferToTexture {0} -> {1}", VKSrc.Name, VKDst.Name) .str()); vkCmdCopyBufferToImage(CB.CmdBuffer, VKSrc.Buffer, VKDst.Image, - VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 0, nullptr); + VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, Regions.size(), + Regions.data()); CB.addImageTransition(VK_ACCESS_TRANSFER_WRITE_BIT, /*SrcAccessMask*/ VK_ACCESS_NONE, /*DstAccessMask*/ @@ -1081,6 +1480,8 @@ class VKComputeEncoder : public offloadtest::ComputeEncoder { }; addDstBarrier(VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT); insertDebugSignpost("copyCounterToBuffer 4B"); + assert(VKSrc.CounterBuffer != nullptr && "Counter buffer was nullptr >:("); + assert(VKDst.Buffer != nullptr && "Dst buffer was nullptr >:("); vkCmdCopyBuffer(CB.CmdBuffer, VKSrc.CounterBuffer, VKDst.Buffer, 1, &Region); return llvm::Error::success(); @@ -1106,9 +1507,20 @@ class VKComputeEncoder : public offloadtest::ComputeEncoder { insertDebugSignpost( llvm::formatv("copyTextureToBuffer {0} -> {1}", VKSrc.Name, VKDst.Name) .str()); + + VkBufferImageCopy Region = {}; + Region.imageSubresource.aspectMask = + VKSrc.FullRange.aspectMask & + ~VK_IMAGE_ASPECT_STENCIL_BIT; // color or depth + Region.imageSubresource.mipLevel = 0; + Region.imageSubresource.baseArrayLayer = 0; + Region.imageSubresource.layerCount = 1; + Region.imageExtent.width = VKSrc.Desc.Width; + Region.imageExtent.height = VKSrc.Desc.Height; + Region.imageExtent.depth = 1; vkCmdCopyImageToBuffer(CB.CmdBuffer, VKSrc.Image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VKDst.Buffer, - 0, nullptr); + 1, &Region); CB.addImageTransition(VK_ACCESS_TRANSFER_READ_BIT, /*SrcAccessMask*/ VK_ACCESS_NONE, /*DstAccessMask*/ @@ -1254,6 +1666,18 @@ class VulkanRenderEncoder : public offloadtest::RenderEncoder { vkCmdBindVertexBuffers(CB.CmdBuffer, Slot, 1, &NullBuf, &Zero); } } + void bindDescriptorSets(const PipelineState &PSO, + const DescriptorSets &DSets) override { + const VulkanPipelineState &VulkanPipeline = + llvm::cast(PSO); + const VulkanDescriptorSets &VulkanDSets = + llvm::cast(DSets); + + if (!VulkanDSets.Sets.empty()) + vkCmdBindDescriptorSets(CB.CmdBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, + VulkanPipeline.Layout, 0, VulkanDSets.Sets.size(), + VulkanDSets.Sets.data(), 0, 0); + } llvm::Error drawInstanced(const offloadtest::PipelineState &PSO, uint32_t VertexCount, uint32_t InstanceCount, @@ -1576,33 +2000,34 @@ class VulkanDevice : public offloadtest::Device { llvm::SmallVector CounterResourceRefs; }; - struct InvocationState { - std::unique_ptr CB; - VkDescriptorPool Pool = VK_NULL_HANDLE; - - std::unique_ptr Pipeline; - // Lifetime-tied to the pipeline; only set for RT pipelines. - std::unique_ptr SBT; - - std::unique_ptr RenderPass; - std::unique_ptr RenderTarget; - std::unique_ptr RTReadback; - std::unique_ptr DepthStencil; - std::unique_ptr VB; - - llvm::SmallVector Resources; - llvm::SmallVector DescriptorSets; - llvm::SmallVector BufferViews; - llvm::SmallVector ImageViews; - - // Parallel-indexed to `P.AccelStructs.BLAS`. - llvm::SmallVector> - BLASes; - // Keyed by `TLASDesc::Name`. - llvm::StringMap> TLASes; - // Vertex/index buffers consumed during AS builds; must outlive submission. - llvm::SmallVector> ASInputBuffers; - }; + // struct InvocationState { + // std::unique_ptr CB; + + // std::unique_ptr Pipeline; + // // Lifetime-tied to the pipeline; only set for RT pipelines. + // std::unique_ptr SBT; + + // std::unique_ptr RenderPass; + // std::unique_ptr RenderTarget; + // std::unique_ptr RTReadback; + // std::unique_ptr DepthStencil; + // std::unique_ptr VB; + + // llvm::SmallVector Resources; + // llvm::SmallVector DescriptorSets; + // llvm::SmallVector BufferViews; + // llvm::SmallVector ImageViews; + + // // Parallel-indexed to `P.AccelStructs.BLAS`. + // llvm::SmallVector> + // BLASes; + // // Keyed by `TLASDesc::Name`. + // llvm::StringMap> + // TLASes; + // // Vertex/index buffers consumed during AS builds; must outlive + // submission. llvm::SmallVector> + // ASInputBuffers; + // }; public: static llvm::Expected> @@ -2016,12 +2441,12 @@ class VulkanDevice : public offloadtest::Device { Queue &getGraphicsQueue() override { return GraphicsQueue; } - llvm::Error - createPipelineLayout(const BindingsDesc &BindingsDesc, - VkShaderStageFlags StageFlags, - llvm::SmallVectorImpl &SetLayouts, - VkPipelineLayout &PipelineLayout) { + llvm::Error createPipelineLayout( + const BindingsDesc &BindingsDesc, VkShaderStageFlags StageFlags, + llvm::SmallVectorImpl &SetLayouts, + VkPipelineLayout &PipelineLayout, DescriptorCounts &DescCounts) { assert(SetLayouts.empty() && "Output vector SetLayouts must be empty."); + DescCounts = {}; // Build descriptor set layouts from BindingsDesc. for (const DescriptorSetLayoutDesc &SetDesc : @@ -2030,6 +2455,35 @@ class VulkanDevice : public offloadtest::Device { for (const ResourceBindingDesc &RB : SetDesc.ResourceBindings) { const VulkanBinding VKBinding = RB.VKBinding.value(); + switch (RB.Kind) { + case ResourceKind::Buffer: + case ResourceKind::RWBuffer: + DescCounts.BufferViewCount += RB.DescriptorCount; + break; + case ResourceKind::StructuredBuffer: + case ResourceKind::RWStructuredBuffer: + case ResourceKind::ByteAddressBuffer: + case ResourceKind::RWByteAddressBuffer: + case ResourceKind::ConstantBuffer: + DescCounts.BufferInfoCount += RB.DescriptorCount; + break; + case ResourceKind::Texture2D: + case ResourceKind::RWTexture2D: + case ResourceKind::Sampler: + case ResourceKind::SampledTexture2D: + DescCounts.ImageInfoCount += RB.DescriptorCount; + break; + case ResourceKind::AccelerationStructure: + DescCounts.ASHandleCount += RB.DescriptorCount; + + // This will end up being to high if the user uses the + // `DescriptorSetsBuilder` efficiently, but will always be correct + // this way. Alternatively, the `DescriptorSetsBuilder` could patch + // pointers on reallocation instead. + DescCounts.ASInfoCount += RB.DescriptorCount; + break; + } + VkDescriptorSetLayoutBinding B = {}; B.binding = VKBinding.Binding; B.descriptorType = getDescriptorType(RB.Kind); @@ -2038,6 +2492,8 @@ class VulkanDevice : public offloadtest::Device { Binds.push_back(B); if (VKBinding.CounterBinding) { + DescCounts.BufferInfoCount += RB.DescriptorCount; + VkDescriptorSetLayoutBinding CB = {}; CB.binding = *VKBinding.CounterBinding; CB.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; @@ -2046,6 +2502,9 @@ class VulkanDevice : public offloadtest::Device { Binds.push_back(CB); } } + + DescCounts.DescriptorWriteHint += static_cast(Binds.size()); + VkDescriptorSetLayoutCreateInfo SetCI = {}; SetCI.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO; SetCI.bindingCount = static_cast(Binds.size()); @@ -2121,9 +2580,10 @@ class VulkanDevice : public offloadtest::Device { llvm::SmallVector SetLayouts; VkPipelineLayout PipelineLayout = VK_NULL_HANDLE; + DescriptorCounts DescCounts = {}; if (auto Err = createPipelineLayout(BindingsDesc, VK_SHADER_STAGE_COMPUTE_BIT, - SetLayouts, PipelineLayout)) + SetLayouts, PipelineLayout, DescCounts)) return Err; auto CleanupState = llvm::scope_exit([&]() { @@ -2153,7 +2613,8 @@ class VulkanDevice : public offloadtest::Device { } return std::make_unique( - Name, Device, Pipeline, PipelineLayout, std::move(SetLayouts)); + Name, Device, Pipeline, PipelineLayout, std::move(SetLayouts), + DescCounts); } llvm::Expected> @@ -2334,8 +2795,9 @@ class VulkanDevice : public offloadtest::Device { llvm::SmallVector SetLayouts; VkPipelineLayout PipelineLayout = VK_NULL_HANDLE; + DescriptorCounts DescCounts; if (auto Err = createPipelineLayout(BindingsDesc, GraphicsFlags, SetLayouts, - PipelineLayout)) + PipelineLayout, DescCounts)) return Err; // Build vertex input attribute and binding descriptions from InputLayout. @@ -2458,7 +2920,8 @@ class VulkanDevice : public offloadtest::Device { } return std::make_unique( - Name, Device, Pipeline, PipelineLayout, std::move(SetLayouts)); + Name, Device, Pipeline, PipelineLayout, std::move(SetLayouts), + DescCounts); } llvm::Expected> createMeshShaderRasterPipeline( @@ -2575,8 +3038,9 @@ class VulkanDevice : public offloadtest::Device { llvm::SmallVector SetLayouts; VkPipelineLayout PipelineLayout = VK_NULL_HANDLE; + DescriptorCounts DescCounts = {}; if (auto Err = createPipelineLayout(BindingsDesc, GraphicsFlags, SetLayouts, - PipelineLayout)) + PipelineLayout, DescCounts)) return Err; VkPipelineViewportStateCreateInfo ViewportCI = {}; @@ -2648,7 +3112,8 @@ class VulkanDevice : public offloadtest::Device { } return std::make_unique( - Name, Device, Pipeline, PipelineLayout, std::move(SetLayouts)); + Name, Device, Pipeline, PipelineLayout, std::move(SetLayouts), + DescCounts); } // Defined out-of-line below — needs VKRayTracingPipelineState's full type @@ -2713,6 +3178,10 @@ class VulkanDevice : public offloadtest::Device { VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR; + if (Desc.AccessType == BufferShaderAccessType::Typed) + BufInfo.usage |= VK_BUFFER_USAGE_UNIFORM_TEXEL_BUFFER_BIT | + VK_BUFFER_USAGE_STORAGE_TEXEL_BUFFER_BIT; + VkBuffer BufferObject; if (auto Err = VK::toError( vkCreateBuffer(Device, &BufInfo, nullptr, &BufferObject), @@ -2743,7 +3212,6 @@ class VulkanDevice : public offloadtest::Device { VkMemoryRequirements CounterMemReqs = {}; VkDeviceSize CounterOffsetInBytes = 0; if (Desc.HasCounter) { - VkBuffer CounterBuffer; VkBufferCreateInfo CounterBufferInfo = {}; CounterBufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO; CounterBufferInfo.size = sizeof(uint32_t); @@ -2808,9 +3276,28 @@ class VulkanDevice : public offloadtest::Device { DevAddr = vkGetBufferDeviceAddress(Device, &AddrInfo); } + VkBufferView View = VK_NULL_HANDLE; + if (Desc.AccessType == BufferShaderAccessType::Typed) { + // Create buffer view + VkBufferViewCreateInfo BufferViewCI = {}; + BufferViewCI.sType = VK_STRUCTURE_TYPE_BUFFER_VIEW_CREATE_INFO; + BufferViewCI.buffer = BufferObject; + BufferViewCI.format = getVulkanFormat(Desc.AccessTypeParams.Fmt); + BufferViewCI.range = VK_WHOLE_SIZE; + + if (auto Err = VK::toError( + vkCreateBufferView(Device, &BufferViewCI, nullptr, &View), + "Failed to create buffer view.")) { + if (CounterBuffer) + vkDestroyBuffer(Device, CounterBuffer, nullptr); + vkDestroyBuffer(Device, BufferObject, nullptr); + vkFreeMemory(Device, DeviceMemory, nullptr); + } + } + return std::make_unique(Device, BufferObject, CounterBuffer, DeviceMemory, DevAddr, Name, Desc, - SizeInBytes); + SizeInBytes, View); } llvm::Expected> @@ -2882,44 +3369,71 @@ class VulkanDevice : public offloadtest::Device { 1, /*layerCount*/ }; + VkImageViewCreateInfo ViewCi = {}; + ViewCi.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO; + ViewCi.image = Image; + ViewCi.viewType = VK_IMAGE_VIEW_TYPE_2D; + ViewCi.format = getVulkanFormat(Desc.Fmt); + ViewCi.components = {VK_COMPONENT_SWIZZLE_R, VK_COMPONENT_SWIZZLE_G, + VK_COMPONENT_SWIZZLE_B, VK_COMPONENT_SWIZZLE_A}; + ViewCi.subresourceRange.baseMipLevel = 0; + ViewCi.subresourceRange.levelCount = Desc.MipLevels; + ViewCi.subresourceRange.baseArrayLayer = 0; + ViewCi.subresourceRange.layerCount = 1; + ViewCi.subresourceRange.aspectMask = FullAspectMask; + + VkImageView View = VK_NULL_HANDLE; + if (auto Err = + VK::toError(vkCreateImageView(Device, &ViewCi, nullptr, &View), + "Failed to create image view.")) { + vkDestroyImage(Device, Image, nullptr); + vkFreeMemory(Device, DeviceMemory, nullptr); + return Err; + } + VkImageLayout PreferredLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL; if ((Desc.Usage & TextureUsage::Storage)) PreferredLayout = VK_IMAGE_LAYOUT_GENERAL; auto Tex = std::make_unique( - Device, Image, DeviceMemory, Name, Desc, PreferredLayout, FullRange, - ImageInfo.tiling, MemReqs.size); - - const bool IsRT = (Desc.Usage & TextureUsage::RenderTarget) != 0; - const bool IsDS = (Desc.Usage & TextureUsage::DepthStencil) != 0; - if (IsRT || IsDS) { - VkImageViewCreateInfo ViewCi = {}; - ViewCi.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO; - ViewCi.viewType = VK_IMAGE_VIEW_TYPE_2D; - ViewCi.format = getVulkanFormat(Desc.Fmt); - ViewCi.subresourceRange.baseMipLevel = 0; - ViewCi.subresourceRange.levelCount = 1; - ViewCi.subresourceRange.baseArrayLayer = 0; - ViewCi.subresourceRange.layerCount = 1; - ViewCi.image = Image; - if (IsRT) { - ViewCi.components = {VK_COMPONENT_SWIZZLE_R, VK_COMPONENT_SWIZZLE_G, - VK_COMPONENT_SWIZZLE_B, VK_COMPONENT_SWIZZLE_A}; - ViewCi.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT; - } else { - ViewCi.subresourceRange.aspectMask = - VK_IMAGE_ASPECT_DEPTH_BIT | VK_IMAGE_ASPECT_STENCIL_BIT; - } - // Tex destructor will clean up Image + Memory on failure. - if (auto Err = VK::toError( - vkCreateImageView(Device, &ViewCi, nullptr, &Tex->View), - "Failed to create image view.")) - return Err; - } + Device, Image, DeviceMemory, View, Name, Desc, PreferredLayout, + FullRange, ImageInfo.tiling, MemReqs.size); return Tex; } + llvm::Expected> + createSampler(std::string Name, const SamplerCreateDesc &Desc) override { + + VkSamplerCreateInfo SamplerInfo = {}; + SamplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO; + SamplerInfo.magFilter = getVKFilter(Desc.MagFilter); + SamplerInfo.minFilter = getVKFilter(Desc.MinFilter); + SamplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_NEAREST; + SamplerInfo.addressModeU = getVKAddressMode(Desc.Address); + SamplerInfo.addressModeV = getVKAddressMode(Desc.Address); + SamplerInfo.addressModeW = getVKAddressMode(Desc.Address); + SamplerInfo.mipLodBias = Desc.MipLODBias; + SamplerInfo.anisotropyEnable = VK_FALSE; + SamplerInfo.maxAnisotropy = 1.0f; + SamplerInfo.compareEnable = + Desc.Kind == SamplerKind::SamplerComparison ? VK_TRUE : VK_FALSE; + SamplerInfo.compareOp = getVKCompareOp(Desc.ComparisonOp); + SamplerInfo.minLod = Desc.MinLOD; + SamplerInfo.maxLod = Desc.MaxLOD; + SamplerInfo.borderColor = VK_BORDER_COLOR_FLOAT_TRANSPARENT_BLACK; + SamplerInfo.unnormalizedCoordinates = VK_FALSE; + + VkSampler Sampler; + if (auto Err = VK::toError( + vkCreateSampler(Device, &SamplerInfo, nullptr, &Sampler), + "Failed to create sampler.")) + return Err; + + return std::make_unique(std::move(Name), Desc, Sampler, + Device); + } + uint32_t getTextureUploadRowStrideInBytes( const TextureCreateDesc &Desc) const override { const uint64_t TightRow = @@ -2928,6 +3442,12 @@ class VulkanDevice : public offloadtest::Device { TightRow, Props.limits.optimalBufferCopyRowPitchAlignment)); } + TextureUploadLayout + getTextureUploadLayout(const TextureCreateDesc &Desc) const override { + // copyBufferToTexture consumes a tightly-packed staging buffer. + return computeTightTextureUploadLayout(Desc); + } + const Capabilities &getCapabilities() override { if (Caps.empty()) queryCapabilities(); @@ -3317,11 +3837,46 @@ class VulkanDevice : public offloadtest::Device { } llvm::Expected> - createTLAS(const AccelerationStructureSizes &Sizes) override { + createTLAS(const AccelerationStructureSizes &Sizes, + uint32_t /*InstanceCount*/) override { return allocateAS(Sizes, VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR, "TLAS"); } + llvm::Expected> + createDescriptorPool() override { + return VulkanDescriptorPool::create(Device); + } + + llvm::Expected> + createDescriptorSetsBuilder(DescriptorPool &Pool, + const PipelineState &PSO) override { + const VulkanDescriptorPool &PoolVK = llvm::cast(Pool); + const VulkanPipelineState &PipelineVK = + llvm::cast(PSO); + + llvm::SmallVector DescriptorSets; + if (!PipelineVK.SetLayouts.empty()) { + VkDescriptorSetAllocateInfo DSAllocInfo = {}; + DSAllocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO; + DSAllocInfo.descriptorPool = PoolVK.Pool; + DSAllocInfo.descriptorSetCount = PipelineVK.SetLayouts.size(); + DSAllocInfo.pSetLayouts = PipelineVK.SetLayouts.data(); + assert(DescriptorSets.empty()); + + DescriptorSets.insert(DescriptorSets.begin(), + PipelineVK.SetLayouts.size(), VkDescriptorSet()); + if (auto Err = + VK::toError(vkAllocateDescriptorSets(Device, &DSAllocInfo, + DescriptorSets.data()), + "Failed to allocate descriptor sets.")) + return Err; + } + + return std::make_unique( + Device, std::move(DescriptorSets), PipelineVK.DescCounts); + } + llvm::Expected createBuffer(VkBufferUsageFlags Usage, VkMemoryPropertyFlags MemoryFlags, size_t Size, void *Data = nullptr, @@ -3388,539 +3943,6 @@ class VulkanDevice : public offloadtest::Device { return BufferRef{Buffer, Memory}; } - llvm::Expected createImage(Resource &R, BufferRef &Host, - int UsageOverride = 0) { - const offloadtest::CPUBuffer &B = *R.BufferPtr; - if (B.Format == DataFormat::Depth32 && R.isReadWrite()) - return llvm::createStringError(std::errc::invalid_argument, - "Image memory allocation failed."); - VkImageCreateInfo ImageCreateInfo = {}; - ImageCreateInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO; - ImageCreateInfo.imageType = getVKImageType(R.Kind); - ImageCreateInfo.format = getVKFormat(B.Format, B.Channels); - ImageCreateInfo.mipLevels = B.OutputProps.MipLevels; - ImageCreateInfo.arrayLayers = 1; - ImageCreateInfo.samples = VK_SAMPLE_COUNT_1_BIT; - ImageCreateInfo.tiling = VK_IMAGE_TILING_OPTIMAL; - ImageCreateInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE; - // Set initial layout of the image to undefined - ImageCreateInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; - ImageCreateInfo.extent = {static_cast(B.OutputProps.Width), - static_cast(B.OutputProps.Height), 1}; - if (UsageOverride == 0) { - ImageCreateInfo.usage = - VK_IMAGE_USAGE_TRANSFER_DST_BIT | - (R.isReadWrite() - ? (VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT) - : VK_IMAGE_USAGE_SAMPLED_BIT); - } else { - ImageCreateInfo.usage = UsageOverride; - } - - VkImage Image; - if (auto Err = VK::toError( - vkCreateImage(Device, &ImageCreateInfo, nullptr, &Image), - "Failed to create image.")) - return Err; - - VkSampler Sampler = 0; - - VkMemoryRequirements MemReqs; - vkGetImageMemoryRequirements(Device, Image, &MemReqs); - VkMemoryAllocateInfo AllocInfo = {}; - AllocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO; - AllocInfo.allocationSize = MemReqs.size; - - VkDeviceMemory Memory; - if (auto Err = - VK::toError(vkAllocateMemory(Device, &AllocInfo, nullptr, &Memory), - "Image memory allocation failed.")) - return Err; - if (auto Err = VK::toError(vkBindImageMemory(Device, Image, Memory, 0), - "Image memory binding failed.")) - return Err; - - return ResourceRef(Host, ImageRef{Image, Sampler, Memory}); - } - - llvm::Expected createSampler(Resource &R, BufferRef &Host) { - VkSamplerCreateInfo SamplerInfo = {}; - SamplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO; - const Sampler &S = *R.SamplerPtr; - SamplerInfo.magFilter = getVKFilter(S.MagFilter); - SamplerInfo.minFilter = getVKFilter(S.MinFilter); - SamplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_NEAREST; - SamplerInfo.addressModeU = getVKAddressMode(S.Address); - SamplerInfo.addressModeV = getVKAddressMode(S.Address); - SamplerInfo.addressModeW = getVKAddressMode(S.Address); - SamplerInfo.mipLodBias = S.MipLODBias; - SamplerInfo.anisotropyEnable = VK_FALSE; - SamplerInfo.maxAnisotropy = 1.0f; - SamplerInfo.compareEnable = - S.Kind == SamplerKind::SamplerComparison ? VK_TRUE : VK_FALSE; - SamplerInfo.compareOp = getVKCompareOp(S.ComparisonOp); - SamplerInfo.minLod = S.MinLOD; - SamplerInfo.maxLod = S.MaxLOD; - SamplerInfo.borderColor = VK_BORDER_COLOR_FLOAT_TRANSPARENT_BLACK; - SamplerInfo.unnormalizedCoordinates = VK_FALSE; - - VkSampler Sampler; - if (auto Err = VK::toError( - vkCreateSampler(Device, &SamplerInfo, nullptr, &Sampler), - "Failed to create sampler.")) - return Err; - - return ResourceRef(Host, ImageRef{0, Sampler, 0}); - } - - llvm::Expected> createAS(Resource &R) { - assert(R.TLASPtr && "AS resource must be resolved to a TLAS"); - assert(R.getArraySize() == 1 && "AS arrays not yet supported"); - auto SizesOrErr = - getTLASBuildSizes(static_cast(R.TLASPtr->Instances.size())); - if (!SizesOrErr) - return SizesOrErr.takeError(); - return createTLAS(*SizesOrErr); - } - - llvm::Error createResource(Resource &R, InvocationState &IS) { - // Samplers don't have backing data buffers, so handle them separately - if (R.isSampler()) { - ResourceBundle Bundle{getDescriptorType(R.Kind), 0, nullptr}; - BufferRef HostBuf = {0, 0}; - auto ExSamplerRef = createSampler(R, HostBuf); - if (!ExSamplerRef) - return ExSamplerRef.takeError(); - Bundle.ResourceRefs.push_back(*ExSamplerRef); - IS.Resources.push_back(Bundle); - return llvm::Error::success(); - } - - ResourceBundle Bundle{getDescriptorType(R.Kind), R.size(), R.BufferPtr}; - for (auto &ResData : R.BufferPtr->Data) { - auto ExHostBuf = createBuffer( - VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, - VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, R.size(), ResData.get()); - if (!ExHostBuf) - return ExHostBuf.takeError(); - - if (R.isTexture()) { - auto ExImageRef = createImage(R, *ExHostBuf); - if (!ExImageRef) - return ExImageRef.takeError(); - - // Sampled textures use combined-image-sampler descriptors and need - // both valid image and sampler handles. - if (R.isSampledTexture()) { - BufferRef NullHost = {0, 0}; - auto ExSamplerRef = createSampler(R, NullHost); - if (!ExSamplerRef) - return ExSamplerRef.takeError(); - ExImageRef->Image.Sampler = ExSamplerRef->Image.Sampler; - } - - Bundle.ResourceRefs.push_back(*ExImageRef); - } else { - auto ExDeviceBuf = createBuffer( - getFlagBits(R.Kind) | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | - VK_BUFFER_USAGE_TRANSFER_DST_BIT, - VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, R.size()); - if (!ExDeviceBuf) - return ExDeviceBuf.takeError(); - VkBufferCopy Copy = {}; - Copy.size = R.size(); - vkCmdCopyBuffer(IS.CB->CmdBuffer, ExHostBuf->Buffer, - ExDeviceBuf->Buffer, 1, &Copy); - Bundle.ResourceRefs.emplace_back(*ExHostBuf, *ExDeviceBuf); - } - } - if (R.HasCounter) { - for (uint32_t I = 0; I < R.getArraySize(); ++I) { - uint32_t CounterValue = 0; - auto ExHostBuf = createBuffer(VK_BUFFER_USAGE_TRANSFER_SRC_BIT | - VK_BUFFER_USAGE_TRANSFER_DST_BIT, - VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, - sizeof(uint32_t), &CounterValue); - if (!ExHostBuf) - return ExHostBuf.takeError(); - - auto ExDeviceBuf = createBuffer( - getFlagBits(R.Kind) | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | - VK_BUFFER_USAGE_TRANSFER_DST_BIT, - VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, sizeof(uint32_t)); - if (!ExDeviceBuf) - return ExDeviceBuf.takeError(); - VkBufferCopy Copy = {}; - Copy.size = sizeof(uint32_t); - vkCmdCopyBuffer(IS.CB->CmdBuffer, ExHostBuf->Buffer, - ExDeviceBuf->Buffer, 1, &Copy); - Bundle.CounterResourceRefs.emplace_back(*ExHostBuf, *ExDeviceBuf); - } - } - IS.Resources.push_back(Bundle); - return llvm::Error::success(); - } - - llvm::Error createRenderTarget(Pipeline &P, InvocationState &IS) { - if (!P.Bindings.RTargetBufferPtr) - return llvm::createStringError( - std::errc::invalid_argument, - "No render target bound for graphics pipeline."); - const CPUBuffer &RTBuf = *P.Bindings.RTargetBufferPtr; - - auto TexOrErr = offloadtest::createRenderTargetFromCPUBuffer(*this, RTBuf); - if (!TexOrErr) - return TexOrErr.takeError(); - - IS.RenderTarget = std::move(*TexOrErr); - - // Create a host-visible staging buffer for readback. - BufferCreateDesc BufDesc = {}; - BufDesc.Location = MemoryLocation::GpuToCpu; - BufDesc.Usage = BufferUsage::Storage; - auto BufOrErr = createBuffer("RTReadback", BufDesc, RTBuf.size()); - if (!BufOrErr) - return BufOrErr.takeError(); - IS.RTReadback = std::move(*BufOrErr); - - return llvm::Error::success(); - } - - llvm::Error createDepthStencil(Pipeline &P, InvocationState &IS) { - auto TexOrErr = offloadtest::createDefaultDepthStencilTarget( - *this, P.Bindings.RTargetBufferPtr->OutputProps.Width, - P.Bindings.RTargetBufferPtr->OutputProps.Height); - if (!TexOrErr) - return TexOrErr.takeError(); - IS.DepthStencil = std::move(*TexOrErr); - return llvm::Error::success(); - } - - llvm::Error createResources(Pipeline &P, InvocationState &IS) { - for (auto &D : P.Sets) { - for (auto &R : D.Resources) { - if (R.isAccelerationStructure()) { - auto ASOrErr = createAS(R); - if (!ASOrErr) - return ASOrErr.takeError(); - auto *VkAS = llvm::cast(ASOrErr->get()); - ResourceBundle Bundle{getDescriptorType(R.Kind), 0, nullptr}; - Bundle.ResourceRefs.push_back(ResourceRef{VkAS}); - IS.Resources.push_back(std::move(Bundle)); - auto Inserted = - IS.TLASes.try_emplace(R.TLASPtr->Name, std::move(*ASOrErr)); - assert(Inserted.second && "TLAS bound to multiple resources NYI"); - (void)Inserted; - continue; - } - if (auto Err = createResource(R, IS)) - return Err; - } - } - - if (P.isRaster()) { - if (auto Err = createRenderTarget(P, IS)) - return Err; - // TODO: Always created for graphics pipelines. Consider making this - // conditional on the pipeline definition. - if (auto Err = createDepthStencil(P, IS)) - return Err; - } - - if (P.isTraditionalRaster() && P.Bindings.VertexBufferPtr) { - const CPUBuffer *VBuffer = P.Bindings.VertexBufferPtr; - - BufferCreateDesc BufDesc = {}; - BufDesc.Location = MemoryLocation::CpuToGpu; - BufDesc.Usage = BufferUsage::VertexBuffer; - auto BufOrErr = createBufferWithData(*this, "VertexBuffer", BufDesc, - VBuffer->Data[0].get(), - VBuffer->size(), nullptr, nullptr); - if (!BufOrErr) - return BufOrErr.takeError(); - IS.VB = std::move(*BufOrErr); - llvm::outs() << "Vertex buffer created.\n"; - } - - return llvm::Error::success(); - } - - llvm::Error createDescriptorPool(Pipeline &P, InvocationState &IS) { - - constexpr VkDescriptorType DescriptorTypes[] = { - VK_DESCRIPTOR_TYPE_SAMPLER, - VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, - VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, - VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, - VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, - VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, - VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, - VK_DESCRIPTOR_TYPE_STORAGE_BUFFER}; - constexpr size_t DescriptorTypesSize = - sizeof(DescriptorTypes) / sizeof(VkDescriptorType); - uint32_t DescriptorCounts[DescriptorTypesSize] = {0}; - // VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR has an out-of-range enum - // value (1000150000), so it can't share the indexed array above. - uint32_t ASDescriptorCount = 0; - for (const auto &S : P.Sets) { - for (const auto &R : S.Resources) { - if (R.isAccelerationStructure()) { - ASDescriptorCount += R.getArraySize(); - continue; - } - DescriptorCounts[getDescriptorType(R.Kind)] += R.getArraySize(); - if (R.HasCounter) - DescriptorCounts[VK_DESCRIPTOR_TYPE_STORAGE_BUFFER] += - R.getArraySize(); - } - } - llvm::SmallVector PoolSizes; - for (const VkDescriptorType Type : DescriptorTypes) { - if (DescriptorCounts[Type] > 0) { - llvm::outs() << "Descriptors: { type = " << Type - << ", count = " << DescriptorCounts[Type] << " }\n"; - VkDescriptorPoolSize PoolSize = {}; - PoolSize.type = Type; - PoolSize.descriptorCount = DescriptorCounts[Type]; - PoolSizes.push_back(PoolSize); - } - } - if (ASDescriptorCount > 0) { - llvm::outs() << "Descriptors: { type = ACCELERATION_STRUCTURE_KHR" - << ", count = " << ASDescriptorCount << " }\n"; - PoolSizes.push_back( - {VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR, ASDescriptorCount}); - } - - if (P.Sets.size() > 0) { - VkDescriptorPoolCreateInfo PoolCreateInfo = {}; - PoolCreateInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO; - PoolCreateInfo.poolSizeCount = PoolSizes.size(); - PoolCreateInfo.pPoolSizes = PoolSizes.data(); - PoolCreateInfo.maxSets = P.Sets.size(); - if (auto Err = VK::toError(vkCreateDescriptorPool(Device, &PoolCreateInfo, - nullptr, &IS.Pool), - "Failed to create descriptor pool.")) - return Err; - } - return llvm::Error::success(); - } - - llvm::Error createDescriptorSets(Pipeline &P, InvocationState &IS) { - if (P.Sets.size() == 0) - return llvm::Error::success(); - - const VulkanPipelineState &VulkanPipeline = - llvm::cast(*IS.Pipeline.get()); - - VkDescriptorSetAllocateInfo DSAllocInfo = {}; - DSAllocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO; - DSAllocInfo.descriptorPool = IS.Pool; - DSAllocInfo.descriptorSetCount = VulkanPipeline.SetLayouts.size(); - DSAllocInfo.pSetLayouts = VulkanPipeline.SetLayouts.data(); - assert(IS.DescriptorSets.empty()); - IS.DescriptorSets.insert(IS.DescriptorSets.begin(), - VulkanPipeline.SetLayouts.size(), - VkDescriptorSet()); - llvm::outs() << "Num Descriptor sets: " << VulkanPipeline.SetLayouts.size() - << "\n"; - if (auto Err = - VK::toError(vkAllocateDescriptorSets(Device, &DSAllocInfo, - IS.DescriptorSets.data()), - "Failed to allocate descriptor sets.")) - return Err; - - // Calculate the number of infos/views we are going to need for each type - uint32_t ImageInfoCount = 0; - uint32_t BufferInfoCount = 0; - uint32_t BufferViewCount = 0; - uint32_t ASInfoCount = 0; - uint32_t ASHandleCount = 0; - for (auto &D : P.Sets) { - for (auto &R : D.Resources) { - if (R.isAccelerationStructure()) { - ASInfoCount += 1; - ASHandleCount += R.getArraySize(); - continue; - } - if (R.isSampler()) { - ImageInfoCount += 1; - continue; - } - const uint32_t Count = R.getArraySize(); - if (R.isTexture()) - ImageInfoCount += Count; - else if (R.isRaw()) - BufferInfoCount += Count; - else - BufferViewCount += Count; - if (R.HasCounter) - BufferInfoCount += Count; - } - } - - // reserve enough space for the descriptor infos so it never needs to be - // resized (we need the memory fixed in place) - llvm::SmallVector ImageInfos; - llvm::SmallVector BufferInfos; - llvm::SmallVector BufferViews; - llvm::SmallVector ASInfos; - llvm::SmallVector ASHandles; - ImageInfos.reserve(ImageInfoCount); - BufferInfos.reserve(BufferInfoCount); - BufferViews.reserve(BufferViewCount); - ASInfos.reserve(ASInfoCount); - ASHandles.reserve(ASHandleCount); - - llvm::SmallVector WriteDescriptors; - WriteDescriptors.reserve(ImageInfoCount + BufferInfoCount + - BufferViewCount + ASInfoCount); - assert(IS.BufferViews.empty()); - - uint32_t OverallResIdx = 0; - for (uint32_t SetIdx = 0; SetIdx < P.Sets.size(); ++SetIdx) { - for (uint32_t RIdx = 0; RIdx < P.Sets[SetIdx].Resources.size(); - ++RIdx, ++OverallResIdx) { - const Resource &R = P.Sets[SetIdx].Resources[RIdx]; - if (VulkanAccelerationStructure *VkAS = - IS.Resources[OverallResIdx].ResourceRefs[0].AS) { - const size_t HandleStart = ASHandles.size(); - ASHandles.push_back(VkAS->AccelStruct); - VkWriteDescriptorSetAccelerationStructureKHR ASWrite = {}; - ASWrite.sType = - VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET_ACCELERATION_STRUCTURE_KHR; - ASWrite.accelerationStructureCount = 1; - ASWrite.pAccelerationStructures = &ASHandles[HandleStart]; - ASInfos.push_back(ASWrite); - - VkWriteDescriptorSet WDS = {}; - WDS.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; - WDS.pNext = &ASInfos.back(); - WDS.dstSet = IS.DescriptorSets[SetIdx]; - WDS.dstBinding = R.VKBinding->Binding; - WDS.descriptorCount = 1; - WDS.descriptorType = VK_DESCRIPTOR_TYPE_ACCELERATION_STRUCTURE_KHR; - llvm::outs() << "Updating AS Descriptor [" << OverallResIdx << "] { " - << SetIdx << ", " << RIdx << " }\n"; - WriteDescriptors.push_back(WDS); - continue; - } - uint32_t IndexOfFirstBufferDataInArray; - if (R.isSampler()) { - IndexOfFirstBufferDataInArray = ImageInfos.size(); - for (auto &ResRef : IS.Resources[OverallResIdx].ResourceRefs) { - const VkDescriptorImageInfo ImageInfo = { - ResRef.Image.Sampler, 0, - VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL}; - ImageInfos.push_back(ImageInfo); - } - } else if (R.isTexture()) { - VkImageViewCreateInfo ViewCreateInfo = {}; - ViewCreateInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO; - ViewCreateInfo.viewType = getImageViewType(R.Kind); - ViewCreateInfo.format = - getVKFormat(R.BufferPtr->Format, R.BufferPtr->Channels); - ViewCreateInfo.components = { - VK_COMPONENT_SWIZZLE_R, VK_COMPONENT_SWIZZLE_G, - VK_COMPONENT_SWIZZLE_B, VK_COMPONENT_SWIZZLE_A}; - ViewCreateInfo.subresourceRange.aspectMask = - R.BufferPtr->Format == DataFormat::Depth32 - ? VK_IMAGE_ASPECT_DEPTH_BIT - : VK_IMAGE_ASPECT_COLOR_BIT; - ViewCreateInfo.subresourceRange.baseMipLevel = 0; - ViewCreateInfo.subresourceRange.baseArrayLayer = 0; - ViewCreateInfo.subresourceRange.layerCount = 1; - ViewCreateInfo.subresourceRange.levelCount = - R.BufferPtr->OutputProps.MipLevels; - IndexOfFirstBufferDataInArray = ImageInfos.size(); - for (auto &ResRef : IS.Resources[OverallResIdx].ResourceRefs) { - ViewCreateInfo.image = ResRef.Image.Image; - VkImageView View = {0}; - if (auto Err = VK::toError( - vkCreateImageView(Device, &ViewCreateInfo, nullptr, &View), - "Failed to create image view.")) - return Err; - const VkDescriptorImageInfo ImageInfo = {ResRef.Image.Sampler, View, - VK_IMAGE_LAYOUT_GENERAL}; - IS.ImageViews.push_back(View); - ImageInfos.push_back(ImageInfo); - } - } else if (R.isRaw()) { - IndexOfFirstBufferDataInArray = BufferInfos.size(); - for (auto ResRef : IS.Resources[OverallResIdx].ResourceRefs) { - const VkDescriptorBufferInfo BI = {ResRef.Device.Buffer, 0, - VK_WHOLE_SIZE}; - BufferInfos.push_back(BI); - } - } else { - VkBufferViewCreateInfo ViewCreateInfo = {}; - const VkFormat Format = - getVKFormat(R.BufferPtr->Format, R.BufferPtr->Channels); - ViewCreateInfo.sType = VK_STRUCTURE_TYPE_BUFFER_VIEW_CREATE_INFO; - ViewCreateInfo.format = Format; - ViewCreateInfo.range = VK_WHOLE_SIZE; - VkBufferView View = {0}; - IndexOfFirstBufferDataInArray = BufferViews.size(); - for (auto &ResRef : IS.Resources[OverallResIdx].ResourceRefs) { - ViewCreateInfo.buffer = ResRef.Device.Buffer; - if (auto Err = VK::toError( - vkCreateBufferView(Device, &ViewCreateInfo, nullptr, &View), - "Failed to create buffer view.")) - return Err; - IS.BufferViews.push_back(View); - BufferViews.push_back(View); - } - } - - VkWriteDescriptorSet WDS = {}; - WDS.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; - WDS.dstSet = IS.DescriptorSets[SetIdx]; - WDS.dstBinding = R.VKBinding->Binding; - WDS.descriptorCount = R.getArraySize(); - WDS.descriptorType = getDescriptorType(R.Kind); - if (R.isTexture() || R.isSampler()) - WDS.pImageInfo = &ImageInfos[IndexOfFirstBufferDataInArray]; - else if (R.isRaw()) - WDS.pBufferInfo = &BufferInfos[IndexOfFirstBufferDataInArray]; - else - WDS.pTexelBufferView = &BufferViews[IndexOfFirstBufferDataInArray]; - llvm::outs() << "Updating Descriptor [" << OverallResIdx << "] { " - << SetIdx << ", " << RIdx << " }\n"; - WriteDescriptors.push_back(WDS); - - if (R.HasCounter) { - IndexOfFirstBufferDataInArray = BufferInfos.size(); - for (auto ResRef : IS.Resources[OverallResIdx].CounterResourceRefs) { - const VkDescriptorBufferInfo BI = {ResRef.Device.Buffer, 0, - VK_WHOLE_SIZE}; - BufferInfos.push_back(BI); - } - - VkWriteDescriptorSet CounterWDS = {}; - CounterWDS.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; - CounterWDS.dstSet = IS.DescriptorSets[SetIdx]; - CounterWDS.dstBinding = *R.VKBinding->CounterBinding; - CounterWDS.descriptorCount = R.getArraySize(); - CounterWDS.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; - CounterWDS.pBufferInfo = &BufferInfos[IndexOfFirstBufferDataInArray]; - llvm::outs() << "Updating Counter Descriptor [" << OverallResIdx - << "] { " << SetIdx << ", " << RIdx << " }\n"; - llvm::outs() << "Binding = " << CounterWDS.dstBinding << "\n"; - WriteDescriptors.push_back(CounterWDS); - } - } - } - assert(ImageInfos.size() == ImageInfoCount && - BufferInfos.size() == BufferInfoCount && - BufferViews.size() == BufferViewCount && - "size of buffer infos does not match expected count"); - - llvm::outs() << "WriteDescriptors: " << WriteDescriptors.size() << "\n"; - vkUpdateDescriptorSets(Device, WriteDescriptors.size(), - WriteDescriptors.data(), 0, nullptr); - return llvm::Error::success(); - } - static llvm::Expected parseSpecializationConstant(const SpecializationConstant &SpecConst, llvm::SmallVectorImpl &SpecData) { @@ -4054,291 +4076,22 @@ class VulkanDevice : public offloadtest::Device { return llvm::Error::success(); } - void copyResourceDataToDevice(InvocationState &IS, ResourceBundle &R) { - if (R.isSampler() || R.isAccelerationStructure()) - return; - if (R.isImage()) { - const offloadtest::CPUBuffer &B = *R.BufferPtr; - llvm::SmallVector Regions; - uint64_t CurrentOffset = 0; - for (int I = 0; I < B.OutputProps.MipLevels; ++I) { - VkBufferImageCopy Region = {}; - Region.imageSubresource.aspectMask = B.Format == DataFormat::Depth32 - ? VK_IMAGE_ASPECT_DEPTH_BIT - : VK_IMAGE_ASPECT_COLOR_BIT; - Region.imageSubresource.mipLevel = I; - Region.imageSubresource.baseArrayLayer = 0; - Region.imageSubresource.layerCount = 1; - Region.imageExtent.width = - std::max(1u, static_cast(B.OutputProps.Width) >> I); - Region.imageExtent.height = - std::max(1u, static_cast(B.OutputProps.Height) >> I); - Region.imageExtent.depth = - std::max(1u, static_cast(B.OutputProps.Depth) >> I); - Region.bufferOffset = CurrentOffset; - Regions.push_back(Region); - CurrentOffset += static_cast(Region.imageExtent.width) * - Region.imageExtent.height * Region.imageExtent.depth * - B.getElementSize(); - } - - VkImageSubresourceRange SubRange = {}; - SubRange.aspectMask = B.Format == DataFormat::Depth32 - ? VK_IMAGE_ASPECT_DEPTH_BIT - : VK_IMAGE_ASPECT_COLOR_BIT; - SubRange.baseMipLevel = 0; - SubRange.levelCount = B.OutputProps.MipLevels; - SubRange.layerCount = 1; - - VkImageMemoryBarrier ImageBarrier = {}; - ImageBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER; - - ImageBarrier.subresourceRange = SubRange; - ImageBarrier.srcAccessMask = 0; - ImageBarrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - ImageBarrier.oldLayout = R.ImageLayout; - ImageBarrier.newLayout = VK_IMAGE_LAYOUT_GENERAL; - R.ImageLayout = VK_IMAGE_LAYOUT_GENERAL; - - for (auto &ResRef : R.ResourceRefs) { - ImageBarrier.image = ResRef.Image.Image; - vkCmdPipelineBarrier(IS.CB->CmdBuffer, VK_PIPELINE_STAGE_HOST_BIT, - VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, - nullptr, 1, &ImageBarrier); - - vkCmdCopyBufferToImage(IS.CB->CmdBuffer, ResRef.Host.Buffer, - ResRef.Image.Image, VK_IMAGE_LAYOUT_GENERAL, - Regions.size(), Regions.data()); - } - - ImageBarrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - ImageBarrier.dstAccessMask = - VK_ACCESS_SHADER_READ_BIT | - (R.isReadWrite() ? VK_ACCESS_SHADER_WRITE_BIT : 0); - ImageBarrier.oldLayout = VK_IMAGE_LAYOUT_GENERAL; - ImageBarrier.newLayout = VK_IMAGE_LAYOUT_GENERAL; - - for (auto &ResRef : R.ResourceRefs) { - ImageBarrier.image = ResRef.Image.Image; - vkCmdPipelineBarrier(IS.CB->CmdBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 0, - nullptr, 0, nullptr, 1, &ImageBarrier); - } - return; - } - VkBufferMemoryBarrier Barrier = {}; - Barrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - Barrier.size = VK_WHOLE_SIZE; - Barrier.srcAccessMask = VK_ACCESS_HOST_WRITE_BIT; - Barrier.dstAccessMask = 0; - Barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - Barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - for (auto &ResRef : R.ResourceRefs) { - Barrier.buffer = ResRef.Host.Buffer; - vkCmdPipelineBarrier(IS.CB->CmdBuffer, VK_PIPELINE_STAGE_HOST_BIT, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 0, nullptr, - 1, &Barrier, 0, nullptr); - } - } - - // Record commands to copy a texture into a readback buffer. - void copyTextureToReadback(VkCommandBuffer CmdBuffer, - const VulkanTexture &Tex, - const VulkanBuffer &Readback, - VkImageLayout OldLayout, - VkAccessFlags SrcAccessMask, - VkPipelineStageFlags SrcStageMask) { - const VkImageAspectFlags AspectMask = isDepthFormat(Tex.Desc.Fmt) - ? VK_IMAGE_ASPECT_DEPTH_BIT - : VK_IMAGE_ASPECT_COLOR_BIT; - - // Transition texture to transfer source. - VkImageSubresourceRange SubRange = {}; - SubRange.aspectMask = AspectMask; - SubRange.baseMipLevel = 0; - SubRange.levelCount = 1; - SubRange.layerCount = 1; - - VkImageMemoryBarrier ImageBarrier = {}; - ImageBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER; - ImageBarrier.subresourceRange = SubRange; - ImageBarrier.srcAccessMask = SrcAccessMask; - ImageBarrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT; - ImageBarrier.oldLayout = OldLayout; - ImageBarrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL; - ImageBarrier.image = Tex.Image; - vkCmdPipelineBarrier(CmdBuffer, SrcStageMask, - VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, - nullptr, 1, &ImageBarrier); - - // Copy image to readback buffer. - VkBufferImageCopy Region = {}; - Region.imageSubresource.aspectMask = AspectMask; - Region.imageSubresource.mipLevel = 0; - Region.imageSubresource.baseArrayLayer = 0; - Region.imageSubresource.layerCount = 1; - Region.imageExtent.width = Tex.Desc.Width; - Region.imageExtent.height = Tex.Desc.Height; - Region.imageExtent.depth = 1; - vkCmdCopyImageToBuffer(CmdBuffer, Tex.Image, - VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, - Readback.Buffer, 1, &Region); - - // Barrier to make the readback buffer visible to the host. These - // explicit HOST barriers are not managed by the encoder's barrier - // tracking — they are recorded directly on the command buffer. - VkBufferMemoryBarrier BufBarrier = {}; - BufBarrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - BufBarrier.size = VK_WHOLE_SIZE; - BufBarrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - BufBarrier.dstAccessMask = VK_ACCESS_HOST_READ_BIT; - BufBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - BufBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - BufBarrier.buffer = Readback.Buffer; - vkCmdPipelineBarrier(CmdBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_HOST_BIT, 0, 0, nullptr, 1, - &BufBarrier, 0, nullptr); - } - - void copyResourceDataToHost(InvocationState &IS, ResourceBundle &R) { - if (!R.isReadWrite()) - return; - if (R.isImage()) { - const offloadtest::CPUBuffer &B = *R.BufferPtr; - VkImageSubresourceRange SubRange = {}; - SubRange.aspectMask = B.Format == DataFormat::Depth32 - ? VK_IMAGE_ASPECT_DEPTH_BIT - : VK_IMAGE_ASPECT_COLOR_BIT; - SubRange.baseMipLevel = 0; - SubRange.levelCount = B.OutputProps.MipLevels; - SubRange.layerCount = 1; - - VkImageMemoryBarrier ImageBarrier = {}; - ImageBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER; - - ImageBarrier.subresourceRange = SubRange; - ImageBarrier.srcAccessMask = 0; - ImageBarrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT; - ImageBarrier.oldLayout = R.ImageLayout; - ImageBarrier.newLayout = VK_IMAGE_LAYOUT_GENERAL; - R.ImageLayout = VK_IMAGE_LAYOUT_GENERAL; - - for (auto &ResRef : R.ResourceRefs) { - ImageBarrier.image = ResRef.Image.Image; - vkCmdPipelineBarrier(IS.CB->CmdBuffer, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, - VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, - nullptr, 1, &ImageBarrier); - } - - llvm::SmallVector Regions; - uint64_t CurrentOffset = 0; - for (int I = 0; I < B.OutputProps.MipLevels; ++I) { - VkBufferImageCopy Region = {}; - Region.imageSubresource.aspectMask = B.Format == DataFormat::Depth32 - ? VK_IMAGE_ASPECT_DEPTH_BIT - : VK_IMAGE_ASPECT_COLOR_BIT; - Region.imageSubresource.mipLevel = I; - Region.imageSubresource.baseArrayLayer = 0; - Region.imageSubresource.layerCount = 1; - Region.imageExtent.width = - std::max(1u, static_cast(B.OutputProps.Width) >> I); - Region.imageExtent.height = - std::max(1u, static_cast(B.OutputProps.Height) >> I); - Region.imageExtent.depth = - std::max(1u, static_cast(B.OutputProps.Depth) >> I); - Region.bufferOffset = CurrentOffset; - Regions.push_back(Region); - CurrentOffset += static_cast(Region.imageExtent.width) * - Region.imageExtent.height * Region.imageExtent.depth * - B.getElementSize(); - } - - for (auto &ResRef : R.ResourceRefs) - vkCmdCopyImageToBuffer(IS.CB->CmdBuffer, ResRef.Image.Image, - VK_IMAGE_LAYOUT_GENERAL, ResRef.Host.Buffer, - Regions.size(), Regions.data()); - - VkBufferMemoryBarrier Barrier = {}; - Barrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - Barrier.size = VK_WHOLE_SIZE; - Barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - Barrier.dstAccessMask = VK_ACCESS_HOST_READ_BIT; - Barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - Barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - for (auto &ResRef : R.ResourceRefs) { - Barrier.buffer = ResRef.Host.Buffer; - vkCmdPipelineBarrier(IS.CB->CmdBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_HOST_BIT, 0, 0, nullptr, 1, - &Barrier, 0, nullptr); - } - return; - } - VkBufferMemoryBarrier Barrier = {}; - Barrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER; - Barrier.size = VK_WHOLE_SIZE; - Barrier.srcAccessMask = - VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT; - Barrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT; - Barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - Barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - for (auto &ResRef : R.ResourceRefs) { - Barrier.buffer = ResRef.Host.Buffer; - vkCmdPipelineBarrier(IS.CB->CmdBuffer, - VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, - VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 1, - &Barrier, 0, nullptr); - } - VkBufferCopy CopyRegion = {}; - CopyRegion.size = R.size(); - for (auto &ResRef : R.ResourceRefs) - vkCmdCopyBuffer(IS.CB->CmdBuffer, ResRef.Device.Buffer, - ResRef.Host.Buffer, 1, &CopyRegion); - - VkBufferCopy CounterCopyRegion = {}; - CounterCopyRegion.size = sizeof(uint32_t); - for (auto &ResRef : R.CounterResourceRefs) - vkCmdCopyBuffer(IS.CB->CmdBuffer, ResRef.Device.Buffer, - ResRef.Host.Buffer, 1, &CounterCopyRegion); - - Barrier.size = VK_WHOLE_SIZE; - Barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; - Barrier.dstAccessMask = VK_ACCESS_HOST_READ_BIT; - Barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - Barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; - for (auto &ResRef : R.ResourceRefs) { - Barrier.buffer = ResRef.Host.Buffer; - vkCmdPipelineBarrier(IS.CB->CmdBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_HOST_BIT, 0, 0, nullptr, 1, - &Barrier, 0, nullptr); - } - for (auto &ResRef : R.CounterResourceRefs) { - Barrier.buffer = ResRef.Host.Buffer; - vkCmdPipelineBarrier(IS.CB->CmdBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, - VK_PIPELINE_STAGE_HOST_BIT, 0, 0, nullptr, 1, - &Barrier, 0, nullptr); - } - } - - llvm::Error createCommands(Pipeline &P, InvocationState &IS) { - for (auto &R : IS.Resources) - copyResourceDataToDevice(IS, R); + llvm::Error createCommands(Pipeline &P, SharedInvocationState &IS, + DescriptorPool &Pool) { + auto DescSetsOrErr = + buildDescriptorSets(*this, Pool, *IS.Pipeline, IS.DescTables); + if (!DescSetsOrErr) + return DescSetsOrErr.takeError(); + auto DescSets = std::move(*DescSetsOrErr); - const VkPipelineBindPoint BindPoint = - P.isTraditionalRaster() ? VK_PIPELINE_BIND_POINT_GRAPHICS - : P.isRayTracing() ? VK_PIPELINE_BIND_POINT_RAY_TRACING_KHR - : VK_PIPELINE_BIND_POINT_COMPUTE; + VulkanCommandBuffer &VKCB = llvm::cast(*IS.CB); const VulkanPipelineState &VulkanPipeline = llvm::cast(*IS.Pipeline.get()); - if (IS.DescriptorSets.size() > 0) - vkCmdBindDescriptorSets( - IS.CB->CmdBuffer, BindPoint, VulkanPipeline.Layout, 0, - IS.DescriptorSets.size(), IS.DescriptorSets.data(), 0, 0); for (const auto &PCB : P.PushConstants) { llvm::SmallVector Data; PCB.getContent(Data); - vkCmdPushConstants(IS.CB->CmdBuffer, VulkanPipeline.Layout, + vkCmdPushConstants(VKCB.CmdBuffer, VulkanPipeline.Layout, getShaderStageFlag(PCB.Stage), 0, Data.size(), Data.data()); } @@ -4348,6 +4101,9 @@ class VulkanDevice : public offloadtest::Device { if (!EncoderOrErr) return EncoderOrErr.takeError(); auto &Encoder = *EncoderOrErr.get(); + + Encoder.bindComputeDescriptorSets(*IS.Pipeline, *DescSets); + if (auto Err = Encoder.dispatch( *IS.Pipeline.get(), P.DispatchParameters.DispatchGroupCount[0], P.DispatchParameters.DispatchGroupCount[1], @@ -4363,6 +4119,9 @@ class VulkanDevice : public offloadtest::Device { if (!EncoderOrErr) return EncoderOrErr.takeError(); auto &Encoder = *EncoderOrErr.get(); + + Encoder.bindRayTracingDescriptorSets(*IS.Pipeline, *DescSets); + if (auto Err = Encoder.dispatchRays( *IS.Pipeline, *IS.SBT, P.DispatchParameters.DispatchGroupCount[0], P.DispatchParameters.DispatchGroupCount[1], @@ -4396,6 +4155,8 @@ class VulkanDevice : public offloadtest::Device { Scissor.Height = static_cast(VP.Height); Encoder.setScissor(Scissor); + Encoder.bindDescriptorSets(*IS.Pipeline, *DescSets); + if (P.isTraditionalRaster()) { if (IS.VB) Encoder.setVertexBuffer(0, IS.VB.get(), 0, @@ -4413,146 +4174,48 @@ class VulkanDevice : public offloadtest::Device { return Err; } Encoder.endEncoding(); - - copyTextureToReadback(IS.CB->CmdBuffer, - llvm::cast(*IS.RenderTarget), - llvm::cast(*IS.RTReadback), - VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL, - VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT, - VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT); } - for (auto &R : IS.Resources) - copyResourceDataToHost(IS, R); - return llvm::Error::success(); - } + auto EncoderOrErr = IS.CB->createComputeEncoder(); + if (!EncoderOrErr) + return EncoderOrErr.takeError(); + auto ReadbackEncoder = std::move(*EncoderOrErr); - llvm::Error readBackData(Pipeline &P, InvocationState &IS) { - uint32_t BufIdx = 0; - for (auto &S : P.Sets) { - for (int I = 0, E = S.Resources.size(); I < E; ++I, ++BufIdx) { - const Resource &R = S.Resources[I]; - if (!R.isReadWrite()) - continue; - VkMappedMemoryRange Range = {}; - Range.sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE; - Range.offset = 0; - Range.size = VK_WHOLE_SIZE; - auto &ResourceRef = IS.Resources[BufIdx].ResourceRefs; - auto &DataSet = R.BufferPtr->Data; - auto *ResRefIt = ResourceRef.begin(); - auto *DataIt = DataSet.begin(); - for (; ResRefIt != ResourceRef.end() && DataIt != DataSet.end(); - ++ResRefIt, ++DataIt) { - void *Mapped = nullptr; // NOLINT(misc-const-correctness) - vkMapMemory(Device, ResRefIt->Host.Memory, 0, VK_WHOLE_SIZE, 0, - &Mapped); - Range.memory = ResRefIt->Host.Memory; - vkInvalidateMappedMemoryRanges(Device, 1, &Range); - memcpy(DataIt->get(), Mapped, R.size()); - vkUnmapMemory(Device, ResRefIt->Host.Memory); - } - if (R.HasCounter) { - R.BufferPtr->Counters.clear(); - for (uint32_t I = 0; I < R.getArraySize(); ++I) { - uint32_t *Mapped = nullptr; // NOLINT(misc-const-correctness) - auto &CounterRef = IS.Resources[BufIdx].CounterResourceRefs[I]; - vkMapMemory(Device, CounterRef.Host.Memory, 0, VK_WHOLE_SIZE, 0, - (void **)&Mapped); - Range.memory = CounterRef.Host.Memory; - vkInvalidateMappedMemoryRanges(Device, 1, &Range); - R.BufferPtr->Counters.push_back(*Mapped); - vkUnmapMemory(Device, CounterRef.Host.Memory); - } - } - } + if (P.isRaster()) { + if (auto Err = ReadbackEncoder->copyTextureToBuffer(*IS.RenderTarget, + *IS.RTReadback)) + return Err; } - // Copy back the frame buffer data if this was a graphics pipeline. - if (P.isRaster()) { - auto &Readback = llvm::cast(*IS.RTReadback); + for (auto &Table : IS.DescTables) + for (auto &R : Table.Resources) + if (auto Err = copyBackResource(*ReadbackEncoder, R)) + return Err; - VkMappedMemoryRange Range = {}; - Range.sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE; - Range.offset = 0; - Range.size = VK_WHOLE_SIZE; - Range.memory = Readback.Memory; + for (auto &R : IS.RootResources) + if (auto Err = copyBackResource(*ReadbackEncoder, R)) + return Err; - void *Mapped = nullptr; // NOLINT(misc-const-correctness) - vkMapMemory(Device, Readback.Memory, 0, VK_WHOLE_SIZE, 0, &Mapped); - vkInvalidateMappedMemoryRanges(Device, 1, &Range); + ReadbackEncoder->endEncoding(); - auto *RT = P.Bindings.RTargetBufferPtr; - RT->copyFromTexture(Mapped, RT->getImageRowBytes()); - vkUnmapMemory(Device, Readback.Memory); - } return llvm::Error::success(); } - void cleanup(InvocationState &IS) { - // Wait for all in-flight submissions to complete before destroying - // resources. On the happy path the caller already waited, but this - // handles early-return error paths. - llvm::consumeError(GraphicsQueue.SubmitFence->waitForCompletion( - GraphicsQueue.FenceCounter)); - for (auto &V : IS.BufferViews) - vkDestroyBufferView(Device, V, nullptr); - - for (auto &V : IS.ImageViews) - vkDestroyImageView(Device, V, nullptr); - - for (auto &R : IS.Resources) { - // AS resources are owned by `IS.TLASes`; ResourceRef.AS is non-owning. - if (R.isAccelerationStructure()) - continue; - for (auto &ResRef : R.ResourceRefs) { - if (R.isBuffer()) { - vkDestroyBuffer(Device, ResRef.Device.Buffer, nullptr); - vkFreeMemory(Device, ResRef.Device.Memory, nullptr); - } else if (R.isSampler()) { - vkDestroySampler(Device, ResRef.Image.Sampler, nullptr); - } else if (R.DescriptorType == - VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER) { - vkDestroySampler(Device, ResRef.Image.Sampler, nullptr); - vkDestroyImage(Device, ResRef.Image.Image, nullptr); - vkFreeMemory(Device, ResRef.Image.Memory, nullptr); - } else { - assert(R.isImage()); - vkDestroyImage(Device, ResRef.Image.Image, nullptr); - vkFreeMemory(Device, ResRef.Image.Memory, nullptr); - } - vkDestroyBuffer(Device, ResRef.Host.Buffer, nullptr); - vkFreeMemory(Device, ResRef.Host.Memory, nullptr); - } - for (auto &ResRef : R.CounterResourceRefs) { - vkDestroyBuffer(Device, ResRef.Device.Buffer, nullptr); - vkFreeMemory(Device, ResRef.Device.Memory, nullptr); - vkDestroyBuffer(Device, ResRef.Host.Buffer, nullptr); - vkFreeMemory(Device, ResRef.Host.Memory, nullptr); - } - } - - if (IS.Pool) - vkDestroyDescriptorPool(Device, IS.Pool, nullptr); - } - llvm::Error executeProgram(Pipeline &P) override { - InvocationState State; - auto CleanupState = llvm::scope_exit([&]() { - cleanup(State); - llvm::outs() << "Cleanup complete.\n"; - }); + SharedInvocationState State; - auto CBOrErr = VulkanCommandBuffer::create( - Device, GraphicsQueue.QueueFamilyIdx, CmdBeginDebugUtilsLabel, - CmdEndDebugUtilsLabel, CmdInsertDebugUtilsLabel, MeshShaderFns); + auto DescriptorPoolOrErr = createDescriptorPool(); + if (!DescriptorPoolOrErr) + return DescriptorPoolOrErr.takeError(); + auto DescriptorPool = std::move(*DescriptorPoolOrErr); + + auto CBOrErr = createCommandBuffer(); if (!CBOrErr) return CBOrErr.takeError(); State.CB = std::move(*CBOrErr); - State.CB->Dev = this; llvm::outs() << "Command buffer created.\n"; - if (auto Err = createResources(P, State)) + if (auto Err = createResources(*this, P, State)) return Err; if (!P.AccelStructs.BLAS.empty() || !P.AccelStructs.TLAS.empty()) { @@ -4741,28 +4404,7 @@ class VulkanDevice : public offloadtest::Device { "Pipeline was neither Compute nor Traditional Raster"); } - llvm::outs() << "Memory buffers created.\n"; - // No explicit wait: the next submit's GPU-side timeline semaphore - // dependency ensures the copy completes before the dispatch runs. - auto CopyResult = GraphicsQueue.submit(std::move(State.CB)); - if (!CopyResult) - return CopyResult.takeError(); - llvm::outs() << "Executed copy command buffer.\n"; - auto DispatchCBOrErr = VulkanCommandBuffer::create( - Device, GraphicsQueue.QueueFamilyIdx, CmdBeginDebugUtilsLabel, - CmdEndDebugUtilsLabel, CmdInsertDebugUtilsLabel, MeshShaderFns); - if (!DispatchCBOrErr) - return DispatchCBOrErr.takeError(); - State.CB = std::move(*DispatchCBOrErr); - State.CB->Dev = this; - llvm::outs() << "Execute command buffer created.\n"; - if (auto Err = createDescriptorPool(P, State)) - return Err; - llvm::outs() << "Descriptor pool created.\n"; - if (auto Err = createDescriptorSets(P, State)) - return Err; - llvm::outs() << "Descriptor sets created.\n"; - if (auto Err = createCommands(P, State)) + if (auto Err = createCommands(P, State, *DescriptorPool)) return Err; llvm::outs() << "Commands created.\n"; auto DispatchResult = GraphicsQueue.submit(std::move(State.CB)); @@ -4771,9 +4413,9 @@ class VulkanDevice : public offloadtest::Device { llvm::outs() << "Executed compute command buffer.\n"; if (auto Err = DispatchResult->waitForCompletion()) return Err; - if (auto Err = readBackData(P, State)) + if (auto Err = readBack(*this, P, State)) return Err; - llvm::outs() << "Compute pipeline created.\n"; + llvm::outs() << "Read data back.\n"; return llvm::Error::success(); } @@ -5002,7 +4644,9 @@ VulkanDevice::createPipelineRT(llvm::StringRef Name, const BindingsDesc &BD, VK_SHADER_STAGE_INTERSECTION_BIT_KHR | VK_SHADER_STAGE_CALLABLE_BIT_KHR; llvm::SmallVector SetLayouts; VkPipelineLayout Layout = VK_NULL_HANDLE; - if (auto Err = createPipelineLayout(BD, AllRTStages, SetLayouts, Layout)) + DescriptorCounts DescCounts = {}; + if (auto Err = + createPipelineLayout(BD, AllRTStages, SetLayouts, Layout, DescCounts)) return Err; auto LayoutCleanup = llvm::scope_exit([&] { if (Layout != VK_NULL_HANDLE) @@ -5089,7 +4733,7 @@ VulkanDevice::createPipelineRT(llvm::StringRef Name, const BindingsDesc &BD, return Err; auto State = std::make_unique( - Name, Device, Pipeline, Layout, std::move(SetLayouts)); + Name, Device, Pipeline, Layout, std::move(SetLayouts), DescCounts); State->ShaderGroupIndices = std::move(NameToGroup); State->NumRaygenGroups = NumRG; State->NumMissGroups = NumMS; diff --git a/lib/Support/CMakeLists.txt b/lib/Support/CMakeLists.txt index 14bd70ea5..8d1a31b4c 100644 --- a/lib/Support/CMakeLists.txt +++ b/lib/Support/CMakeLists.txt @@ -1 +1 @@ -add_offloadtest_library(Support Pipeline.cpp Check.cpp) +add_offloadtest_library(Support Pipeline.cpp Check.cpp OffloadMigration.cpp) diff --git a/lib/Support/OffloadMigration.cpp b/lib/Support/OffloadMigration.cpp new file mode 100644 index 000000000..70ec6272e --- /dev/null +++ b/lib/Support/OffloadMigration.cpp @@ -0,0 +1,961 @@ +#include "OffloadMigration.h" + +#include "API/Device.h" +#include "API/FormatConversion.h" + + +#ifdef OFFLOADTEST_ENABLE_D3D12 +// Needed for direct Root Signature binding +#include "API/DX/Buffer.h" +#include "API/DX/CommandBuffer.h" +#include "API/DX/Descriptors.h" +#include "API/DX/PipelineState.h" +#endif + +namespace offloadtest { + +static BufferUsage bufferUsageFromResourceKind(ResourceKind Kind) { + // Determine Buffer Usage + switch (Kind) { + case ResourceKind::Buffer: + case ResourceKind::StructuredBuffer: + case ResourceKind::ByteAddressBuffer: + case ResourceKind::RWBuffer: + case ResourceKind::RWStructuredBuffer: + case ResourceKind::RWByteAddressBuffer: + return BufferUsage::Storage; + case ResourceKind::ConstantBuffer: + return BufferUsage::ConstantBuffer; + case ResourceKind::Texture2D: + case ResourceKind::RWTexture2D: + case ResourceKind::Sampler: + case ResourceKind::SampledTexture2D: + case ResourceKind::AccelerationStructure: + llvm_unreachable("Invalid case, ResourceKind is not a buffer."); + } + llvm_unreachable("All ResourceKind cases handled"); +} + +static BufferShaderAccessType bufferShaderAccessTypeFromResourceKind( + const Resource &Resource, BufferShaderAccessTypeParams &OutParams) { + // Determine Buffer Access Type + switch (Resource.Kind) { + case ResourceKind::Buffer: + case ResourceKind::RWBuffer: { + auto FmtOrErr = + toFormat(Resource.BufferPtr->Format, Resource.BufferPtr->Channels); + if (!FmtOrErr) { + printf("Invalid format! FMT: %d, CHANNELS: %d\n", + Resource.BufferPtr->Format, Resource.BufferPtr->Channels); + assert(false && "Invalid format."); + } + OutParams.Fmt = *FmtOrErr; + return BufferShaderAccessType::Typed; + } + case ResourceKind::StructuredBuffer: + case ResourceKind::RWStructuredBuffer: + OutParams.StructureStride = Resource.BufferPtr->getElementSize(); + return BufferShaderAccessType::Structured; + case ResourceKind::ByteAddressBuffer: + case ResourceKind::RWByteAddressBuffer: + case ResourceKind::ConstantBuffer: + return BufferShaderAccessType::Raw; + case ResourceKind::Texture2D: + case ResourceKind::RWTexture2D: + case ResourceKind::Sampler: + case ResourceKind::SampledTexture2D: + case ResourceKind::AccelerationStructure: + llvm_unreachable( + "Invalid case, non-buffers should have been filtered out."); + } + llvm_unreachable("All ResourceKind cases handled"); +} + +static llvm::Expected> +createAS(Device &Dev, Resource &R) { + assert(R.TLASPtr && "AS resource must be resolved to a TLAS"); + assert(R.getArraySize() == 1 && "AS arrays not yet supported"); + const uint32_t InstanceCount = + static_cast(R.TLASPtr->Instances.size()); + auto SizesOrErr = Dev.getTLASBuildSizes(InstanceCount); + if (!SizesOrErr) + return SizesOrErr.takeError(); + return Dev.createTLAS(*SizesOrErr, InstanceCount); +} + +llvm::Error copyBackResource(offloadtest::ComputeEncoder &ReadbackEncoder, + ResourcePair &R) { + if (R.first->isTexture()) { + for (const ResourceSet &RS : R.second) { + if (RS.Readback == nullptr) + continue; + + if (auto Err = + ReadbackEncoder.copyTextureToBuffer(*RS.Texture, *RS.Readback)) + return Err; + } + } else if (R.first->isBuffer()) { + for (const ResourceSet &RS : R.second) { + if (RS.Readback == nullptr) + continue; + + if (auto Err = ReadbackEncoder.copyBufferToBuffer( + *RS.Buffer, 0, *RS.Readback, 0, RS.Buffer->getSizeInBytes())) + return Err; + + if (!RS.Buffer->getDesc().HasCounter) + continue; + + if (auto Err = ReadbackEncoder.copyCounterToBuffer(*RS.Buffer, + *RS.CounterReadback)) + return Err; + } + } + + return llvm::Error::success(); +} + +llvm::Error readBack(Device &Dev, Pipeline &P, SharedInvocationState &IS) { + auto MemCpyBack = [&Dev](ResourcePair &R) -> llvm::Error { + if (!R.first->isReadWrite()) + return llvm::Error::success(); + + auto *RSIt = R.second.begin(); + auto *DataIt = R.first->BufferPtr->Data.begin(); + for (; RSIt != R.second.end() && DataIt != R.first->BufferPtr->Data.end(); + ++RSIt, ++DataIt) { + offloadtest::Buffer &Readback = *RSIt->Readback; + auto DataPtrOrErr = Readback.map(); + if (!DataPtrOrErr) + return DataPtrOrErr.takeError(); + const void *DataPtr = *DataPtrOrErr; + + if (R.first->isTexture()) { + const TextureCreateDesc &Desc = RSIt->Texture->getDesc(); + const uint32_t SrcStrideInBytes = + Dev.getTextureUploadRowStrideInBytes(Desc); + const uint32_t DstStrideInBytes = + Desc.Width * getFormatSizeInBytes(Desc.Fmt); + assert(DstStrideInBytes <= SrcStrideInBytes && + "Destination should not have padding and thus should be <= " + "than SrcStride where we do expect potential padding."); + uint8_t *Dst = (uint8_t *)DataIt->get(); + const uint8_t *Src = (const uint8_t *)DataPtr; + + for (uint32_t Y = 0; Y < Desc.Height; ++Y) { + memcpy(Dst, Src, DstStrideInBytes); + Dst += DstStrideInBytes; + Src += SrcStrideInBytes; + } + } else { + memcpy(DataIt->get(), DataPtr, R.first->size()); + } + + Readback.unmap(); + + if (R.first->HasCounter) { + offloadtest::Buffer &CounterReadback = *RSIt->CounterReadback; + auto CounterPtrOrErr = CounterReadback.map(); + if (!CounterPtrOrErr) + return CounterPtrOrErr.takeError(); + const uint32_t *CounterPtr = (const uint32_t *)*CounterPtrOrErr; + R.first->BufferPtr->Counters.push_back(*CounterPtr); + CounterReadback.unmap(); + } + } + + return llvm::Error::success(); + }; + + for (auto &Table : IS.DescTables) + for (auto &R : Table.Resources) + if (auto Err = MemCpyBack(R)) + return Err; + + for (auto &R : IS.RootResources) + if (auto Err = MemCpyBack(R)) + return Err; + + // If there is no render target, return early. + if (!IS.RTReadback) + return llvm::Error::success(); + + auto DataPtrOrErr = IS.RTReadback->map(); + if (!DataPtrOrErr) + return DataPtrOrErr.takeError(); + const void *Mapped = *DataPtrOrErr; + + const uint32_t SrcStrideInBytes = + Dev.getTextureUploadRowStrideInBytes(IS.RenderTarget->getDesc()); + + P.Bindings.RTargetBufferPtr->copyFromTexture(Mapped, SrcStrideInBytes); + IS.RTReadback->unmap(); + return llvm::Error::success(); +} + +static llvm::Error createRenderTarget(Device &Dev, Pipeline &P, + SharedInvocationState &IS) { + if (!P.Bindings.RTargetBufferPtr) + return llvm::createStringError( + std::errc::invalid_argument, + "No render target bound for graphics pipeline."); + const CPUBuffer &OutBuf = *P.Bindings.RTargetBufferPtr; + + auto TexOrErr = offloadtest::createRenderTargetFromCPUBuffer(Dev, OutBuf); + if (!TexOrErr) + return TexOrErr.takeError(); + + IS.RenderTarget = std::move(*TexOrErr); + + // Create readback buffer sized for the pixel data with row pitch padded + // up to D3D12_TEXTURE_DATA_PITCH_ALIGNMENT, which is what D3D12 requires + // for the placed footprint used by CopyTextureRegion. The compaction + // back to a tight layout happens in readBack() via GetCopyableFootprints. + BufferCreateDesc BufDesc = {}; + BufDesc.Location = MemoryLocation::GpuToCpu; + BufDesc.Usage = BufferUsage::Storage; + auto BufOrErr = Dev.createBuffer( + "RTReadback", BufDesc, IS.RenderTarget->calculateLinearSizeInBytes(Dev)); + if (!BufOrErr) + return BufOrErr.takeError(); + IS.RTReadback = std::move(*BufOrErr); + + return llvm::Error::success(); +} + +static llvm::Error createDepthStencil(Device &Dev, Pipeline &P, + SharedInvocationState &IS) { + auto TexOrErr = offloadtest::createDefaultDepthStencilTarget( + Dev, P.Bindings.RTargetBufferPtr->OutputProps.Width, + P.Bindings.RTargetBufferPtr->OutputProps.Height); + if (!TexOrErr) + return TexOrErr.takeError(); + IS.DepthStencil = std::move(*TexOrErr); + return llvm::Error::success(); +} + +llvm::Error createResources(Device &Dev, Pipeline &P, + SharedInvocationState &IS) { + auto EncOrErr = IS.CB->createComputeEncoder(); + if (!EncOrErr) + return EncOrErr.takeError(); + auto Enc = std::move(*EncOrErr); + + auto CreateBuffer = + [&Dev, &Enc, + &IS](Resource &R, + llvm::SmallVectorImpl &Resources) -> llvm::Error { + ResourceBundle ResBundle; + if (R.isBuffer()) { + BufferCreateDesc CreateDesc = {}; + CreateDesc.Location = MemoryLocation::GpuOnly; + CreateDesc.Backing = + R.IsReserved ? MemoryBacking::Sparse : MemoryBacking::Automatic; + CreateDesc.Usage = bufferUsageFromResourceKind(R.Kind); + CreateDesc.AccessType = bufferShaderAccessTypeFromResourceKind( + R, CreateDesc.AccessTypeParams); + CreateDesc.HasCounter = R.HasCounter; + + for (auto &Data : R.BufferPtr->Data) { + std::unique_ptr UploadBuffer; + std::unique_ptr BackingMemoryHeap; + + std::unique_ptr Buffer; + if (R.IsReserved) { + auto BufferOrErr = createSparseBufferWithData( + Dev, Dev.getGraphicsQueue(), "Sparse Buffer", CreateDesc, + R.size(), R.TilesMapped, Data.get(), R.size(), *Enc.get(), + UploadBuffer, BackingMemoryHeap); + if (!BufferOrErr) + return BufferOrErr.takeError(); + + Buffer = std::move(*BufferOrErr); + } else { + auto BufferOrErr = + createBufferWithData(Dev, "Buffer", CreateDesc, Data.get(), + R.size(), Enc.get(), &UploadBuffer); + if (!BufferOrErr) + return BufferOrErr.takeError(); + + Buffer = std::move(*BufferOrErr); + } + + std::unique_ptr ReadbackBuffer; + std::unique_ptr CounterReadbackBuffer; + if (getDescriptorKind(R.Kind) == DescriptorKind::UAV) { + const BufferCreateDesc ReadbackDesc = + BufferCreateDesc::readbackBuffer(); + auto ReadbackOrErr = Dev.createBuffer("Readback", ReadbackDesc, + Buffer->getSizeInBytes()); + if (!ReadbackOrErr) + return ReadbackOrErr.takeError(); + ReadbackBuffer = std::move(*ReadbackOrErr); + + if (R.HasCounter) { + auto CounterReadbackOrErr = + Dev.createBuffer("Readback", ReadbackDesc, sizeof(uint32_t)); + if (!CounterReadbackOrErr) + return CounterReadbackOrErr.takeError(); + CounterReadbackBuffer = std::move(*CounterReadbackOrErr); + } + } + + IS.KeepAliveBuffers.push_back(std::move(UploadBuffer)); + ResourceSet RSet(std::move(Buffer), std::move(BackingMemoryHeap), + std::move(ReadbackBuffer), + std::move(CounterReadbackBuffer)); + ResBundle.push_back(std::move(RSet)); + } + } else if (R.isTexture()) { + auto FormatOrErr = toFormat(R.BufferPtr->Format, R.BufferPtr->Channels); + if (!FormatOrErr) + return FormatOrErr.takeError(); + + LLVM_ENABLE_BITMASK_ENUMS_IN_NAMESPACE(); + + TextureCreateDesc CreateDesc = {}; + CreateDesc.Location = MemoryLocation::GpuOnly; + CreateDesc.Backing = + R.IsReserved ? MemoryBacking::Sparse : MemoryBacking::Automatic; + CreateDesc.Usage = TextureUsage::Sampled; + if (R.Kind == ResourceKind::RWTexture2D) + CreateDesc.Usage |= TextureUsage::Storage; + CreateDesc.Fmt = *FormatOrErr; + CreateDesc.Width = R.BufferPtr->OutputProps.Width; + CreateDesc.Height = R.BufferPtr->OutputProps.Height; + CreateDesc.MipLevels = R.BufferPtr->OutputProps.MipLevels; + + for (auto &Data : R.BufferPtr->Data) { + std::unique_ptr UploadBuffer; + std::unique_ptr BackingMemoryHeap; + + std::unique_ptr Texture; + if (R.IsReserved) { + auto TextureOrErr = createSparseTextureWithData( + Dev, Dev.getGraphicsQueue(), "Sparse Texture", CreateDesc, + Data.get(), R.size(), *Enc.get(), UploadBuffer, + BackingMemoryHeap); + if (!TextureOrErr) + return TextureOrErr.takeError(); + + Texture = std::move(*TextureOrErr); + } else { + auto TextureOrErr = + createTextureWithData(Dev, R.Name, CreateDesc, Data.get(), + R.size(), Enc.get(), &UploadBuffer); + if (!TextureOrErr) + return TextureOrErr.takeError(); + + Texture = std::move(*TextureOrErr); + } + + std::unique_ptr ReadbackBuffer; + if (getDescriptorKind(R.Kind) == DescriptorKind::UAV) { + const BufferCreateDesc ReadbackDesc = + BufferCreateDesc::readbackBuffer(); + auto ReadbackOrErr = + Dev.createBuffer("Readback", ReadbackDesc, + Texture->calculateLinearSizeInBytes(Dev)); + if (!ReadbackOrErr) + return ReadbackOrErr.takeError(); + ReadbackBuffer = std::move(*ReadbackOrErr); + } + + std::unique_ptr Sampler; + if (R.Kind == ResourceKind::SampledTexture2D) { + SamplerCreateDesc Desc = { + R.SamplerPtr->MinFilter, R.SamplerPtr->MagFilter, + R.SamplerPtr->Address, R.SamplerPtr->MinLOD, + R.SamplerPtr->MaxLOD, R.SamplerPtr->MipLODBias, + R.SamplerPtr->ComparisonOp, R.SamplerPtr->Kind, + }; + + auto SamplerOrErr = Dev.createSampler(R.SamplerPtr->Name, Desc); + if (!SamplerOrErr) + return SamplerOrErr.takeError(); + Sampler = std::move(*SamplerOrErr); + } + + IS.KeepAliveBuffers.push_back(std::move(UploadBuffer)); + ResourceSet RSet(std::move(Texture), std::move(Sampler), + std::move(BackingMemoryHeap), + std::move(ReadbackBuffer)); + ResBundle.push_back(std::move(RSet)); + } + } else if (R.isAccelerationStructure()) { + auto ASOrErr = createAS(Dev, R); + if (!ASOrErr) + return ASOrErr.takeError(); + ResBundle.emplace_back(ASOrErr->get()); + auto Inserted = + IS.TLASes.try_emplace(R.TLASPtr->Name, std::move(*ASOrErr)); + assert(Inserted.second && "TLAS bound to multiple resources NYI"); + (void)Inserted; + } else if (R.isSampler()) { + SamplerCreateDesc Desc = { + R.SamplerPtr->MinFilter, R.SamplerPtr->MagFilter, + R.SamplerPtr->Address, R.SamplerPtr->MinLOD, + R.SamplerPtr->MaxLOD, R.SamplerPtr->MipLODBias, + R.SamplerPtr->ComparisonOp, R.SamplerPtr->Kind, + }; + + auto SamplerOrErr = Dev.createSampler(R.SamplerPtr->Name, Desc); + if (!SamplerOrErr) + return SamplerOrErr.takeError(); + + ResourceSet RSet(std::move(*SamplerOrErr)); + ResBundle.push_back(std::move(RSet)); + } else { + return llvm::createStringError(std::errc::not_supported, + "Unrecognized resource type."); + } + + Resources.push_back(std::make_pair(&R, std::move(ResBundle))); + return llvm::Error::success(); + }; + + if (P.isRaster()) { + // Create render target and depth/stencil + if (auto Err = createRenderTarget(Dev, P, IS)) + return Err; + llvm::outs() << "Render target created.\n"; + // TODO: Always created for graphics pipelines. Consider making this + // conditional on the pipeline definition. + if (auto Err = createDepthStencil(Dev, P, IS)) + return Err; + llvm::outs() << "Depth stencil created.\n"; + } + + for (auto &D : P.Sets) { + IS.DescTables.emplace_back(DescriptorTable()); + DescriptorTable &Table = IS.DescTables.back(); + for (auto &R : D.Resources) + if (auto Err = CreateBuffer(R, Table.Resources)) + return Err; + } + + Enc->endEncoding(); + + // Setup root descriptors + for (auto &R : P.Settings.DX.RootParams) { + if (R.Kind != dx::RootParamKind::RootDescriptor) + continue; + auto &Resource = std::get(R.Data); + if (!Resource.IsReserved && Resource.TilesMapped.has_value()) { + return llvm::createStringError( + std::errc::invalid_argument, + "Error: Cannot define tiles mapped without declaring resource as " + "reserved."); + } + if (auto Err = CreateBuffer(Resource, IS.RootResources)) + return Err; + } + + if (P.isTraditionalRaster() && P.Bindings.VertexBufferPtr) { + const CPUBuffer *VBuffer = P.Bindings.VertexBufferPtr; + + BufferCreateDesc BufDesc = {}; + BufDesc.Location = MemoryLocation::CpuToGpu; + BufDesc.Usage = BufferUsage::VertexBuffer; + auto BufOrErr = createBufferWithData(Dev, "VertexBuffer", BufDesc, + VBuffer->Data[0].get(), + VBuffer->size(), nullptr, nullptr); + if (!BufOrErr) + return BufOrErr.takeError(); + IS.VB = std::move(*BufOrErr); + llvm::outs() << "Vertex buffer created.\n"; + } + + return llvm::Error::success(); +} + +llvm::Expected> +buildDescriptorSets(Device &Dev, DescriptorPool &Pool, PipelineState &Pipeline, + llvm::ArrayRef DescTables) { + auto BuilderOrErr = Dev.createDescriptorSetsBuilder(Pool, Pipeline); + if (!BuilderOrErr) + return BuilderOrErr.takeError(); + auto Builder = std::move(*BuilderOrErr); + + uint32_t SetIndex = 0; + for (auto &T : DescTables) { + for (auto &R : T.Resources) { + VKBind BindVK = {}; + if (R.first->VKBinding.has_value()) { + const VulkanBinding &VKBinding = *R.first->VKBinding; + BindVK.Binding = VKBinding.Binding; + if (VKBinding.CounterBinding) + BindVK.CounterBinding = *VKBinding.CounterBinding; + } + + if (R.first->isAccelerationStructure()) { + assert(R.second[0].AS != nullptr && + "Acceleration structure was nullptr ?!"); + Builder->read(SetIndex, R.second[0].AS, BindVK); + } else if (R.first->isBuffer()) { + llvm::SmallVector Buffers; + for (const auto &Set : R.second) + Buffers.push_back(Set.Buffer.get()); + + if (R.first->Kind == ResourceKind::ConstantBuffer) + Builder->constant(SetIndex, Buffers, BindVK); + else if (R.first->isReadWrite()) + Builder->write(SetIndex, Buffers, BindVK); + else + Builder->read(SetIndex, Buffers, BindVK); + } else if (R.first->isTexture()) { + llvm::SmallVector Textures; + llvm::SmallVector Samplers; + for (const auto &Set : R.second) { + if (Set.Sampler.get() != nullptr) + Samplers.push_back(Set.Sampler.get()); + Textures.push_back(Set.Texture.get()); + } + + if (R.first->isReadWrite()) { + assert(Samplers.empty() && + "Combined image samplers cannot be bound as write"); + Builder->write(SetIndex, Textures, BindVK); + } else { + Builder->read(SetIndex, Textures, Samplers, BindVK); + } + } else if (R.first->isSampler()) { + llvm::SmallVector Samplers; + for (const auto &Set : R.second) + Samplers.push_back(Set.Sampler.get()); + + Builder->sampler(SetIndex, Samplers, BindVK); + } + } + SetIndex += 1; + } + + return Builder->build(); +} + +#ifdef OFFLOADTEST_ENABLE_D3D12 +static llvm::Error bindDXComputeRootSignature(Pipeline &P, + SharedInvocationState &IS, + DescriptorPool &Pool) { + const DXCommandBuffer &DXCB = llvm::cast(*IS.CB); + const DXDescriptorPool &PoolDX = llvm::cast(Pool); + const DXPipelineState &DXPipeline = + llvm::cast(*IS.Pipeline.get()); + + D3D12_GPU_DESCRIPTOR_HANDLE Handle = PoolDX.CSUHandleGPU; + const uint32_t Inc = PoolDX.CSUIncSize; + DXCB.CmdList->SetComputeRootSignature(DXPipeline.RootSig.Get()); + + uint32_t ConstantOffset = 0u; + uint32_t RootParamIndex = 0u; + uint32_t DescriptorTableIndex = 0u; + auto *RootDescIt = IS.RootResources.begin(); + for (const auto &Param : P.Settings.DX.RootParams) { + switch (Param.Kind) { + case dx::RootParamKind::Constant: { + auto &Constant = std::get(Param.Data); + if (Constant.BufferPtr->ArraySize != 1) + return llvm::createStringError( + std::errc::value_too_large, + "Root constant cannot refer to resource arrays."); + const uint32_t NumValues = Constant.BufferPtr->size() / sizeof(uint32_t); + DXCB.CmdList->SetComputeRoot32BitConstants( + RootParamIndex++, NumValues, Constant.BufferPtr->Data.back().get(), + ConstantOffset); + ConstantOffset += NumValues; + break; + } + case dx::RootParamKind::DescriptorTable: + // TODO(manon): Add support for descriptor tables containing samplers + DXCB.CmdList->SetComputeRootDescriptorTable(RootParamIndex++, Handle); + Handle.ptr += P.Sets[DescriptorTableIndex++].Resources.size() * Inc; + break; + case dx::RootParamKind::RootDescriptor: + assert(RootDescIt != IS.RootResources.end()); + if (RootDescIt->first->getArraySize() != 1) + return llvm::createStringError( + std::errc::value_too_large, + "Root descriptor cannot refer to resource arrays."); + + const DXBuffer *BufferDX = llvm::cast_if_present( + RootDescIt->second.back().Buffer.get()); + if (!BufferDX) { + return llvm::createStringError( + std::errc::value_too_large, + "Root descriptor can only refer to buffers."); + } + + const D3D12_GPU_VIRTUAL_ADDRESS VirtualAddress = + BufferDX->Buffer->GetGPUVirtualAddress(); + switch (getDescriptorKind(RootDescIt->first->Kind)) { + case DescriptorKind::SRV: + DXCB.CmdList->SetComputeRootShaderResourceView(RootParamIndex++, + VirtualAddress); + break; + case DescriptorKind::UAV: + DXCB.CmdList->SetComputeRootUnorderedAccessView(RootParamIndex++, + VirtualAddress); + break; + case DescriptorKind::CBV: + DXCB.CmdList->SetComputeRootConstantBufferView(RootParamIndex++, + VirtualAddress); + break; + case DescriptorKind::SAMPLER: + llvm_unreachable( + "Samplers cannot be written directly into the Root Signature."); + } + ++RootDescIt; + break; + } + } + + return llvm::Error::success(); +} +#endif // OFFLOADTEST_ENABLE_D3D12 + +static llvm::Error createComputeCommands(Pipeline &P, SharedInvocationState &IS, + Device &Dev, DescriptorPool &Pool) { + IS.CB->bindPool(Pool); + + auto DescSetsOrErr = + buildDescriptorSets(Dev, Pool, *IS.Pipeline, IS.DescTables); + if (!DescSetsOrErr) + return DescSetsOrErr.takeError(); + auto DescSets = std::move(*DescSetsOrErr); + + if (P.Settings.DX.RootParams.size() > 0) { + #ifdef OFFLOADTEST_ENABLE_D3D12 + if (auto Err = bindDXComputeRootSignature(P, IS, Pool)) + return Err; + #else + return llvm::createStringError(std::errc::not_supported,, "RootParams are only supported on DX12, but API is not compiled in."); + #endif + } + + auto EncoderOrErr = IS.CB->createComputeEncoder(); + if (!EncoderOrErr) + return EncoderOrErr.takeError(); + auto &Encoder = *EncoderOrErr.get(); + + if (P.isRayTracing()) { + if (P.Settings.DX.RootParams.empty()) + Encoder.bindRayTracingDescriptorSets(*IS.Pipeline, *DescSets); + + if (auto Err = Encoder.dispatchRays( + *IS.Pipeline, *IS.SBT, P.DispatchParameters.DispatchGroupCount[0], + P.DispatchParameters.DispatchGroupCount[1], + P.DispatchParameters.DispatchGroupCount[2])) + return Err; + } else { + if (P.Settings.DX.RootParams.empty()) + Encoder.bindComputeDescriptorSets(*IS.Pipeline, *DescSets); + + if (auto Err = Encoder.dispatch(*IS.Pipeline.get(), + P.DispatchParameters.DispatchGroupCount[0], + P.DispatchParameters.DispatchGroupCount[1], + P.DispatchParameters.DispatchGroupCount[2])) + return Err; + } + Encoder.endEncoding(); + + return llvm::Error::success(); +} + +static llvm::Error createGraphicsCommands(Pipeline &P, + SharedInvocationState &IS, + Device &Dev, DescriptorPool &Pool) { + IS.CB->bindPool(Pool); + auto DescSetsOrErr = + buildDescriptorSets(Dev, Pool, *IS.Pipeline, IS.DescTables); + if (!DescSetsOrErr) + return DescSetsOrErr.takeError(); + auto DescSets = std::move(*DescSetsOrErr); + + RenderPassBeginDesc BeginDesc = {}; + BeginDesc.Pass = IS.RenderPass.get(); + BeginDesc.ColorAttachments.push_back(IS.RenderTarget.get()); + BeginDesc.DepthStencil = IS.DepthStencil.get(); + + auto EncOrErr = IS.CB->createRenderEncoder(BeginDesc); + if (!EncOrErr) + return EncOrErr.takeError(); + auto &Encoder = *EncOrErr.get(); + + Encoder.bindDescriptorSets(*IS.Pipeline, *DescSets); + + Viewport VP; + VP.Width = static_cast(P.Bindings.RTargetBufferPtr->OutputProps.Width); + VP.Height = + static_cast(P.Bindings.RTargetBufferPtr->OutputProps.Height); + Encoder.setViewport(VP); + + ScissorRect Scissor; + Scissor.Width = static_cast(VP.Width); + Scissor.Height = static_cast(VP.Height); + Encoder.setScissor(Scissor); + + if (P.isTraditionalRaster()) { + if (IS.VB) + Encoder.setVertexBuffer(0, IS.VB.get(), 0, P.Bindings.getVertexStride()); + + if (auto Err = Encoder.drawInstanced(*IS.Pipeline.get(), P.getVertexCount(), + /*InstanceCount=*/1)) + return Err; + } else { + if (auto Err = Encoder.dispatchMesh( + *IS.Pipeline.get(), P.DispatchParameters.DispatchGroupCount[0], + P.DispatchParameters.DispatchGroupCount[1], + P.DispatchParameters.DispatchGroupCount[2])) + return Err; + } + + Encoder.endEncoding(); + + return llvm::Error::success(); +} + +llvm::Error executeUnitTest(Device &Dev, Pipeline &P) { + llvm::outs() << "Configuring execution on device: " << Dev.Description + << "\n"; + + auto DescriptorPoolOrErr = Dev.createDescriptorPool(); + if (!DescriptorPoolOrErr) + return DescriptorPoolOrErr.takeError(); + auto DescriptorPool = std::move(*DescriptorPoolOrErr); + + SharedInvocationState State; + auto CBOrErr = Dev.createCommandBuffer(); + if (!CBOrErr) + return CBOrErr.takeError(); + State.CB = std::move(*CBOrErr); + llvm::outs() << "Command buffer created.\n"; + + if (auto Err = createResources(Dev, P, State)) + return Err; + llvm::outs() << "Buffers created.\n"; + + if (!P.AccelStructs.BLAS.empty() || !P.AccelStructs.TLAS.empty()) { + auto EncOrErr = State.CB->createComputeEncoder(); + if (!EncOrErr) + return EncOrErr.takeError(); + if (auto Err = offloadtest::buildPipelineAccelerationStructures( + Dev, **EncOrErr, P, State.BLASes, State.TLASes, + State.ASInputBuffers)) + return Err; + (*EncOrErr)->endEncoding(); + } + + BindingsDesc BndDesc = {}; + for (auto &S : P.Sets) { + DescriptorSetLayoutDesc Layout; + for (auto &R : S.Resources) { + ResourceBindingDesc ResourceBinding = {}; + ResourceBinding.Kind = R.Kind; + ResourceBinding.DXBinding.Register = R.DXBinding.Register; + ResourceBinding.DXBinding.Space = R.DXBinding.Space; + ResourceBinding.VKBinding = R.VKBinding; + ResourceBinding.DescriptorCount = R.getArraySize(); + + Layout.ResourceBindings.push_back(ResourceBinding); + } + + BndDesc.DescriptorSetDescs.push_back(Layout); + } + + if (P.isCompute()) { + // This is an arbitrary distinction that we could alter in the future. + if (P.Shaders.size() != 1 || P.Shaders[0].Stage != Stages::Compute) + return llvm::createStringError( + std::errc::invalid_argument, + "Compute pipeline must have exactly one compute shader."); + + ShaderContainer CS = {}; + CS.EntryPoint = P.Shaders[0].Entry; + CS.Shader = P.Shaders[0].Shader.get(); + + auto PipelineStateOrErr = + Dev.createPipelineCs("Compute Pipeline State", BndDesc, CS); + if (!PipelineStateOrErr) + return PipelineStateOrErr.takeError(); + State.Pipeline = std::move(*PipelineStateOrErr); + llvm::outs() << "Compute Pipeline created.\n"; + if (auto Err = createComputeCommands(P, State, Dev, *DescriptorPool)) + return Err; + llvm::outs() << "Compute command list created.\n"; + + } else if (P.isRaster()) { + + // Begin a render pass: bind RT/DSV and clear depth-stencil. Color + // load action is Load — the existing inline code didn't clear color. + ColorAttachmentFormatDesc ColorAttachment = {}; + ColorAttachment.Fmt = State.RenderTarget->getDesc().Fmt; + ColorAttachment.Load = LoadAction::Load; + ColorAttachment.Store = StoreAction::Store; + + DepthStencilAttachmentFormatDesc DSAttachment = {}; + DSAttachment.Fmt = State.DepthStencil->getDesc().Fmt; + DSAttachment.DepthLoad = LoadAction::Clear; + DSAttachment.DepthStore = StoreAction::Store; + DSAttachment.StencilLoad = LoadAction::DontCare; + DSAttachment.StencilStore = StoreAction::DontCare; + + RenderPassDesc PassDesc; + PassDesc.ColorAttachments.push_back(ColorAttachment); + PassDesc.DepthStencil = DSAttachment; + + auto RenderPassOrErr = Dev.createRenderPass(PassDesc); + if (!RenderPassOrErr) + return RenderPassOrErr.takeError(); + State.RenderPass = std::move(*RenderPassOrErr); + llvm::outs() << "Render pass created.\n"; + + if (P.isTraditionalRaster()) { + ShaderContainer VS = {}; + ShaderContainer PS = {}; + for (auto &Shader : P.Shaders) { + if (Shader.Stage == Stages::Vertex) { + VS.EntryPoint = Shader.Entry; + VS.Shader = Shader.Shader.get(); + } else if (Shader.Stage == Stages::Pixel) { + PS.EntryPoint = Shader.Entry; + PS.Shader = Shader.Shader.get(); + } + } + + TraditionalRasterPipelineCreateDesc PipelineDesc = {}; + PipelineDesc.Topology = P.Bindings.Topology; + PipelineDesc.PatchControlPoints = P.Bindings.PatchControlPoints; + PipelineDesc.DSFormat = Format::D32FloatS8Uint; + for (auto &Shader : P.Shaders) { + ShaderContainer SC = {}; + SC.EntryPoint = Shader.Entry; + SC.Shader = Shader.Shader.get(); + PipelineDesc.setShader(Shader.Stage, std::move(SC)); + } + + // Create the input layout based on the vertex attributes. + for (auto &Attr : P.Bindings.VertexAttributes) { + auto FormatOrErr = toFormat(Attr.Format, Attr.Channels); + if (!FormatOrErr) + return FormatOrErr.takeError(); + + InputLayoutDesc Layout = {}; + Layout.Name = Attr.Name; + Layout.Fmt = *FormatOrErr; + Layout.OffsetInBytes = Attr.Offset; + PipelineDesc.InputLayout.push_back(Layout); + } + + auto FormatOrErr = toFormat(P.Bindings.RTargetBufferPtr->Format, + P.Bindings.RTargetBufferPtr->Channels); + if (!FormatOrErr) + return FormatOrErr.takeError(); + PipelineDesc.RTFormats.push_back(*FormatOrErr); + + auto PipelineStateOrErr = Dev.createTraditionalRasterPipeline( + "Graphics Pipeline State", BndDesc, PipelineDesc); + if (!PipelineStateOrErr) + return PipelineStateOrErr.takeError(); + State.Pipeline = std::move(*PipelineStateOrErr); + llvm::outs() << "Traditional Raster Pipeline created.\n"; + + } else if (P.isMeshShaderRaster()) { + MeshShaderRasterPipelineCreateDesc PipelineDesc = {}; + PipelineDesc.Topology = P.Bindings.Topology; + PipelineDesc.DSFormat = Format::D32FloatS8Uint; + for (auto &Shader : P.Shaders) { + ShaderContainer SC = {}; + SC.EntryPoint = Shader.Entry; + SC.Shader = Shader.Shader.get(); + PipelineDesc.setShader(Shader.Stage, std::move(SC)); + } + + auto FormatOrErr = toFormat(P.Bindings.RTargetBufferPtr->Format, + P.Bindings.RTargetBufferPtr->Channels); + if (!FormatOrErr) + return FormatOrErr.takeError(); + PipelineDesc.RTFormats.push_back(*FormatOrErr); + + auto PipelineStateOrErr = Dev.createMeshShaderRasterPipeline( + "Mesh Shader Pipeline State", BndDesc, PipelineDesc); + + if (!PipelineStateOrErr) + return PipelineStateOrErr.takeError(); + State.Pipeline = std::move(*PipelineStateOrErr); + llvm::outs() << "Mesh Shader Pipeline created.\n"; + } + + if (auto Err = createGraphicsCommands(P, State, Dev, *DescriptorPool)) + return Err; + llvm::outs() << "Graphics command list created complete.\n"; + } else if (P.isRayTracing()) { + if (P.Shaders.empty() || !P.SBT || !P.RTConfig) + return llvm::createStringError( + std::errc::invalid_argument, + "RayTracing pipeline requires Shaders, " + "ShaderBindingTable, and RayTracingPipelineConfig."); + + RayTracingPipelineCreateDesc RTDesc{}; + RTDesc.Library = P.Shaders.front().Shader.get(); + RTDesc.HitGroups = P.HitGroups; + RTDesc.Config = *P.RTConfig; + RTDesc.Shaders.reserve(P.Shaders.size()); + for (const auto &Sh : P.Shaders) + RTDesc.Shaders.push_back({Sh.Stage, Sh.Entry}); + + auto PSOOrErr = + Dev.createPipelineRT("RayTracing Pipeline State", BndDesc, RTDesc); + if (!PSOOrErr) + return PSOOrErr.takeError(); + State.Pipeline = std::move(*PSOOrErr); + llvm::outs() << "RayTracing Pipeline created.\n"; + + auto SBTOrErr = Dev.createShaderBindingTable(*State.Pipeline, *P.SBT); + if (!SBTOrErr) + return SBTOrErr.takeError(); + State.SBT = std::move(*SBTOrErr); + llvm::outs() << "Shader Binding Table created.\n"; + + if (auto Err = createComputeCommands(P, State, Dev, *DescriptorPool)) + return Err; + llvm::outs() << "RayTracing command list created.\n"; + } else { + return llvm::createStringError("Pipeline was neither Compute nor Raster"); + } + + auto EncoderOrErr = State.CB->createComputeEncoder(); + if (!EncoderOrErr) + return EncoderOrErr.takeError(); + auto ReadbackEncoder = std::move(*EncoderOrErr); + + if (State.RenderTarget) { + if (auto Err = ReadbackEncoder->copyTextureToBuffer(*State.RenderTarget, + *State.RTReadback)) + return Err; + } + + for (auto &Table : State.DescTables) + for (auto &R : Table.Resources) + if (auto Err = copyBackResource(*ReadbackEncoder, R)) + return Err; + + for (auto &R : State.RootResources) + if (auto Err = copyBackResource(*ReadbackEncoder, R)) + return Err; + + ReadbackEncoder->endEncoding(); + + auto SubmitResult = Dev.getGraphicsQueue().submit(std::move(State.CB)); + if (!SubmitResult) + return SubmitResult.takeError(); + llvm::outs() << "Compute commands executed.\n"; + if (auto Err = SubmitResult->waitForCompletion()) + return Err; + if (auto Err = readBack(Dev, P, State)) + return Err; + llvm::outs() << "Read data back.\n"; + + return llvm::Error::success(); +} + +} // namespace offloadtest diff --git a/lib/Support/OffloadMigration.h b/lib/Support/OffloadMigration.h new file mode 100644 index 000000000..49d8a19ce --- /dev/null +++ b/lib/Support/OffloadMigration.h @@ -0,0 +1,132 @@ +//===- Texture.h - Offload API Texture ------------------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Code that is shared between backends, which should eventually live in th +// offloader tool, but cannot at the moment because the graphics backend layer +// is not yet finished. +// +//===----------------------------------------------------------------------===// + +#ifndef OFFLOADTEST_OFFLOAD_MIGRATION_H +#define OFFLOADTEST_OFFLOAD_MIGRATION_H + +#include "Support/Pipeline.h" +#include "llvm/ADT/SmallVector.h" +#include + +namespace offloadtest { + +class AccelerationStructure; +class Buffer; +class CommandBuffer; +class ComputeEncoder; +class DescriptorPool; +class DescriptorSets; +class Device; +class MemoryHeap; +class PipelineState; +class RenderPass; +class ShaderBindingTable; +class Texture; + +struct ResourceSet { + std::unique_ptr BackingMemory; + std::unique_ptr Buffer; + std::unique_ptr Texture; + std::unique_ptr Sampler; + std::unique_ptr Readback; + std::unique_ptr CounterReadback; + + // AS-only; mutually exclusive with the buffer/texture fields above. + AccelerationStructure *AS = nullptr; + + ResourceSet(std::unique_ptr Buffer, + std::unique_ptr BackingMemory, + std::unique_ptr Readback, + std::unique_ptr CounterReadback) + : BackingMemory(std::move(BackingMemory)), Buffer(std::move(Buffer)), + Readback(std::move(Readback)), + CounterReadback(std::move(CounterReadback)) {} + ResourceSet(std::unique_ptr Texture, + std::unique_ptr Sampler, + std::unique_ptr BackingMemory, + std::unique_ptr Readback) + : BackingMemory(std::move(BackingMemory)), Texture(std::move(Texture)), + Readback(std::move(Readback)) {} + ResourceSet(std::unique_ptr Sampler) + : Sampler(std::move(Sampler)) {} + explicit ResourceSet(AccelerationStructure *AS) : AS(AS) {} + + ResourceSet(const ResourceSet &) = delete; + ResourceSet &operator=(const ResourceSet &) = delete; + + ResourceSet(ResourceSet &&A) + : BackingMemory(std::move(A.BackingMemory)), Buffer(std::move(A.Buffer)), + Texture(std::move(A.Texture)), Sampler(std::move(A.Sampler)), + Readback(std::move(A.Readback)), + CounterReadback(std::move(A.CounterReadback)), AS(A.AS) {} + ResourceSet &operator=(ResourceSet &&A) { + BackingMemory = std::move(A.BackingMemory); + Buffer = std::move(A.Buffer); + Texture = std::move(A.Texture); + Sampler = std::move(A.Sampler); + Readback = std::move(A.Readback); + CounterReadback = std::move(A.CounterReadback); + AS = A.AS; + return *this; + } +}; + +// ResourceBundle will contain one ResourceSet for a singular resource +// or multiple ResourceSets for resource array. +using ResourceBundle = llvm::SmallVector; +using ResourcePair = std::pair; + +struct DescriptorTable { + llvm::SmallVector Resources; +}; + +struct SharedInvocationState { + std::unique_ptr CB; + std::unique_ptr Pipeline; + // Lifetime-tied to the pipeline; only set for RT pipelines. + std::unique_ptr SBT; + + // Resources for graphics pipelines. + std::unique_ptr RenderPass; + std::unique_ptr RenderTarget; + std::unique_ptr RTReadback; + std::unique_ptr DepthStencil; + std::unique_ptr VB; + + llvm::SmallVector> KeepAliveBuffers; + + llvm::SmallVector DescTables; + llvm::SmallVector RootResources; + + // Parallel-indexed to `P.AccelStructs.BLAS`. + llvm::SmallVector> BLASes; + // Keyed by `TLASDesc::Name`. + llvm::StringMap> TLASes; + // Vertex/index buffers consumed during AS builds; must outlive submission. + llvm::SmallVector> ASInputBuffers; +}; + +llvm::Error copyBackResource(offloadtest::ComputeEncoder &ReadbackEncoder, + ResourcePair &R); +llvm::Error readBack(Device &Dev, Pipeline &P, SharedInvocationState &IS); +llvm::Error createResources(Device &Dev, Pipeline &P, + SharedInvocationState &IS); +llvm::Expected> +buildDescriptorSets(Device &Dev, DescriptorPool &Pool, PipelineState &Pipeline, + llvm::ArrayRef DescTables); +llvm::Error executeUnitTest(Device &Dev, Pipeline &P); + +} // namespace offloadtest + +#endif // OFFLOADTEST_OFFLOAD_MIGRATION_H diff --git a/lib/Support/Pipeline.cpp b/lib/Support/Pipeline.cpp index 41fa3f7d6..19552cc47 100644 --- a/lib/Support/Pipeline.cpp +++ b/lib/Support/Pipeline.cpp @@ -331,8 +331,8 @@ static void setCounters(IO &I, offloadtest::CPUBuffer &B) { } } -void MappingTraits::mapping(IO &I, - offloadtest::Sampler &S) { +void MappingTraits::mapping( + IO &I, offloadtest::YAMLSampler &S) { I.mapRequired("Name", S.Name); I.mapOptional("Kind", S.Kind); I.mapOptional("MinFilter", S.MinFilter); diff --git a/test/Feature/Textures/Texture2D.CalculateLevelOfDetail.test.yaml b/test/Feature/Textures/Texture2D.CalculateLevelOfDetail.test.yaml index 74d3302fc..bd104af5a 100644 --- a/test/Feature/Textures/Texture2D.CalculateLevelOfDetail.test.yaml +++ b/test/Feature/Textures/Texture2D.CalculateLevelOfDetail.test.yaml @@ -126,8 +126,8 @@ Results: ... #--- end -# Unimplemented: Clang + DX: https://github.com/llvm/llvm-project/issues/101558 -# XFAIL: DirectX || Metal +# Unimplemented: Metal: https://github.com/llvm/llvm-project/issues/101558 +# XFAIL: Metal # RUN: split-file %s %t # RUN: %dxc_target -T vs_6_6 -E mainVS -Fo %t-vs.o %t/vertex.hlsl diff --git a/test/Feature/Textures/Texture2D.Gather.test.yaml b/test/Feature/Textures/Texture2D.Gather.test.yaml index 538471766..7bf8bf2c2 100644 --- a/test/Feature/Textures/Texture2D.Gather.test.yaml +++ b/test/Feature/Textures/Texture2D.Gather.test.yaml @@ -117,8 +117,8 @@ Results: ... #--- end -# Unimplemented: Clang + DX: https://github.com/llvm/llvm-project/issues/101558 -# XFAIL: DirectX || Metal +# Unimplemented: Clang: https://github.com/llvm/llvm-project/issues/101558 +# XFAIL: Metal # RUN: split-file %s %t # RUN: %dxc_target -T cs_6_0 -Fo %t.o %t/source.hlsl diff --git a/test/Feature/Textures/Texture2D.GatherCmp.test.yaml b/test/Feature/Textures/Texture2D.GatherCmp.test.yaml index 9a950a6c8..ffd15395f 100644 --- a/test/Feature/Textures/Texture2D.GatherCmp.test.yaml +++ b/test/Feature/Textures/Texture2D.GatherCmp.test.yaml @@ -138,8 +138,8 @@ Results: ... #--- end -# Unimplemented: Clang + DX: https://github.com/llvm/llvm-project/issues/101558 -# XFAIL: DirectX || Metal +# Unimplemented: Clang: https://github.com/llvm/llvm-project/issues/101558 +# XFAIL: Metal # RUN: split-file %s %t # RUN: %dxc_target -T cs_6_0 -Fo %t.o %t/source.hlsl diff --git a/test/Feature/Textures/Texture2D.GetDimensions.test.yaml b/test/Feature/Textures/Texture2D.GetDimensions.test.yaml index 484fcc93e..14d2aa74f 100644 --- a/test/Feature/Textures/Texture2D.GetDimensions.test.yaml +++ b/test/Feature/Textures/Texture2D.GetDimensions.test.yaml @@ -114,8 +114,8 @@ Results: ... #--- end -# Unimplemented: Clang + DX: https://github.com/llvm/llvm-project/issues/101558 -# XFAIL: DirectX || Metal +# Unimplemented: Metal: https://github.com/llvm/llvm-project/issues/101558 +# XFAIL: Metal # Bug https://github.com/llvm/llvm-project/issues/197837 # XFAIL: Clang && Vulkan diff --git a/test/Feature/Textures/Texture2D.OperatorIndex.test.yaml b/test/Feature/Textures/Texture2D.OperatorIndex.test.yaml index e57df3748..80d5d8b11 100644 --- a/test/Feature/Textures/Texture2D.OperatorIndex.test.yaml +++ b/test/Feature/Textures/Texture2D.OperatorIndex.test.yaml @@ -60,9 +60,6 @@ Results: ... #--- end -# Unimplemented: Clang + DX: https://github.com/llvm/llvm-project/issues/101558 -# XFAIL: DirectX - # RUN: split-file %s %t # RUN: %dxc_target -T cs_6_0 -Fo %t.o %t/source.hlsl # RUN: %offloader %t/pipeline.yaml %t.o diff --git a/test/Feature/Textures/Texture2D.Sample.test.yaml b/test/Feature/Textures/Texture2D.Sample.test.yaml index 34ba1bff0..5818f2013 100644 --- a/test/Feature/Textures/Texture2D.Sample.test.yaml +++ b/test/Feature/Textures/Texture2D.Sample.test.yaml @@ -160,7 +160,7 @@ Results: #--- end # Unimplemented: https://github.com/llvm/offload-test-suite/issues/664 -# XFAIL: DirectX || Metal || Vulkan && Darwin +# XFAIL: Metal || Vulkan && Darwin # XFAIL: Clang && !Vulkan diff --git a/test/Feature/Textures/Texture2D.SampleCmp.test.yaml b/test/Feature/Textures/Texture2D.SampleCmp.test.yaml index 52a06c512..2efcb2b56 100644 --- a/test/Feature/Textures/Texture2D.SampleCmp.test.yaml +++ b/test/Feature/Textures/Texture2D.SampleCmp.test.yaml @@ -193,7 +193,7 @@ Results: #--- end # Unimplemented: https://github.com/llvm/offload-test-suite/issues/664 -# XFAIL: DirectX || Metal || Vulkan && Darwin +# XFAIL: Metal || Vulkan && Darwin # RUN: split-file %s %t # RUN: %dxc_target -T vs_6_0 -E mainVS -Fo %t-vs.o %t/vertex.hlsl diff --git a/test/Feature/Textures/Texture2D.SampleGrad.test.yaml b/test/Feature/Textures/Texture2D.SampleGrad.test.yaml index 6ea23f4d3..e53e5470c 100644 --- a/test/Feature/Textures/Texture2D.SampleGrad.test.yaml +++ b/test/Feature/Textures/Texture2D.SampleGrad.test.yaml @@ -102,8 +102,8 @@ Results: ... #--- end -# Unimplemented: Clang + DX: https://github.com/llvm/llvm-project/issues/101558 -# XFAIL: DirectX || Metal +# Unimplemented: Metal: https://github.com/llvm/llvm-project/issues/101558 +# XFAIL: Metal # XFAIL: Clang && !Vulkan # RUN: split-file %s %t diff --git a/test/Feature/Textures/Texture2D.Sampler.address.test.yaml b/test/Feature/Textures/Texture2D.Sampler.address.test.yaml index 844b8f5e7..ec67a3d5d 100644 --- a/test/Feature/Textures/Texture2D.Sampler.address.test.yaml +++ b/test/Feature/Textures/Texture2D.Sampler.address.test.yaml @@ -160,7 +160,7 @@ Results: #--- end # Unimplemented: https://github.com/llvm/offload-test-suite/issues/664 -# XFAIL: DirectX || Metal +# XFAIL: Metal # RUN: split-file %s %t # RUN: %dxc_target -T vs_6_0 -E mainVS -Fo %t-vs.o %t/vertex.hlsl diff --git a/test/Feature/Textures/Texture2D.Sampler.filter.test.yaml b/test/Feature/Textures/Texture2D.Sampler.filter.test.yaml index 5dfd8255b..b7b1ed5a3 100644 --- a/test/Feature/Textures/Texture2D.Sampler.filter.test.yaml +++ b/test/Feature/Textures/Texture2D.Sampler.filter.test.yaml @@ -131,7 +131,7 @@ Results: #--- end # Unimplemented: https://github.com/llvm/offload-test-suite/issues/664 -# XFAIL: DirectX || Metal +# XFAIL: Metal # RUN: split-file %s %t # RUN: %dxc_target -T vs_6_0 -E mainVS -Fo %t-vs.o %t/vertex.hlsl diff --git a/test/Feature/Textures/Texture2D.Sampler.mips.test.yaml b/test/Feature/Textures/Texture2D.Sampler.mips.test.yaml index f11b76940..4e0b58122 100644 --- a/test/Feature/Textures/Texture2D.Sampler.mips.test.yaml +++ b/test/Feature/Textures/Texture2D.Sampler.mips.test.yaml @@ -93,8 +93,8 @@ Results: ... #--- end -# Unimplemented: Clang + DX: https://github.com/llvm/llvm-project/issues/101558 -# XFAIL: DirectX || Metal +# Unimplemented: Metal: https://github.com/llvm/llvm-project/issues/101558 +# XFAIL: Metal # XFAIL: Clang && !Vulkan # RUN: split-file %s %t diff --git a/test/Feature/Textures/Texture2D.mips.OperatorIndex.test.yaml b/test/Feature/Textures/Texture2D.mips.OperatorIndex.test.yaml index 49f8ee4ba..67c1f2bab 100644 --- a/test/Feature/Textures/Texture2D.mips.OperatorIndex.test.yaml +++ b/test/Feature/Textures/Texture2D.mips.OperatorIndex.test.yaml @@ -61,8 +61,6 @@ Results: #--- end # Unimplemented: https://github.com/llvm/offload-test-suite/issues/1039 -# XFAIL: DirectX - # XFAIL: Metal # RUN: split-file %s %t