HaxeJolt

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/BufferVK.h

@@ -0,0 +1,42 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2024 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <Jolt/Compute/VK/IncludeVK.h>
+#include <Jolt/Core/Reference.h>
+#include <Jolt/Core/NonCopyable.h>
+
+JPH_NAMESPACE_BEGIN
+
+/// Simple wrapper class to manage a Vulkan memory block
+class MemoryVK : public RefTarget<MemoryVK>, public NonCopyable
+{
+public:
+								~MemoryVK()
+	{
+		// We should have unmapped and freed the block before destruction
+		JPH_ASSERT(mMappedCount == 0);
+		JPH_ASSERT(mMemory == VK_NULL_HANDLE);
+	}
+
+	VkDeviceMemory				mMemory = VK_NULL_HANDLE;		///< The Vulkan memory handle
+	VkDeviceSize				mSize = 0;						///< Size of the memory block
+	VkDeviceSize				mBufferSize = 0;				///< Size of each of the buffers that this memory block has been divided into
+	VkMemoryPropertyFlags		mProperties = 0;				///< Vulkan memory properties used to allocate this block
+	int							mMappedCount = 0;				///< How often buffers using this memory block were mapped
+	void *						mMappedPtr = nullptr;			///< The CPU address of the memory block when mapped
+};
+
+/// Simple wrapper class to manage a Vulkan buffer
+class BufferVK
+{
+public:
+	Ref<MemoryVK>				mMemory;						///< The memory block that contains the buffer (note that filling this in is optional if you do your own buffer allocation)
+	VkBuffer					mBuffer = VK_NULL_HANDLE;		///< The Vulkan buffer handle
+	VkDeviceSize				mOffset = 0;					///< Offset in the memory block where the buffer starts
+	VkDeviceSize				mSize = 0;						///< Real size of the buffer
+};
+
+JPH_NAMESPACE_END

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeBufferVK.cpp

@@ -0,0 +1,140 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/ComputeBufferVK.h>
+#include <Jolt/Compute/VK/ComputeSystemVK.h>
+
+JPH_NAMESPACE_BEGIN
+
+ComputeBufferVK::ComputeBufferVK(ComputeSystemVK *inComputeSystem, EType inType, uint64 inSize, uint inStride) :
+	ComputeBuffer(inType, inSize, inStride),
+	mComputeSystem(inComputeSystem)
+{
+}
+
+bool ComputeBufferVK::Initialize(const void *inData)
+{
+	VkDeviceSize buffer_size = VkDeviceSize(mSize * mStride);
+
+	switch (mType)
+	{
+	case EType::Buffer:
+		JPH_ASSERT(inData != nullptr);
+		[[fallthrough]];
+
+	case EType::UploadBuffer:
+	case EType::RWBuffer:
+		if (!mComputeSystem->CreateBuffer(buffer_size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT | VK_MEMORY_PROPERTY_HOST_CACHED_BIT, mBufferCPU))
+			return false;
+		if (!mComputeSystem->CreateBuffer(buffer_size, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, mBufferGPU))
+			return false;
+		if (inData != nullptr)
+		{
+			void *data = mComputeSystem->MapBuffer(mBufferCPU);
+			memcpy(data, inData, size_t(buffer_size));
+			mComputeSystem->UnmapBuffer(mBufferCPU);
+			mNeedsSync = true;
+		}
+		break;
+
+	case EType::ConstantBuffer:
+		if (!mComputeSystem->CreateBuffer(buffer_size, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT | VK_MEMORY_PROPERTY_HOST_CACHED_BIT, mBufferCPU))
+			return false;
+		if (inData != nullptr)
+		{
+			void* data = mComputeSystem->MapBuffer(mBufferCPU);
+			memcpy(data, inData, size_t(buffer_size));
+			mComputeSystem->UnmapBuffer(mBufferCPU);
+		}
+		break;
+
+	case EType::ReadbackBuffer:
+		JPH_ASSERT(inData == nullptr, "Can't upload data to a readback buffer");
+		if (!mComputeSystem->CreateBuffer(buffer_size, VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT | VK_MEMORY_PROPERTY_HOST_CACHED_BIT, mBufferCPU))
+			return false;
+		break;
+	}
+
+	return true;
+}
+
+ComputeBufferVK::~ComputeBufferVK()
+{
+	mComputeSystem->FreeBuffer(mBufferGPU);
+	mComputeSystem->FreeBuffer(mBufferCPU);
+}
+
+void ComputeBufferVK::Barrier(VkCommandBuffer inCommandBuffer, VkPipelineStageFlags inToStage, VkAccessFlagBits inToFlags, bool inForce) const
+{
+	if (mAccessStage == inToStage && mAccessFlagBits == inToFlags && !inForce)
+		return;
+
+	VkBufferMemoryBarrier b = {};
+	b.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
+	b.srcAccessMask = mAccessFlagBits;
+	b.dstAccessMask = inToFlags;
+	b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
+	b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
+	b.buffer = mBufferGPU.mBuffer != VK_NULL_HANDLE? mBufferGPU.mBuffer : mBufferCPU.mBuffer;
+	b.offset = 0;
+	b.size = VK_WHOLE_SIZE;
+	vkCmdPipelineBarrier(inCommandBuffer, mAccessStage, inToStage, 0, 0, nullptr, 1, &b, 0, nullptr);
+
+	mAccessStage = inToStage;
+	mAccessFlagBits = inToFlags;
+}
+
+bool ComputeBufferVK::SyncCPUToGPU(VkCommandBuffer inCommandBuffer) const
+{
+	if (!mNeedsSync)
+		return false;
+
+	// Barrier before write
+	Barrier(inCommandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, false);
+
+	// Copy from CPU to GPU
+	VkBufferCopy copy = {};
+	copy.srcOffset = 0;
+	copy.dstOffset = 0;
+	copy.size = GetSize() * GetStride();
+	vkCmdCopyBuffer(inCommandBuffer, mBufferCPU.mBuffer, mBufferGPU.mBuffer, 1, &copy);
+
+	mNeedsSync = false;
+	return true;
+}
+
+void *ComputeBufferVK::MapInternal(EMode inMode)
+{
+	switch (inMode)
+	{
+	case EMode::Read:
+		JPH_ASSERT(mType == EType::ReadbackBuffer);
+		break;
+
+	case EMode::Write:
+		JPH_ASSERT(mType == EType::UploadBuffer || mType == EType::ConstantBuffer);
+		mNeedsSync = true;
+		break;
+	}
+
+	return mComputeSystem->MapBuffer(mBufferCPU);
+}
+
+void ComputeBufferVK::UnmapInternal()
+{
+	mComputeSystem->UnmapBuffer(mBufferCPU);
+}
+
+ComputeBufferResult ComputeBufferVK::CreateReadBackBuffer() const
+{
+	return mComputeSystem->CreateComputeBuffer(ComputeBuffer::EType::ReadbackBuffer, mSize, mStride);
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeBufferVK.h

@@ -0,0 +1,52 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <Jolt/Compute/ComputeBuffer.h>
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/BufferVK.h>
+
+JPH_NAMESPACE_BEGIN
+
+class ComputeSystemVK;
+
+/// Buffer that can be read from / written to by a compute shader
+class JPH_EXPORT ComputeBufferVK final : public ComputeBuffer
+{
+public:
+	JPH_OVERRIDE_NEW_DELETE
+
+	/// Constructor
+									ComputeBufferVK(ComputeSystemVK *inComputeSystem, EType inType, uint64 inSize, uint inStride);
+	virtual							~ComputeBufferVK() override;
+
+	bool							Initialize(const void *inData);
+
+	virtual ComputeBufferResult		CreateReadBackBuffer() const override;
+
+	VkBuffer						GetBufferCPU() const									{ return mBufferCPU.mBuffer; }
+	VkBuffer						GetBufferGPU() const									{ return mBufferGPU.mBuffer; }
+	BufferVK						ReleaseBufferCPU() const								{ BufferVK tmp = mBufferCPU; mBufferCPU = BufferVK(); return tmp; }
+
+	void							Barrier(VkCommandBuffer inCommandBuffer, VkPipelineStageFlags inToStage, VkAccessFlagBits inToFlags, bool inForce) const;
+	bool							SyncCPUToGPU(VkCommandBuffer inCommandBuffer) const;
+
+private:
+	virtual void *					MapInternal(EMode inMode) override;
+	virtual void					UnmapInternal() override;
+
+	ComputeSystemVK *				mComputeSystem;
+	mutable BufferVK				mBufferCPU;
+	BufferVK						mBufferGPU;
+	mutable bool					mNeedsSync = false;										///< If this buffer needs to be synced from CPU to GPU
+	mutable VkAccessFlagBits		mAccessFlagBits = VK_ACCESS_SHADER_READ_BIT;			///< Access flags of the last usage, used for barriers
+	mutable VkPipelineStageFlags	mAccessStage = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;	///< Pipeline stage of the last usage, used for barriers
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeQueueVK.cpp

@@ -0,0 +1,304 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/ComputeQueueVK.h>
+#include <Jolt/Compute/VK/ComputeBufferVK.h>
+#include <Jolt/Compute/VK/ComputeSystemVK.h>
+
+JPH_NAMESPACE_BEGIN
+
+ComputeQueueVK::~ComputeQueueVK()
+{
+	Wait();
+
+	VkDevice device = mComputeSystem->GetDevice();
+
+	if (mCommandBuffer != VK_NULL_HANDLE)
+		vkFreeCommandBuffers(device, mCommandPool, 1, &mCommandBuffer);
+
+	if (mCommandPool != VK_NULL_HANDLE)
+		vkDestroyCommandPool(device, mCommandPool, nullptr);
+
+	if (mDescriptorPool != VK_NULL_HANDLE)
+		vkDestroyDescriptorPool(device, mDescriptorPool, nullptr);
+
+	if (mFence != VK_NULL_HANDLE)
+		vkDestroyFence(device, mFence, nullptr);
+}
+
+bool ComputeQueueVK::Initialize(uint32 inComputeQueueIndex, ComputeQueueResult &outResult)
+{
+	// Get the queue
+	VkDevice device = mComputeSystem->GetDevice();
+	vkGetDeviceQueue(device, inComputeQueueIndex, 0, &mQueue);
+
+	// Create a command pool
+	VkCommandPoolCreateInfo pool_info = {};
+	pool_info.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
+	pool_info.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
+	pool_info.queueFamilyIndex = inComputeQueueIndex;
+	if (VKFailed(vkCreateCommandPool(device, &pool_info, nullptr, &mCommandPool), outResult))
+		return false;
+
+	// Create descriptor pool
+	VkDescriptorPoolSize descriptor_pool_sizes[] = {
+		{ VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1024 },
+		{ VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 16 * 1024 },
+	};
+	VkDescriptorPoolCreateInfo descriptor_info = {};
+	descriptor_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
+	descriptor_info.poolSizeCount = (uint32)std::size(descriptor_pool_sizes);
+	descriptor_info.pPoolSizes = descriptor_pool_sizes;
+	descriptor_info.maxSets = 256;
+	if (VKFailed(vkCreateDescriptorPool(device, &descriptor_info, nullptr, &mDescriptorPool), outResult))
+		return false;
+
+	// Create a command buffer
+	VkCommandBufferAllocateInfo alloc_info = {};
+	alloc_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
+	alloc_info.commandPool = mCommandPool;
+	alloc_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
+	alloc_info.commandBufferCount = 1;
+	if (VKFailed(vkAllocateCommandBuffers(device, &alloc_info, &mCommandBuffer), outResult))
+		return false;
+
+	// Create a fence
+	VkFenceCreateInfo fence_info = {};
+	fence_info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
+	if (VKFailed(vkCreateFence(device, &fence_info, nullptr, &mFence), outResult))
+		return false;
+
+	return true;
+}
+
+bool ComputeQueueVK::BeginCommandBuffer()
+{
+	if (!mCommandBufferRecording)
+	{
+		VkCommandBufferBeginInfo begin_info = {};
+		begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
+		begin_info.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
+		if (VKFailed(vkBeginCommandBuffer(mCommandBuffer, &begin_info)))
+			return false;
+		mCommandBufferRecording = true;
+	}
+	return true;
+}
+
+void ComputeQueueVK::SetShader(const ComputeShader *inShader)
+{
+	mShader = static_cast<const ComputeShaderVK *>(inShader);
+	mBufferInfos = mShader->GetBufferInfos();
+}
+
+void ComputeQueueVK::SetConstantBuffer(const char *inName, const ComputeBuffer *inBuffer)
+{
+	if (inBuffer == nullptr)
+		return;
+	JPH_ASSERT(inBuffer->GetType() == ComputeBuffer::EType::ConstantBuffer);
+
+	if (!BeginCommandBuffer())
+		return;
+
+	const ComputeBufferVK *buffer = static_cast<const ComputeBufferVK *>(inBuffer);
+	buffer->Barrier(mCommandBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_UNIFORM_READ_BIT, false);
+
+	uint index = mShader->NameToBufferInfoIndex(inName);
+	JPH_ASSERT(mShader->GetLayoutBindings()[index].descriptorType == VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER);
+	mBufferInfos[index].buffer = buffer->GetBufferCPU();
+
+	mUsedBuffers.insert(buffer);
+}
+
+void ComputeQueueVK::SyncCPUToGPU(const ComputeBufferVK *inBuffer)
+{
+	// Ensure that any CPU writes are visible to the GPU
+	if (inBuffer->SyncCPUToGPU(mCommandBuffer)
+		&& (inBuffer->GetType() == ComputeBuffer::EType::Buffer || inBuffer->GetType()  == ComputeBuffer::EType::RWBuffer))
+	{
+		// After the first upload, the CPU buffer is no longer needed for Buffer and RWBuffer types
+		mDelayedFreedBuffers.push_back(inBuffer->ReleaseBufferCPU());
+	}
+}
+
+void ComputeQueueVK::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);
+
+	if (!BeginCommandBuffer())
+		return;
+
+	const ComputeBufferVK *buffer = static_cast<const ComputeBufferVK *>(inBuffer);
+	SyncCPUToGPU(buffer);
+	buffer->Barrier(mCommandBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_SHADER_READ_BIT, false);
+
+	uint index = mShader->NameToBufferInfoIndex(inName);
+	JPH_ASSERT(mShader->GetLayoutBindings()[index].descriptorType == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
+	mBufferInfos[index].buffer = buffer->GetBufferGPU();
+
+	mUsedBuffers.insert(buffer);
+}
+
+void ComputeQueueVK::SetRWBuffer(const char *inName, ComputeBuffer *inBuffer, EBarrier inBarrier)
+{
+	if (inBuffer == nullptr)
+		return;
+	JPH_ASSERT(inBuffer->GetType() == ComputeBuffer::EType::RWBuffer);
+
+	if (!BeginCommandBuffer())
+		return;
+
+	const ComputeBufferVK *buffer = static_cast<const ComputeBufferVK *>(inBuffer);
+	SyncCPUToGPU(buffer);
+	buffer->Barrier(mCommandBuffer, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VkAccessFlagBits(VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT), inBarrier == EBarrier::Yes);
+
+	uint index = mShader->NameToBufferInfoIndex(inName);
+	JPH_ASSERT(mShader->GetLayoutBindings()[index].descriptorType == VK_DESCRIPTOR_TYPE_STORAGE_BUFFER);
+	mBufferInfos[index].buffer = buffer->GetBufferGPU();
+
+	mUsedBuffers.insert(buffer);
+}
+
+void ComputeQueueVK::ScheduleReadback(ComputeBuffer *inDst, const ComputeBuffer *inSrc)
+{
+	if (inDst == nullptr || inSrc == nullptr)
+		return;
+	JPH_ASSERT(inDst->GetType() == ComputeBuffer::EType::ReadbackBuffer);
+
+	if (!BeginCommandBuffer())
+		return;
+
+	const ComputeBufferVK *src_vk = static_cast<const ComputeBufferVK *>(inSrc);
+	const ComputeBufferVK *dst_vk = static_cast<ComputeBufferVK *>(inDst);
+
+	// Barrier to start reading from GPU buffer and writing to CPU buffer
+	src_vk->Barrier(mCommandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_READ_BIT, false);
+	dst_vk->Barrier(mCommandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, false);
+
+	// Copy
+	VkBufferCopy copy = {};
+	copy.srcOffset = 0;
+	copy.dstOffset = 0;
+	copy.size = src_vk->GetSize() * src_vk->GetStride();
+	vkCmdCopyBuffer(mCommandBuffer, src_vk->GetBufferGPU(), dst_vk->GetBufferCPU(), 1, &copy);
+
+	// Barrier to indicate that CPU can read from the buffer
+	dst_vk->Barrier(mCommandBuffer, VK_PIPELINE_STAGE_HOST_BIT, VK_ACCESS_HOST_READ_BIT, false);
+
+	mUsedBuffers.insert(src_vk);
+	mUsedBuffers.insert(dst_vk);
+}
+
+void ComputeQueueVK::Dispatch(uint inThreadGroupsX, uint inThreadGroupsY, uint inThreadGroupsZ)
+{
+	if (!BeginCommandBuffer())
+		return;
+
+	vkCmdBindPipeline(mCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, mShader->GetPipeline());
+
+	VkDevice device = mComputeSystem->GetDevice();
+	const Array<VkDescriptorSetLayoutBinding> &ds_bindings = mShader->GetLayoutBindings();
+	if (!ds_bindings.empty())
+	{
+		// Create a descriptor set
+		VkDescriptorSetAllocateInfo alloc_info = {};
+		alloc_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
+		alloc_info.descriptorPool = mDescriptorPool;
+		alloc_info.descriptorSetCount = 1;
+		VkDescriptorSetLayout ds_layout = mShader->GetDescriptorSetLayout();
+		alloc_info.pSetLayouts = &ds_layout;
+		VkDescriptorSet descriptor_set;
+		if (VKFailed(vkAllocateDescriptorSets(device, &alloc_info, &descriptor_set)))
+			return;
+
+		// Write the values to the descriptor set
+		Array<VkWriteDescriptorSet> writes;
+		writes.reserve(ds_bindings.size());
+		for (uint32 i = 0; i < (uint32)ds_bindings.size(); ++i)
+		{
+			VkWriteDescriptorSet w = {};
+			w.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
+			w.dstSet = descriptor_set;
+			w.dstBinding = ds_bindings[i].binding;
+			w.dstArrayElement = 0;
+			w.descriptorCount = ds_bindings[i].descriptorCount;
+			w.descriptorType = ds_bindings[i].descriptorType;
+			w.pBufferInfo = &mBufferInfos[i];
+			writes.push_back(w);
+		}
+		vkUpdateDescriptorSets(device, (uint32)writes.size(), writes.data(), 0, nullptr);
+
+		// Bind the descriptor set
+		vkCmdBindDescriptorSets(mCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, mShader->GetPipelineLayout(), 0, 1, &descriptor_set, 0, nullptr);
+	}
+
+	vkCmdDispatch(mCommandBuffer, inThreadGroupsX, inThreadGroupsY, inThreadGroupsZ);
+}
+
+void ComputeQueueVK::Execute()
+{
+	// End command buffer
+	if (!mCommandBufferRecording)
+		return;
+	if (VKFailed(vkEndCommandBuffer(mCommandBuffer)))
+		return;
+	mCommandBufferRecording = false;
+
+	// Reset fence
+	VkDevice device = mComputeSystem->GetDevice();
+	if (VKFailed(vkResetFences(device, 1, &mFence)))
+		return;
+
+	// Submit
+	VkSubmitInfo submit = {};
+	submit.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
+	submit.commandBufferCount = 1;
+	submit.pCommandBuffers = &mCommandBuffer;
+	if (VKFailed(vkQueueSubmit(mQueue, 1, &submit, mFence)))
+		return;
+
+	// Clear the current shader
+	mShader = nullptr;
+
+	// Mark that we're executing
+	mIsExecuting = true;
+}
+
+void ComputeQueueVK::Wait()
+{
+	if (!mIsExecuting)
+		return;
+
+	// Wait for the work to complete
+	VkDevice device = mComputeSystem->GetDevice();
+	if (VKFailed(vkWaitForFences(device, 1, &mFence, VK_TRUE, UINT64_MAX)))
+		return;
+
+	// Reset command buffer so it can be reused
+	if (mCommandBuffer != VK_NULL_HANDLE)
+		vkResetCommandBuffer(mCommandBuffer, 0);
+
+	// Allow reusing the descriptors for next run
+	vkResetDescriptorPool(device, mDescriptorPool, 0);
+
+	// Buffers can be freed now
+	mUsedBuffers.clear();
+
+	// Free delayed buffers
+	for (BufferVK &buffer : mDelayedFreedBuffers)
+		mComputeSystem->FreeBuffer(buffer);
+	mDelayedFreedBuffers.clear();
+
+	mIsExecuting = false;
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeQueueVK.h

@@ -0,0 +1,66 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <Jolt/Compute/ComputeQueue.h>
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/ComputeShaderVK.h>
+#include <Jolt/Compute/VK/BufferVK.h>
+#include <Jolt/Core/UnorderedMap.h>
+#include <Jolt/Core/UnorderedSet.h>
+
+JPH_NAMESPACE_BEGIN
+
+class ComputeSystemVK;
+class ComputeBufferVK;
+
+/// A command queue for Vulkan for executing compute workloads on the GPU.
+class JPH_EXPORT ComputeQueueVK final : public ComputeQueue
+{
+public:
+	JPH_OVERRIDE_NEW_DELETE
+
+	/// Constructor / Destructor
+	explicit							ComputeQueueVK(ComputeSystemVK *inComputeSystem) : mComputeSystem(inComputeSystem) { }
+	virtual								~ComputeQueueVK() override;
+
+	/// Initialize the queue
+	bool								Initialize(uint32 inComputeQueueIndex, ComputeQueueResult &outResult);
+
+	// 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:
+	bool								BeginCommandBuffer();
+
+	// Copy the CPU buffer to the GPU buffer if needed
+	void								SyncCPUToGPU(const ComputeBufferVK *inBuffer);
+
+	ComputeSystemVK *					mComputeSystem;
+	VkQueue								mQueue = VK_NULL_HANDLE;
+	VkCommandPool						mCommandPool = VK_NULL_HANDLE;
+	VkDescriptorPool					mDescriptorPool = VK_NULL_HANDLE;
+	VkCommandBuffer						mCommandBuffer = VK_NULL_HANDLE;
+	bool								mCommandBufferRecording = false;				///< If we are currently recording commands into the command buffer
+	VkFence								mFence = VK_NULL_HANDLE;
+	bool								mIsExecuting = false;							///< If Execute has been called and we are waiting for it to finish
+	RefConst<ComputeShaderVK>			mShader;										///< Shader that has been activated
+	Array<VkDescriptorBufferInfo>		mBufferInfos;									///< List of parameters that will be sent to the current shader
+	UnorderedSet<RefConst<ComputeBuffer>> mUsedBuffers;									///< Buffers that are in use by the current execution, these will be retained until execution is finished so that we don't free buffers that are in use
+	Array<BufferVK>						mDelayedFreedBuffers;							///< Hardware buffers that need to be freed after execution is done
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeShaderVK.cpp

@@ -0,0 +1,232 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/ComputeShaderVK.h>
+
+JPH_NAMESPACE_BEGIN
+
+ComputeShaderVK::~ComputeShaderVK()
+{
+	if (mShaderModule != VK_NULL_HANDLE)
+		vkDestroyShaderModule(mDevice, mShaderModule, nullptr);
+
+	if (mDescriptorSetLayout != VK_NULL_HANDLE)
+		vkDestroyDescriptorSetLayout(mDevice, mDescriptorSetLayout, nullptr);
+
+	if (mPipelineLayout != VK_NULL_HANDLE)
+		vkDestroyPipelineLayout(mDevice, mPipelineLayout, nullptr);
+
+	if (mPipeline != VK_NULL_HANDLE)
+		vkDestroyPipeline(mDevice, mPipeline, nullptr);
+}
+
+bool ComputeShaderVK::Initialize(const Array<uint8> &inSPVCode, VkBuffer inDummyBuffer, ComputeShaderResult &outResult)
+{
+	const uint32 *spv_words = reinterpret_cast<const uint32 *>(inSPVCode.data());
+	size_t spv_word_count = inSPVCode.size() / sizeof(uint32);
+
+	// Minimal SPIR-V parser to extract name to binding info
+	UnorderedMap<uint32, String> id_to_name;
+	UnorderedMap<uint32, uint32> id_to_binding;
+	UnorderedMap<uint32, VkDescriptorType> id_to_descriptor_type;
+	UnorderedMap<uint32, uint32> pointer_to_pointee;
+	UnorderedMap<uint32, uint32> var_to_ptr_type;
+	size_t i = 5; // Skip 5 word header
+	while (i < spv_word_count)
+	{
+		// Parse next word
+		uint32 word = spv_words[i];
+		uint16 opcode = uint16(word & 0xffff);
+		uint16 word_count = uint16(word >> 16);
+		if (word_count == 0 || i + word_count > spv_word_count)
+			break;
+
+		switch (opcode)
+		{
+		case 5: // OpName
+			if (word_count >= 2)
+			{
+				uint32 target_id = spv_words[i + 1];
+				const char* name = reinterpret_cast<const char*>(&spv_words[i + 2]);
+				if (*name != 0)
+					id_to_name.insert({ target_id, name });
+			}
+			break;
+
+		case 16: // OpExecutionMode
+			if (word_count >= 6)
+			{
+				uint32 execution_mode = spv_words[i + 2];
+				if (execution_mode == 17) // LocalSize
+				{
+					// Assert that the group size provided matches the one in the shader
+					JPH_ASSERT(GetGroupSizeX() == spv_words[i + 3], "Group size X mismatch");
+					JPH_ASSERT(GetGroupSizeY() == spv_words[i + 4], "Group size Y mismatch");
+					JPH_ASSERT(GetGroupSizeZ() == spv_words[i + 5], "Group size Z mismatch");
+				}
+			}
+			break;
+
+		case 32: // OpTypePointer
+			if (word_count >= 4)
+			{
+				uint32 result_id = spv_words[i + 1];
+				uint32 type_id = spv_words[i + 3];
+				pointer_to_pointee.insert({ result_id, type_id });
+			}
+			break;
+
+		case 59: // OpVariable
+			if (word_count >= 3)
+			{
+				uint32 ptr_type_id = spv_words[i + 1];
+				uint32 result_id = spv_words[i + 2];
+				var_to_ptr_type.insert({ result_id, ptr_type_id });
+			}
+			break;
+
+		case 71: // OpDecorate
+			if (word_count >= 3)
+			{
+				uint32 target_id = spv_words[i + 1];
+				uint32 decoration = spv_words[i + 2];
+				if (decoration == 2) // Block
+				{
+					id_to_descriptor_type.insert({ target_id, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER });
+				}
+				else if (decoration == 3) // BufferBlock
+				{
+					id_to_descriptor_type.insert({ target_id, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER });
+				}
+				else if (decoration == 33 && word_count >= 4) // Binding
+				{
+					uint32 binding = spv_words[i + 3];
+					id_to_binding.insert({ target_id, binding });
+				}
+			}
+			break;
+
+		default:
+			break;
+		}
+
+		i += word_count;
+	}
+
+	// Build name to binding map
+	UnorderedMap<String, std::pair<uint32, VkDescriptorType>> name_to_binding;
+	for (const UnorderedMap<uint32, uint32>::value_type &entry : id_to_binding)
+	{
+		uint32 target_id = entry.first;
+		uint32 binding = entry.second;
+
+		// Get the name of the variable
+		UnorderedMap<uint32, String>::const_iterator it_name = id_to_name.find(target_id);
+		if (it_name != id_to_name.end())
+		{
+			// Find variable that links to the target
+			UnorderedMap<uint32, uint32>::const_iterator it_var_ptr = var_to_ptr_type.find(target_id);
+			if (it_var_ptr != var_to_ptr_type.end())
+			{
+				// Find type pointed at
+				uint32 ptr_type = it_var_ptr->second;
+				UnorderedMap<uint32, uint32>::const_iterator it_pointee = pointer_to_pointee.find(ptr_type);
+				if (it_pointee != pointer_to_pointee.end())
+				{
+					uint32 pointee_type = it_pointee->second;
+
+					// Find descriptor type
+					UnorderedMap<uint32, VkDescriptorType>::iterator it_descriptor_type = id_to_descriptor_type.find(pointee_type);
+					VkDescriptorType descriptor_type = it_descriptor_type != id_to_descriptor_type.end() ? it_descriptor_type->second : VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
+
+					name_to_binding.insert({ it_name->second, { binding, descriptor_type } });
+					continue;
+				}
+			}
+		}
+	}
+
+	// Create layout bindings and buffer infos
+	if (!name_to_binding.empty())
+	{
+		mLayoutBindings.reserve(name_to_binding.size());
+		mBufferInfos.reserve(name_to_binding.size());
+
+		mBindingNames.reserve(name_to_binding.size());
+		for (const UnorderedMap<String, std::pair<uint32, VkDescriptorType>>::value_type &b : name_to_binding)
+		{
+			const String &name = b.first;
+			uint binding = b.second.first;
+			VkDescriptorType descriptor_type = b.second.second;
+
+			VkDescriptorSetLayoutBinding l = {};
+			l.binding = binding;
+			l.descriptorCount = 1;
+			l.stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;
+			l.descriptorType = descriptor_type;
+			mLayoutBindings.push_back(l);
+
+			mBindingNames.push_back(name); // Add all strings to a pool to keep them alive
+			mNameToBufferInfoIndex[string_view(mBindingNames.back())] = (uint32)mBufferInfos.size();
+
+			VkDescriptorBufferInfo bi = {};
+			bi.offset = 0;
+			bi.range = VK_WHOLE_SIZE;
+			bi.buffer = inDummyBuffer; // Avoid: The Vulkan spec states: If the nullDescriptor feature is not enabled, buffer must not be VK_NULL_HANDLE
+			mBufferInfos.push_back(bi);
+		}
+
+		// Create descriptor set layout
+		VkDescriptorSetLayoutCreateInfo layout_info = {};
+		layout_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
+		layout_info.bindingCount = (uint32)mLayoutBindings.size();
+		layout_info.pBindings = mLayoutBindings.data();
+		if (VKFailed(vkCreateDescriptorSetLayout(mDevice, &layout_info, nullptr, &mDescriptorSetLayout), outResult))
+			return false;
+	}
+
+	// Create pipeline layout
+	VkPipelineLayoutCreateInfo pl_info = {};
+	pl_info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
+	pl_info.setLayoutCount = mDescriptorSetLayout != VK_NULL_HANDLE ? 1 : 0;
+	pl_info.pSetLayouts = mDescriptorSetLayout != VK_NULL_HANDLE ? &mDescriptorSetLayout : nullptr;
+	if (VKFailed(vkCreatePipelineLayout(mDevice, &pl_info, nullptr, &mPipelineLayout), outResult))
+		return false;
+
+	// Create shader module
+	VkShaderModuleCreateInfo create_info = {};
+	create_info.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
+	create_info.codeSize = inSPVCode.size();
+	create_info.pCode = spv_words;
+	if (VKFailed(vkCreateShaderModule(mDevice, &create_info, nullptr, &mShaderModule), outResult))
+		return false;
+
+	// Create compute pipeline
+	VkComputePipelineCreateInfo pipe_info = {};
+	pipe_info.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
+	pipe_info.stage.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
+	pipe_info.stage.stage = VK_SHADER_STAGE_COMPUTE_BIT;
+	pipe_info.stage.module = mShaderModule;
+	pipe_info.stage.pName = "main";
+	pipe_info.layout = mPipelineLayout;
+	if (VKFailed(vkCreateComputePipelines(mDevice, VK_NULL_HANDLE, 1, &pipe_info, nullptr, &mPipeline), outResult))
+		return false;
+
+	return true;
+}
+
+uint32 ComputeShaderVK::NameToBufferInfoIndex(const char *inName) const
+{
+	UnorderedMap<string_view, uint>::const_iterator it = mNameToBufferInfoIndex.find(inName);
+	JPH_ASSERT(it != mNameToBufferInfoIndex.end());
+	return it->second;
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeShaderVK.h

@@ -0,0 +1,53 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <Jolt/Compute/ComputeShader.h>
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/IncludeVK.h>
+#include <Jolt/Core/UnorderedMap.h>
+
+JPH_NAMESPACE_BEGIN
+
+/// Compute shader handle for Vulkan
+class JPH_EXPORT ComputeShaderVK : public ComputeShader
+{
+public:
+	JPH_OVERRIDE_NEW_DELETE
+
+	/// Constructor / destructor
+										ComputeShaderVK(VkDevice inDevice, uint32 inGroupSizeX, uint32 inGroupSizeY, uint32 inGroupSizeZ) : ComputeShader(inGroupSizeX, inGroupSizeY, inGroupSizeZ), mDevice(inDevice) { }
+	virtual								~ComputeShaderVK() override;
+
+	/// Initialize from SPIR-V code
+	bool								Initialize(const Array<uint8> &inSPVCode, VkBuffer inDummyBuffer, ComputeShaderResult &outResult);
+
+	/// Get index of parameter in buffer infos
+	uint32								NameToBufferInfoIndex(const char *inName) const;
+
+	/// Getters
+	VkPipeline							GetPipeline() const							{ return mPipeline; }
+	VkPipelineLayout					GetPipelineLayout() const					{ return mPipelineLayout; }
+	VkDescriptorSetLayout				GetDescriptorSetLayout() const				{ return mDescriptorSetLayout; }
+	const Array<VkDescriptorSetLayoutBinding> &GetLayoutBindings() const			{ return mLayoutBindings; }
+	const Array<VkDescriptorBufferInfo> &GetBufferInfos() const						{ return mBufferInfos; }
+
+private:
+	VkDevice							mDevice;
+	VkShaderModule						mShaderModule = VK_NULL_HANDLE;
+	VkPipelineLayout					mPipelineLayout = VK_NULL_HANDLE;
+	VkPipeline							mPipeline = VK_NULL_HANDLE;
+	VkDescriptorSetLayout				mDescriptorSetLayout = VK_NULL_HANDLE;
+	Array<String>						mBindingNames;								///< A list of binding names, mNameToBufferInfoIndex points to these strings
+	UnorderedMap<string_view, uint32>	mNameToBufferInfoIndex;						///< Binding name to buffer index, using a string_view so we can do find() without an allocation
+	Array<VkDescriptorSetLayoutBinding>	mLayoutBindings;
+	Array<VkDescriptorBufferInfo>		mBufferInfos;
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeSystemVK.cpp

@@ -0,0 +1,118 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/ComputeSystemVK.h>
+#include <Jolt/Compute/VK/ComputeShaderVK.h>
+#include <Jolt/Compute/VK/ComputeBufferVK.h>
+#include <Jolt/Compute/VK/ComputeQueueVK.h>
+
+JPH_NAMESPACE_BEGIN
+
+JPH_IMPLEMENT_RTTI_ABSTRACT(ComputeSystemVK)
+{
+	JPH_ADD_BASE_CLASS(ComputeSystemVK, ComputeSystem)
+}
+
+bool ComputeSystemVK::Initialize(VkPhysicalDevice inPhysicalDevice, VkDevice inDevice, uint32 inComputeQueueIndex, ComputeSystemResult &outResult)
+{
+	mPhysicalDevice = inPhysicalDevice;
+	mDevice = inDevice;
+	mComputeQueueIndex = inComputeQueueIndex;
+
+	// Get function to set a debug name
+	mVkSetDebugUtilsObjectNameEXT = reinterpret_cast<PFN_vkSetDebugUtilsObjectNameEXT>(reinterpret_cast<void *>(vkGetDeviceProcAddr(mDevice, "vkSetDebugUtilsObjectNameEXT")));
+
+	if (!InitializeMemory())
+	{
+		outResult.SetError("Failed to initialize memory subsystem");
+		return false;
+	}
+
+	// Create the dummy buffer. This is used to bind to shaders for which we have no buffer. We can't rely on VK_EXT_robustness2 being available to set nullDescriptor = VK_TRUE (it is unavailable on macOS).
+	if (!CreateBuffer(1024, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, mDummyBuffer))
+	{
+		outResult.SetError("Failed to create dummy buffer");
+		return false;
+	}
+
+	return true;
+}
+
+void ComputeSystemVK::Shutdown()
+{
+	if (mDevice != VK_NULL_HANDLE)
+		vkDeviceWaitIdle(mDevice);
+
+	// Free the dummy buffer
+	FreeBuffer(mDummyBuffer);
+
+	ShutdownMemory();
+}
+
+ComputeShaderResult ComputeSystemVK::CreateComputeShader(const char *inName, uint32 inGroupSizeX, uint32 inGroupSizeY, uint32 inGroupSizeZ)
+{
+	ComputeShaderResult result;
+
+	// Read shader source file
+	Array<uint8> data;
+	String file_name = String(inName) + ".spv";
+	String error;
+	if (!mShaderLoader(file_name.c_str(), data, error))
+	{
+		result.SetError(error);
+		return result;
+	}
+
+	Ref<ComputeShaderVK> shader = new ComputeShaderVK(mDevice, inGroupSizeX, inGroupSizeY, inGroupSizeZ);
+	if (!shader->Initialize(data, mDummyBuffer.mBuffer, result))
+		return result;
+
+	// Name the pipeline so we can easily find it in a profile
+	if (mVkSetDebugUtilsObjectNameEXT != nullptr)
+	{
+		VkDebugUtilsObjectNameInfoEXT info = {};
+		info.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_OBJECT_NAME_INFO_EXT;
+		info.pNext = nullptr;
+		info.objectType = VK_OBJECT_TYPE_PIPELINE;
+		info.objectHandle = (uint64)shader->GetPipeline();
+		info.pObjectName = inName;
+		mVkSetDebugUtilsObjectNameEXT(mDevice, &info);
+	}
+
+	result.Set(shader.GetPtr());
+	return result;
+}
+
+ComputeBufferResult ComputeSystemVK::CreateComputeBuffer(ComputeBuffer::EType inType, uint64 inSize, uint inStride, const void *inData)
+{
+	ComputeBufferResult result;
+
+	Ref<ComputeBufferVK> buffer = new ComputeBufferVK(this, inType, inSize, inStride);
+	if (!buffer->Initialize(inData))
+	{
+		result.SetError("Failed to create compute buffer");
+		return result;
+	}
+
+	result.Set(buffer.GetPtr());
+	return result;
+}
+
+ComputeQueueResult ComputeSystemVK::CreateComputeQueue()
+{
+	ComputeQueueResult result;
+	Ref<ComputeQueueVK> q = new ComputeQueueVK(this);
+	if (!q->Initialize(mComputeQueueIndex, result))
+		return result;
+	result.Set(q.GetPtr());
+	return result;
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeSystemVK.h

@@ -0,0 +1,57 @@
+// 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_VK
+
+#include <Jolt/Compute/VK/ComputeQueueVK.h>
+
+JPH_NAMESPACE_BEGIN
+
+/// Interface to run a workload on the GPU using Vulkan.
+/// Minimal implementation that can integrate with your own Vulkan setup.
+class JPH_EXPORT ComputeSystemVK : public ComputeSystem
+{
+public:
+	JPH_DECLARE_RTTI_ABSTRACT(JPH_EXPORT, ComputeSystemVK)
+
+	// Initialize / shutdown the compute system
+	bool							Initialize(VkPhysicalDevice inPhysicalDevice, VkDevice inDevice, uint32 inComputeQueueIndex, ComputeSystemResult &outResult);
+	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;
+
+	/// Access to the Vulkan device
+	VkDevice						GetDevice() const												{ return mDevice; }
+
+	/// Allow the application to override buffer creation and memory mapping in case it uses its own allocator
+	virtual bool					CreateBuffer(VkDeviceSize inSize, VkBufferUsageFlags inUsage, VkMemoryPropertyFlags inProperties, BufferVK &outBuffer) = 0;
+	virtual void					FreeBuffer(BufferVK &ioBuffer) = 0;
+	virtual void *					MapBuffer(BufferVK &ioBuffer) = 0;
+	virtual void					UnmapBuffer(BufferVK &ioBuffer) = 0;
+
+protected:
+	/// Initialize / shutdown the memory subsystem
+	virtual bool					InitializeMemory() = 0;
+	virtual void					ShutdownMemory() = 0;
+
+	VkPhysicalDevice				mPhysicalDevice = VK_NULL_HANDLE;
+	VkDevice						mDevice = VK_NULL_HANDLE;
+	uint32							mComputeQueueIndex = 0;
+	PFN_vkSetDebugUtilsObjectNameEXT mVkSetDebugUtilsObjectNameEXT = nullptr;
+
+private:
+	// Buffer that can be bound when we have no buffer
+	BufferVK						mDummyBuffer;
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeSystemVKImpl.cpp

@@ -0,0 +1,330 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/ComputeSystemVKImpl.h>
+#include <Jolt/Core/QuickSort.h>
+
+JPH_NAMESPACE_BEGIN
+
+JPH_IMPLEMENT_RTTI_VIRTUAL(ComputeSystemVKImpl)
+{
+	JPH_ADD_BASE_CLASS(ComputeSystemVKImpl, ComputeSystemVKWithAllocator)
+}
+
+#ifdef JPH_DEBUG
+
+static VKAPI_ATTR VkBool32 VKAPI_CALL sVulkanDebugCallback(VkDebugUtilsMessageSeverityFlagBitsEXT inSeverity, [[maybe_unused]] VkDebugUtilsMessageTypeFlagsEXT inType, const VkDebugUtilsMessengerCallbackDataEXT *inCallbackData, [[maybe_unused]] void *inUserData)
+{
+	if (inSeverity & (VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT))
+		Trace("VK: %s", inCallbackData->pMessage);
+	JPH_ASSERT((inSeverity & VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT) == 0);
+	return VK_FALSE;
+}
+
+#endif // JPH_DEBUG
+
+ComputeSystemVKImpl::~ComputeSystemVKImpl()
+{
+	ComputeSystemVK::Shutdown();
+
+	if (mDevice != VK_NULL_HANDLE)
+		vkDestroyDevice(mDevice, nullptr);
+
+#ifdef JPH_DEBUG
+	PFN_vkDestroyDebugUtilsMessengerEXT vkDestroyDebugUtilsMessengerEXT = (PFN_vkDestroyDebugUtilsMessengerEXT)(void *)vkGetInstanceProcAddr(mInstance, "vkDestroyDebugUtilsMessengerEXT");
+	if (mInstance != VK_NULL_HANDLE && mDebugMessenger != VK_NULL_HANDLE && vkDestroyDebugUtilsMessengerEXT != nullptr)
+		vkDestroyDebugUtilsMessengerEXT(mInstance, mDebugMessenger, nullptr);
+#endif
+
+	if (mInstance != VK_NULL_HANDLE)
+		vkDestroyInstance(mInstance, nullptr);
+}
+
+bool ComputeSystemVKImpl::Initialize(ComputeSystemResult &outResult)
+{
+	// Required instance extensions
+	Array<const char *> required_instance_extensions;
+	required_instance_extensions.push_back(VK_KHR_SURFACE_EXTENSION_NAME);
+	required_instance_extensions.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);
+#ifdef JPH_PLATFORM_MACOS
+	required_instance_extensions.push_back("VK_KHR_portability_enumeration");
+	required_instance_extensions.push_back("VK_KHR_get_physical_device_properties2");
+#endif
+	GetInstanceExtensions(required_instance_extensions);
+
+	// Required device extensions
+	Array<const char *> required_device_extensions;
+	required_device_extensions.push_back(VK_EXT_SCALAR_BLOCK_LAYOUT_EXTENSION_NAME);
+#ifdef JPH_PLATFORM_MACOS
+	required_device_extensions.push_back("VK_KHR_portability_subset"); // VK_KHR_PORTABILITY_SUBSET_EXTENSION_NAME
+#endif
+	GetDeviceExtensions(required_device_extensions);
+
+	// Query supported instance extensions
+	uint32 instance_extension_count = 0;
+	if (VKFailed(vkEnumerateInstanceExtensionProperties(nullptr, &instance_extension_count, nullptr), outResult))
+		return false;
+	Array<VkExtensionProperties> instance_extensions;
+	instance_extensions.resize(instance_extension_count);
+	if (VKFailed(vkEnumerateInstanceExtensionProperties(nullptr, &instance_extension_count, instance_extensions.data()), outResult))
+		return false;
+
+	// Query supported validation layers
+	uint32 validation_layer_count;
+	vkEnumerateInstanceLayerProperties(&validation_layer_count, nullptr);
+	Array<VkLayerProperties> validation_layers(validation_layer_count);
+	vkEnumerateInstanceLayerProperties(&validation_layer_count, validation_layers.data());
+
+	VkApplicationInfo app_info = {};
+	app_info.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
+	app_info.apiVersion = VK_API_VERSION_1_1;
+
+	// Create Vulkan instance
+	VkInstanceCreateInfo instance_create_info = {};
+	instance_create_info.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
+#ifdef JPH_PLATFORM_MACOS
+	instance_create_info.flags = VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR;
+#endif
+	instance_create_info.pApplicationInfo = &app_info;
+
+#ifdef JPH_DEBUG
+	// Enable validation layer if supported
+	const char *desired_validation_layers[] = { "VK_LAYER_KHRONOS_validation" };
+	for (const VkLayerProperties &p : validation_layers)
+		if (strcmp(desired_validation_layers[0], p.layerName) == 0)
+		{
+			instance_create_info.enabledLayerCount = 1;
+			instance_create_info.ppEnabledLayerNames = desired_validation_layers;
+			break;
+		}
+
+	// Setup debug messenger callback if the extension is supported
+	VkDebugUtilsMessengerCreateInfoEXT messenger_create_info = {};
+	for (const VkExtensionProperties &ext : instance_extensions)
+		if (strcmp(VK_EXT_DEBUG_UTILS_EXTENSION_NAME, ext.extensionName) == 0)
+		{
+			messenger_create_info.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT;
+			messenger_create_info.messageSeverity = VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_INFO_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT;
+			messenger_create_info.messageType = VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT;
+			messenger_create_info.pfnUserCallback = sVulkanDebugCallback;
+			instance_create_info.pNext = &messenger_create_info;
+			required_instance_extensions.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);
+			break;
+		}
+#endif
+
+	instance_create_info.enabledExtensionCount = (uint32)required_instance_extensions.size();
+	instance_create_info.ppEnabledExtensionNames = required_instance_extensions.data();
+	if (VKFailed(vkCreateInstance(&instance_create_info, nullptr, &mInstance), outResult))
+		return false;
+
+#ifdef JPH_DEBUG
+	// Finalize debug messenger callback
+	PFN_vkCreateDebugUtilsMessengerEXT vkCreateDebugUtilsMessengerEXT = (PFN_vkCreateDebugUtilsMessengerEXT)(std::uintptr_t)vkGetInstanceProcAddr(mInstance, "vkCreateDebugUtilsMessengerEXT");
+	if (vkCreateDebugUtilsMessengerEXT != nullptr)
+		if (VKFailed(vkCreateDebugUtilsMessengerEXT(mInstance, &messenger_create_info, nullptr, &mDebugMessenger), outResult))
+			return false;
+#endif
+
+	// Notify that instance has been created
+	OnInstanceCreated();
+
+	// Select device
+	uint32 device_count = 0;
+	if (VKFailed(vkEnumeratePhysicalDevices(mInstance, &device_count, nullptr), outResult))
+		return false;
+	Array<VkPhysicalDevice> devices;
+	devices.resize(device_count);
+	if (VKFailed(vkEnumeratePhysicalDevices(mInstance, &device_count, devices.data()), outResult))
+		return false;
+	struct Device
+	{
+		VkPhysicalDevice		mPhysicalDevice;
+		String					mName;
+		VkSurfaceFormatKHR		mFormat;
+		uint32					mGraphicsQueueIndex;
+		uint32					mPresentQueueIndex;
+		uint32					mComputeQueueIndex;
+		int						mScore;
+	};
+	Array<Device> available_devices;
+	for (VkPhysicalDevice device : devices)
+	{
+		// Get device properties
+		VkPhysicalDeviceProperties properties;
+		vkGetPhysicalDeviceProperties(device, &properties);
+
+		// Test if it is an appropriate type
+		int score = 0;
+		switch (properties.deviceType)
+		{
+		case VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU:
+			score = 30;
+			break;
+		case VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU:
+			score = 20;
+			break;
+		case VK_PHYSICAL_DEVICE_TYPE_VIRTUAL_GPU:
+			score = 10;
+			break;
+		case VK_PHYSICAL_DEVICE_TYPE_CPU:
+			score = 5;
+			break;
+		case VK_PHYSICAL_DEVICE_TYPE_OTHER:
+		case VK_PHYSICAL_DEVICE_TYPE_MAX_ENUM:
+			continue;
+		}
+
+		// Check if the device supports all our required extensions
+		uint32 device_extension_count;
+		vkEnumerateDeviceExtensionProperties(device, nullptr, &device_extension_count, nullptr);
+		Array<VkExtensionProperties> available_extensions;
+		available_extensions.resize(device_extension_count);
+		vkEnumerateDeviceExtensionProperties(device, nullptr, &device_extension_count, available_extensions.data());
+		int found_extensions = 0;
+		for (const char *required_device_extension : required_device_extensions)
+			for (const VkExtensionProperties &ext : available_extensions)
+				if (strcmp(required_device_extension, ext.extensionName) == 0)
+				{
+					found_extensions++;
+					break;
+				}
+		if (found_extensions != int(required_device_extensions.size()))
+			continue;
+
+		// Find the right queues
+		uint32 queue_family_count = 0;
+		vkGetPhysicalDeviceQueueFamilyProperties(device, &queue_family_count, nullptr);
+		Array<VkQueueFamilyProperties> queue_families;
+		queue_families.resize(queue_family_count);
+		vkGetPhysicalDeviceQueueFamilyProperties(device, &queue_family_count, queue_families.data());
+		uint32 graphics_queue = ~uint32(0);
+		uint32 present_queue = ~uint32(0);
+		uint32 compute_queue = ~uint32(0);
+		for (uint32 i = 0; i < uint32(queue_families.size()); ++i)
+		{
+			if (queue_families[i].queueFlags & VK_QUEUE_GRAPHICS_BIT)
+			{
+				graphics_queue = i;
+
+				if (queue_families[i].queueFlags & VK_QUEUE_COMPUTE_BIT)
+					compute_queue = i;
+			}
+
+			if (HasPresentSupport(device, i))
+				present_queue = i;
+
+			if (graphics_queue != ~uint32(0) && present_queue != ~uint32(0) && compute_queue != ~uint32(0))
+				break;
+		}
+		if (graphics_queue == ~uint32(0) || present_queue == ~uint32(0) || compute_queue == ~uint32(0))
+			continue;
+
+		// Select surface format
+		VkSurfaceFormatKHR selected_format = SelectFormat(device);
+		if (selected_format.format == VK_FORMAT_UNDEFINED)
+			continue;
+
+		// Add the device
+		available_devices.push_back({ device, properties.deviceName, selected_format, graphics_queue, present_queue, compute_queue, score });
+	}
+	if (available_devices.empty())
+	{
+		outResult.SetError("No suitable Vulkan device found");
+		return false;
+	}
+
+	// Sort the devices by score
+	QuickSort(available_devices.begin(), available_devices.end(), [](const Device &inLHS, const Device &inRHS) {
+		return inLHS.mScore > inRHS.mScore;
+	});
+	const Device &selected_device = available_devices[0];
+
+	// Create device
+	float queue_priority = 1.0f;
+	VkDeviceQueueCreateInfo queue_create_info[3] = {};
+	for (VkDeviceQueueCreateInfo &q : queue_create_info)
+	{
+		q.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
+		q.queueCount = 1;
+		q.pQueuePriorities = &queue_priority;
+	}
+	uint32 num_queues = 0;
+	queue_create_info[num_queues++].queueFamilyIndex = selected_device.mGraphicsQueueIndex;
+	bool found = false;
+	for (uint32 i = 0; i < num_queues; ++i)
+		if (queue_create_info[i].queueFamilyIndex == selected_device.mPresentQueueIndex)
+		{
+			found = true;
+			break;
+		}
+	if (!found)
+		queue_create_info[num_queues++].queueFamilyIndex = selected_device.mPresentQueueIndex;
+	found = false;
+	for (uint32 i = 0; i < num_queues; ++i)
+		if (queue_create_info[i].queueFamilyIndex == selected_device.mComputeQueueIndex)
+		{
+			found = true;
+			break;
+		}
+	if (!found)
+		queue_create_info[num_queues++].queueFamilyIndex = selected_device.mComputeQueueIndex;
+
+	VkPhysicalDeviceScalarBlockLayoutFeatures enable_scalar_block = {};
+	enable_scalar_block.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SCALAR_BLOCK_LAYOUT_FEATURES;
+	enable_scalar_block.scalarBlockLayout = VK_TRUE;
+
+	VkPhysicalDeviceFeatures2 enabled_features2 = {};
+	enabled_features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2;
+	GetEnabledFeatures(enabled_features2);
+	enable_scalar_block.pNext = enabled_features2.pNext;
+	enabled_features2.pNext = &enable_scalar_block;
+
+	VkDeviceCreateInfo device_create_info = {};
+	device_create_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
+	device_create_info.queueCreateInfoCount = num_queues;
+	device_create_info.pQueueCreateInfos = queue_create_info;
+	device_create_info.enabledLayerCount = instance_create_info.enabledLayerCount;
+	device_create_info.ppEnabledLayerNames = instance_create_info.ppEnabledLayerNames;
+	device_create_info.enabledExtensionCount = uint32(required_device_extensions.size());
+	device_create_info.ppEnabledExtensionNames = required_device_extensions.data();
+	device_create_info.pNext = &enabled_features2;
+	device_create_info.pEnabledFeatures = nullptr;
+
+	VkDevice device = VK_NULL_HANDLE;
+	if (VKFailed(vkCreateDevice(selected_device.mPhysicalDevice, &device_create_info, nullptr, &device), outResult))
+		return false;
+
+	// Get the queues
+	mGraphicsQueueIndex = selected_device.mGraphicsQueueIndex;
+	mPresentQueueIndex = selected_device.mPresentQueueIndex;
+	vkGetDeviceQueue(device, mGraphicsQueueIndex, 0, &mGraphicsQueue);
+	vkGetDeviceQueue(device, mPresentQueueIndex, 0, &mPresentQueue);
+
+	// Store selected format
+	mSelectedFormat = selected_device.mFormat;
+
+	// Initialize the compute system
+	return ComputeSystemVK::Initialize(selected_device.mPhysicalDevice, device, selected_device.mComputeQueueIndex, outResult);
+}
+
+ComputeSystemResult CreateComputeSystemVK()
+{
+	ComputeSystemResult result;
+
+	Ref<ComputeSystemVKImpl> compute = new ComputeSystemVKImpl;
+	if (!compute->Initialize(result))
+		return result;
+
+	result.Set(compute.GetPtr());
+	return result;
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeSystemVKImpl.h

@@ -0,0 +1,57 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/ComputeSystemVKWithAllocator.h>
+
+JPH_NAMESPACE_BEGIN
+
+/// Implementation of ComputeSystemVK that fully initializes Vulkan
+class JPH_EXPORT ComputeSystemVKImpl : public ComputeSystemVKWithAllocator
+{
+public:
+	JPH_DECLARE_RTTI_VIRTUAL(JPH_EXPORT, ComputeSystemVKImpl)
+
+	/// Destructor
+	virtual							~ComputeSystemVKImpl() override;
+
+	/// Initialize the compute system
+	bool							Initialize(ComputeSystemResult &outResult);
+
+protected:
+	/// Override to perform actions once the instance has been created
+	virtual void					OnInstanceCreated()																{ /* Do nothing */ }
+
+	/// Override to add platform specific instance extensions
+	virtual void					GetInstanceExtensions(Array<const char *> &outExtensions)						{ /* Add nothing */ }
+
+	/// Override to add platform specific device extensions
+	virtual void					GetDeviceExtensions(Array<const char *> &outExtensions)							{ /* Add nothing */ }
+
+	/// Override  to enable specific features
+	virtual void					GetEnabledFeatures(VkPhysicalDeviceFeatures2 &ioFeatures)						{ /* Add nothing */ }
+
+	/// Override to check for present support on a given device and queue family
+	virtual bool					HasPresentSupport(VkPhysicalDevice inDevice, uint32 inQueueFamilyIndex)			{ return true; }
+
+	/// Override to select the surface format
+	virtual VkSurfaceFormatKHR		SelectFormat(VkPhysicalDevice inDevice)											{ return { VK_FORMAT_B8G8R8A8_UNORM, VK_COLOR_SPACE_SRGB_NONLINEAR_KHR }; }
+
+	VkInstance						mInstance = VK_NULL_HANDLE;
+#ifdef JPH_DEBUG
+	VkDebugUtilsMessengerEXT		mDebugMessenger = VK_NULL_HANDLE;
+#endif
+	uint32							mGraphicsQueueIndex = 0;
+	uint32							mPresentQueueIndex = 0;
+	VkQueue							mGraphicsQueue = VK_NULL_HANDLE;
+	VkQueue							mPresentQueue = VK_NULL_HANDLE;
+	VkSurfaceFormatKHR				mSelectedFormat;
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeSystemVKWithAllocator.cpp

@@ -0,0 +1,172 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#include <Jolt/Jolt.h>
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/ComputeSystemVKWithAllocator.h>
+#include <Jolt/Compute/VK/ComputeShaderVK.h>
+#include <Jolt/Compute/VK/ComputeBufferVK.h>
+#include <Jolt/Compute/VK/ComputeQueueVK.h>
+
+JPH_NAMESPACE_BEGIN
+
+JPH_IMPLEMENT_RTTI_VIRTUAL(ComputeSystemVKWithAllocator)
+{
+	JPH_ADD_BASE_CLASS(ComputeSystemVKWithAllocator, ComputeSystemVK)
+}
+
+bool ComputeSystemVKWithAllocator::InitializeMemory()
+{
+	// Get memory properties
+	vkGetPhysicalDeviceMemoryProperties(mPhysicalDevice, &mMemoryProperties);
+
+	return true;
+}
+
+void ComputeSystemVKWithAllocator::ShutdownMemory()
+{
+	// Free all memory
+	for (const MemoryCache::value_type &mc : mMemoryCache)
+		for (const Memory &m : mc.second)
+			if (m.mOffset == 0)
+				FreeMemory(*m.mMemory);
+	mMemoryCache.clear();
+}
+
+uint32 ComputeSystemVKWithAllocator::FindMemoryType(uint32 inTypeFilter, VkMemoryPropertyFlags inProperties) const
+{
+	for (uint32 i = 0; i < mMemoryProperties.memoryTypeCount; i++)
+		if ((inTypeFilter & (1 << i))
+			&& (mMemoryProperties.memoryTypes[i].propertyFlags & inProperties) == inProperties)
+			return i;
+
+	JPH_ASSERT(false, "Failed to find memory type!");
+	return 0;
+}
+
+void ComputeSystemVKWithAllocator::AllocateMemory(VkDeviceSize inSize, uint32 inMemoryTypeBits, VkMemoryPropertyFlags inProperties, MemoryVK &ioMemory)
+{
+	JPH_ASSERT(ioMemory.mMemory == VK_NULL_HANDLE);
+
+	ioMemory.mSize = inSize;
+	ioMemory.mProperties = inProperties;
+
+	VkMemoryAllocateInfo alloc_info = {};
+	alloc_info.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
+	alloc_info.allocationSize = inSize;
+	alloc_info.memoryTypeIndex = FindMemoryType(inMemoryTypeBits, inProperties);
+	vkAllocateMemory(mDevice, &alloc_info, nullptr, &ioMemory.mMemory);
+}
+
+void ComputeSystemVKWithAllocator::FreeMemory(MemoryVK &ioMemory)
+{
+	vkFreeMemory(mDevice, ioMemory.mMemory, nullptr);
+	ioMemory.mMemory = VK_NULL_HANDLE;
+}
+
+bool ComputeSystemVKWithAllocator::CreateBuffer(VkDeviceSize inSize, VkBufferUsageFlags inUsage, VkMemoryPropertyFlags inProperties, BufferVK &outBuffer)
+{
+	// Create a new buffer
+	outBuffer.mSize = inSize;
+
+	VkBufferCreateInfo create_info = {};
+	create_info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
+	create_info.size = inSize;
+	create_info.usage = inUsage;
+	create_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
+	if (VKFailed(vkCreateBuffer(mDevice, &create_info, nullptr, &outBuffer.mBuffer)))
+	{
+		outBuffer.mBuffer = VK_NULL_HANDLE;
+		return false;
+	}
+
+	VkMemoryRequirements mem_requirements;
+	vkGetBufferMemoryRequirements(mDevice, outBuffer.mBuffer, &mem_requirements);
+
+	if (mem_requirements.size > cMaxAllocSize)
+	{
+		// Allocate block directly
+		Ref<MemoryVK> memory_vk = new MemoryVK();
+		memory_vk->mBufferSize = mem_requirements.size;
+		AllocateMemory(mem_requirements.size, mem_requirements.memoryTypeBits, inProperties, *memory_vk);
+		outBuffer.mMemory = memory_vk;
+		outBuffer.mOffset = 0;
+	}
+	else
+	{
+		// Round allocation to the next power of 2 so that we can use a simple block based allocator
+		VkDeviceSize buffer_size = max(VkDeviceSize(GetNextPowerOf2(uint32(mem_requirements.size))), cMinAllocSize);
+
+		// Ensure that we have memory available from the right pool
+		Array<Memory> &mem_array = mMemoryCache[{ buffer_size, inProperties }];
+		if (mem_array.empty())
+		{
+			// Allocate a bigger block
+			Ref<MemoryVK> memory_vk = new MemoryVK();
+			memory_vk->mBufferSize = buffer_size;
+			AllocateMemory(cBlockSize, mem_requirements.memoryTypeBits, inProperties, *memory_vk);
+
+			// Divide into sub blocks
+			for (VkDeviceSize offset = 0; offset < cBlockSize; offset += buffer_size)
+				mem_array.push_back({ memory_vk, offset });
+		}
+
+		// Claim memory from the pool
+		Memory &memory = mem_array.back();
+		outBuffer.mMemory = memory.mMemory;
+		outBuffer.mOffset = memory.mOffset;
+		mem_array.pop_back();
+	}
+
+	// Bind the memory to the buffer
+	vkBindBufferMemory(mDevice, outBuffer.mBuffer, outBuffer.mMemory->mMemory, outBuffer.mOffset);
+	return true;
+}
+
+void ComputeSystemVKWithAllocator::FreeBuffer(BufferVK &ioBuffer)
+{
+	if (ioBuffer.mBuffer != VK_NULL_HANDLE)
+	{
+		// Destroy the buffer
+		vkDestroyBuffer(mDevice, ioBuffer.mBuffer, nullptr);
+		ioBuffer.mBuffer = VK_NULL_HANDLE;
+
+		// Hand the memory back to the cache
+		VkDeviceSize buffer_size = ioBuffer.mMemory->mBufferSize;
+		if (buffer_size > cMaxAllocSize)
+			FreeMemory(*ioBuffer.mMemory);
+		else
+			mMemoryCache[{ buffer_size, ioBuffer.mMemory->mProperties }].push_back({ ioBuffer.mMemory, ioBuffer.mOffset });
+
+		ioBuffer = BufferVK();
+	}
+}
+
+void *ComputeSystemVKWithAllocator::MapBuffer(BufferVK& ioBuffer)
+{
+	if (++ioBuffer.mMemory->mMappedCount == 1
+		&& VKFailed(vkMapMemory(mDevice, ioBuffer.mMemory->mMemory, 0, VK_WHOLE_SIZE, 0, &ioBuffer.mMemory->mMappedPtr)))
+	{
+		ioBuffer.mMemory->mMappedCount = 0;
+		return nullptr;
+	}
+
+	return static_cast<uint8 *>(ioBuffer.mMemory->mMappedPtr) + ioBuffer.mOffset;
+}
+
+void ComputeSystemVKWithAllocator::UnmapBuffer(BufferVK& ioBuffer)
+{
+	JPH_ASSERT(ioBuffer.mMemory->mMappedCount > 0);
+	if (--ioBuffer.mMemory->mMappedCount == 0)
+	{
+		vkUnmapMemory(mDevice, ioBuffer.mMemory->mMemory);
+		ioBuffer.mMemory->mMappedPtr = nullptr;
+	}
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/ComputeSystemVKWithAllocator.h

@@ -0,0 +1,70 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#ifdef JPH_USE_VK
+
+#include <Jolt/Compute/VK/ComputeSystemVK.h>
+#include <Jolt/Core/UnorderedMap.h>
+
+JPH_NAMESPACE_BEGIN
+
+/// This extends ComputeSystemVK to provide a default implementation for memory allocation and mapping.
+/// It uses a simple block based allocator to reduce the number of allocations done to Vulkan.
+class JPH_EXPORT ComputeSystemVKWithAllocator : public ComputeSystemVK
+{
+public:
+	JPH_DECLARE_RTTI_VIRTUAL(JPH_EXPORT, ComputeSystemVKWithAllocator)
+
+	/// Allow the application to override buffer creation and memory mapping in case it uses its own allocator
+	virtual bool					CreateBuffer(VkDeviceSize inSize, VkBufferUsageFlags inUsage, VkMemoryPropertyFlags inProperties, BufferVK &outBuffer) override;
+	virtual void					FreeBuffer(BufferVK &ioBuffer) override;
+	virtual void *					MapBuffer(BufferVK &ioBuffer) override;
+	virtual void					UnmapBuffer(BufferVK &ioBuffer) override;
+
+protected:
+	virtual bool					InitializeMemory() override;
+	virtual void					ShutdownMemory() override;
+
+	uint32							FindMemoryType(uint32 inTypeFilter, VkMemoryPropertyFlags inProperties) const;
+	void							AllocateMemory(VkDeviceSize inSize, uint32 inMemoryTypeBits, VkMemoryPropertyFlags inProperties, MemoryVK &ioMemory);
+	void							FreeMemory(MemoryVK &ioMemory);
+
+	VkPhysicalDeviceMemoryProperties mMemoryProperties;
+
+private:
+	// Smaller allocations (from cMinAllocSize to cMaxAllocSize) will be done in blocks of cBlockSize bytes.
+	// We do this because there is a limit to the number of allocations that we can make in Vulkan.
+	static constexpr VkDeviceSize	cMinAllocSize = 512;
+	static constexpr VkDeviceSize	cMaxAllocSize = 65536;
+	static constexpr VkDeviceSize	cBlockSize = 524288;
+
+	struct MemoryKey
+	{
+		bool						operator == (const MemoryKey &inRHS) const
+		{
+			return mSize == inRHS.mSize && mProperties == inRHS.mProperties;
+		}
+
+		VkDeviceSize				mSize;
+		VkMemoryPropertyFlags		mProperties;
+	};
+
+	JPH_MAKE_HASH_STRUCT(MemoryKey, MemoryKeyHasher, t.mProperties, t.mSize)
+
+	struct Memory
+	{
+		Ref<MemoryVK>				mMemory;
+		VkDeviceSize				mOffset;
+	};
+
+	using MemoryCache = UnorderedMap<MemoryKey, Array<Memory>, MemoryKeyHasher>;
+
+	MemoryCache						mMemoryCache;
+};
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK

lib/haxejolt/JoltPhysics/Jolt/Compute/VK/IncludeVK.h

@@ -0,0 +1,44 @@
+// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
+// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <Jolt/Core/StringTools.h>
+
+#ifdef JPH_USE_VK
+
+JPH_SUPPRESS_WARNINGS_STD_BEGIN
+JPH_CLANG_SUPPRESS_WARNING("-Wc++98-compat-pedantic")
+
+#include <vulkan/vulkan.h>
+
+JPH_SUPPRESS_WARNINGS_STD_END
+
+JPH_NAMESPACE_BEGIN
+
+inline bool VKFailed(VkResult inResult)
+{
+	if (inResult == VK_SUCCESS)
+		return false;
+
+	Trace("Vulkan call failed with error code: %d", (int)inResult);
+	JPH_ASSERT(false);
+	return true;
+}
+
+template <class Result>
+inline bool VKFailed(VkResult inResult, Result &outResult)
+{
+	if (inResult == VK_SUCCESS)
+		return false;
+
+	String error = StringFormat("Vulkan call failed with error code: %d", (int)inResult);
+	outResult.SetError(error);
+	JPH_ASSERT(false);
+	return true;
+}
+
+JPH_NAMESPACE_END
+
+#endif // JPH_USE_VK