HaxeJolt

lib/haxejolt/JoltPhysics/Jolt/Compute/MTL/ComputeBufferMTL.h

@@ -0,0 +1,39 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#ifdef JPH_USE_MTL
+
+#include <Jolt/Compute/MTL/ComputeSystemMTL.h>
+
+JPH_NAMESPACE_BEGIN
+
+/// Buffer that can be read from / written to by a compute shader
+class JPH_EXPORT ComputeBufferMTL final : public ComputeBuffer
+{
+public:
+	JPH_OVERRIDE_NEW_DELETE
+
+	/// Constructor
+									ComputeBufferMTL(ComputeSystemMTL *inComputeSystem, EType inType, uint64 inSize, uint inStride);
+	virtual							~ComputeBufferMTL() override;
+
+	bool							Initialize(const void *inData);
+
+	virtual ComputeBufferResult		CreateReadBackBuffer() const override;
+
+	id<MTLBuffer>					GetBuffer() const							{ return mBuffer; }
+
+private:
+	virtual void *					MapInternal(EMode inMode) override;
+	virtual void					UnmapInternal() override;
+
+	ComputeSystemMTL *				mComputeSystem;
+	id<MTLBuffer>					mBuffer;
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_MTL

lib/haxejolt/JoltPhysics/Jolt/Compute/MTL/ComputeBufferMTL.mm

@@ -0,0 +1,52 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_MTL
+
+#include <Jolt/Compute/MTL/ComputeBufferMTL.h>
+
+JPH_NAMESPACE_BEGIN
+
+ComputeBufferMTL::ComputeBufferMTL(ComputeSystemMTL *inComputeSystem, EType inType, uint64 inSize, uint inStride) :
+	ComputeBuffer(inType, inSize, inStride),
+	mComputeSystem(inComputeSystem)
+{
+}
+
+bool ComputeBufferMTL::Initialize(const void *inData)
+{
+	NSUInteger size = NSUInteger(mSize) * mStride;
+	if (inData != nullptr)
+		mBuffer = [mComputeSystem->GetDevice() newBufferWithBytes: inData length: size options: MTLResourceCPUCacheModeDefaultCache | MTLResourceStorageModeShared | MTLResourceHazardTrackingModeTracked];
+	else
+		mBuffer = [mComputeSystem->GetDevice() newBufferWithLength: size options: MTLResourceCPUCacheModeDefaultCache | MTLResourceStorageModeShared | MTLResourceHazardTrackingModeTracked];
+	return mBuffer != nil;
+}
+
+ComputeBufferMTL::~ComputeBufferMTL()
+{
+	[mBuffer release];
+}
+
+void *ComputeBufferMTL::MapInternal(EMode inMode)
+{
+	return mBuffer.contents;
+}
+
+void ComputeBufferMTL::UnmapInternal()
+{
+}
+
+ComputeBufferResult ComputeBufferMTL::CreateReadBackBuffer() const
+{
+	ComputeBufferResult result;
+	result.Set(const_cast<ComputeBufferMTL *>(this));
+	return result;
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_MTL

lib/haxejolt/JoltPhysics/Jolt/Compute/MTL/ComputeQueueMTL.h

@@ -0,0 +1,49 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#ifdef JPH_USE_MTL
+
+#include <MetalKit/MetalKit.h>
+
+#include <Jolt/Compute/ComputeQueue.h>
+
+JPH_NAMESPACE_BEGIN
+
+class ComputeShaderMTL;
+
+/// A command queue for Metal for executing compute workloads on the GPU.
+class JPH_EXPORT ComputeQueueMTL final : public ComputeQueue
+{
+public:
+	JPH_OVERRIDE_NEW_DELETE
+
+	/// Constructor / destructor
+										ComputeQueueMTL(id<MTLDevice> inDevice);
+	virtual								~ComputeQueueMTL() override;
+
+	// See: ComputeQueue
+	virtual void						SetShader(const ComputeShader *inShader) override;
+	virtual void						SetConstantBuffer(const char *inName, const ComputeBuffer *inBuffer) override;
+	virtual void						SetBuffer(const char *inName, const ComputeBuffer *inBuffer) override;
+	virtual void 						SetRWBuffer(const char *inName, ComputeBuffer *inBuffer, EBarrier inBarrier = EBarrier::Yes) override;
+	virtual void						ScheduleReadback(ComputeBuffer *inDst, const ComputeBuffer *inSrc) override;
+	virtual void						Dispatch(uint inThreadGroupsX, uint inThreadGroupsY, uint inThreadGroupsZ) override;
+	virtual void						Execute() override;
+	virtual void						Wait() override;
+
+private:
+	void								BeginCommandBuffer();
+
+	id<MTLCommandQueue>					mCommandQueue;
+	id<MTLCommandBuffer> 				mCommandBuffer;
+	id<MTLComputeCommandEncoder>		mComputeEncoder;
+	RefConst<ComputeShaderMTL>			mShader;
+	bool								mIsExecuting = false;
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_MTL

lib/haxejolt/JoltPhysics/Jolt/Compute/MTL/ComputeQueueMTL.mm

@@ -0,0 +1,123 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_MTL
+
+#include <Jolt/Compute/MTL/ComputeQueueMTL.h>
+#include <Jolt/Compute/MTL/ComputeShaderMTL.h>
+#include <Jolt/Compute/MTL/ComputeBufferMTL.h>
+#include <Jolt/Compute/MTL/ComputeSystemMTL.h>
+
+JPH_NAMESPACE_BEGIN
+
+ComputeQueueMTL::~ComputeQueueMTL()
+{
+	Wait();
+
+	[mCommandQueue release];
+}
+
+ComputeQueueMTL::ComputeQueueMTL(id<MTLDevice> inDevice)
+{
+	// Create the command queue
+	mCommandQueue = [inDevice newCommandQueue];
+}
+
+void ComputeQueueMTL::BeginCommandBuffer()
+{
+	if (mCommandBuffer == nil)
+	{
+		// Start a new command buffer
+		mCommandBuffer = [mCommandQueue commandBuffer];
+		mComputeEncoder = [mCommandBuffer computeCommandEncoder];
+	}
+}
+
+void ComputeQueueMTL::SetShader(const ComputeShader *inShader)
+{
+	BeginCommandBuffer();
+
+	mShader = static_cast<const ComputeShaderMTL *>(inShader);
+
+	[mComputeEncoder setComputePipelineState: mShader->GetPipelineState()];
+}
+
+void ComputeQueueMTL::SetConstantBuffer(const char *inName, const ComputeBuffer *inBuffer)
+{
+	if (inBuffer == nullptr)
+		return;
+	JPH_ASSERT(inBuffer->GetType() == ComputeBuffer::EType::ConstantBuffer);
+
+	BeginCommandBuffer();
+
+	const ComputeBufferMTL *buffer = static_cast<const ComputeBufferMTL *>(inBuffer);
+	[mComputeEncoder setBuffer: buffer->GetBuffer() offset: 0 atIndex: mShader->NameToBindingIndex(inName)];
+}
+
+void ComputeQueueMTL::SetBuffer(const char *inName, const ComputeBuffer *inBuffer)
+{
+	if (inBuffer == nullptr)
+		return;
+	JPH_ASSERT(inBuffer->GetType() == ComputeBuffer::EType::UploadBuffer || inBuffer->GetType() == ComputeBuffer::EType::Buffer || inBuffer->GetType() == ComputeBuffer::EType::RWBuffer);
+
+	BeginCommandBuffer();
+
+	const ComputeBufferMTL *buffer = static_cast<const ComputeBufferMTL *>(inBuffer);
+	[mComputeEncoder setBuffer: buffer->GetBuffer() offset: 0 atIndex: mShader->NameToBindingIndex(inName)];
+}
+
+void ComputeQueueMTL::SetRWBuffer(const char *inName, ComputeBuffer *inBuffer, EBarrier inBarrier)
+{
+	if (inBuffer == nullptr)
+		return;
+	JPH_ASSERT(inBuffer->GetType() == ComputeBuffer::EType::RWBuffer);
+
+	BeginCommandBuffer();
+
+	const ComputeBufferMTL *buffer = static_cast<const ComputeBufferMTL *>(inBuffer);
+	[mComputeEncoder setBuffer: buffer->GetBuffer() offset: 0 atIndex: mShader->NameToBindingIndex(inName)];
+}
+
+void ComputeQueueMTL::ScheduleReadback(ComputeBuffer *inDst, const ComputeBuffer *inSrc)
+{
+	JPH_ASSERT(inDst == inSrc); // Since ComputeBuffer::CreateReadBackBuffer returns the same buffer, we don't need to copy
+}
+
+void ComputeQueueMTL::Dispatch(uint inThreadGroupsX, uint inThreadGroupsY, uint inThreadGroupsZ)
+{
+	BeginCommandBuffer();
+
+	MTLSize thread_groups = MTLSizeMake(inThreadGroupsX, inThreadGroupsY, inThreadGroupsZ);
+	MTLSize group_size = MTLSizeMake(mShader->GetGroupSizeX(), mShader->GetGroupSizeY(), mShader->GetGroupSizeZ());
+	[mComputeEncoder dispatchThreadgroups: thread_groups threadsPerThreadgroup: group_size];
+}
+
+void ComputeQueueMTL::Execute()
+{
+	// End command buffer
+	if (mCommandBuffer == nil)
+		return;
+
+	[mComputeEncoder endEncoding];
+	[mCommandBuffer commit];
+	mShader = nullptr;
+	mIsExecuting = true;
+}
+
+void ComputeQueueMTL::Wait()
+{
+	if (!mIsExecuting)
+		return;
+
+	[mCommandBuffer waitUntilCompleted];
+	mComputeEncoder = nil;
+	mCommandBuffer = nil;
+	mIsExecuting = false;
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_MTL

lib/haxejolt/JoltPhysics/Jolt/Compute/MTL/ComputeShaderMTL.h

@@ -0,0 +1,39 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#ifdef JPH_USE_MTL
+
+#include <MetalKit/MetalKit.h>
+
+#include <Jolt/Compute/ComputeShader.h>
+#include <Jolt/Core/UnorderedMap.h>
+
+JPH_NAMESPACE_BEGIN
+
+/// Compute shader handle for Metal
+class JPH_EXPORT ComputeShaderMTL : public ComputeShader
+{
+public:
+	JPH_OVERRIDE_NEW_DELETE
+
+	/// Constructor
+								ComputeShaderMTL(id<MTLComputePipelineState> inPipelineState, MTLComputePipelineReflection *inReflection, uint32 inGroupSizeX, uint32 inGroupSizeY, uint32 inGroupSizeZ);
+	virtual						~ComputeShaderMTL() override 					{ [mPipelineState release]; }
+
+	/// Access to the function
+	id<MTLComputePipelineState>	GetPipelineState() const						{ return mPipelineState; }
+
+	/// Get index of buffer name
+	uint						NameToBindingIndex(const char *inName) const;
+
+private:
+	id<MTLComputePipelineState>	mPipelineState;
+	UnorderedMap<String, uint>	mNameToBindingIndex;
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_MTL

lib/haxejolt/JoltPhysics/Jolt/Compute/MTL/ComputeShaderMTL.mm

@@ -0,0 +1,34 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_MTL
+
+#include <Jolt/Compute/MTL/ComputeShaderMTL.h>
+
+JPH_NAMESPACE_BEGIN
+
+ComputeShaderMTL::ComputeShaderMTL(id<MTLComputePipelineState> inPipelineState, MTLComputePipelineReflection *inReflection, uint32 inGroupSizeX, uint32 inGroupSizeY, uint32 inGroupSizeZ) :
+	ComputeShader(inGroupSizeX, inGroupSizeY, inGroupSizeZ),
+	mPipelineState(inPipelineState)
+{
+	for (id<MTLBinding> binding in inReflection.bindings)
+	{
+		const char *name = [binding.name UTF8String];
+		uint index = uint(binding.index);
+		mNameToBindingIndex[name] = index;
+	}
+}
+
+uint ComputeShaderMTL::NameToBindingIndex(const char *inName) const
+{
+	UnorderedMap<String, uint>::const_iterator it = mNameToBindingIndex.find(inName);
+	JPH_ASSERT(it != mNameToBindingIndex.end());
+	return it->second;
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_MTL

lib/haxejolt/JoltPhysics/Jolt/Compute/MTL/ComputeSystemMTL.h

@@ -0,0 +1,40 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <Jolt/Compute/ComputeSystem.h>
+
+#ifdef JPH_USE_MTL
+
+#include <MetalKit/MetalKit.h>
+
+JPH_NAMESPACE_BEGIN
+
+/// Interface to run a workload on the GPU
+class JPH_EXPORT ComputeSystemMTL : public ComputeSystem
+{
+public:
+	JPH_DECLARE_RTTI_VIRTUAL(JPH_EXPORT, ComputeSystemMTL)
+
+	// Initialize / shutdown the compute system
+	bool							Initialize(id<MTLDevice> inDevice);
+	void							Shutdown();
+
+	// See: ComputeSystem
+	virtual ComputeShaderResult		CreateComputeShader(const char *inName, uint32 inGroupSizeX, uint32 inGroupSizeY, uint32 inGroupSizeZ) override;
+	virtual ComputeBufferResult		CreateComputeBuffer(ComputeBuffer::EType inType, uint64 inSize, uint inStride, const void *inData = nullptr) override;
+	virtual ComputeQueueResult		CreateComputeQueue() override;
+
+	/// Get the metal device
+	id<MTLDevice>					GetDevice() const						{ return mDevice; }
+
+private:
+	id<MTLDevice>					mDevice;
+	id<MTLLibrary>					mShaderLibrary;
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_MTL

lib/haxejolt/JoltPhysics/Jolt/Compute/MTL/ComputeSystemMTL.mm

@@ -0,0 +1,110 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_MTL
+
+#include <Jolt/Compute/MTL/ComputeSystemMTL.h>
+#include <Jolt/Compute/MTL/ComputeBufferMTL.h>
+#include <Jolt/Compute/MTL/ComputeShaderMTL.h>
+#include <Jolt/Compute/MTL/ComputeQueueMTL.h>
+
+JPH_NAMESPACE_BEGIN
+
+JPH_IMPLEMENT_RTTI_VIRTUAL(ComputeSystemMTL)
+{
+	JPH_ADD_BASE_CLASS(ComputeSystemMTL, ComputeSystem)
+}
+
+bool ComputeSystemMTL::Initialize(id<MTLDevice> inDevice)
+{
+	mDevice = [inDevice retain];
+
+	return true;
+}
+
+void ComputeSystemMTL::Shutdown()
+{
+	[mShaderLibrary release];
+	[mDevice release];
+}
+
+ComputeShaderResult ComputeSystemMTL::CreateComputeShader(const char *inName, uint32 inGroupSizeX, uint32 inGroupSizeY, uint32 inGroupSizeZ)
+{
+	ComputeShaderResult result;
+
+	if (mShaderLibrary == nil)
+	{
+		// Load the shader library containing all shaders
+		Array<uint8> *data = new Array<uint8>();
+		String error;
+		if (!mShaderLoader("Jolt.metallib", *data, error))
+		{
+			result.SetError(error);
+			delete data;
+			return result;
+		}
+
+		// Convert to dispatch data
+		dispatch_data_t data_dispatch = dispatch_data_create(data->data(), data->size(), nullptr, ^{ delete data; });
+
+		// Create the library
+		NSError *ns_error = nullptr;
+		mShaderLibrary = [mDevice newLibraryWithData: data_dispatch error: &ns_error];
+		if (ns_error != nil)
+		{
+			result.SetError("Failed to laod shader library");
+			return result;
+		}
+	}
+
+	// Get the shader function
+	id<MTLFunction> function = [mShaderLibrary newFunctionWithName: [NSString stringWithCString: inName encoding: NSUTF8StringEncoding]];
+	if (function == nil)
+	{
+		result.SetError("Failed to instantiate compute shader");
+		return result;
+	}
+
+	// Create the pipeline
+	NSError *error = nil;
+	MTLComputePipelineReflection *reflection = nil;
+	id<MTLComputePipelineState> pipeline_state = [mDevice newComputePipelineStateWithFunction: function options: MTLPipelineOptionBindingInfo | MTLPipelineOptionBufferTypeInfo reflection: &reflection error: &error];
+	if (error != nil || pipeline_state == nil)
+	{
+		result.SetError("Failed to create compute pipeline");
+		[function release];
+		return result;
+	}
+
+	result.Set(new ComputeShaderMTL(pipeline_state, reflection, inGroupSizeX, inGroupSizeY, inGroupSizeZ));
+	return result;
+}
+
+ComputeBufferResult ComputeSystemMTL::CreateComputeBuffer(ComputeBuffer::EType inType, uint64 inSize, uint inStride, const void *inData)
+{
+	ComputeBufferResult result;
+
+	Ref<ComputeBufferMTL> buffer = new ComputeBufferMTL(this, inType, inSize, inStride);
+	if (!buffer->Initialize(inData))
+	{
+		result.SetError("Failed to create compute buffer");
+		return result;
+	}
+
+	result.Set(buffer.GetPtr());
+	return result;
+}
+
+ComputeQueueResult ComputeSystemMTL::CreateComputeQueue()
+{
+	ComputeQueueResult result;
+	result.Set(new ComputeQueueMTL(mDevice));
+	return result;
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_MTL

lib/haxejolt/JoltPhysics/Jolt/Compute/MTL/ComputeSystemMTLImpl.h

@@ -0,0 +1,28 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#ifdef JPH_USE_MTL
+
+#include <Jolt/Compute/MTL/ComputeSystemMTL.h>
+
+JPH_NAMESPACE_BEGIN
+
+/// Interface to run a workload on the GPU that fully initializes Metal.
+class JPH_EXPORT ComputeSystemMTLImpl : public ComputeSystemMTL
+{
+public:
+	JPH_DECLARE_RTTI_VIRTUAL(JPH_EXPORT, ComputeSystemMTLImpl)
+
+	/// Destructor
+	virtual							~ComputeSystemMTLImpl() override;
+
+	/// Initialize / shutdown the compute system
+	bool							Initialize();
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_MTL

lib/haxejolt/JoltPhysics/Jolt/Compute/MTL/ComputeSystemMTLImpl.mm

@@ -0,0 +1,49 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_MTL
+
+#include <Jolt/Compute/MTL/ComputeSystemMTLImpl.h>
+
+JPH_NAMESPACE_BEGIN
+
+JPH_IMPLEMENT_RTTI_VIRTUAL(ComputeSystemMTLImpl)
+{
+	JPH_ADD_BASE_CLASS(ComputeSystemMTLImpl, ComputeSystemMTL)
+}
+
+ComputeSystemMTLImpl::~ComputeSystemMTLImpl()
+{
+	Shutdown();
+
+	[GetDevice() release];
+}
+
+bool ComputeSystemMTLImpl::Initialize()
+{
+	id<MTLDevice> device = MTLCreateSystemDefaultDevice();
+
+	return ComputeSystemMTL::Initialize(device);
+}
+
+ComputeSystemResult CreateComputeSystemMTL()
+{
+	ComputeSystemResult result;
+
+	Ref<ComputeSystemMTLImpl> compute = new ComputeSystemMTLImpl;
+	if (!compute->Initialize())
+	{
+		result.SetError("Failed to initialize compute system");
+		return result;
+	}
+
+	result.Set(compute.GetPtr());
+	return result;
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_MTL