wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

d4r: official NVIDIA DLSS on AMD GPUs under Linux and Proton

Browse Source
This commit is contained in:
ayo committed 2026-09-27 15:26:56 -04:00
commit 118173ecd7
128 files changed
+31333
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

patches/vkd3d-proton/0001-split-command-list-for-external-wait.patch
+304
View File
@@ -0,0 +1,304 @@
diff --git a/include/vkd3d_device_vkd3d_ext.idl b/include/vkd3d_device_vkd3d_ext.idl
index 680f2c96..1fddc28e 100644
--- a/include/vkd3d_device_vkd3d_ext.idl
+++ b/include/vkd3d_device_vkd3d_ext.idl
@@ -104,6 +104,21 @@ interface ID3D12DXVKInteropDevice3 : ID3D12DXVKInteropDevice2
UINT32 *vk_memory_type);
}
+[
+ uuid(5a7c8b3e-2f61-4d0e-9c1a-7e3b52d4a901),
+ object,
+ local,
+ pointer_default(unique)
+]
+interface ID3D12DXVKInteropDeviceD4R : ID3D12DXVKInteropDevice3
+{
+ /* d4r: ends the command list's current Vulkan command buffer. Everything
+ * recorded afterwards is submitted separately and waits (as a queue
+ * submission wait, not a GPU-side wait) until the timeline semaphore
+ * vk_semaphore reaches value. */
+ HRESULT SplitCommandListForExternalWait(ID3D12CommandList *list, UINT64 vk_semaphore, UINT64 value);
+}
+
[
uuid(f3112584-41f9-348d-a59b-00b7e1d285d6),
object,
diff --git a/libs/vkd3d/command.c b/libs/vkd3d/command.c
index 611e373b..23af78a5 100644
--- a/libs/vkd3d/command.c
+++ b/libs/vkd3d/command.c
@@ -2557,6 +2557,38 @@ static void d3d12_command_list_consider_new_sequence(struct d3d12_command_list *
}
}
+HRESULT d3d12_command_list_split_for_external_wait(struct d3d12_command_list *list,
+ VkSemaphore semaphore, uint64_t value)
+{
+ struct d3d12_command_list_iteration *iteration;
+ unsigned int count;
+
+ if (!list->is_recording)
+ return E_INVALIDARG;
+ if (!d3d12_command_list_allows_new_sequence(list))
+ {
+ WARN("Cannot split command list for external wait with a running query.\n");
+ return E_FAIL;
+ }
+
+ /* Flush everything recorded so far into the current command buffer. */
+ d3d12_command_list_decay_tracked_state(list);
+ d3d12_command_list_invalidate_all_state(list);
+
+ count = list->cmd.iteration_count;
+ d3d12_command_list_begin_new_sequence(list, false);
+ if (list->cmd.iteration_count == count)
+ {
+ WARN("Out of command list sequences for external wait.\n");
+ return E_FAIL;
+ }
+
+ iteration = &list->cmd.iterations[list->cmd.iteration_count - 1];
+ iteration->external_wait_semaphore = semaphore;
+ iteration->external_wait_value = value;
+ return S_OK;
+}
+
static void d3d12_command_list_copy_queue_fallback(struct d3d12_command_list *list)
{
if (!list->cmd.iterations[list->cmd.iteration_count - 1].fallback)
@@ -22769,6 +22801,7 @@ static void STDMETHODCALLTYPE d3d12_command_queue_ExecuteCommandLists(ID3D12Comm
struct vkd3d_initial_transition *transitions;
size_t num_transitions, num_command_buffers;
VkCommandBufferSubmitInfo *buffers, *buffer;
+ VkSemaphoreSubmitInfo *external_waits = NULL;
struct d3d12_command_allocator **allocators;
struct d3d12_command_queue_submission sub;
unsigned int indirect_barrier_hoist_index;
@@ -22954,6 +22987,7 @@ static void STDMETHODCALLTYPE d3d12_command_queue_ExecuteCommandLists(ID3D12Comm
vkd3d_free(allocators);
vkd3d_free(buffers);
vkd3d_free(cmd_cost);
+ vkd3d_free(external_waits);
#ifdef VKD3D_ENABLE_BREADCRUMBS
vkd3d_free(breadcrumb_indices);
#endif
@@ -23036,6 +23070,21 @@ static void STDMETHODCALLTYPE d3d12_command_queue_ExecuteCommandLists(ID3D12Comm
buffer->commandBuffer = cmd_list->cmd.iterations[iter].vk_command_buffer;
/* Pilfer this for purposes of signalling aux data. No need to allocate more stuff just to signal this. */
buffer->deviceMask = cmd_list->cmd.iterations[iter].fallback ? VKD3D_COMMAND_BUFFER_SUBMIT_INFO_DEVICE_MASK_FALLBACK_QUEUE : 0;
+
+ if (cmd_list->cmd.iterations[iter].external_wait_semaphore)
+ {
+ if (!external_waits)
+ external_waits = vkd3d_calloc(num_command_buffers, sizeof(*external_waits));
+ if (external_waits)
+ {
+ VkSemaphoreSubmitInfo *wait = &external_waits[cmd_submit_count - 1];
+ wait->sType = VK_STRUCTURE_TYPE_SEMAPHORE_SUBMIT_INFO;
+ wait->semaphore = cmd_list->cmd.iterations[iter].external_wait_semaphore;
+ wait->value = cmd_list->cmd.iterations[iter].external_wait_value;
+ wait->stageMask = VK_PIPELINE_STAGE_2_ALL_COMMANDS_BIT;
+ buffer->deviceMask |= VKD3D_COMMAND_BUFFER_SUBMIT_INFO_DEVICE_MASK_EXTERNAL_WAIT;
+ }
+ }
}
need_suspend_fixup = cmd_list->cmd.suspend_resume.suspend.vk_fixup_cmd_buffer && (i + 1 == command_list_count ||
@@ -23155,6 +23204,7 @@ static void STDMETHODCALLTYPE d3d12_command_queue_ExecuteCommandLists(ID3D12Comm
sub.type = VKD3D_SUBMISSION_EXECUTE;
sub.execute.cmd = buffers;
sub.execute.cmd_cost = cmd_cost;
+ sub.execute.external_waits = external_waits;
sub.execute.cmd_count = cmd_submit_count;
sub.execute.command_allocators = allocators;
sub.execute.num_command_allocators = command_list_count;
@@ -24377,6 +24427,8 @@ static void d3d12_command_queue_execute(struct d3d12_command_queue *command_queu
VkSemaphoreSubmitInfo *binary_semaphore_info;
bool stagger_submissions, is_first, is_last;
VkSemaphoreSubmitInfo wait_semaphore_info;
+ VkSemaphoreSubmitInfo combined_wait_infos[2];
+ bool external_wait;
uint32_t cmd_index, cmd_count, total_cost;
struct vkd3d_fence_wait_info fence_info;
VkSubmitInfo2 submit_desc[2], *submit;
@@ -24516,7 +24568,15 @@ static void d3d12_command_queue_execute(struct d3d12_command_queue *command_queu
total_cost = 0;
for (i = cmd_index + cmd_count; i < exec->cmd_count && total_cost < VKD3D_COMMAND_COST_MERGE_THRESHOLD; i++)
+ {
+ /* d4r: never merge across an external wait. */
+ if (exec->cmd[i].deviceMask & VKD3D_COMMAND_BUFFER_SUBMIT_INFO_DEVICE_MASK_EXTERNAL_WAIT)
+ {
+ total_cost = VKD3D_COMMAND_COST_MERGE_THRESHOLD;
+ break;
+ }
total_cost += exec->cmd_cost[i];
+ }
if (total_cost < VKD3D_COMMAND_COST_MERGE_THRESHOLD)
cmd_count = exec->cmd_count - cmd_index;
@@ -24525,8 +24585,14 @@ static void d3d12_command_queue_execute(struct d3d12_command_queue *command_queu
{
/* Submit everything at once, until we hit a fallback queue submission. Pick that up next iteration. */
for (cmd_count = 0; cmd_index + cmd_count < exec->cmd_count; cmd_count++)
+ {
if (exec->cmd[cmd_index + cmd_count].deviceMask & VKD3D_COMMAND_BUFFER_SUBMIT_INFO_DEVICE_MASK_FALLBACK_QUEUE)
break;
+ /* d4r: an external wait starts a new submit. */
+ if (cmd_count && (exec->cmd[cmd_index + cmd_count].deviceMask &
+ VKD3D_COMMAND_BUFFER_SUBMIT_INFO_DEVICE_MASK_EXTERNAL_WAIT))
+ break;
+ }
}
submit = &submit_desc[num_submits++];
@@ -24536,6 +24602,9 @@ static void d3d12_command_queue_execute(struct d3d12_command_queue *command_queu
submit->signalSemaphoreInfoCount = 1;
submit->pSignalSemaphoreInfos = signal_semaphore_infos;
+ external_wait = cmd_count && exec->external_waits &&
+ (exec->cmd[cmd_index].deviceMask & VKD3D_COMMAND_BUFFER_SUBMIT_INFO_DEVICE_MASK_EXTERNAL_WAIT);
+
/* Clear out any shenanigans we added due to fallback submit. */
for (i = 0; i < cmd_count; i++)
exec->cmd[cmd_index + i].deviceMask = 0;
@@ -24555,6 +24624,23 @@ static void d3d12_command_queue_execute(struct d3d12_command_queue *command_queu
submit->pWaitSemaphoreInfos = &wait_semaphore_info;
}
+ if (external_wait)
+ {
+ /* d4r: a queue-level wait-before-signal; the host signals the value later. */
+ if (submit->waitSemaphoreInfoCount)
+ {
+ combined_wait_infos[0] = submit->pWaitSemaphoreInfos[0];
+ combined_wait_infos[1] = exec->external_waits[cmd_index];
+ submit->waitSemaphoreInfoCount = 2;
+ submit->pWaitSemaphoreInfos = combined_wait_infos;
+ }
+ else
+ {
+ submit->waitSemaphoreInfoCount = 1;
+ submit->pWaitSemaphoreInfos = &exec->external_waits[cmd_index];
+ }
+ }
+
cmd_index += cmd_count;
is_last = cmd_index == exec->cmd_count;
@@ -25347,6 +25433,7 @@ static void *d3d12_command_queue_submission_worker_main(void *userdata)
* On error, the counters are freed early, so there is no risk of leak. */
vkd3d_free(submission.execute.cmd);
vkd3d_free(submission.execute.cmd_cost);
+ vkd3d_free(submission.execute.external_waits);
vkd3d_free(submission.execute.transitions);
#ifdef VKD3D_ENABLE_BREADCRUMBS
for (i = 0; i < submission.execute.breadcrumb_indices_count; i++)
diff --git a/libs/vkd3d/device.c b/libs/vkd3d/device.c
index f6db6954..06955941 100644
--- a/libs/vkd3d/device.c
+++ b/libs/vkd3d/device.c
@@ -4657,7 +4657,8 @@ HRESULT STDMETHODCALLTYPE d3d12_device_QueryInterface(d3d12_device_iface *iface,
if (IsEqualGUID(riid, &IID_ID3D12DXVKInteropDevice)
|| IsEqualGUID(riid, &IID_ID3D12DXVKInteropDevice1)
|| IsEqualGUID(riid, &IID_ID3D12DXVKInteropDevice2)
- || IsEqualGUID(riid, &IID_ID3D12DXVKInteropDevice3))
+ || IsEqualGUID(riid, &IID_ID3D12DXVKInteropDevice3)
+ || IsEqualGUID(riid, &IID_ID3D12DXVKInteropDeviceD4R))
{
d3d12_dxvk_interop_device_AddRef(&device->ID3D12DXVKInteropDevice_iface);
*object = &device->ID3D12DXVKInteropDevice_iface;
@@ -11379,7 +11380,7 @@ static void d3d12_device_replace_vtable(struct d3d12_device *device)
}
extern CONST_VTBL struct ID3D12DeviceExt5Vtbl d3d12_device_vkd3d_ext_vtbl;
-extern CONST_VTBL struct ID3D12DXVKInteropDevice3Vtbl d3d12_dxvk_interop_device_vtbl;
+extern CONST_VTBL struct ID3D12DXVKInteropDeviceD4RVtbl d3d12_dxvk_interop_device_vtbl;
extern CONST_VTBL struct ID3DLowLatencyDeviceVtbl d3d_low_latency_device_vtbl;
extern CONST_VTBL struct IAmdExtAntiLagApiVtbl d3d_amd_ext_anti_lag_vtbl;
diff --git a/libs/vkd3d/device_vkd3d_ext.c b/libs/vkd3d/device_vkd3d_ext.c
index 94818b9a..bd5dd7a5 100644
--- a/libs/vkd3d/device_vkd3d_ext.c
+++ b/libs/vkd3d/device_vkd3d_ext.c
@@ -1145,7 +1145,19 @@ static HRESULT STDMETHODCALLTYPE d3d12_dxvk_interop_device_GetVulkanHeapInfo(d3d
return S_OK;
}
-CONST_VTBL struct ID3D12DXVKInteropDevice3Vtbl d3d12_dxvk_interop_device_vtbl =
+static HRESULT STDMETHODCALLTYPE d3d12_dxvk_interop_device_SplitCommandListForExternalWait(
+ d3d12_dxvk_interop_device_iface *iface, ID3D12CommandList *list, UINT64 vk_semaphore, UINT64 value)
+{
+ struct d3d12_command_list *cmd_list = d3d12_command_list_from_iface(list);
+
+ TRACE("iface %p, list %p, vk_semaphore 0x%"PRIx64", value %"PRIu64".\n", iface, list, vk_semaphore, value);
+
+ if (!cmd_list || !vk_semaphore)
+ return E_INVALIDARG;
+ return d3d12_command_list_split_for_external_wait(cmd_list, (VkSemaphore)vk_semaphore, value);
+}
+
+CONST_VTBL struct ID3D12DXVKInteropDeviceD4RVtbl d3d12_dxvk_interop_device_vtbl =
{
/* IUnknown methods */
d3d12_dxvk_interop_device_QueryInterface,
@@ -1177,6 +1189,9 @@ CONST_VTBL struct ID3D12DXVKInteropDevice3Vtbl d3d12_dxvk_interop_device_vtbl =
/* ID3D12DXVKInteropDevice3 methods */
d3d12_dxvk_interop_device_GetVulkanHeapInfo,
+
+ /* ID3D12DXVKInteropDeviceD4R methods */
+ d3d12_dxvk_interop_device_SplitCommandListForExternalWait,
};
static inline struct d3d12_device *d3d12_device_from_ID3DLowLatencyDevice(d3d_low_latency_device_iface *iface)
diff --git a/libs/vkd3d/vkd3d_private.h b/libs/vkd3d/vkd3d_private.h
index a2a78000..64338585 100644
--- a/libs/vkd3d/vkd3d_private.h
+++ b/libs/vkd3d/vkd3d_private.h
@@ -3207,6 +3207,9 @@ struct d3d12_command_list_iteration
uint32_t estimated_cost;
bool fallback; /* Only applies to vk_command_buffer. */
struct d3d12_command_list_iteration_indirect_meta indirect_meta;
+ /* d4r: vk_command_buffer is submitted separately, after this timeline value is reached. */
+ VkSemaphore external_wait_semaphore;
+ uint64_t external_wait_value;
};
enum
@@ -3216,9 +3219,12 @@ enum
* without having to add more allocations. */
VKD3D_COMMAND_BUFFER_SUBMIT_INFO_DEVICE_MASK_DEFAULT = 0,
VKD3D_COMMAND_BUFFER_SUBMIT_INFO_DEVICE_MASK_FALLBACK_QUEUE = 1u << 0,
+ /* d4r: starts a new submit that waits for submission_execute::external_waits[index]. */
+ VKD3D_COMMAND_BUFFER_SUBMIT_INFO_DEVICE_MASK_EXTERNAL_WAIT = 1u << 1,
};
-#define VKD3D_MAX_COMMAND_LIST_SEQUENCES 2
+/* d4r: 4 so an external-wait split fits next to DGC and fallback splits. */
+#define VKD3D_MAX_COMMAND_LIST_SEQUENCES 4
#define VKD3D_COMMAND_COST_LOW (1u)
#define VKD3D_COMMAND_COST_HIGH (16u)
@@ -3737,6 +3743,8 @@ struct d3d12_command_queue_submission_execute
{
VkCommandBufferSubmitInfo *cmd;
uint32_t *cmd_cost;
+ /* d4r: indexed like cmd, only allocated when a command buffer waits externally. */
+ VkSemaphoreSubmitInfo *external_waits;
struct d3d12_command_allocator **command_allocators;
UINT cmd_count;
UINT num_command_allocators;
@@ -5463,7 +5471,10 @@ struct vkd3d_timestamp_profiler;
typedef ID3D12DeviceExt5 d3d12_device_vkd3d_ext_iface;
/* ID3D12DXVKInteropDevice */
-typedef ID3D12DXVKInteropDevice3 d3d12_dxvk_interop_device_iface;
+HRESULT d3d12_command_list_split_for_external_wait(struct d3d12_command_list *list,
+ VkSemaphore semaphore, uint64_t value);
+
+typedef ID3D12DXVKInteropDeviceD4R d3d12_dxvk_interop_device_iface;
/* ID3DLowLatencyDevice */
typedef ID3DLowLatencyDevice d3d_low_latency_device_iface;
patches/zluda/0002-dlss-runtime-cumulative.patch
+2768
View File
File diff suppressed because it is too large. Load diff
patches/zluda/0003-perf-implicit-bound-gfx11-wmma.patch
+772
View File
@@ -0,0 +1,772 @@
diff --git a/ptx/src/pass/mod.rs b/ptx/src/pass/mod.rs
index 15dd9ba..cff04dd 100644
--- a/ptx/src/pass/mod.rs
+++ b/ptx/src/pass/mod.rs
@@ -61,6 +61,42 @@ pub struct Attributes {
pub clock_rate: u32,
}
+/// Work-group size bound applied to kernels without `.maxntid`/`.reqntid`
+/// (D4R_ZLUDA_IMPLICIT_MAX_BLOCK, 1..=1024; unset keeps LLVM's 1024 default).
+///
+/// Without a bound, LLVM must assume 1024-thread work-groups, which on gfx11
+/// in CU mode caps each wave at 96 VGPRs and makes register-heavy kernels
+/// (e.g. DLSS's tensor-core convolutions) spill kilobytes to scratch. CUDA's
+/// own compiler gives unbounded kernels up to 255 registers, which likewise
+/// limits them to 256-thread blocks; a launch exceeding the bound fails
+/// (as it would on NVIDIA) instead of misbehaving.
+/// Optional AMDGPU machine-scheduler strategy for kernels
+/// (D4R_ZLUDA_SCHED_STRATEGY: max-ilp, max-memory-clause, iterative-ilp,
+/// iterative-minreg or iterative-maxocc); unset keeps LLVM's default.
+pub fn sched_strategy() -> Option<&'static str> {
+ static VALUE: std::sync::OnceLock<Option<String>> = std::sync::OnceLock::new();
+ VALUE
+ .get_or_init(|| {
+ std::env::var("D4R_ZLUDA_SCHED_STRATEGY").ok().filter(|value| {
+ matches!(
+ value.as_str(),
+ "max-ilp" | "max-memory-clause" | "iterative-ilp" | "iterative-minreg" | "iterative-maxocc"
+ )
+ })
+ })
+ .as_deref()
+}
+
+pub fn implicit_max_block() -> Option<u32> {
+ static VALUE: std::sync::OnceLock<Option<u32>> = std::sync::OnceLock::new();
+ *VALUE.get_or_init(|| {
+ std::env::var("D4R_ZLUDA_IMPLICIT_MAX_BLOCK")
+ .ok()
+ .and_then(|value| value.trim().parse::<u32>().ok())
+ .filter(|value| (1..=1024).contains(value))
+ })
+}
+
pub fn to_llvm_module<'input>(
ast: ast::Module<'input>,
attributes: Attributes,
diff --git a/ptx/src/lib.rs b/ptx/src/lib.rs
index 8daa637..79ed812 100644
--- a/ptx/src/lib.rs
+++ b/ptx/src/lib.rs
@@ -4,4 +4,6 @@ mod test;
pub use pass::to_llvm_module;
pub use pass::Attributes;
+pub use pass::implicit_max_block;
+pub use pass::sched_strategy;
pub use pass::TranslateError;
diff --git a/ptx/src/pass/llvm/emit.rs b/ptx/src/pass/llvm/emit.rs
index 5ab8985..e891c34 100644
--- a/ptx/src/pass/llvm/emit.rs
+++ b/ptx/src/pass/llvm/emit.rs
@@ -147,6 +147,23 @@ impl<'a, 'input> ModuleEmitContext<'a, 'input> {
}
self.emit_target_features(fn_);
self.emit_tuning(fn_, &method.tuning);
+ if method.is_kernel()
+ && !method.tuning.iter().any(|tuning| {
+ matches!(
+ tuning,
+ ptx_parser::TuningDirective::MaxNtid(..) | ptx_parser::TuningDirective::ReqNtid(..)
+ )
+ })
+ {
+ if let Some(size) = crate::pass::implicit_max_block() {
+ self.emit_fn_attribute_string(fn_, "amdgpu-flat-work-group-size", &format!("1,{size}"));
+ }
+ }
+ if method.is_kernel() {
+ if let Some(strategy) = crate::pass::sched_strategy() {
+ self.emit_fn_attribute_string(fn_, "amdgpu-sched-strategy", strategy);
+ }
+ }
if let Some(kernel_attrs) = &method.kernel_attributes {
self.emit_fn_attribute_string(
fn_,
diff --git a/zluda/src/impl/module.rs b/zluda/src/impl/module.rs
index a6adeef..0fbc7f7 100644
--- a/zluda/src/impl/module.rs
+++ b/zluda/src/impl/module.rs
@@ -260,6 +260,9 @@ fn get_best_ptx_and_compile(
let attributes = ExtraCacheAttributes {
clock_rate: hip_properties.clockRate as u32,
is_debug: cfg!(debug_assertions),
+ implicit_max_block: ptx::implicit_max_block(),
+ sched_strategy: ptx::sched_strategy(),
+ wmma: llvm_zluda::compile::wmma_rewrite_requested(),
};
let mut cache_with_key = match (text, global_state.cache_path.as_ref()) {
(Some(text), Some(p)) => (|| {
@@ -296,6 +299,9 @@ pub(crate) fn load_hip_module(
struct ExtraCacheAttributes {
is_debug: bool,
clock_rate: u32,
+ implicit_max_block: Option<u32>,
+ sched_strategy: Option<&'static str>,
+ wmma: bool,
}
fn get_hip_properties<'a>() -> Result<hipDeviceProp_tR0600, CUerror> {
diff --git a/llvm_zluda/src/compile.rs b/llvm_zluda/src/compile.rs
index b80262b..1284cbc 100644
--- a/llvm_zluda/src/compile.rs
+++ b/llvm_zluda/src/compile.rs
@@ -87,6 +87,18 @@ fn create_oclc_constants(ctx: &Context, gcn_arch: &str) -> Result<Module, String
Ok(module)
}
+/// D4R_ZLUDA_WMMA=1 lowers f16 m16n8k16 MMAs to gfx11 WMMA (see
+/// LLVMZludaRewriteMmaF16Wmma); anything else keeps the dot-product fallback.
+pub fn wmma_rewrite_requested() -> bool {
+ static VALUE: OnceLock<bool> = OnceLock::new();
+ *VALUE.get_or_init(|| std::env::var("D4R_ZLUDA_WMMA").map_or(false, |value| value.trim() == "1"))
+}
+
+fn wmma_rewrite_enabled(gcn_arch: &str) -> bool {
+ // The helpers assume RDNA3's wave32 WMMA register layout.
+ wmma_rewrite_requested() && gcn_arch.starts_with("gfx11")
+}
+
pub fn compile(
ctx: &Context,
gcn_arch: &str,
@@ -106,6 +118,9 @@ pub fn compile(
let oclc_constants = create_oclc_constants(ctx, gcn_arch)?;
+ if wmma_rewrite_enabled(gcn_arch) {
+ unsafe { crate::ffi::LLVMZludaRewriteMmaF16Wmma(main.get()) };
+ }
linked.link(main)?;
linked.link(attributes)?;
linked.link(oclc_constants)?;
diff --git a/llvm_zluda/src/ffi.rs b/llvm_zluda/src/ffi.rs
index 7ca4289..6bf0d30 100644
--- a/llvm_zluda/src/ffi.rs
+++ b/llvm_zluda/src/ffi.rs
@@ -42,6 +42,9 @@ pub const LLVMZludaFastMathAll: ::std::ffi::c_uint = LLVMZludaFastMathAllowReass
pub type LLVMZludaFastMathFlags = std::ffi::c_uint;
extern "C" {
+ /// See LLVMZludaRewriteMmaF16Wmma in lib.cpp.
+ pub fn LLVMZludaRewriteMmaF16Wmma(M: LLVMModuleRef) -> u32;
+
pub fn LLVMZludaBuildAlloca(
B: LLVMBuilderRef,
Ty: LLVMTypeRef,
diff --git a/llvm_zluda/src/lib.cpp b/llvm_zluda/src/lib.cpp
index 14b4f40..30d5490 100644
--- a/llvm_zluda/src/lib.cpp
+++ b/llvm_zluda/src/lib.cpp
@@ -9,6 +9,11 @@
#include <llvm/Support/raw_ostream.h>
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/raw_ostream.h"
+#include <llvm/ADT/DenseMap.h>
+#include <llvm/ADT/DenseSet.h>
+#include <llvm/ADT/SmallVector.h>
+#include <llvm/IR/Module.h>
+#include <llvm/IR/InstIterator.h>
#pragma GCC diagnostic pop
#include <mutex>
@@ -255,6 +260,367 @@ int LLVMZludaLinkWithLLD(uint32_t inputs, const char **input_paths, const char *
return 0;
}
+// gfx11 WMMA lowering of mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16,
+// run on ZLUDA's emitted module before linking, while PTX registers are still
+// allocas. Each MMA becomes
+// d = d4r_wmma_t_to_nv(d4r_wmma_t_mma_f16(a_from_b(B), b_from_a(A), C'))
+// (see zluda_ptx_impl.cpp for the layouts). Register pairs used as MMA
+// accumulators get a shadow alloca holding the pair in WMMA layout: an MMA
+// whose C is a plain load of such a pair reads the shadow, an MMA whose D is
+// stored to it writes the shadow, and any other store to either register
+// refreshes the shadow with d4r_wmma_t_from_nv. The invariant "shadow ==
+// WMMA layout of the pair" thus holds at every program point, and once
+// mem2reg promotes the allocas the layout conversions between chained MMAs
+// become dead code. Returns the number of MMAs rewritten.
+namespace
+{
+constexpr const char *kMmaF16 = "__zluda_ptx_impl_mma_sync_aligned_m16n8k16_row_col_f16_f16_f16_f16";
+
+struct RegisterPair
+{
+ AllocaInst *Low = nullptr;
+ AllocaInst *High = nullptr;
+ bool operator==(const RegisterPair &Other) const { return Low == Other.Low && High == Other.High; }
+};
+
+// Only allocas used purely as the address of loads and stores are tracked.
+static bool isPlainRegister(AllocaInst *Alloca)
+{
+ for (User *U : Alloca->users())
+ {
+ if (isa<LoadInst>(U))
+ continue;
+ auto *Store = dyn_cast<StoreInst>(U);
+ if (Store == nullptr || Store->getPointerOperand() != Alloca)
+ return false;
+ }
+ return true;
+}
+
+// Matches insertelement(insertelement(undef, load Low, 0), load High, 1) where
+// the loads sit in the call's block with no store to either register between
+// them and the call.
+static bool matchLoadedPair(Value *Vector, CallInst *Call, RegisterPair &Pair)
+{
+ auto *Second = dyn_cast<InsertElementInst>(Vector);
+ if (Second == nullptr)
+ return false;
+ auto *First = dyn_cast<InsertElementInst>(Second->getOperand(0));
+ auto *SecondIndex = dyn_cast<ConstantInt>(Second->getOperand(2));
+ if (First == nullptr || SecondIndex == nullptr || SecondIndex->getZExtValue() != 1)
+ return false;
+ auto *FirstIndex = dyn_cast<ConstantInt>(First->getOperand(2));
+ if (FirstIndex == nullptr || FirstIndex->getZExtValue() != 0 || !isa<UndefValue>(First->getOperand(0)))
+ return false;
+ auto *LowLoad = dyn_cast<LoadInst>(First->getOperand(1));
+ auto *HighLoad = dyn_cast<LoadInst>(Second->getOperand(1));
+ if (LowLoad == nullptr || HighLoad == nullptr || LowLoad->getParent() != Call->getParent() ||
+ HighLoad->getParent() != Call->getParent())
+ return false;
+ Pair.Low = dyn_cast<AllocaInst>(LowLoad->getPointerOperand());
+ Pair.High = dyn_cast<AllocaInst>(HighLoad->getPointerOperand());
+ if (Pair.Low == nullptr || Pair.High == nullptr || Pair.Low == Pair.High ||
+ !Pair.Low->getAllocatedType()->isIntegerTy(32) || !Pair.High->getAllocatedType()->isIntegerTy(32) ||
+ !isPlainRegister(Pair.Low) || !isPlainRegister(Pair.High))
+ return false;
+ Instruction *Earliest = LowLoad->comesBefore(HighLoad) ? LowLoad : HighLoad;
+ for (Instruction *I = Earliest->getNextNode(); I != Call; I = I->getNextNode())
+ if (auto *Store = dyn_cast<StoreInst>(I))
+ if (Store->getPointerOperand() == Pair.Low || Store->getPointerOperand() == Pair.High)
+ return false;
+ return true;
+}
+
+// Finds the stores of both lanes of an MMA result into a register pair.
+static bool matchStoredPair(CallInst *Call, RegisterPair &Pair, StoreInst *Stores[2])
+{
+ Stores[0] = Stores[1] = nullptr;
+ for (User *U : Call->users())
+ {
+ auto *Extract = dyn_cast<ExtractElementInst>(U);
+ auto *Index = Extract != nullptr ? dyn_cast<ConstantInt>(Extract->getIndexOperand()) : nullptr;
+ if (Index == nullptr || Index->getZExtValue() > 1 || !Extract->hasOneUse())
+ return false;
+ auto *Store = dyn_cast<StoreInst>(Extract->user_back());
+ if (Store == nullptr || Store->getValueOperand() != Extract || Stores[Index->getZExtValue()] != nullptr)
+ return false;
+ Stores[Index->getZExtValue()] = Store;
+ }
+ if (Stores[0] == nullptr || Stores[1] == nullptr)
+ return false;
+ Pair.Low = dyn_cast<AllocaInst>(Stores[0]->getPointerOperand());
+ Pair.High = dyn_cast<AllocaInst>(Stores[1]->getPointerOperand());
+ return Pair.Low != nullptr && Pair.High != nullptr && Pair.Low != Pair.High;
+}
+
+// Matches an N-element vector built as a chain of insertelements (indices
+// 0..N-1 into undef) of loads from plain i32 registers; returns the registers.
+static bool matchLoadedVector(Value *Vector, unsigned Elements, SmallVectorImpl<AllocaInst *> &Registers)
+{
+ Registers.assign(Elements, nullptr);
+ for (unsigned Remaining = Elements; Remaining > 0; --Remaining)
+ {
+ auto *Insert = dyn_cast<InsertElementInst>(Vector);
+ auto *Index = Insert != nullptr ? dyn_cast<ConstantInt>(Insert->getOperand(2)) : nullptr;
+ if (Index == nullptr || Index->getZExtValue() != Remaining - 1)
+ return false;
+ auto *Load = dyn_cast<LoadInst>(Insert->getOperand(1));
+ auto *Register = Load != nullptr ? dyn_cast<AllocaInst>(Load->getPointerOperand()) : nullptr;
+ if (Register == nullptr || !Register->getAllocatedType()->isIntegerTy(32) || !isPlainRegister(Register))
+ return false;
+ Registers[Remaining - 1] = Register;
+ Vector = Insert->getOperand(0);
+ }
+ return isa<UndefValue>(Vector);
+}
+
+// The registers are plain allocas (only loaded and stored directly), so only
+// explicit stores can change them.
+static bool storesBetween(Instruction *From, Instruction *To, ArrayRef<AllocaInst *> Registers)
+{
+ for (Instruction *I = From->getNextNode(); I != nullptr && I != To; I = I->getNextNode())
+ if (auto *Store = dyn_cast<StoreInst>(I))
+ if (is_contained(Registers, Store->getPointerOperand()))
+ return true;
+ return false;
+}
+
+struct Helpers
+{
+ FunctionCallee AFromB, BFromA, Mma, FromNv, ToNv, AFromBPair, MmaPair;
+};
+
+static Value *loadPair(IRBuilder<> &Builder, const RegisterPair &Pair)
+{
+ Type *I32 = Builder.getInt32Ty();
+ Value *Vector = UndefValue::get(FixedVectorType::get(I32, 2));
+ Vector = Builder.CreateInsertElement(Vector, Builder.CreateLoad(I32, Pair.Low), Builder.getInt8(0));
+ return Builder.CreateInsertElement(Vector, Builder.CreateLoad(I32, Pair.High), Builder.getInt8(1));
+}
+
+static unsigned rewriteFunction(Function &F, Function *Mma, const Helpers &H)
+{
+ SmallVector<CallInst *, 64> Calls;
+ for (Instruction &I : instructions(F))
+ if (auto *Call = dyn_cast<CallInst>(&I))
+ if (Call->getCalledFunction() == Mma && Call->arg_size() == 3)
+ Calls.push_back(Call);
+ if (Calls.empty())
+ return 0;
+
+ struct Site
+ {
+ CallInst *Call;
+ bool HasC;
+ RegisterPair C;
+ bool HasD;
+ RegisterPair D;
+ StoreInst *DStores[2];
+ };
+ SmallVector<Site, 64> Sites;
+ for (CallInst *Call : Calls)
+ {
+ Site S{Call, false, {}, false, {}, {nullptr, nullptr}};
+ S.HasC = matchLoadedPair(Call->getArgOperand(2), Call, S.C);
+ S.HasD = matchStoredPair(Call, S.D, S.DStores);
+ Sites.push_back(S);
+ }
+
+ // Accumulator pairs whose registers belong to no other pair are replaced
+ // outright by a shadow in WMMA layout.
+ SmallVector<RegisterPair, 32> Shadowed;
+ DenseMap<AllocaInst *, int> Owner; // register -> shadow index, -1 if ambiguous
+ auto consider = [&](const RegisterPair &Pair) {
+ for (const RegisterPair &Known : Shadowed)
+ if (Known == Pair)
+ return;
+ int Index = static_cast<int>(Shadowed.size());
+ for (AllocaInst *Member : {Pair.Low, Pair.High})
+ {
+ auto Found = Owner.find(Member);
+ if (Found != Owner.end())
+ Found->second = -1;
+ else
+ Owner[Member] = Index;
+ }
+ Shadowed.push_back(Pair);
+ };
+ for (const Site &S : Sites)
+ if (S.HasC)
+ consider(S.C);
+ auto shadowOf = [&](const RegisterPair &Pair) -> int {
+ auto Low = Owner.find(Pair.Low);
+ auto High = Owner.find(Pair.High);
+ if (Low == Owner.end() || High == Owner.end() || Low->second < 0 || Low->second != High->second)
+ return -1;
+ return Shadowed[Low->second] == Pair ? Low->second : -1;
+ };
+
+ Type *Pair32 = FixedVectorType::get(Type::getInt32Ty(F.getContext()), 2);
+ SmallVector<AllocaInst *, 32> Shadows;
+ {
+ IRBuilder<> Builder(&*F.getEntryBlock().getFirstInsertionPt());
+ for (size_t Index = 0; Index < Shadowed.size(); ++Index)
+ Shadows.push_back(
+ shadowOf(Shadowed[Index]) == static_cast<int>(Index)
+ ? Builder.CreateAlloca(Pair32, Shadowed[Index].Low->getAddressSpace(), nullptr, "wmma.shadow")
+ : nullptr);
+ }
+ auto shadowFor = [&](const RegisterPair &Pair) -> AllocaInst * {
+ int Index = shadowOf(Pair);
+ return Index >= 0 ? Shadows[Index] : nullptr;
+ };
+
+ // Snapshot every access to shadowed registers before rewriting anything.
+ SmallVector<LoadInst *, 64> Loads;
+ SmallVector<StoreInst *, 64> Stores;
+ for (auto &[Register, Index] : Owner)
+ {
+ if (Index < 0 || Shadows[Index] == nullptr)
+ continue;
+ for (User *U : Register->users())
+ {
+ if (auto *Load = dyn_cast<LoadInst>(U))
+ Loads.push_back(Load);
+ else
+ Stores.push_back(cast<StoreInst>(U));
+ }
+ }
+ DenseSet<StoreInst *> ResultStores; // MMA result stores replaced by shadow stores
+
+ auto accumulator = [&](const Site &S, IRBuilder<> &Builder) -> Value * {
+ if (S.HasC)
+ if (AllocaInst *Shadow = shadowFor(S.C))
+ return Builder.CreateLoad(Pair32, Shadow);
+ return Builder.CreateCall(H.FromNv, {S.Call->getArgOperand(2)});
+ };
+ auto finish = [&](const Site &S, Value *D, IRBuilder<> &Builder) {
+ if (S.HasD)
+ if (AllocaInst *Shadow = shadowFor(S.D))
+ {
+ // The result goes to the shadow where its register stores were.
+ StoreInst *Last = S.DStores[0]->comesBefore(S.DStores[1]) ? S.DStores[1] : S.DStores[0];
+ IRBuilder<> At(Last);
+ At.CreateStore(D, Shadow);
+ ResultStores.insert(S.DStores[0]);
+ ResultStores.insert(S.DStores[1]);
+ }
+ return Builder.CreateCall(H.ToNv, {D});
+ };
+ // Two MMAs that share A (same registers, no store in between to anything
+ // the second one reads) run as one 16x16x16 WMMA placed at the first.
+ auto pairable = [&](const Site &First, const Site &Second) {
+ SmallVector<AllocaInst *, 4> AFirst, ASecond, BSecond;
+ if (First.Call->getParent() != Second.Call->getParent() || !Second.HasC ||
+ shadowFor(Second.C) == nullptr || !First.Call->comesBefore(Second.Call) ||
+ !matchLoadedVector(First.Call->getArgOperand(0), 4, AFirst) ||
+ !matchLoadedVector(Second.Call->getArgOperand(0), 4, ASecond) || AFirst != ASecond ||
+ !matchLoadedVector(Second.Call->getArgOperand(1), 2, BSecond))
+ return false;
+ // The second B is reloaded at the first MMA, so it must not live in a
+ // shadow (those registers are no longer kept up to date).
+ for (AllocaInst *Register : BSecond)
+ {
+ auto Found = Owner.find(Register);
+ if (Found != Owner.end() && Found->second >= 0 && Shadows[Found->second] != nullptr)
+ return false;
+ }
+ SmallVector<AllocaInst *, 8> Read(AFirst.begin(), AFirst.end());
+ Read.append(BSecond.begin(), BSecond.end());
+ Read.push_back(Second.C.Low);
+ Read.push_back(Second.C.High);
+ return !storesBetween(First.Call, Second.Call, Read);
+ };
+ for (size_t Index = 0; Index < Sites.size(); ++Index)
+ {
+ const Site &S = Sites[Index];
+ CallInst *Call = S.Call;
+ IRBuilder<> Builder(Call);
+ Value *BTransposed = Builder.CreateCall(H.BFromA, {Call->getArgOperand(0)});
+ if (Index + 1 < Sites.size() && pairable(S, Sites[Index + 1]))
+ {
+ const Site &Next = Sites[Index + 1];
+ SmallVector<AllocaInst *, 2> BNext;
+ matchLoadedVector(Next.Call->getArgOperand(1), 2, BNext);
+ Value *B1 = loadPair(Builder, RegisterPair{BNext[0], BNext[1]});
+ Value *ATransposed = Builder.CreateCall(H.AFromBPair, {Call->getArgOperand(1), B1});
+ Value *C0 = accumulator(S, Builder);
+ Value *C1 = accumulator(Next, Builder);
+ Value *D = Builder.CreateCall(H.MmaPair, {ATransposed, BTransposed, C0, C1});
+ Value *D0 = Builder.CreateShuffleVector(D, ArrayRef<int>{0, 1});
+ Value *D1 = Builder.CreateShuffleVector(D, ArrayRef<int>{2, 3});
+ Call->replaceAllUsesWith(finish(S, D0, Builder));
+ Next.Call->replaceAllUsesWith(finish(Next, D1, Builder));
+ Call->eraseFromParent();
+ Next.Call->eraseFromParent();
+ ++Index;
+ continue;
+ }
+ Value *ATransposed = Builder.CreateCall(H.AFromB, {Call->getArgOperand(1)});
+ Value *D = Builder.CreateCall(H.Mma, {ATransposed, BTransposed, accumulator(S, Builder)});
+ Call->replaceAllUsesWith(finish(S, D, Builder));
+ Call->eraseFromParent();
+ }
+
+ // Every other access now goes through the shadow.
+ for (LoadInst *Load : Loads)
+ {
+ auto *Register = cast<AllocaInst>(Load->getPointerOperand());
+ const RegisterPair &Pair = Shadowed[Owner[Register]];
+ IRBuilder<> Builder(Load);
+ Value *Nv = Builder.CreateCall(H.ToNv, {Builder.CreateLoad(Pair32, shadowFor(Pair))});
+ Load->replaceAllUsesWith(Builder.CreateExtractElement(Nv, Builder.getInt8(Register == Pair.High)));
+ Load->eraseFromParent();
+ }
+ for (StoreInst *Store : Stores)
+ {
+ if (!ResultStores.contains(Store))
+ {
+ auto *Register = cast<AllocaInst>(Store->getPointerOperand());
+ const RegisterPair &Pair = Shadowed[Owner[Register]];
+ AllocaInst *Shadow = shadowFor(Pair);
+ IRBuilder<> Builder(Store);
+ Value *Nv = Builder.CreateCall(H.ToNv, {Builder.CreateLoad(Pair32, Shadow)});
+ Nv = Builder.CreateInsertElement(Nv, Store->getValueOperand(), Builder.getInt8(Register == Pair.High));
+ Builder.CreateStore(Builder.CreateCall(H.FromNv, {Nv}), Shadow);
+ }
+ auto *Operand = dyn_cast<Instruction>(Store->getValueOperand());
+ Store->eraseFromParent();
+ if (Operand != nullptr && Operand->use_empty() && !Operand->mayHaveSideEffects())
+ Operand->eraseFromParent();
+ }
+ return Calls.size();
+}
+} // namespace
+
+unsigned LLVMZludaRewriteMmaF16Wmma(LLVMModuleRef ModuleRef)
+{
+ Module *M = unwrap(ModuleRef);
+ Function *Mma = M->getFunction(kMmaF16);
+ if (Mma == nullptr || Mma->use_empty())
+ return 0;
+ LLVMContext &Context = M->getContext();
+ Type *I32 = Type::getInt32Ty(Context);
+ Type *V2 = FixedVectorType::get(I32, 2);
+ Type *V4 = FixedVectorType::get(I32, 4);
+ Type *V8 = FixedVectorType::get(I32, 8);
+ AttributeList Attributes = AttributeList::get(Context, AttributeList::FunctionIndex,
+ AttrBuilder(Context, Mma->getAttributes().getFnAttrs()));
+ Helpers H;
+ H.AFromB = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_a_from_b", Attributes, V8, V2);
+ H.BFromA = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_b_from_a", Attributes, V8, V4);
+ H.Mma = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_mma_f16", Attributes, V2, V8, V8, V2);
+ H.FromNv = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_from_nv", Attributes, V2, V2);
+ H.ToNv = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_to_nv", Attributes, V2, V2);
+ H.AFromBPair = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_a_from_b_pair", Attributes, V8, V2, V2);
+ H.MmaPair = M->getOrInsertFunction("__zluda_ptx_impl_d4r_wmma_t_mma_f16_pair", Attributes, V4, V8, V8, V2, V2);
+ unsigned Rewritten = 0;
+ for (Function &F : *M)
+ if (!F.isDeclaration())
+ Rewritten += rewriteFunction(F, Mma, H);
+ return Rewritten;
+}
+
bool LLVMZludaParseCommandLineOptions(int argc, const char *const *argv,
char **ErrorMessage)
{
diff --git a/ptx/lib/zluda_ptx_impl.cpp b/ptx/lib/zluda_ptx_impl.cpp
index 69d1faa..c11e06a 100644
--- a/ptx/lib/zluda_ptx_impl.cpp
+++ b/ptx/lib/zluda_ptx_impl.cpp
@@ -1235,6 +1235,184 @@ extern "C"
}
}
+// gfx11 WMMA lowering of mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16.
+//
+// llvm_zluda's MMA rewrite (LLVMZludaRewriteMmaF16Wmma) replaces each such MMA
+// with these helpers. It computes D^T = B^T A^T + C^T with
+// v_wmma_f32_16x16x16_f16, so the 16x8 accumulator occupies WMMA result
+// rows 0..7, i.e. VGPRs 0..3 of the f32 result: lane l holds D[m][n] for
+// m = l % 16 and n = 2v + l / 16, v = 0..3 (the "transposed WMMA layout").
+// Those four values are kept as two f16x2 words (v0 | v1 << 16,
+// v2 | v3 << 16), the same register footprint as NVIDIA's own fragment.
+// Converting between the two layouts is pure 16-bit data movement, so the
+// rewrite keeps accumulators in WMMA layout across chained MMAs and converts
+// only where other code reads or writes them. Each MMA still rounds its result
+// to f16, as the .f16 accumulator requires; only the f32 summation order
+// within one MMA differs from the dot-product fallback.
+typedef _Float16 wmma_half16 __attribute__((ext_vector_type(16)));
+typedef float wmma_float8 __attribute__((ext_vector_type(8)));
+typedef uint32_t wmma_uint8 __attribute__((ext_vector_type(8)));
+__device__ wmma_float8 __zluda_wmma_f32_16x16x16_f16(wmma_half16 a, wmma_half16 b, wmma_float8 c) __asm("llvm.amdgcn.wmma.f32.16x16x16.f16.v8f32.v16f16");
+
+__device__ static uint32_t wmma_lane()
+{
+ return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
+}
+
+__device__ static uint32_t wmma_from_lane(uint32_t lane, uint32_t value)
+{
+ return uint32_t(__builtin_amdgcn_ds_bpermute(int(lane << 2), int(value)));
+}
+
+// Gathers WMMA operand dword q for every lane from two NVIDIA fragment
+// registers: rows r < 8 take `low` from lane 4r+q, rows r >= 8 take `high`
+// from lane 4(r-8)+q (r = lane % 16). Swapping `high` between quad halves
+// (DPP quad_perm [2,3,0,1]) parks the value rows r >= 8 need in lane
+// 4(r-8)+(q^2), so one bpermute serves both halves.
+__device__ static uint32_t wmma_gather_rows(uint32_t low, uint32_t high_swapped, uint32_t q)
+{
+ uint32_t lane = wmma_lane();
+ uint32_t row = lane % 16;
+ uint32_t source = lane % 4 == q ? low : high_swapped;
+ uint32_t from = row < 8 ? row * 4 + q : (row - 8) * 4 + (q ^ 2);
+ return wmma_from_lane(from, source);
+}
+
+__device__ static uint32_t wmma_swap_quad_halves(uint32_t value)
+{
+ return uint32_t(__builtin_amdgcn_mov_dpp(int(value), 0b01'00'11'10, 0xf, 0xf, true));
+}
+
+extern "C"
+{
+ // WMMA A operand (lane l holds row l % 16 over k) = B^T, from an NVIDIA B
+ // fragment (lane 4n+q: B[2q..2q+1][n] in .x, B[2q+8..2q+9][n] in .y).
+ // Rows 8..15 (padding columns of the 16x8 B) are zero.
+ wmma_uint8 FUNC(d4r_wmma_t_a_from_b)(uint2::Native_vec_ b)
+ {
+ // Keep the scheduler from hoisting these shuffles (and the second
+ // candidate of each select) far ahead of the WMMA that consumes them,
+ // which otherwise spills them to scratch.
+ __builtin_amdgcn_sched_barrier(0);
+ uint32_t n = wmma_lane() % 16;
+ uint32_t source = (n % 8) * 4;
+ wmma_uint8 result;
+ for (uint32_t q = 0; q < 4; q++)
+ {
+ uint32_t low = wmma_from_lane(source + q, b.x);
+ uint32_t high = wmma_from_lane(source + q, b.y);
+ result[q] = n < 8 ? low : 0u;
+ result[q + 4] = n < 8 ? high : 0u;
+ }
+ return result;
+ }
+
+ // WMMA B operand (lane l holds column l % 16 over k) = A^T, i.e. row
+ // m = l % 16 of the NVIDIA A fragment (lane 4g+t: rows g/g+8, k 2t..2t+1 in
+ // .x/.y and k 2t+8..2t+9 in .z/.w).
+ wmma_uint8 FUNC(d4r_wmma_t_b_from_a)(uint4::Native_vec_ a)
+ {
+ // Keep the scheduler from hoisting these shuffles (and the second
+ // candidate of each select) far ahead of the WMMA that consumes them,
+ // which otherwise spills them to scratch.
+ __builtin_amdgcn_sched_barrier(0);
+ uint32_t ay = wmma_swap_quad_halves(a.y);
+ uint32_t aw = wmma_swap_quad_halves(a.w);
+ wmma_uint8 result;
+ for (uint32_t p = 0; p < 4; p++)
+ {
+ result[p] = wmma_gather_rows(a.x, ay, p);
+ result[p + 4] = wmma_gather_rows(a.z, aw, p);
+ }
+ return result;
+ }
+
+ // Paired form: the WMMA A operand rows 0..7 are B0^T and rows 8..15 are
+ // B1^T, so one WMMA computes two m16n8k16 MMAs that share A.
+ wmma_uint8 FUNC(d4r_wmma_t_a_from_b_pair)(uint2::Native_vec_ b0, uint2::Native_vec_ b1)
+ {
+ // Keep the scheduler from hoisting these shuffles (and the second
+ // candidate of each select) far ahead of the WMMA that consumes them,
+ // which otherwise spills them to scratch.
+ __builtin_amdgcn_sched_barrier(0);
+ uint32_t b1x = wmma_swap_quad_halves(b1.x);
+ uint32_t b1y = wmma_swap_quad_halves(b1.y);
+ wmma_uint8 result;
+ for (uint32_t q = 0; q < 4; q++)
+ {
+ result[q] = wmma_gather_rows(b0.x, b1x, q);
+ result[q + 4] = wmma_gather_rows(b0.y, b1y, q);
+ }
+ return result;
+ }
+
+ // Two accumulators (for B0 and B1) in transposed WMMA layout: .xy is the
+ // first, .zw the second; they are WMMA result VGPRs 0..3 and 4..7.
+ uint4::Native_vec_ FUNC(d4r_wmma_t_mma_f16_pair)(wmma_uint8 a_t, wmma_uint8 b_t, uint2::Native_vec_ c0_t, uint2::Native_vec_ c1_t)
+ {
+ f16x2 c01 = std::bit_cast<f16x2>(c0_t.x);
+ f16x2 c23 = std::bit_cast<f16x2>(c0_t.y);
+ f16x2 c45 = std::bit_cast<f16x2>(c1_t.x);
+ f16x2 c67 = std::bit_cast<f16x2>(c1_t.y);
+ wmma_float8 c = {float(c01.x), float(c01.y), float(c23.x), float(c23.y),
+ float(c45.x), float(c45.y), float(c67.x), float(c67.y)};
+ wmma_float8 d = __zluda_wmma_f32_16x16x16_f16(std::bit_cast<wmma_half16>(a_t), std::bit_cast<wmma_half16>(b_t), c);
+ f16x2 d01{static_cast<f16>(d[0]), static_cast<f16>(d[1])};
+ f16x2 d23{static_cast<f16>(d[2]), static_cast<f16>(d[3])};
+ f16x2 d45{static_cast<f16>(d[4]), static_cast<f16>(d[5])};
+ f16x2 d67{static_cast<f16>(d[6]), static_cast<f16>(d[7])};
+ return {std::bit_cast<uint32_t>(d01), std::bit_cast<uint32_t>(d23),
+ std::bit_cast<uint32_t>(d45), std::bit_cast<uint32_t>(d67)};
+ }
+
+ // D = round_f16(A B + C) with C and D in the transposed WMMA layout.
+ uint2::Native_vec_ FUNC(d4r_wmma_t_mma_f16)(wmma_uint8 a_t, wmma_uint8 b_t, uint2::Native_vec_ c_t)
+ {
+ f16x2 c01 = std::bit_cast<f16x2>(c_t.x);
+ f16x2 c23 = std::bit_cast<f16x2>(c_t.y);
+ wmma_float8 c = {float(c01.x), float(c01.y), float(c23.x), float(c23.y), 0.0f, 0.0f, 0.0f, 0.0f};
+ wmma_float8 d = __zluda_wmma_f32_16x16x16_f16(std::bit_cast<wmma_half16>(a_t), std::bit_cast<wmma_half16>(b_t), c);
+ f16x2 d01{static_cast<f16>(d[0]), static_cast<f16>(d[1])};
+ f16x2 d23{static_cast<f16>(d[2]), static_cast<f16>(d[3])};
+ return {std::bit_cast<uint32_t>(d01), std::bit_cast<uint32_t>(d23)};
+ }
+
+ // NVIDIA accumulator fragment (lane 4g+t: D[g][2t..2t+1] in .x,
+ // D[g+8][2t..2t+1] in .y) -> transposed WMMA layout.
+ uint2::Native_vec_ FUNC(d4r_wmma_t_from_nv)(uint2::Native_vec_ c)
+ {
+ uint32_t lane = wmma_lane();
+ uint32_t m = lane % 16;
+ uint32_t shift = (lane / 16) * 16;
+ uint32_t source = (m % 8) * 4;
+ uint32_t halves[4];
+ for (uint32_t v = 0; v < 4; v++)
+ {
+ uint32_t low_rows = wmma_from_lane(source + v, c.x);
+ uint32_t high_rows = wmma_from_lane(source + v, c.y);
+ halves[v] = ((m >= 8 ? high_rows : low_rows) >> shift) & 0xffffu;
+ }
+ return {halves[0] | (halves[1] << 16), halves[2] | (halves[3] << 16)};
+ }
+
+ // Transposed WMMA layout -> NVIDIA accumulator fragment.
+ uint2::Native_vec_ FUNC(d4r_wmma_t_to_nv)(uint2::Native_vec_ d)
+ {
+ uint32_t lane = wmma_lane();
+ uint32_t g = lane / 4;
+ uint32_t t = lane % 4;
+ uint32_t sources[4] = {g, g + 16, g + 8, g + 24}; // D[g][2t], D[g][2t+1], D[g+8][2t], D[g+8][2t+1]
+ uint32_t halves[4];
+ for (uint32_t i = 0; i < 4; i++)
+ {
+ uint32_t v01 = wmma_from_lane(sources[i], d.x);
+ uint32_t v23 = wmma_from_lane(sources[i], d.y);
+ halves[i] = ((t >= 2 ? v23 : v01) >> ((t % 2) * 16)) & 0xffffu;
+ }
+ return {halves[0] | (halves[1] << 16), halves[2] | (halves[3] << 16)};
+ }
+}
+
template <typename T, typename A, typename B, int BOffset, typename OcklFunc>
__device__ T dp2a_impl(T a, T b, T c, OcklFunc ockl_func)
{
diff --git a/zluda_common/examples/d4r_fatbin_ptx.rs b/zluda_common/examples/d4r_fatbin_ptx.rs
new file mode 100644
index 0000000..50f7a09
--- /dev/null
+++ b/zluda_common/examples/d4r_fatbin_ptx.rs
@@ -0,0 +1,36 @@
+// Extracts (and decompresses) every PTX file from captured CUDA fatbin images.
+// usage: d4r_fatbin_ptx <outdir> <image>...
+use cuda_types::dark_api::FatbinFileHeader;
+use std::{env, fs, path::PathBuf};
+use zluda_common::{CodeLibraryRef, CodeModuleRef};
+
+fn main() {
+ let mut args = env::args().skip(1);
+ let outdir = PathBuf::from(args.next().expect("output directory"));
+ fs::create_dir_all(&outdir).expect("create output directory");
+ for path in args {
+ let mut bytes = fs::read(&path).expect("read image");
+ bytes.extend_from_slice(&[0; 16]);
+ let stem = PathBuf::from(&path).file_stem().unwrap().to_string_lossy().into_owned();
+ let library = unsafe { CodeLibraryRef::try_load(bytes.as_ptr().cast()) }.expect("load image");
+ let mut count = 0;
+ unsafe {
+ library.iterate_modules(|_, module| {
+ let text = match module {
+ Ok(CodeModuleRef::Text(ptx)) => Some(ptx.to_string()),
+ Ok(CodeModuleRef::File(file)) if file.header.kind == FatbinFileHeader::HEADER_KIND_PTX => file
+ .get_or_decompress_content(true)
+ .ok()
+ .map(|content| String::from_utf8_lossy(&content).trim_end_matches('\0').to_string()),
+ _ => None,
+ };
+ if let Some(text) = text {
+ let name = outdir.join(format!("{stem}-{count:02}.ptx"));
+ fs::write(&name, text).expect("write PTX");
+ count += 1;
+ }
+ })
+ };
+ println!("{path}: {count} PTX modules");
+ }
+}
patches/zluda/0004-native-overrides-and-gpu-handoffs.patch
+1672
View File
File diff suppressed because it is too large. Load diff