Implement sparse texture support.

Author: manon-traverse

include/API/Buffer.h

@@ -20,6 +20,8 @@
 
 namespace offloadtest {
 
+class Device;
+
 enum class BufferShaderAccessType {
   Raw,
   Typed,
@@ -85,11 +87,8 @@ class Buffer {
   virtual ~Buffer();
   virtual size_t getSizeInBytes() const = 0;
 
-  // The granularity, in bytes, of a single sparse tile for this buffer. Tile
-  // counts are computed as ceil(sizeInBytes / granularity). Only meaningful for
-  // buffers created with MemoryBacking::Sparse. The value is backend-defined
-  // (e.g. 64 KiB on DX12/Vulkan, commonly 16 KiB on Metal/Apple Silicon).
-  virtual size_t querySparseTileSizeInBytes() const = 0;
+  // The granularity, in bytes, of a single sparse tile for this buffer.
+  virtual size_t querySparseTileSizeInBytes(const Device &Dev) const = 0;
 
   // Maps the buffer's memory for host access. Only valid for CpuToGpu and
   // GpuToCpu buffers; returns an error for GpuOnly. Each successful map() must

include/API/Device.h

@@ -395,6 +395,13 @@ llvm::Expected<std::unique_ptr<offloadtest::Buffer>> createSparseBufferWithData(
     std::unique_ptr<offloadtest::Buffer> &OutUploadBuffer,
     std::unique_ptr<offloadtest::MemoryHeap> &OutBackingMemoryHeap);
 
+llvm::Expected<std::unique_ptr<offloadtest::Texture>>
+createSparseTextureWithData(
+    Device &Dev, Queue &Q, std::string Name, const TextureCreateDesc &Desc,
+    const void *Data, size_t SizeInBytes, ComputeEncoder &Encoder,
+    std::unique_ptr<offloadtest::Buffer> &OutUploadBuffer,
+    std::unique_ptr<offloadtest::MemoryHeap> &OutBackingMemoryHeap);
+
 // TLAS handles come in pre-allocated because the caller's binding loop
 // stamps the AS pointer into descriptor bundles before this helper runs;
 // BLAS handles are allocated inline since BLASes aren't user-bindable.

include/API/Texture.h

@@ -68,12 +68,13 @@ using ClearValue = std::variant<ClearColor, ClearDepthStencil>;
 // cube, or array textures, add a TextureType enum and validation between usage
 // and type (e.g. 3D textures cannot be used as DepthStencil).
 struct TextureCreateDesc {
-  MemoryLocation Location;
-  TextureUsage Usage;
-  Format Fmt;
-  uint32_t Width;
-  uint32_t Height;
-  uint32_t MipLevels;
+  MemoryLocation Location = MemoryLocation::GpuOnly;
+  MemoryBacking Backing = MemoryBacking::Automatic;
+  TextureUsage Usage = {};
+  Format Fmt = Format::RGBA32Float;
+  uint32_t Width = 1;
+  uint32_t Height = 1;
+  uint32_t MipLevels = 1;
   // Clear value for render target or depth/stencil textures.
   // How and when this is applied depends on the backend:
   // - DX uses it as an optimized clear hint at resource creation time
@@ -144,6 +145,15 @@ inline llvm::Error validateTextureCreateDesc(const TextureCreateDesc &Desc) {
   return llvm::Error::success();
 }
 
+// The texel dimensions of a single sparse tile. A tile is a fixed byte size
+// (~64 KiB) laid out as this WxHxD box of texels; Depth is 1 for 2D textures.
+// Tile counts are computed per dimension as ceil(extent / tileExtent).
+struct TileShape {
+  uint32_t Width = 1;
+  uint32_t Height = 1;
+  uint32_t Depth = 1;
+};
+
 class Texture {
   GPUAPI API;
 
@@ -154,7 +164,12 @@ class Texture {
 
   // 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 calculateLinearSizeInBytes(Device &Dev) const;
+  size_t calculateLinearSizeInBytes(const Device &Dev) const;
+
+  // The texel dimensions of a single sparse tile for this texture. The shape
+  // varies by format, dimensionality, and sample count. Only meaningful for a
+  // texture created with MemoryBacking::Sparse.
+  virtual TileShape querySparseTileShape(const Device &Dev) const = 0;
 
   GPUAPI getAPI() const { return API; }
   virtual const TextureCreateDesc &getDesc() const = 0;

lib/API/DX/Device.cpp

@@ -193,6 +193,7 @@ static BufferShaderAccessType bufferShaderAccessTypeFromResourceKind(
 }
 
 namespace {
+class DXDevice;
 
 class DXBuffer : public offloadtest::Buffer {
 public:
@@ -231,7 +232,7 @@ class DXBuffer : public offloadtest::Buffer {
 
   size_t getSizeInBytes() const override { return SizeInBytes; }
 
-  size_t querySparseTileSizeInBytes() const override {
+  size_t querySparseTileSizeInBytes(const Device & /*Dev*/) const override {
     return D3D12_TILED_RESOURCE_TILE_SIZE_IN_BYTES;
   }
 
@@ -284,6 +285,8 @@ class DXTexture : public offloadtest::Texture {
         PreferredState(PreferredState), SRVHandle(SRVHandle),
         UAVHandle(UAVHandle), Name(Name), Desc(Desc) {}
 
+  TileShape querySparseTileShape(const Device &Dev) const override;
+
   const TextureCreateDesc &getDesc() const override { return Desc; }
 
   static bool classof(const offloadtest::Texture *T) {
@@ -1149,7 +1152,7 @@ class DXDevice : public offloadtest::Device {
   // raw ID3D12Device for scratch buffer allocation.
   friend class DXComputeEncoder;
 
-private:
+public:
   ComPtr<IDXCoreAdapter> Adapter;
   ComPtr<ID3D12DeviceX> Device;
   DXQueue GraphicsQueue;
@@ -1175,31 +1178,36 @@ class DXDevice : public offloadtest::Device {
                 std::unique_ptr<MemoryHeap> BackingMemory,
                 std::unique_ptr<offloadtest::Buffer> Readback,
                 std::unique_ptr<offloadtest::Buffer> CounterReadback)
-        : UploadBuffer(std::move(UploadBuffer)), Buffer(std::move(Buffer)),
-          BackingMemory(std::move(BackingMemory)),
+        : UploadBuffer(std::move(UploadBuffer)),
+          BackingMemory(std::move(BackingMemory)), Buffer(std::move(Buffer)),
+
           Readback(std::move(Readback)),
           CounterReadback(std::move(CounterReadback)) {}
     ResourceSet(std::unique_ptr<offloadtest::Buffer> UploadBuffer,
                 std::unique_ptr<offloadtest::Texture> Texture,
+                std::unique_ptr<MemoryHeap> BackingMemory,
                 std::unique_ptr<offloadtest::Buffer> Readback)
-        : UploadBuffer(std::move(UploadBuffer)), Texture(std::move(Texture)),
+        : UploadBuffer(std::move(UploadBuffer)),
+          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)
-        : UploadBuffer(std::move(A.UploadBuffer)), Buffer(std::move(A.Buffer)),
-          Texture(std::move(A.Texture)),
+        : UploadBuffer(std::move(A.UploadBuffer)),
           BackingMemory(std::move(A.BackingMemory)),
-          Readback(std::move(A.Readback)),
+          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) {
       UploadBuffer = std::move(A.UploadBuffer);
+      BackingMemory = std::move(A.BackingMemory);
       Buffer = std::move(A.Buffer);
       Texture = std::move(A.Texture);
-      BackingMemory = std::move(A.BackingMemory);
       Readback = std::move(A.Readback);
       CounterReadback = std::move(A.CounterReadback);
       AS = A.AS;
@@ -1284,6 +1292,10 @@ class DXDevice : public offloadtest::Device {
   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
@@ -1770,9 +1782,14 @@ class DXDevice : public offloadtest::Device {
     TexDesc.MipLevels = static_cast<UINT16>(Desc.MipLevels);
     TexDesc.Format = getDXGIFormat(Desc.Fmt);
     TexDesc.SampleDesc.Count = 1;
-    TexDesc.Layout = Desc.Location == MemoryLocation::GpuOnly
-                         ? D3D12_TEXTURE_LAYOUT_UNKNOWN
-                         : D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
+    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;
@@ -1803,12 +1820,20 @@ class DXDevice : public offloadtest::Device {
       InitialState = D3D12_RESOURCE_STATE_UNORDERED_ACCESS;
 
     ComPtr<ID3D12Resource> TextureObject;
-    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;
+    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 = {};
     {
@@ -2250,7 +2275,8 @@ class DXDevice : public offloadtest::Device {
       if (R.isBuffer()) {
         BufferCreateDesc CreateDesc = {};
         CreateDesc.Location = MemoryLocation::GpuOnly;
-        CreateDesc.Backing = R.IsReserved ? MemoryBacking::Sparse : MemoryBacking::Automatic;
+        CreateDesc.Backing =
+            R.IsReserved ? MemoryBacking::Sparse : MemoryBacking::Automatic;
         CreateDesc.Usage = bufferUsageFromResourceKind(R.Kind);
         CreateDesc.AccessType = bufferShaderAccessTypeFromResourceKind(
             R, CreateDesc.AccessTypeParams);
@@ -2319,6 +2345,8 @@ class DXDevice : public offloadtest::Device {
 
         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;
@@ -2329,12 +2357,26 @@ class DXDevice : public offloadtest::Device {
 
         for (auto &Data : R.BufferPtr->Data) {
           std::unique_ptr<offloadtest::Buffer> UploadBuffer;
-          auto TextureOrErr =
-              createTextureWithData(*this, "Texture", CreateDesc, Data.get(),
-                                    R.size(), Enc.get(), &UploadBuffer);
-          if (!TextureOrErr)
-            return TextureOrErr.takeError();
-          auto Texture = std::move(*TextureOrErr);
+          std::unique_ptr<offloadtest::MemoryHeap> BackingMemoryHeap;
+
+          std::unique_ptr<offloadtest::Texture> 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<Buffer> ReadbackBuffer;
           if (getDescriptorKind(R.Kind) == DescriptorKind::UAV) {
@@ -2349,6 +2391,7 @@ class DXDevice : public offloadtest::Device {
           }
 
           ResourceSet RSet(std::move(UploadBuffer), std::move(Texture),
+                           std::move(BackingMemoryHeap),
                            std::move(ReadbackBuffer));
           ResBundle.push_back(std::move(RSet));
         }
@@ -3048,6 +3091,18 @@ class DXDevice : public offloadtest::Device {
   }
 };
 
+TileShape DXTexture::querySparseTileShape(const Device &Dev) const {
+  const DXDevice &DeviceDX = llvm::cast<DXDevice>(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};
+}
+
 llvm::Error DXComputeEncoder::batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) {
   if (Items.empty())
     return llvm::Error::success();