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> = 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 { + static VALUE: std::sync::OnceLock> = std::sync::OnceLock::new(); + *VALUE.get_or_init(|| { + std::env::var("D4R_ZLUDA_IMPLICIT_MAX_BLOCK") + .ok() + .and_then(|value| value.trim().parse::().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, + sched_strategy: Option<&'static str>, + wmma: bool, } fn get_hip_properties<'a>() -> Result { 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 bool { + static VALUE: OnceLock = 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 #include "llvm/Support/CommandLine.h" #include "llvm/Support/raw_ostream.h" +#include +#include +#include +#include +#include #pragma GCC diagnostic pop #include @@ -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(U)) + continue; + auto *Store = dyn_cast(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(Vector); + if (Second == nullptr) + return false; + auto *First = dyn_cast(Second->getOperand(0)); + auto *SecondIndex = dyn_cast(Second->getOperand(2)); + if (First == nullptr || SecondIndex == nullptr || SecondIndex->getZExtValue() != 1) + return false; + auto *FirstIndex = dyn_cast(First->getOperand(2)); + if (FirstIndex == nullptr || FirstIndex->getZExtValue() != 0 || !isa(First->getOperand(0))) + return false; + auto *LowLoad = dyn_cast(First->getOperand(1)); + auto *HighLoad = dyn_cast(Second->getOperand(1)); + if (LowLoad == nullptr || HighLoad == nullptr || LowLoad->getParent() != Call->getParent() || + HighLoad->getParent() != Call->getParent()) + return false; + Pair.Low = dyn_cast(LowLoad->getPointerOperand()); + Pair.High = dyn_cast(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(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(U); + auto *Index = Extract != nullptr ? dyn_cast(Extract->getIndexOperand()) : nullptr; + if (Index == nullptr || Index->getZExtValue() > 1 || !Extract->hasOneUse()) + return false; + auto *Store = dyn_cast(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(Stores[0]->getPointerOperand()); + Pair.High = dyn_cast(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 &Registers) +{ + Registers.assign(Elements, nullptr); + for (unsigned Remaining = Elements; Remaining > 0; --Remaining) + { + auto *Insert = dyn_cast(Vector); + auto *Index = Insert != nullptr ? dyn_cast(Insert->getOperand(2)) : nullptr; + if (Index == nullptr || Index->getZExtValue() != Remaining - 1) + return false; + auto *Load = dyn_cast(Insert->getOperand(1)); + auto *Register = Load != nullptr ? dyn_cast(Load->getPointerOperand()) : nullptr; + if (Register == nullptr || !Register->getAllocatedType()->isIntegerTy(32) || !isPlainRegister(Register)) + return false; + Registers[Remaining - 1] = Register; + Vector = Insert->getOperand(0); + } + return isa(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 Registers) +{ + for (Instruction *I = From->getNextNode(); I != nullptr && I != To; I = I->getNextNode()) + if (auto *Store = dyn_cast(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 Calls; + for (Instruction &I : instructions(F)) + if (auto *Call = dyn_cast(&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 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 Shadowed; + DenseMap 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(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 Shadows; + { + IRBuilder<> Builder(&*F.getEntryBlock().getFirstInsertionPt()); + for (size_t Index = 0; Index < Shadowed.size(); ++Index) + Shadows.push_back( + shadowOf(Shadowed[Index]) == static_cast(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 Loads; + SmallVector Stores; + for (auto &[Register, Index] : Owner) + { + if (Index < 0 || Shadows[Index] == nullptr) + continue; + for (User *U : Register->users()) + { + if (auto *Load = dyn_cast(U)) + Loads.push_back(Load); + else + Stores.push_back(cast(U)); + } + } + DenseSet 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 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 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 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{0, 1}); + Value *D1 = Builder.CreateShuffleVector(D, ArrayRef{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(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(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(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(c0_t.x); + f16x2 c23 = std::bit_cast(c0_t.y); + f16x2 c45 = std::bit_cast(c1_t.x); + f16x2 c67 = std::bit_cast(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(a_t), std::bit_cast(b_t), c); + f16x2 d01{static_cast(d[0]), static_cast(d[1])}; + f16x2 d23{static_cast(d[2]), static_cast(d[3])}; + f16x2 d45{static_cast(d[4]), static_cast(d[5])}; + f16x2 d67{static_cast(d[6]), static_cast(d[7])}; + return {std::bit_cast(d01), std::bit_cast(d23), + std::bit_cast(d45), std::bit_cast(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(c_t.x); + f16x2 c23 = std::bit_cast(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(a_t), std::bit_cast(b_t), c); + f16x2 d01{static_cast(d[0]), static_cast(d[1])}; + f16x2 d23{static_cast(d[2]), static_cast(d[3])}; + return {std::bit_cast(d01), std::bit_cast(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 __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 ... +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"); + } +}