[Metal] Make dispatchRays a self-contained encoder operation

createRayTracingCommands was a backend-private device method that bound
the descriptor heap and top-level argument buffer, synthesized and bound
the IRDispatchRaysArgument, marked RT resources resident, and only then
called the encoder's thin dispatchRays(). The DX and Vulkan backends
instead create an encoder and call Encoder.dispatchRays() directly.

Mirror that here. Since Metal lowers raygen to a compute kernel, route
ray tracing through the shared createComputeCommands path, which already
binds the descriptor heap / argument buffer and marks the common
resources resident, and now branches on P.isRayTracing(). The
IRDispatchRaysArgument construction, upload, binding, and RT-specific
residency (SBT buffer, visible/intersection function tables, the
argument buffer itself) move into MTLComputeEncoder::dispatchRays, which
is defined out-of-line so it can allocate via CB->Dev and keep the
buffer alive in CB->KeepAliveOwned, like batchBuildAS.

IRDispatchRaysArgument::ResDescHeap needs the resource descriptor heap
address, which is not in the cross-backend dispatchRays signature, so
the encoder remembers the heap bound to it via bindResourceHeap() and
reads its GPU address back when dispatching. createRayTracingCommands is
deleted.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: MarijnS95

lib/API/MTL/MTLDevice.cpp

@@ -609,6 +609,10 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder {
   /// the next encoder transition (via endEncodingImpl).
   MTL::AccelerationStructureCommandEncoder *ASEnc = nullptr;
 
+  /// Resource descriptor heap bound to this encoder, if any. Remembered so a
+  /// ray dispatch can point IRDispatchRaysArgument::ResDescHeap at it.
+  MTLDescriptorHeap *BoundResourceHeap = nullptr;
+
   /// Accumulated barrier scope from commands recorded since the last barrier.
   MTL::BarrierScope PendingScope = MTL::BarrierScope(0);
 
@@ -665,6 +669,16 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder {
 
   MTL::ComputeCommandEncoder *getNative() const { return ComputeEnc; }
 
+  /// Bind \p Heap as the global resource descriptor heap and remember it so a
+  /// subsequent dispatchRays can reference it from IRDispatchRaysArgument.
+  llvm::Error bindResourceHeap(MTLDescriptorHeap &Heap) {
+    if (auto Err = ensureComputeEncoder())
+      return Err;
+    Heap.bind(ComputeEnc);
+    BoundResourceHeap = &Heap;
+    return llvm::Error::success();
+  }
+
   MTL::CommandEncoder *getActiveEncoder() const {
     if (ComputeEnc)
       return ComputeEnc;
@@ -787,40 +801,17 @@ class MTLComputeEncoder : public offloadtest::ComputeEncoder {
   // MTL::Device handle (used to allocate scratch and instance buffers).
   llvm::Error batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) override;
 
-  // Dispatch threads using a raygen compute kernel synthesized by the
-  // irconverter. All bindings (descriptor heap, top-level argument buffer,
-  // IRDispatchRaysArgument at slot 3, visible/intersection function tables,
-  // 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<MTLRayTracingPipelineState>(&PSO))
-      return llvm::createStringError(
-          std::errc::invalid_argument,
-          "dispatchRays requires a RayTracing PipelineState.");
-    const auto &RTPSO = llvm::cast<MTLRayTracingPipelineState>(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();
-  }
+  // Issue a ray dispatch: synthesizes the per-dispatch IRDispatchRaysArgument
+  // (SBT region addresses, GRS + resource-heap pointers, visible/intersection
+  // function-table IDs), binds it at kIRRayDispatchArgumentsBindPoint, marks
+  // the RT-specific resources resident, and dispatches the irconverter raygen
+  // compute kernel. The caller must have already bound the descriptor heap and
+  // top-level argument buffer (the compute-dispatch path does the same).
+  // Defined out-of-line below — allocating the argument buffer needs
+  // MTLDevice's full type.
+  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.
@@ -1742,7 +1733,8 @@ class MTLDevice : public offloadtest::Device {
     const auto &PS = llvm::cast<MTLPipelineState>(IS.Pipeline.get());
     MTLGPUDescriptorHandle Handle = {};
     if (IS.DescHeap) {
-      IS.DescHeap->bind(NativeEncoder);
+      if (auto Err = Encoder.bindResourceHeap(*IS.DescHeap))
+        return Err;
       Handle = IS.DescHeap->getGPUDescriptorHandleForHeapStart();
     }
 
@@ -1772,7 +1764,18 @@ class MTLDevice : public offloadtest::Device {
       NativeEncoder->useResource(llvm::cast<MTLBuffer>(B.get())->Buf,
                                  MTL::ResourceUsageRead);
 
-    if (auto Err = Encoder.dispatch(*IS.Pipeline.get(),
+    // Metal compiles raygen as a compute kernel, so a ray-tracing pipeline
+    // dispatches through the same compute encoder as a plain compute shader;
+    // dispatchRays adds the RT-specific argument buffer and residency.
+    if (P.isRayTracing()) {
+      if (auto Err =
+              Encoder.dispatchRays(*IS.Pipeline.get(), *IS.SBT.get(),
+                                   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]))
@@ -1781,112 +1784,6 @@ class MTLDevice : public offloadtest::Device {
     return llvm::Error::success();
   }
 
-  llvm::Error createRayTracingCommands(Pipeline &P, InvocationState &IS) {
-    auto EncoderOrErr = IS.CB->createComputeEncoder();
-    if (!EncoderOrErr)
-      return EncoderOrErr.takeError();
-    auto &Encoder = llvm::cast<MTLComputeEncoder>(*EncoderOrErr.get());
-    MTL::ComputeCommandEncoder *NativeEncoder = Encoder.getNative();
-
-    const auto &RTPSO =
-        llvm::cast<MTLRayTracingPipelineState>(*IS.Pipeline.get());
-    const auto &SBT = llvm::cast<MTLShaderBindingTable>(*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<MTLBuffer>(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<offloadtest::AccelerationStructure> &AS) {
-          auto *MTLAS = llvm::cast<MetalAccelerationStructure>(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<MTLBuffer>(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);
-
-    if (auto Err =
-            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(
@@ -3094,7 +2991,9 @@ class MTLDevice : public offloadtest::Device {
       IS.SBT = std::move(*SBTOrErr);
       llvm::outs() << "Shader Binding Table created.\n";
 
-      if (auto Err = createRayTracingCommands(P, IS))
+      // Metal lowers raygen to a compute kernel, so ray tracing records its
+      // dispatch through the shared compute-command path.
+      if (auto Err = createComputeCommands(P, IS))
         return Err;
     }
 
@@ -3303,6 +3202,88 @@ llvm::Error MTLComputeEncoder::batchBuildAS(llvm::ArrayRef<ASBuildItem> Items) {
 
   return llvm::Error::success();
 }
+
+llvm::Error MTLComputeEncoder::dispatchRays(const PipelineState &PSO,
+                                            const ShaderBindingTable &SBT,
+                                            uint32_t Width, uint32_t Height,
+                                            uint32_t Depth) {
+  if (!llvm::isa<MTLRayTracingPipelineState>(&PSO))
+    return llvm::createStringError(
+        std::errc::invalid_argument,
+        "dispatchRays requires a RayTracing PipelineState.");
+  if (!llvm::isa<MTLShaderBindingTable>(&SBT))
+    return llvm::createStringError(
+        std::errc::invalid_argument,
+        "dispatchRays requires a Metal ShaderBindingTable.");
+  const auto &RTPSO = llvm::cast<MTLRayTracingPipelineState>(PSO);
+  const auto &MTLSBT = llvm::cast<MTLShaderBindingTable>(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());
+
+  // 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 kIRRayDispatchArgumentsBindPoint and any
+  // visible-function callees inherit it through the same pointer.
+  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 =
+      BoundResourceHeap
+          ? BoundResourceHeap->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(
+      *CB->Dev, "MTL Dispatch Rays Arguments", ArgsBufDesc, &Args,
+      sizeof(IRDispatchRaysArgument), nullptr, nullptr);
+  if (!ArgsBufOrErr)
+    return ArgsBufOrErr.takeError();
+  auto *MTLArgsBuf = llvm::cast<MTLBuffer>(ArgsBufOrErr->get());
+  CB->KeepAliveOwned.push_back(std::move(*ArgsBufOrErr));
+
+  ComputeEnc->setBuffer(MTLArgsBuf->Buf, 0, kIRRayDispatchArgumentsBindPoint);
+  ComputeEnc->useResource(MTLArgsBuf->Buf, MTL::ResourceUsageRead);
+
+  // RT-specific residency beyond what the shared compute path marks: the SBT
+  // buffer (the raygen kernel dereferences SBT addresses) and the visible /
+  // intersection function tables.
+  if (MTLSBT.Buffer)
+    ComputeEnc->useResource(MTLSBT.Buffer, MTL::ResourceUsageRead);
+  if (RTPSO.VFT)
+    ComputeEnc->useResource(RTPSO.VFT, MTL::ResourceUsageRead);
+  if (RTPSO.IFT)
+    ComputeEnc->useResource(RTPSO.IFT, MTL::ResourceUsageRead);
+
+  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();
+}
 } // namespace
 
 llvm::Error offloadtest::initializeMetalDevices(