[gfx] Metal improvements (#4337)

* fix binding resources for render pass

* compute vertex buffer binding offset based on root layout

* fix result code

* cleanup

* implement readTextureResource

* more getNativeHandle

---------

Co-authored-by: Yong He <yonghe@outlook.com>

6 files changed, 65 insertions, 36 deletions

diff --git a/tools/gfx/metal/metal-command-encoder.cpp b/tools/gfx/metal/metal-command-encoder.cpp
index fa1df263e..2447295c4 100644
--- a/tools/gfx/metal/metal-command-encoder.cpp
+++ b/tools/gfx/metal/metal-command-encoder.cpp
@@ -360,8 +360,8 @@ Result RenderCommandEncoder::prepareDraw(MTL::RenderCommandEncoder*& encoder)
 
     encoder = m_commandBuffer->getMetalRenderCommandEncoder(m_renderPassDesc.get());
     encoder->setRenderPipelineState(pipeline->m_renderPipelineState.get());
-    // TODO only vertex stage bindings are set for now
-    VertexBindingContext bindingContext;
+
+    RenderBindingContext bindingContext;
     bindingContext.init(m_commandBuffer->m_device, encoder);
     auto program = static_cast<ShaderProgramImpl*>(m_currentPipeline->m_program.get());
     m_commandBuffer->m_rootObject.bindAsRoot(&bindingContext, program->m_rootObjectLayout);
@@ -406,7 +406,7 @@ Result RenderCommandEncoder::drawIndexed(
     SLANG_RETURN_ON_FAIL(prepareDraw(encoder));
     // TODO baseVertex is not supported by Metal
     encoder->drawIndexedPrimitives(m_primitiveType, indexCount, m_indexBufferType, m_indexBuffer, m_indexBufferOffset);
-    return SLANG_E_NOT_IMPLEMENTED;
+    return SLANG_OK;
 }
 
 Result RenderCommandEncoder::drawIndirect(
@@ -483,10 +483,12 @@ Result ComputeCommandEncoder::dispatchCompute(int x, int y, int z)
 
     MTL::ComputeCommandEncoder* encoder = m_commandBuffer->getMetalComputeCommandEncoder();
     encoder->setComputePipelineState(pipeline->m_computePipelineState.get());
+
     ComputeBindingContext bindingContext;
     bindingContext.init(m_commandBuffer->m_device, encoder);
     auto program = static_cast<ShaderProgramImpl*>(m_currentPipeline->m_program.get());
     m_commandBuffer->m_rootObject.bindAsRoot(&bindingContext, program->m_rootObjectLayout);
+
     encoder->dispatchThreadgroups(MTL::Size(x, y, z), pipeline->m_threadGroupSize);
 
     return SLANG_OK;
diff --git a/tools/gfx/metal/metal-device.cpp b/tools/gfx/metal/metal-device.cpp
index 1be9d1019..4a1c02480 100644
--- a/tools/gfx/metal/metal-device.cpp
+++ b/tools/gfx/metal/metal-device.cpp
@@ -189,7 +189,53 @@ SlangResult DeviceImpl::readTextureResource(
 {
     AUTORELEASEPOOL
 
+    TextureResourceImpl* textureImpl = static_cast<TextureResourceImpl*>(texture);
+
+    if (textureImpl->getDesc()->sampleDesc.numSamples > 1)
+    {
     return SLANG_E_NOT_IMPLEMENTED;
+    }
+
+    NS::SharedPtr<MTL::Texture> srcTexture = textureImpl->m_texture;
+
+    const ITextureResource::Desc& desc = *textureImpl->getDesc();
+    Count width = Math::Max(desc.size.width, 1);
+    Count height = Math::Max(desc.size.height, 1);
+    Count depth = Math::Max(desc.size.depth, 1);
+    FormatInfo formatInfo;
+    gfxGetFormatInfo(desc.format, &formatInfo);
+    Size bytesPerPixel = formatInfo.blockSizeInBytes / formatInfo.pixelsPerBlock;
+    Size bytesPerRow = Size(width) * bytesPerPixel;
+    Size bytesPerSlice = Size(height) * bytesPerRow;
+    Size bufferSize = Size(depth) * bytesPerSlice;
+    if (outRowPitch)
+        *outRowPitch = bytesPerRow;
+    if (outPixelSize)
+        *outPixelSize = bytesPerPixel;
+
+    // create staging buffer
+    NS::SharedPtr<MTL::Buffer> stagingBuffer = NS::TransferPtr(m_device->newBuffer(bufferSize, MTL::StorageModeShared));
+    if (!stagingBuffer)
+    {
+        return SLANG_FAIL;
+    }
+
+    MTL::CommandBuffer* commandBuffer = m_commandQueue->commandBuffer();
+    MTL::BlitCommandEncoder* encoder = commandBuffer->blitCommandEncoder();
+    encoder->copyFromTexture(
+        srcTexture.get(), 0, 0, MTL::Origin(0, 0, 0), MTL::Size(width, height, depth),
+        stagingBuffer.get(), 0, bytesPerRow, bytesPerSlice);
+    encoder->endEncoding();
+    commandBuffer->commit();
+    commandBuffer->waitUntilCompleted();
+
+    List<uint8_t> blobData;
+    blobData.setCount(bufferSize);
+    ::memcpy(blobData.getBuffer(), stagingBuffer->contents(), bufferSize);
+    auto blob = ListBlob::moveCreate(blobData);
+
+    returnComPtr(outBlob, blob);
+    return SLANG_OK;
 }
 
 SlangResult DeviceImpl::readBufferResource(
diff --git a/tools/gfx/metal/metal-fence.cpp b/tools/gfx/metal/metal-fence.cpp
index c8c318d7f..a908a8493 100644
--- a/tools/gfx/metal/metal-fence.cpp
+++ b/tools/gfx/metal/metal-fence.cpp
@@ -1,4 +1,4 @@
-// metal-gfence.cpp
+// metal-fence.cpp
 #include "metal-fence.h"
 #include "metal-device.h"
 
diff --git a/tools/gfx/metal/metal-helper-functions.h b/tools/gfx/metal/metal-helper-functions.h
index 9fe4c9ff9..bdc1f4e88 100644
--- a/tools/gfx/metal/metal-helper-functions.h
+++ b/tools/gfx/metal/metal-helper-functions.h
@@ -85,7 +85,7 @@ struct ComputeBindingContext : public BindingContext
     }
 };
 
-struct VertexBindingContext : public BindingContext
+struct RenderBindingContext : public BindingContext
 {
     MTL::RenderCommandEncoder* encoder;
 
@@ -99,42 +99,18 @@ struct VertexBindingContext : public BindingContext
     void setBuffer(MTL::Buffer* buffer, NS::UInteger index) override
     {
         encoder->setVertexBuffer(buffer, 0, index);
-    }
-
-    void setTexture(MTL::Texture* texture, NS::UInteger index) override
-    {
-        encoder->setVertexTexture(texture, index);
-    }
-
-    void setSampler(MTL::SamplerState* sampler, NS::UInteger index) override
-    {
-        encoder->setVertexSamplerState(sampler, index);
-    }
-};
-
-struct FragmentBindingContext : public BindingContext
-{
-    MTL::RenderCommandEncoder* encoder;
-
-    Result init(DeviceImpl* device, MTL::RenderCommandEncoder* encoder)
-    {
-        this->device = device;
-        this->encoder = encoder;
-        return SLANG_OK;
-    }
-
-    void setBuffer(MTL::Buffer* buffer, NS::UInteger index) override
-    {
         encoder->setFragmentBuffer(buffer, 0, index);
     }
 
     void setTexture(MTL::Texture* texture, NS::UInteger index) override
     {
+        encoder->setVertexTexture(texture, index);
         encoder->setFragmentTexture(texture, index);
     }
 
     void setSampler(MTL::SamplerState* sampler, NS::UInteger index) override
     {
+        encoder->setVertexSamplerState(sampler, index);
         encoder->setFragmentSamplerState(sampler, index);
     }
 };
diff --git a/tools/gfx/metal/metal-pipeline-state.cpp b/tools/gfx/metal/metal-pipeline-state.cpp
index b23aab2b9..d8ca3f793 100644
--- a/tools/gfx/metal/metal-pipeline-state.cpp
+++ b/tools/gfx/metal/metal-pipeline-state.cpp
@@ -76,9 +76,10 @@ Result PipelineStateImpl::createMetalRenderPipelineState()
         }
     }
 
-    // Offset vertex buffer indices in vertex layout.
-    // They need to be in a range not used by other buffer bindings.
-    m_vertexBufferOffset = 10; // TODO get from layout
+    // Create a vertex descriptor with the vertex buffer binding indices being offset.
+    // They need to be in a range not used by any buffers in the root object layout.
+    // The +1 is to account for a potential constant buffer at index 0.
+    m_vertexBufferOffset = programImpl->m_rootObjectLayout->getBufferCount() + 1;
     auto inputLayoutImpl = static_cast<InputLayoutImpl*>(desc.graphics.inputLayout);
     NS::SharedPtr<MTL::VertexDescriptor> vertexDescriptor = inputLayoutImpl->createVertexDescriptor(m_vertexBufferOffset);
     pd->setVertexDescriptor(vertexDescriptor.get());
diff --git a/tools/gfx/metal/metal-resource-views.cpp b/tools/gfx/metal/metal-resource-views.cpp
index 76c0f21ed..ea5c23c10 100644
--- a/tools/gfx/metal/metal-resource-views.cpp
+++ b/tools/gfx/metal/metal-resource-views.cpp
@@ -15,7 +15,9 @@ TextureResourceViewImpl::~TextureResourceViewImpl()
 
 Result TextureResourceViewImpl::getNativeHandle(InteropHandle* outHandle)
 {
-    return SLANG_E_NOT_IMPLEMENTED;
+    outHandle->api = InteropHandleAPI::Metal;
+    outHandle->handleValue = reinterpret_cast<uintptr_t>(m_textureView.get());
+    return SLANG_OK;
 }
 
 BufferResourceViewImpl::~BufferResourceViewImpl()
@@ -24,7 +26,9 @@ BufferResourceViewImpl::~BufferResourceViewImpl()
 
 Result BufferResourceViewImpl::getNativeHandle(InteropHandle* outHandle)
 {
-    return SLANG_E_NOT_IMPLEMENTED;
+    outHandle->api = InteropHandleAPI::Metal;
+    outHandle->handleValue = reinterpret_cast<uintptr_t>(m_buffer->m_buffer.get());
+    return SLANG_OK;
 }
 
 TexelBufferResourceViewImpl::TexelBufferResourceViewImpl(DeviceImpl* device)