From 9f8894c933a3835044da5b5c91e702c7ccfbe195 Mon Sep 17 00:00:00 2001 From: Sanket Pandit Date: Tue, 31 Mar 2026 03:55:10 +0000 Subject: [PATCH 1/5] Add auto-select SALU/VALU emit helpers for scalar arithmetic promotion Add Handlers.h with emitAdd, emitSub, emitMul, emitLshr, emitLshl, emitAnd, emitOr, emitXor, emitMulHi, emitMinI32, emitMaxI32, and emitScalarCmp helpers that auto-select SALU when operands are scalar. Add isScalarOrImm, ensureVGPR, ensureBothVGPR utilities. Extend getConstantValue to recognize S_MOV_B32(ConstantOp) chains. Made-with: Cursor Signed-off-by: Sanket Pandit --- waveasm/include/waveasm/Transforms/Utils.h | 7 +- waveasm/lib/Transforms/handlers/Handlers.h | 273 +++++++++++++++++++++ 2 files changed, 279 insertions(+), 1 deletion(-) diff --git a/waveasm/include/waveasm/Transforms/Utils.h b/waveasm/include/waveasm/Transforms/Utils.h index 5876202885..9dea62eac5 100644 --- a/waveasm/include/waveasm/Transforms/Utils.h +++ b/waveasm/include/waveasm/Transforms/Utils.h @@ -15,7 +15,8 @@ namespace waveasm { /// Extract an integer constant from a Value. -/// Handles ConstantOp directly, or V_MOV_B32(ConstantOp), or ImmType. +/// Handles ConstantOp directly, V_MOV_B32(ConstantOp), S_MOV_B32(ConstantOp), +/// or ImmType. inline std::optional getConstantValue(mlir::Value v) { if (auto constOp = v.getDefiningOp()) return constOp.getValue(); @@ -23,6 +24,10 @@ inline std::optional getConstantValue(mlir::Value v) { if (auto constOp = movOp.getSrc().getDefiningOp()) return constOp.getValue(); } + if (auto movOp = v.getDefiningOp()) { + if (auto constOp = movOp.getSrc().getDefiningOp()) + return constOp.getValue(); + } if (auto immType = llvm::dyn_cast(v.getType())) return immType.getValue(); return std::nullopt; diff --git a/waveasm/lib/Transforms/handlers/Handlers.h b/waveasm/lib/Transforms/handlers/Handlers.h index c3d2265f9c..117d5ac46e 100644 --- a/waveasm/lib/Transforms/handlers/Handlers.h +++ b/waveasm/lib/Transforms/handlers/Handlers.h @@ -26,6 +26,7 @@ #include "waveasm/Transforms/TranslateFromMLIR.h" +#include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Operation.h" #include "mlir/Support/LogicalResult.h" @@ -183,6 +184,8 @@ mlir::LogicalResult handleMemRefAtomicRMW(mlir::Operation *op, mlir::LogicalResult handleReadFirstLane(mlir::Operation *op, TranslationContext &ctx); +mlir::LogicalResult handleROCDLSchedBarrier(mlir::Operation *op, + TranslationContext &ctx); mlir::LogicalResult handleSWaitcnt(mlir::Operation *op, TranslationContext &ctx); @@ -278,6 +281,276 @@ template inline To bitCast(From value) { return result; } +//===----------------------------------------------------------------------===// +// Scalar / VGPR Helpers +//===----------------------------------------------------------------------===// + +/// Check whether a value is scalar (SGPR) or an inline constant. +inline bool isScalarOrImm(mlir::Value v) { + return isSGPRType(v.getType()) || isImmType(v.getType()); +} + +/// If \p v is an SGPR, emit a v_mov_b32 to coerce it into a VGPR so it can +/// be used by VALU-only instructions (v_cvt_*, v_rcp_*, v_mul_f32, etc.). +/// Returns \p v unchanged when it is already a VGPR or immediate. +inline mlir::Value ensureVGPR(mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx, mlir::Value v) { + if (isSGPRType(v.getType())) { + auto vregType = ctx.createVRegType(); + return V_MOV_B32::create(builder, loc, vregType, v); + } + return v; +} + +//===----------------------------------------------------------------------===// +// Auto-select SALU/VALU Emit Helpers +//===----------------------------------------------------------------------===// + +/// Emit add: S_ADD_U32 when both operands are scalar, V_ADD_U32 otherwise. +/// Commutative: swaps to put immediate in src1 (SALU src0 must be SGPR). +inline mlir::Value emitAdd(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, + mlir::Location loc, TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && + !(isImmType(a.getType()) && isImmType(b.getType()))) { + if (isImmType(a.getType())) + std::swap(a, b); + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_ADD_U32::create(builder, loc, sregType, sccType, a, b).getDst(); + } + auto vregType = ctx.createVRegType(); + return V_ADD_U32::create(builder, loc, vregType, a, b); +} + +/// Emit sub: S_SUB_U32 when both operands are scalar, V_SUB_U32 otherwise. +/// Not commutative: src0 (minuend) must be SGPR. +inline mlir::Value emitSub(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, + mlir::Location loc, TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && isSGPRType(a.getType())) { + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_SUB_U32::create(builder, loc, sregType, sccType, a, b).getDst(); + } + auto vregType = ctx.createVRegType(); + return V_SUB_U32::create(builder, loc, vregType, a, b); +} + +/// Emit mul: S_MUL_I32 when both operands are scalar, V_MUL_LO_U32 otherwise. +/// Commutative: swaps to put immediate in src1. +inline mlir::Value emitMul(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, + mlir::Location loc, TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && + !(isImmType(a.getType()) && isImmType(b.getType()))) { + if (isImmType(a.getType())) + std::swap(a, b); + auto sregType = ctx.createSRegType(); + return S_MUL_I32::create(builder, loc, sregType, a, b); + } + auto vregType = ctx.createVRegType(); + return V_MUL_LO_U32::create(builder, loc, vregType, a, b); +} + +/// Emit logical shift right: S_LSHR_B32 when scalar, V_LSHRREV_B32 otherwise. +/// Not commutative: src0 (value) must be SGPR. +inline mlir::Value emitLshr(mlir::Value value, mlir::Value shiftAmt, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { + if (isScalarOrImm(value) && isScalarOrImm(shiftAmt) && + isSGPRType(value.getType())) { + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_LSHR_B32::create(builder, loc, sregType, sccType, value, shiftAmt) + .getDst(); + } + auto vregType = ctx.createVRegType(); + value = ensureVGPR(builder, loc, ctx, value); + return V_LSHRREV_B32::create(builder, loc, vregType, shiftAmt, value); +} + +/// Emit logical shift left: S_LSHL_B32 when scalar, V_LSHLREV_B32 otherwise. +/// Not commutative: src0 (value) must be SGPR. +inline mlir::Value emitLshl(mlir::Value value, mlir::Value shiftAmt, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { + if (isScalarOrImm(value) && isScalarOrImm(shiftAmt) && + isSGPRType(value.getType())) { + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_LSHL_B32::create(builder, loc, sregType, sccType, value, shiftAmt) + .getDst(); + } + auto vregType = ctx.createVRegType(); + value = ensureVGPR(builder, loc, ctx, value); + return V_LSHLREV_B32::create(builder, loc, vregType, shiftAmt, value); +} + +/// Emit bitwise AND: S_AND_B32 when both scalar, V_AND_B32 otherwise. +/// Commutative: swaps to put immediate in src1. +inline mlir::Value emitAnd(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, + mlir::Location loc, TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && + !(isImmType(a.getType()) && isImmType(b.getType()))) { + if (isImmType(a.getType())) + std::swap(a, b); + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_AND_B32::create(builder, loc, sregType, sccType, a, b).getDst(); + } + auto vregType = ctx.createVRegType(); + return V_AND_B32::create(builder, loc, vregType, a, b); +} + +/// Emit bitwise OR: S_OR_B32 when both scalar, V_OR_B32 otherwise. +/// Commutative: swaps to put immediate in src1. +inline mlir::Value emitOr(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, + mlir::Location loc, TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && + !(isImmType(a.getType()) && isImmType(b.getType()))) { + if (isImmType(a.getType())) + std::swap(a, b); + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_OR_B32::create(builder, loc, sregType, sccType, a, b).getDst(); + } + auto vregType = ctx.createVRegType(); + return V_OR_B32::create(builder, loc, vregType, a, b); +} + +/// Emit bitwise XOR: S_XOR_B32 when both scalar, V_XOR_B32 otherwise. +/// Commutative: swaps to put immediate in src1. +inline mlir::Value emitXor(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, + mlir::Location loc, TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && + !(isImmType(a.getType()) && isImmType(b.getType()))) { + if (isImmType(a.getType())) + std::swap(a, b); + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_XOR_B32::create(builder, loc, sregType, sccType, a, b).getDst(); + } + auto vregType = ctx.createVRegType(); + return V_XOR_B32::create(builder, loc, vregType, a, b); +} + +/// Emit mulhi: S_MUL_HI_U32 when both scalar, V_MUL_HI_U32 otherwise. +/// Commutative: swaps to put immediate in src1. +inline mlir::Value emitMulHi(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && + !(isImmType(a.getType()) && isImmType(b.getType()))) { + if (isImmType(a.getType())) + std::swap(a, b); + auto sregType = ctx.createSRegType(); + return S_MUL_HI_U32::create(builder, loc, sregType, a, b); + } + auto vregType = ctx.createVRegType(); + return V_MUL_HI_U32::create(builder, loc, vregType, a, b); +} + +/// Emit signed max: S_MAX_I32 when both scalar, V_MAX_I32 otherwise. +/// Not commutative for SCC semantics (SCC = src0 was selected) but +/// the result value is commutative. +inline mlir::Value emitMaxI32(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && + !(isImmType(a.getType()) && isImmType(b.getType()))) { + if (isImmType(a.getType())) + std::swap(a, b); + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_MAX_I32::create(builder, loc, sregType, sccType, a, b).getDst(); + } + auto vregType = ctx.createVRegType(); + return V_MAX_I32::create(builder, loc, vregType, a, b); +} + +/// Emit signed min: S_MIN_I32 when both scalar, V_MIN_I32 otherwise. +inline mlir::Value emitMinI32(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && + !(isImmType(a.getType()) && isImmType(b.getType()))) { + if (isImmType(a.getType())) + std::swap(a, b); + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_MIN_I32::create(builder, loc, sregType, sccType, a, b).getDst(); + } + auto vregType = ctx.createVRegType(); + return V_MIN_I32::create(builder, loc, vregType, a, b); +} + +/// Emit unsigned max: S_MAX_U32 when both scalar, V_MAX_U32 otherwise. +inline mlir::Value emitMaxU32(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && + !(isImmType(a.getType()) && isImmType(b.getType()))) { + if (isImmType(a.getType())) + std::swap(a, b); + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_MAX_U32::create(builder, loc, sregType, sccType, a, b).getDst(); + } + auto vregType = ctx.createVRegType(); + return V_MAX_U32::create(builder, loc, vregType, a, b); +} + +/// Emit unsigned min: S_MIN_U32 when both scalar, V_MIN_U32 otherwise. +inline mlir::Value emitMinU32(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { + if (isScalarOrImm(a) && isScalarOrImm(b) && + !(isImmType(a.getType()) && isImmType(b.getType()))) { + if (isImmType(a.getType())) + std::swap(a, b); + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_MIN_U32::create(builder, loc, sregType, sccType, a, b).getDst(); + } + auto vregType = ctx.createVRegType(); + return V_MIN_U32::create(builder, loc, vregType, a, b); +} + +//===----------------------------------------------------------------------===// +// Scalar Comparison Helper +//===----------------------------------------------------------------------===// + +/// Emit S_CMP_* for the given predicate with SGPR operands (sets SCC). +/// Both lhs and rhs must be SGPRs (not immediates) before calling. +/// Returns the S_CMP result value (phantom SCC). +inline mlir::Value emitScalarCmp(mlir::OpBuilder &builder, mlir::Location loc, + mlir::arith::CmpIPredicate pred, + mlir::Value lhs, mlir::Value rhs, + TranslationContext &ctx) { + auto sccType = ctx.createSCCType(); + switch (pred) { + case mlir::arith::CmpIPredicate::eq: + return S_CMP_EQ_U32::create(builder, loc, sccType, lhs, rhs); + case mlir::arith::CmpIPredicate::ne: + return S_CMP_NE_U32::create(builder, loc, sccType, lhs, rhs); + case mlir::arith::CmpIPredicate::slt: + return S_CMP_LT_I32::create(builder, loc, sccType, lhs, rhs); + case mlir::arith::CmpIPredicate::sle: + return S_CMP_LE_I32::create(builder, loc, sccType, lhs, rhs); + case mlir::arith::CmpIPredicate::sgt: + return S_CMP_GT_I32::create(builder, loc, sccType, lhs, rhs); + case mlir::arith::CmpIPredicate::sge: + return S_CMP_GE_I32::create(builder, loc, sccType, lhs, rhs); + case mlir::arith::CmpIPredicate::ult: + return S_CMP_LT_U32::create(builder, loc, sccType, lhs, rhs); + case mlir::arith::CmpIPredicate::ule: + return S_CMP_LE_U32::create(builder, loc, sccType, lhs, rhs); + case mlir::arith::CmpIPredicate::ugt: + return S_CMP_GT_U32::create(builder, loc, sccType, lhs, rhs); + case mlir::arith::CmpIPredicate::uge: + return S_CMP_GE_U32::create(builder, loc, sccType, lhs, rhs); + } + llvm_unreachable("unhandled CmpIPredicate"); +} + //===----------------------------------------------------------------------===// // Error Handling Helpers //===----------------------------------------------------------------------===// From e61e36824935dd92ff25cc497b031b7a48fc64d3 Mon Sep 17 00:00:00 2001 From: Sanket Pandit Date: Tue, 31 Mar 2026 18:14:45 +0000 Subject: [PATCH 2/5] Rewrite MLIR-to-waveasm handlers for SALU promotion of scalar arithmetic Rewrite Arith/Affine/AMDGPU/MemRef handlers to use auto-select emit helpers (emitAdd, emitMul, emitScalarCmp, etc.) that route scalar operands through SALU instructions. Keep scalar kernel arguments in SGPRs and promote uniform address arithmetic to SALU, reducing VGPR pressure. Fuse arith.cmpi + arith.select into s_cmp + s_cselect. Rewrite SRD prologue with all-or-nothing preloading strategy and overflow SGPR slots. Modify SRDs in-place to reduce SGPR pressure. Add handleROCDLSchedBarrier and maskedload without mask support. Replace SRD prologue RawOps with typed S_LOAD_DWORDX2, S_MOV_B64, S_MOV_B32, and ExtractOp via SSA value reuse pattern. Adapt upstream PackOp-based SRD construction (PR #1188) to use explicit SCCType for S_AND_B32, S_OR_B32, and S_ADDC_U32 call sites. Made-with: Cursor Signed-off-by: Sanket Pandit --- .../waveasm/Transforms/TranslateFromMLIR.h | 29 +- waveasm/lib/Transforms/MetadataEmitter.cpp | 5 +- waveasm/lib/Transforms/TranslateFromMLIR.cpp | 324 ++++++++++-------- .../Transforms/handlers/AMDGPUHandlers.cpp | 50 ++- .../Transforms/handlers/AffineHandlers.cpp | 143 +++++--- .../lib/Transforms/handlers/ArithHandlers.cpp | 272 +++++++++++---- .../lib/Transforms/handlers/HandlerUtils.cpp | 35 +- .../Transforms/handlers/MemRefHandlers.cpp | 165 ++++++--- 8 files changed, 707 insertions(+), 316 deletions(-) diff --git a/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h b/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h index cbfe19176e..84ac089da0 100644 --- a/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h +++ b/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h @@ -377,6 +377,8 @@ class TranslationContext { /// Get next available SRD index int64_t getNextSRDIndex(); + + /// Update buffer size for a pending SRD (called when we see reinterpret_cast) void updateSRDBufferSize(mlir::Value memref, int64_t bufferSize); @@ -385,6 +387,26 @@ class TranslationContext { return pendingSRDs.size() + pendingScalarArgs.size(); } + /// Get the number of args that fit in hardware preload SGPRs. + /// Uses all-or-nothing strategy: preloads ALL args when they fit, + /// otherwise preloads NOTHING. Partial preloading (SRDs only) is + /// avoided because mixing HW preloading with explicit s_load from + /// the same kernarg buffer produces incorrect values on GFX950. + static constexpr int64_t kMaxPreloadSGPRs = 16; + size_t getNumPreloadedArgs() const { + size_t total = pendingSRDs.size() + pendingScalarArgs.size(); + if (total * 2 <= static_cast(kMaxPreloadSGPRs)) + return total; + return 0; + } + + /// Whether scalar kernel args are hardware-preloaded (true when all args fit + /// in the preload limit) or must be loaded via explicit s_load instructions. + bool areScalarsPreloaded() const { + size_t total = pendingSRDs.size() + pendingScalarArgs.size(); + return total * 2 <= static_cast(kMaxPreloadSGPRs); + } + //===--------------------------------------------------------------------===// // Split Vector Result Tracking //===--------------------------------------------------------------------===// @@ -610,12 +632,11 @@ class TranslationContext { /// Get the number of user SGPRs (for computing system SGPR offsets) /// User SGPRs include: kernarg ptr (2) + preloaded args (on gfx950+) + /// Hardware limits user SGPRs to 16 on gfx950. Args that don't fit are + /// loaded via explicit s_load into overflow SGPR slots after the SRDs. int64_t getUserSgprCount() const { - // Base: 2 SGPRs for kernarg_segment_ptr - int64_t count = 2; - // On gfx950+ with kernarg preloading, add preloaded args + int64_t count = 2; // kernarg_segment_ptr if (llvm::isa(target)) { - // Each kernel arg uses 2 SGPRs, capped at 14 (hardware max 16 total). count += std::min(size_t(14), getNumKernelArgs() * 2); } return count; diff --git a/waveasm/lib/Transforms/MetadataEmitter.cpp b/waveasm/lib/Transforms/MetadataEmitter.cpp index db3ff8de72..e3c97b65d6 100644 --- a/waveasm/lib/Transforms/MetadataEmitter.cpp +++ b/waveasm/lib/Transforms/MetadataEmitter.cpp @@ -148,7 +148,6 @@ static void scanSystemRegisterUsage(ProgramOp program, bool &usesWorkgroupIdX, int64_t userSgprCount = 2; if (isGfx950) { - // Hardware limits user SGPRs to 16 on gfx950. userSgprCount = std::min(int64_t(16), 2 + numArgs * 2); } @@ -206,17 +205,15 @@ MetadataEmitter::emitKernelDescriptor(int64_t peakVGPRs, int64_t peakSGPRs, auto targetAttr = program.getTarget(); auto targetKind = targetAttr.getTargetKind(); - int64_t preloadLength = program.getKernargPreloadLength(); bool usePreloading = llvm::isa(targetKind); + int64_t preloadLength = program.getKernargPreloadLength(); if (usePreloading && preloadLength == 0) { int64_t numArgs = 2; if (auto numArgsAttr = program->getAttrOfType("num_kernel_args")) { numArgs = numArgsAttr.getInt(); } - // Hardware limits user SGPRs to 16 on gfx950 (2 for kernarg ptr + 14 max - // preloaded). Overflow args are loaded via explicit s_load in the prologue. preloadLength = std::min(int64_t(14), numArgs * 2); } diff --git a/waveasm/lib/Transforms/TranslateFromMLIR.cpp b/waveasm/lib/Transforms/TranslateFromMLIR.cpp index fdea6b1d63..5b0b1bcf07 100644 --- a/waveasm/lib/Transforms/TranslateFromMLIR.cpp +++ b/waveasm/lib/Transforms/TranslateFromMLIR.cpp @@ -162,97 +162,116 @@ void TranslationContext::emitSRDPrologue() { srdPrologueEmitted = true; auto loc = builder.getUnknownLoc(); - // Check if this is a gfx95* target (requires preload pattern with - // branch+alignment) bool isGFX95 = llvm::isa(target); + bool usePreloading = isGFX95; // Recompute SRD base indices now that we know the total number of args. // SRDs must start after: user SGPRs + system SGPRs (workgroup IDs). - size_t numPreloadedArgs = getNumKernelArgs(); - int64_t userSgprCount = 2; // kernarg ptr - if (isGFX95) { - userSgprCount += std::min(int64_t(14), (int64_t)getNumKernelArgs() * 2); - } + int64_t userSgprCount = getUserSgprCount(); int64_t systemSgprCount = 3; // workgroup_id_x, y, z int64_t srdStartIndex = (userSgprCount + systemSgprCount + 3) & ~3; // Align to 4 - // Update pending SRDs with correct indices and update srdIndexMap + // Update pending SRDs with correct indices and update srdIndexMap. for (size_t i = 0; i < pendingSRDs.size(); ++i) { int64_t newSrdBase = srdStartIndex + i * 4; pendingSRDs[i].srdBaseIndex = newSrdBase; srdIndexMap[pendingSRDs[i].memref] = newSrdBase; } - // Count overflow scalar args (those that exceed the 16-SGPR preload limit). + // Compute overflow scalar count: args whose preload slot (2+argIndex*2) + // exceeds the 16-SGPR hardware preload window. int64_t numOverflowScalars = 0; for (const auto &pending : pendingScalarArgs) { if (2 + pending.argIndex * 2 >= 16) ++numOverflowScalars; } - // Emit comment for prologue + // Overflow scalars are placed in SGPRs after all SRDs. + int64_t afterLastSrd = srdStartIndex + pendingSRDs.size() * 4; + int64_t overflowSgprBase = (afterLastSrd + 3) & ~3; + // Dedicated SGPR slots for scalar kernel args (placed after overflow slots). + int64_t scalarArgSgprBase = + (overflowSgprBase + numOverflowScalars * 2 + 3) & ~3; + + // Emit comment for prologue. CommentOp::create(builder, loc, "SRD setup prologue"); - // Emit precolored SGPR ops for ABI registers used by the raw assembly - // prologue below. This makes the implicit physical-register usage visible - // to the register allocator so it won't allocate over these indices. - // - // s[0:1] = kernarg segment pointer (used as source in s_load_dwordx2). + // s[0:1] = kernarg segment pointer. auto kernargPtrType = createSRegType(2, 2); PrecoloredSRegOp::create(builder, loc, kernargPtrType, /*index=*/0, /*size=*/2); - if (isGFX95) { - // On gfx95*, the prologue loads kernarg data into preload locations - // s[2:3], s[4:5], etc. Reserve those pairs so regalloc doesn't use them. - // Hardware limits user SGPRs to 16 (s[0:15]), so only reserve preload - // slots for args that fit within the limit. Overflow args are loaded - // via explicit s_load from the kernarg buffer at runtime. - // Reserve all arg positions, not just pointer args with SRDs. - for (size_t i = 0; i < numPreloadedArgs; ++i) { - int64_t preloadBase = 2 + i * 2; + if (usePreloading) { + // Reserve preload SGPR pairs within the 16-SGPR hardware window + // so regalloc doesn't use them. + llvm::DenseSet reservedPreloadBases; + for (const auto &pending : pendingSRDs) { + int64_t preloadBase = 2 + pending.argIndex * 2; + if (preloadBase >= 16) + continue; + if (reservedPreloadBases.insert(preloadBase).second) { + auto preloadType = createSRegType(2, 2); + PrecoloredSRegOp::create(builder, loc, preloadType, preloadBase, + /*size=*/2); + } + } + for (const auto &pending : pendingScalarArgs) { + int64_t preloadBase = 2 + pending.argIndex * 2; if (preloadBase >= 16) continue; - auto preloadType = createSRegType(2, 2); - PrecoloredSRegOp::create(builder, loc, preloadType, preloadBase, - /*size=*/2); + if (reservedPreloadBases.insert(preloadBase).second) { + auto preloadType = createSRegType(2, 2); + PrecoloredSRegOp::create(builder, loc, preloadType, preloadBase, + /*size=*/2); + } } } - if (isGFX95) { - // GFX95* path: Use preload pattern with intermediate locations and - // s_mov_b64 copies. This matches the Python backend behavior for gfx950. - // - // Step 1: Load base addresses into preload locations s[2:3], s[4:5], etc. - // using typed WaveASM ops instead of raw strings. + if (usePreloading) { + // GFX95* preloading path with partial preloading: args that fit in + // s[2:15] are hardware-preloaded, overflow args are loaded via + // explicit s_load after the aligned entry point. + + // Step 1: Load base addresses into preload locations. + // Capture SSA results so typed copy ops below can reference them. auto kernargSRegType = createSRegType(2, 2); auto kernargBase = PrecoloredSRegOp::create(builder, loc, kernargSRegType, 0, 2); - // Load all kernel args (pointers and scalars) into preload positions. - // Capture SSA results so S_MOV_B64 ops below can reference them, - // keeping the loads live and preventing the register allocator from - // aliasing their destination registers. llvm::DenseMap argLoadResults; - for (size_t i = 0; i < numPreloadedArgs; ++i) { - int64_t loadBase = 2 + i * 2; + for (const auto &pending : pendingSRDs) { + int64_t loadBase = 2 + pending.argIndex * 2; if (loadBase >= 16) - continue; // Overflow arg: loaded via s_load_dword path below. - int64_t kernargOffset = i * 8; + continue; + int64_t kernargOffset = pending.argIndex * 8; auto loadDstType = createSRegType(2, loadBase); auto offsetImm = builder.getType(kernargOffset); auto offsetConst = ConstantOp::create(builder, loc, offsetImm, kernargOffset); - auto loadOp = S_LOAD_DWORDX2::create(builder, loc, TypeRange{loadDstType}, + auto loadOp = S_LOAD_DWORDX2::create(builder, loc, + TypeRange{loadDstType}, kernargBase, offsetConst); - argLoadResults[i] = loadOp->getResult(0); + argLoadResults[pending.argIndex] = loadOp->getResult(0); } - // Overflow scalar args: loaded after the aligned entry point below. - int64_t overflowSgprBase = - (srdStartIndex + (int64_t)pendingSRDs.size() * 4 + 3) & ~3; + // Also load scalar args that fit in the preload window. + for (const auto &pending : pendingScalarArgs) { + int64_t loadBase = 2 + pending.argIndex * 2; + if (loadBase >= 16) + continue; + int64_t kernargOffset = pending.argIndex * 8; + + auto loadDstType = createSRegType(2, loadBase); + auto offsetImm = builder.getType(kernargOffset); + auto offsetConst = + ConstantOp::create(builder, loc, offsetImm, kernargOffset); + auto loadOp = S_LOAD_DWORDX2::create(builder, loc, + TypeRange{loadDstType}, + kernargBase, offsetConst); + argLoadResults[pending.argIndex] = loadOp->getResult(0); + } // Step 2: Branch to aligned entry point (gfx95* requirement). // Keep any high-SGPR overflow loads after the aligned entry; LLVM does the @@ -266,29 +285,31 @@ void TranslationContext::emitSRDPrologue() { RawOp::create(builder, loc, ".p2align 8"); RawOp::create(builder, loc, mainLabel + ":"); - // Step 3: Load scalar overflow args after the aligned entry point. - // Reserve overflow SGPR slots so the register allocator avoids them. + // Step 3: Load overflow scalar args after the aligned entry point. for (int64_t i = 0; i < numOverflowScalars; ++i) { int64_t base = overflowSgprBase + i * 2; auto ovfType = createSRegType(2, 2); PrecoloredSRegOp::create(builder, loc, ovfType, base, /*size=*/2); } + llvm::DenseMap overflowLoadResults; int64_t overflowIdx = 0; for (const auto &pending : pendingScalarArgs) { int64_t preloadBase = 2 + pending.argIndex * 2; if (preloadBase < 16) continue; int64_t loadBase = overflowSgprBase + overflowIdx * 2; - overflowIdx++; int64_t kernargOffset = pending.argIndex * 8; auto loadDstType = createSRegType(2, loadBase); auto offsetImm = builder.getType(kernargOffset); auto offsetConst = ConstantOp::create(builder, loc, offsetImm, kernargOffset); - S_LOAD_DWORDX2::create(builder, loc, TypeRange{loadDstType}, kernargBase, - offsetConst); + auto loadOp = S_LOAD_DWORDX2::create(builder, loc, + TypeRange{loadDstType}, + kernargBase, offsetConst); + overflowLoadResults[overflowIdx] = loadOp->getResult(0); + overflowIdx++; } // Step 4: Wait for all scalar loads to complete. @@ -310,13 +331,12 @@ void TranslationContext::emitSRDPrologue() { auto srdType = createSRegType(4, 4); auto srdReg = PrecoloredSRegOp::create(builder, loc, srdType, srdBase, 4); - // Copy base address with s_mov_b64. Value preloadSrc; auto loadIt = argLoadResults.find(pending.argIndex); - if (preloadBase < 16 && loadIt != argLoadResults.end()) { + if (loadIt != argLoadResults.end()) { preloadSrc = loadIt->second; } else { - auto preloadType = createSRegType(2, preloadBase); + auto preloadType = createSRegType(2, 2); preloadSrc = PrecoloredSRegOp::create(builder, loc, preloadType, preloadBase, 2); } @@ -340,84 +360,94 @@ void TranslationContext::emitSRDPrologue() { auto movStride = S_MOV_B32::create(builder, loc, dstW3Type, strideVal); DCEProtectOp::create(builder, loc, movStride); + // Prevent DCE from removing this PrecoloredSRegOp. The SRD registers + // are written to by S_MOV_B64/S_MOV_B32 targeting PSRegType, but + // canonicalize may still remove the PrecoloredSRegOp if it has no + // SSA users. Without the reservation, the register allocator doesn't + // know s[srdBase:srdBase+3] are occupied and may allocate temps there. + DCEProtectOp::create(builder, loc, srdReg); + mapper.mapValue(pending.memref, srdReg); } - // Move scalar args from preload SGPRs to VGPRs. - // Lower 32 bits of the preload pair hold the value (little-endian). - // Overflow args were loaded into overflowSgprBase positions above. + // Copy scalar args from preload/overflow SGPRs to dedicated SGPRs. + // Extract the low 32-bit word from the 64-bit load result. int64_t ovfIdx = 0; - for (const auto &pending : pendingScalarArgs) { + for (size_t i = 0; i < pendingScalarArgs.size(); ++i) { + const auto &pending = pendingScalarArgs[i]; int64_t preloadBase = 2 + pending.argIndex * 2; - int64_t sgprSrc; + int64_t sgprDst = scalarArgSgprBase + i; + + Value srcPairVal; if (preloadBase >= 16) { - sgprSrc = overflowSgprBase + ovfIdx * 2; + auto ovfIt = overflowLoadResults.find(ovfIdx); + if (ovfIt != overflowLoadResults.end()) + srcPairVal = ovfIt->second; ovfIdx++; } else { - sgprSrc = preloadBase; + auto loadIt = argLoadResults.find(pending.argIndex); + if (loadIt != argLoadResults.end()) + srcPairVal = loadIt->second; } - auto vregType = createVRegType(); - auto vreg = - PrecoloredVRegOp::create(builder, loc, vregType, pending.argIndex, 1); - RawOp::create(builder, loc, - "v_mov_b32 v" + std::to_string(pending.argIndex) + ", s" + - std::to_string(sgprSrc)); - mapper.mapValue(pending.blockArg, vreg); + + auto dstType = createSRegType(1, 1); + auto dstSreg = + PrecoloredSRegOp::create(builder, loc, dstType, sgprDst, 1); + + if (srcPairVal) { + auto sregTy = createSRegType(1, 1); + Value lowWord = ExtractOp::create(builder, loc, sregTy, srcPairVal, 0); + auto dstPsType = PSRegType::get(builder.getContext(), sgprDst, 1); + auto movOp = S_MOV_B32::create(builder, loc, dstPsType, lowWord); + DCEProtectOp::create(builder, loc, movOp); + } + + DCEProtectOp::create(builder, loc, dstSreg); + mapper.mapValue(pending.blockArg, dstSreg); } } else { - // Non-GFX95* path (e.g., gfx942): Load directly into SRD positions - // This eliminates the s_mov_b64 copies by loading args directly into the - // SRD base addresses (SRD[0:1]), then only filling size/stride with - // s_mov_b32. - // - // Step 1: Load base addresses directly into SRD[0:1] positions - // using typed WaveASM ops. + // Direct-load path (non-GFX950 targets): + // Load base addresses into virtual registers, then copy to SRD positions. + auto kernargSRegType = createSRegType(2, 2); auto kernargBase = PrecoloredSRegOp::create(builder, loc, kernargSRegType, 0, 2); + llvm::DenseMap srdLoadResults; for (const auto &pending : pendingSRDs) { int64_t srdBase = pending.srdBaseIndex; int64_t kernargOffset = pending.argIndex * 8; - // Load directly into SRD base: s[srdBase:srdBase+1] - auto loadDstType = createSRegType(2, srdBase); + auto loadDstType = createSRegType(2, 2); auto offsetImm = builder.getType(kernargOffset); auto offsetConst = ConstantOp::create(builder, loc, offsetImm, kernargOffset); - S_LOAD_DWORDX2::create(builder, loc, TypeRange{loadDstType}, kernargBase, - offsetConst); + auto loadOp = S_LOAD_DWORDX2::create(builder, loc, + TypeRange{loadDstType}, + kernargBase, offsetConst); + srdLoadResults[srdBase] = loadOp->getResult(0); } - // Load scalar kernel arguments (index types) into pinned SGPRs after - // all SRDs. Must use RawOp + PrecoloredSRegOp so the register - // allocator does not move the load destinations away from the SGPRs - // that the subsequent RawOp v_mov_b32 references. - int64_t scalarSgprBase = - (srdStartIndex + (int64_t)pendingSRDs.size() * 4 + 3) & - ~3; // Align to 4, after all SRDs + llvm::DenseMap scalarLoadResults; for (size_t i = 0; i < pendingScalarArgs.size(); ++i) { const auto &pending = pendingScalarArgs[i]; - int64_t sgprIdx = scalarSgprBase + (int64_t)i; int64_t kernargOffset = pending.argIndex * 8; - // Pin the destination SGPR so regalloc doesn't reuse it. - PrecoloredSRegOp::create(builder, loc, createSRegType(1, 1), sgprIdx, 1); - - // Use RawOp to load directly into the pinned SGPR. - RawOp::create(builder, loc, - "s_load_dword s" + std::to_string(sgprIdx) + ", s[0:1], " + - std::to_string(kernargOffset)); + auto loadDstType = createSRegType(2, 2); + auto offsetImm = builder.getType(kernargOffset); + auto offsetConst = + ConstantOp::create(builder, loc, offsetImm, kernargOffset); + auto loadOp = S_LOAD_DWORDX2::create(builder, loc, + TypeRange{loadDstType}, + kernargBase, offsetConst); + scalarLoadResults[(int64_t)i] = loadOp->getResult(0); } - // Step 2: Wait for all scalar loads to complete auto i32Type = builder.getI32Type(); auto lgkmcntAttr = IntegerAttr::get(i32Type, 0); S_WAITCNT::create(builder, loc, /*vmcnt=*/IntegerAttr{}, lgkmcntAttr, /*expcnt=*/IntegerAttr{}); - // Step 3: Fill SRD[2:3] with size and stride. - // Use typed ops targeting precolored registers with DCEProtectOp. for (size_t i = 0; i < pendingSRDs.size(); ++i) { const auto &pending = pendingSRDs[i]; int64_t srdBase = pending.srdBaseIndex; @@ -425,6 +455,15 @@ void TranslationContext::emitSRDPrologue() { auto srdType = createSRegType(4, 4); auto srdReg = PrecoloredSRegOp::create(builder, loc, srdType, srdBase, 4); + // Copy base address (words 0-1) from load result to SRD position. + auto loadIt = srdLoadResults.find(srdBase); + if (loadIt != srdLoadResults.end()) { + auto dstB64Type = PSRegType::get(builder.getContext(), srdBase, 2); + auto movB64 = S_MOV_B64::create(builder, loc, dstB64Type, + loadIt->second); + DCEProtectOp::create(builder, loc, movB64); + } + // Set num_records (word 2). int64_t clampedSize = std::min(pending.bufferSize, kMaxNumRecords32); auto sizeImm = createImmType(clampedSize); @@ -444,18 +483,27 @@ void TranslationContext::emitSRDPrologue() { mapper.mapValue(pending.memref, srdReg); } - // Move scalar args from SGPRs to VGPRs + // Copy scalar args from load positions to dedicated SGPRs. for (size_t i = 0; i < pendingScalarArgs.size(); ++i) { const auto &pending = pendingScalarArgs[i]; - int64_t sgprIdx = scalarSgprBase + (int64_t)i; + int64_t sgprIdx = scalarArgSgprBase + (int64_t)i; + + auto dstType = createSRegType(1, 1); + auto dstSreg = + PrecoloredSRegOp::create(builder, loc, dstType, sgprIdx, 1); + + auto loadIt = scalarLoadResults.find((int64_t)i); + if (loadIt != scalarLoadResults.end()) { + auto sregTy = createSRegType(1, 1); + Value lowWord = ExtractOp::create(builder, loc, sregTy, + loadIt->second, 0); + auto dstPsType = PSRegType::get(builder.getContext(), sgprIdx, 1); + auto movOp = S_MOV_B32::create(builder, loc, dstPsType, lowWord); + DCEProtectOp::create(builder, loc, movOp); + } - auto vregType = createVRegType(); - auto vreg = - PrecoloredVRegOp::create(builder, loc, vregType, pending.argIndex, 1); - RawOp::create(builder, loc, - "v_mov_b32 v" + std::to_string(pending.argIndex) + ", s" + - std::to_string(sgprIdx)); - mapper.mapValue(pending.blockArg, vreg); + DCEProtectOp::create(builder, loc, dstSreg); + mapper.mapValue(pending.blockArg, dstSreg); } } @@ -601,6 +649,10 @@ VOffsetResult computeVOffsetFromIndices(MemRefType memrefType, auto immType = ctx.createImmType(0); voffset = ConstantOp::create(builder, loc, immType, 0); } + // buffer_load/store with offen mode requires a VGPR for voffset. + if (!llvm::isa(voffset.getType())) { + voffset = V_MOV_B32::create(builder, loc, ctx.createVRegType(), voffset); + } return {voffset, instOffset}; } @@ -1033,12 +1085,19 @@ LogicalResult handleVectorTransferWrite(Operation *op, voffset = ConstantOp::create(builder, loc, immType, 0); } + Value storeData = *data; + if (isAGPRType(storeData.getType())) { + auto vregType = + ctx.createVRegType(numDwords, numDwords > 1 ? numDwords : 1); + storeData = V_ACCVGPR_READ_B32::create(builder, loc, vregType, storeData); + } + if (numDwords == 1) { - BUFFER_STORE_DWORD::create(builder, loc, *data, srd, voffset); + BUFFER_STORE_DWORD::create(builder, loc, storeData, srd, voffset); } else if (numDwords == 2) { - BUFFER_STORE_DWORDX2::create(builder, loc, *data, srd, voffset); + BUFFER_STORE_DWORDX2::create(builder, loc, storeData, srd, voffset); } else { - BUFFER_STORE_DWORDX4::create(builder, loc, *data, srd, voffset); + BUFFER_STORE_DWORDX4::create(builder, loc, storeData, srd, voffset); } } @@ -1245,9 +1304,12 @@ static void emitExecRestore(Value savedExec, TranslationContext &ctx, S_MOV_B64_EXEC::create(builder, loc, savedExec); } -/// Handle vector.maskedload with exec masking. -/// Pre-fills the result register with the passthrough value, then exec-masks -/// the buffer_load so only active lanes overwrite their result. +/// Handle vector.maskedload with optional exec masking. +/// When the mask is available, saves exec and AND-masks it so only active lanes +/// load. When the mask is unavailable (e.g. the mask computation chain uses +/// unsupported vector ops), falls back to an unconditional +/// buffer_load. This is safe because buffer_load returns zero for out-of-bounds +/// accesses, matching the zero passthrough value used by the frontend. LogicalResult handleVectorMaskedLoad(Operation *op, TranslationContext &ctx) { auto maskedLoadOp = cast(op); auto &builder = ctx.getBuilder(); @@ -1260,35 +1322,20 @@ LogicalResult handleVectorMaskedLoad(Operation *op, TranslationContext &ctx) { if (numBytes <= 0) return op->emitError("vector.maskedload with zero-size vector is invalid"); - // Pre-fill result with passthrough value so inactive lanes get the right data - auto passthrough = ctx.getMapper().getMapped(maskedLoadOp.getPassThru()); - Value resultReg; - if (passthrough) { - resultReg = *passthrough; - } else { - auto immZero = ctx.createImmType(0); - resultReg = ConstantOp::create(builder, loc, immZero, 0); - } - - // Get mask VGPR (the mask is broadcast from a scalar i1 comparison) auto mask = ctx.getMapper().getMapped(maskedLoadOp.getMask()); - if (!mask) - return op->emitError("mask operand not mapped"); + Value savedExec; + if (mask) + savedExec = emitExecMask(*mask, ctx, builder, loc); - // Save exec and apply mask - Value savedExec = emitExecMask(*mask, ctx, builder, loc); - - // Emit buffer_load (only active lanes execute) auto [voffset, instOffset] = computeVOffsetFromIndices( memrefType, maskedLoadOp.getIndices(), ctx, loc, maskedLoadOp.getBase()); Value srd = lookupSRD(maskedLoadOp.getBase(), ctx, loc); auto loadResults = emitBufferLoads(srd, voffset, instOffset, numBytes, ctx, loc); - // Restore exec - emitExecRestore(savedExec, ctx, builder, loc); + if (mask) + emitExecRestore(savedExec, ctx, builder, loc); - // Map result if (!loadResults.empty()) { ctx.getMapper().mapValue(maskedLoadOp.getResult(), loadResults[0]); if (loadResults.size() > 1) @@ -1323,6 +1370,14 @@ convertF32ToBF16ForStore(Value srcData, int64_t numElems, auto elemType = PVRegType::get(builder.getContext(), baseIdx, 1); return PrecoloredVRegOp::create(builder, loc, elemType, baseIdx, 1); } + if (isAGPRType(srcType)) { + auto aregType = ctx.createARegType(1, 1); + auto aregElem = ExtractOp::create(builder, loc, aregType, srcData, + builder.getI64IntegerAttr(i)); + auto vregType = ctx.createVRegType(1, 1); + return V_ACCVGPR_READ_B32::create(builder, loc, vregType, + aregElem.getResult()); + } auto elemType = ctx.createVRegType(1, 1); return ExtractOp::create(builder, loc, elemType, srcData, builder.getI64IntegerAttr(i)); @@ -1866,6 +1921,7 @@ LogicalResult handleRawBufferStore(Operation *op, TranslationContext &ctx); LogicalResult handleMemRefAtomicRMW(Operation *op, TranslationContext &ctx); LogicalResult handleReadFirstLane(Operation *op, TranslationContext &ctx); LogicalResult handleROCDLSBarrier(Operation *op, TranslationContext &ctx); +LogicalResult handleROCDLSchedBarrier(Operation *op, TranslationContext &ctx); LogicalResult handleROCDLSetPrio(Operation *op, TranslationContext &ctx); LogicalResult handleSWaitcnt(Operation *op, TranslationContext &ctx); @@ -2029,6 +2085,7 @@ void OpHandlerRegistry::registerDefaultHandlers(mlir::MLIRContext *ctx) { // ROCDL dialect REGISTER_HANDLER(ROCDL::ReadfirstlaneOp, handleReadFirstLane); REGISTER_HANDLER(ROCDL::SBarrierOp, handleROCDLSBarrier); + REGISTER_HANDLER(ROCDL::SchedBarrier, handleROCDLSchedBarrier); REGISTER_HANDLER(ROCDL::SetPrioOp, handleROCDLSetPrio); REGISTER_HANDLER(ROCDL::SWaitcntOp, handleSWaitcnt); @@ -2152,7 +2209,6 @@ LogicalResult translateModule(ModuleOp module, StringRef targetId) { program->setAttr( "num_kernel_args", builder.getI64IntegerAttr(static_cast(numKernelArgs))); - // Set the LDS size attribute if any LDS was allocated int64_t ldsSize = ctx.getTotalLDSSize(); if (ldsSize > 0) { @@ -2252,7 +2308,6 @@ LogicalResult translateModule(ModuleOp module, StringRef targetId) { program->setAttr( "num_kernel_args", builder.getI64IntegerAttr(static_cast(numKernelArgs))); - // Set the LDS size attribute if any LDS was allocated int64_t ldsSize = ctx.getTotalLDSSize(); if (ldsSize > 0) { @@ -2437,7 +2492,6 @@ LogicalResult translateModule(ModuleOp module, program->setAttr( "num_kernel_args", builder.getI64IntegerAttr(static_cast(numKernelArgs))); - // Set LDS size if used int64_t ldsSize = transCtx.getTotalLDSSize(); if (ldsSize > 0) { diff --git a/waveasm/lib/Transforms/handlers/AMDGPUHandlers.cpp b/waveasm/lib/Transforms/handlers/AMDGPUHandlers.cpp index 6d99671326..4286d81c8e 100644 --- a/waveasm/lib/Transforms/handlers/AMDGPUHandlers.cpp +++ b/waveasm/lib/Transforms/handlers/AMDGPUHandlers.cpp @@ -83,6 +83,17 @@ LogicalResult handleROCDLSetPrio(Operation *op, TranslationContext &ctx) { return success(); } +/// Handle rocdl.sched.barrier — compiler scheduling hint, not a HW instruction. +/// The WaveASM backend does its own scheduling, so we emit a comment for +/// traceability rather than a real instruction (the LLVM assembler would +/// reject s_sched_barrier in raw assembly). +LogicalResult handleROCDLSchedBarrier(Operation *op, TranslationContext &ctx) { + auto schedBarrierOp = cast(op); + int32_t mask = schedBarrierOp.getMask(); + ctx.emitComment("sched_barrier mask=" + std::to_string(mask)); + return success(); +} + LogicalResult handleAMDGPUMfma(Operation *op, TranslationContext &ctx) { auto mfmaOp = cast(op); auto &builder = ctx.getBuilder(); @@ -405,6 +416,41 @@ static void patchSrdWord2InPlace(OpBuilder &builder, Location loc, return; } + // Detect arith.select(arith.cmpi(scalar, scalar), scalar, scalar) + // and emit s_cmp + s_cselect directly into SRD word 2. + if (auto selectOp = validBytesVal.getDefiningOp()) { + Value condVal = selectOp.getCondition(); + if (auto cmpOp = condVal.getDefiningOp()) { + Value cmpLhs = cmpOp.getLhs(); + Value cmpRhs = cmpOp.getRhs(); + Value trueVal = selectOp.getTrueValue(); + Value falseVal = selectOp.getFalseValue(); + + auto cmpLhsMapped = ctx.getMapper().getMapped(cmpLhs); + auto cmpRhsMapped = ctx.getMapper().getMapped(cmpRhs); + auto trueMapped = ctx.getMapper().getMapped(trueVal); + auto falseMapped = ctx.getMapper().getMapped(falseVal); + + if (cmpLhsMapped && cmpRhsMapped && trueMapped && falseMapped && + isScalarOrImm(*cmpLhsMapped) && isScalarOrImm(*cmpRhsMapped) && + isScalarOrImm(*trueMapped) && isScalarOrImm(*falseMapped)) { + Value sccVal = emitScalarCmp(builder, loc, cmpOp.getPredicate(), + *cmpLhsMapped, *cmpRhsMapped, ctx); + + auto dstType = PSRegType::get(builder.getContext(), srdBase + 2, 1); + Value trueV = *trueMapped; + Value falseV = *falseMapped; + auto sregType = ctx.createSRegType(); + if (isImmType(trueV.getType())) + trueV = S_MOV_B32::create(builder, loc, sregType, trueV); + auto result = + S_CSELECT_B32::create(builder, loc, dstType, sccVal, trueV, falseV); + DCEProtectOp::create(builder, loc, result); + return; + } + } + } + auto mapped = ctx.getMapper().getMapped(validBytesVal); if (mapped) { Value src = *mapped; @@ -412,7 +458,7 @@ static void patchSrdWord2InPlace(OpBuilder &builder, Location loc, auto result = V_READFIRSTLANE_B32::create(builder, loc, dstType, src); DCEProtectOp::create(builder, loc, result); } else { - auto sccType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); auto zeroImm = ctx.createImmType(0); auto zeroConst = ConstantOp::create(builder, loc, zeroImm, 0); S_ADD_U32::create(builder, loc, dstType, sccType, src, zeroConst); @@ -480,7 +526,7 @@ LogicalResult handleFatRawBufferCast(Operation *op, TranslationContext &ctx) { // When adj exists AND the cast doesn't need non-max validBytes, emit // the SRD adjustment eagerly. The fat_raw_buffer_cast is typically // outside any scf.for loop while the loads are inside; deferring to - // load-time would place the SRD init RawOps inside the loop body. + // load-time would place the SRD init ops inside the loop body. if (adj && !hasNonMaxValidBytes) { Value srd = emitSRDBaseAdjustment(*adj, op->getResult(0), ctx, loc); ctx.getMapper().mapValue(op->getResult(0), srd); diff --git a/waveasm/lib/Transforms/handlers/AffineHandlers.cpp b/waveasm/lib/Transforms/handlers/AffineHandlers.cpp index 83caf692fe..5f61c04dc7 100644 --- a/waveasm/lib/Transforms/handlers/AffineHandlers.cpp +++ b/waveasm/lib/Transforms/handlers/AffineHandlers.cpp @@ -254,13 +254,16 @@ static AffineExpr normalizeExpr(AffineExpr e, MLIRContext *ctx) { } } -// 32-bit unsigned Barrett reduction: floordiv(x, d) exact for all uint32. -// Matches LLVM's __udivsi3 lowering: -// 1. Float rcp seed → Newton-Raphson refinement → exact reciprocal -// 2. q = mulhi(x, recip) -// 3. Two remainder-based correction steps -Value emitUnsignedFloordiv(Value x, Value d, OpBuilder &builder, Location loc, - TranslationContext &ctx) { +// Compute the Barrett reciprocal of divisor d (Steps 1-2 of the Barrett +// reduction). The result is cached by divisor SSA value so that multiple +// divisions by the same d reuse the same reciprocal (~8 VALU saved per hit). +static Value getOrComputeBarrettReciprocal(Value d, OpBuilder &builder, + Location loc, + TranslationContext &ctx) { + // Check the expression cache first. + if (auto cached = ctx.getCachedExpr("barrett_recip", {d})) + return *cached; + auto vregType = ctx.createVRegType(); Value zeroConst = createImmConst(0, builder, loc, ctx); @@ -280,6 +283,29 @@ Value emitUnsignedFloordiv(Value x, Value d, OpBuilder &builder, Location loc, Value err = V_MUL_HI_U32::create(builder, loc, vregType, recip, e); recip = V_ADD_U32::create(builder, loc, vregType, recip, err); + ctx.cacheExpr("barrett_recip", {d}, {}, recip); + return recip; +} + +// 32-bit unsigned Barrett reduction: floordiv(x, d) exact for all uint32. +// Matches LLVM's __udivsi3 lowering: +// 1. Float rcp seed → Newton-Raphson refinement → exact reciprocal +// 2. q = mulhi(x, recip) +// 3. Two remainder-based correction steps +// When both inputs are scalar, extracts the result back to SGPR via +// v_readfirstlane_b32 so downstream arithmetic stays in SALU. +Value emitUnsignedFloordiv(Value x, Value d, OpBuilder &builder, Location loc, + TranslationContext &ctx) { + bool bothScalar = isScalarOrImm(x) && isScalarOrImm(d); + auto vregType = ctx.createVRegType(); + x = ensureVGPR(builder, x.getLoc(), ctx, x); + d = ensureVGPR(builder, d.getLoc(), ctx, d); + + Value zeroConst = createImmConst(0, builder, loc, ctx); + + // Steps 1-2 (reciprocal + Newton-Raphson) cached by divisor. + Value recip = getOrComputeBarrettReciprocal(d, builder, loc, ctx); + // --- Step 3: quotient --- Value q = V_MUL_HI_U32::create(builder, loc, vregType, x, recip); @@ -305,6 +331,11 @@ Value emitUnsignedFloordiv(Value x, Value d, OpBuilder &builder, Location loc, zeroConst); q = V_ADD_U32::create(builder, loc, vregType, q, inc2); + if (bothScalar) { + auto sregType = ctx.createSRegType(); + q = V_READFIRSTLANE_B32::create(builder, loc, sregType, q); + } + return q; } @@ -330,7 +361,6 @@ Value emitConstantUnsignedFloordiv(Value x, int64_t divisor, OpBuilder &builder, assert(static_cast(divisor) <= 0xFFFFFFFFULL && "divisor exceeds 32 bits -- should have been normalized"); - auto vregType = ctx.createVRegType(); llvm::APInt divisorAPInt(32, static_cast(divisor)); auto mag = llvm::UnsignedDivisionByConstantInfo::get( divisorAPInt, /*LeadingZeros=*/0, @@ -345,23 +375,21 @@ Value emitConstantUnsignedFloordiv(Value x, int64_t divisor, OpBuilder &builder, int64_t magicVal = static_cast(mag.Magic.getZExtValue()); Value magicConst = createImmConst(magicVal, builder, loc, ctx); - Value q = V_MUL_HI_U32::create(builder, loc, vregType, x, magicConst); + Value q = emitMulHi(x, magicConst, builder, loc, ctx); auto emitShiftRight = [&](Value val, unsigned amount) -> Value { if (amount == 0) return val; Value shiftConst = createImmConst(static_cast(amount), builder, loc, ctx); - return V_LSHRREV_B32::create(builder, loc, vregType, shiftConst, val); + return emitLshr(val, shiftConst, builder, loc, ctx); }; if (mag.IsAdd) { - // add form: (mulhi(x,m) + ((x - mulhi(x,m)) >> 1)) >> PostShift - Value xSubQ = V_SUB_U32::create(builder, loc, vregType, x, q); + Value xSubQ = emitSub(x, q, builder, loc, ctx); Value oneConst = createImmConst(1, builder, loc, ctx); - Value halfDiff = - V_LSHRREV_B32::create(builder, loc, vregType, oneConst, xSubQ); - Value sum = V_ADD_U32::create(builder, loc, vregType, q, halfDiff); + Value halfDiff = emitLshr(xSubQ, oneConst, builder, loc, ctx); + Value sum = emitAdd(q, halfDiff, builder, loc, ctx); return emitShiftRight(sum, mag.PostShift); } @@ -382,7 +410,25 @@ Value emitConstantUnsignedFloordiv(Value x, int64_t divisor, OpBuilder &builder, static Value emitCeilFromFloorQuotient(Value q, Value x, Value d, OpBuilder &builder, Location loc, TranslationContext &ctx) { + bool allScalar = isScalarOrImm(q) && isScalarOrImm(x) && isScalarOrImm(d); + if (allScalar) { + // Fully scalar path: rem = x - q*d; SCC = (rem != 0); q += SCC + Value qd = emitMul(q, d, builder, loc, ctx); + Value rem = emitSub(x, qd, builder, loc, ctx); + Value zeroConst = createImmConst(0, builder, loc, ctx); + // s_cmp_lg_u32 sets SCC = (rem != 0) + Value sccVal = + S_CMP_NE_U32::create(builder, loc, ctx.createSCCType(), rem, zeroConst); + // s_addc_u32: dst = q + 0 + SCC (carry-in from SCC) + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + return S_ADDC_U32::create(builder, loc, sregType, sccType, sccVal, q, zeroConst) + .getDst(); + } auto vregType = ctx.createVRegType(); + q = ensureVGPR(builder, loc, ctx, q); + x = ensureVGPR(builder, loc, ctx, x); + d = ensureVGPR(builder, loc, ctx, d); Value qd = V_MUL_LO_U32::create(builder, loc, vregType, q, d); Value rem = V_SUB_U32::create(builder, loc, vregType, x, qd); Value zeroConst = createImmConst(0, builder, loc, ctx); @@ -553,6 +599,12 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { BitRange lhsRange = lhsResult.range; BitRange rhsRange = rhsResult.range; + // Helper: coerce both operands to VGPR for VALU binary ops. + auto ensureBothVGPR = [&]() { + lhs = ensureVGPR(builder, loc, ctx, lhs); + rhs = ensureVGPR(builder, loc, ctx, rhs); + }; + switch (binExpr.getKind()) { case AffineExprKind::Add: { if (!lhsRange.overlaps(rhsRange)) { @@ -577,8 +629,11 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { auto shiftConst = ConstantOp::create(builder, loc, shiftImm, shiftAmount); // v_lshl_or_b32: dst = (src << shift) | orend + Value base = ensureVGPR(builder, loc, ctx, + baseResult.value); + orend = ensureVGPR(builder, loc, ctx, orend); Value fusedResult = V_LSHL_OR_B32::create( - builder, loc, vregType, baseResult.value, shiftConst, + builder, loc, vregType, base, shiftConst, orend); BitRange shiftedRange = baseResult.range.shiftLeft(shiftAmount); @@ -597,8 +652,11 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { auto shiftImm = ctx.createImmType(shiftAmount); auto shiftConst = ConstantOp::create(builder, loc, shiftImm, shiftAmount); + Value base2 = ensureVGPR(builder, loc, ctx, + baseResult.value); + orend = ensureVGPR(builder, loc, ctx, orend); Value fusedResult = V_LSHL_OR_B32::create( - builder, loc, vregType, baseResult.value, shiftConst, + builder, loc, vregType, base2, shiftConst, orend); BitRange shiftedRange = baseResult.range.shiftLeft(shiftAmount); @@ -621,14 +679,14 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { return *result; } - // No fusion possible, emit regular v_or_b32 - Value orResult = V_OR_B32::create(builder, loc, vregType, lhs, rhs); + // No fusion possible — emitOr picks S_OR_B32 when both scalar. + Value orResult = emitOr(lhs, rhs, builder, loc, ctx); BitRange resultRange = lhsRange.merge(rhsRange); ctx.setBitRange(orResult, resultRange); return ExprResult(orResult, resultRange); } // Overlapping ranges - must use ADD - Value addResult = V_ADD_U32::create(builder, loc, vregType, lhs, rhs); + Value addResult = emitAdd(lhs, rhs, builder, loc, ctx); BitRange resultRange = lhsRange.extendForAdd(rhsRange); ctx.setBitRange(addResult, resultRange); return ExprResult(addResult, resultRange); @@ -658,10 +716,8 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { int64_t mask = ~(N - 1) & 0xFFFFFFFF; auto maskImm = ctx.createImmType(mask); auto maskConst = ConstantOp::create(builder, loc, maskImm, mask); - // NOTE: constant must be src0 (first operand) for VOP2 encoding. - // src1 must be a VGPR on AMDGCN. - Value andResult = V_AND_B32::create(builder, loc, vregType, maskConst, - innerResult.value); + Value andResult = + emitAnd(maskConst, innerResult.value, builder, loc, ctx); // Result has same bit range as inner, but low bits cleared BitRange resultRange = innerResult.range; // Clear bits below log2(N) -- conservative: use inner range @@ -696,14 +752,12 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { auto shiftConst = ConstantOp::create(builder, loc, shiftAmt, shiftAmount); Value shiftResult = - V_LSHLREV_B32::create(builder, loc, vregType, shiftConst, lhs); - // Shift the bit range left by shiftAmount + emitLshl(lhs, shiftConst, builder, loc, ctx); BitRange resultRange = lhsRange.shiftLeft(shiftAmount); ctx.setBitRange(shiftResult, resultRange); return ExprResult(shiftResult, resultRange); } } - // Also check LHS for power of 2 multiply if (auto constLhs = dyn_cast(binExpr.getLHS())) { int64_t val = constLhs.getValue(); if (isPowerOf2(val)) { @@ -712,14 +766,13 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { auto shiftConst = ConstantOp::create(builder, loc, shiftAmt, shiftAmount); Value shiftResult = - V_LSHLREV_B32::create(builder, loc, vregType, shiftConst, rhs); + emitLshl(rhs, shiftConst, builder, loc, ctx); BitRange resultRange = rhsRange.shiftLeft(shiftAmount); ctx.setBitRange(shiftResult, resultRange); return ExprResult(shiftResult, resultRange); } } - Value mulResult = - V_MUL_LO_U32::create(builder, loc, vregType, lhs, rhs); + Value mulResult = emitMul(lhs, rhs, builder, loc, ctx); return ExprResult(mulResult, BitRange()); // Conservative: full range } @@ -747,7 +800,7 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { auto shiftConst = ConstantOp::create(builder, loc, shiftAmt, shiftAmount); Value shiftResult = - V_LSHRREV_B32::create(builder, loc, vregType, shiftConst, lhs); + emitLshr(lhs, shiftConst, builder, loc, ctx); BitRange resultRange = lhsRange.shiftRight(shiftAmount); ctx.setBitRange(shiftResult, resultRange); return ExprResult(shiftResult, resultRange); @@ -776,13 +829,22 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { if (isPowerOf2(divisor)) { int64_t shiftAmount = log2(divisor); Value shiftConst = createImmConst(shiftAmount, builder, loc, ctx); - Value q = - V_LSHRREV_B32::create(builder, loc, vregType, shiftConst, lhs); + Value q = emitLshr(lhs, shiftConst, builder, loc, ctx); int64_t mask = divisor - 1; Value maskConst = createImmConst(mask, builder, loc, ctx); - Value rem = - V_AND_B32::create(builder, loc, vregType, maskConst, lhs); + Value rem = emitAnd(lhs, maskConst, builder, loc, ctx); Value zeroConst = createImmConst(0, builder, loc, ctx); + if (isScalarOrImm(rem)) { + Value sccVal = S_CMP_NE_U32::create(builder, loc, + ctx.createSCCType(), rem, zeroConst); + auto sregType = ctx.createSRegType(); + auto sccType = ctx.createSCCType(); + Value result = + S_ADDC_U32::create(builder, loc, sregType, sccType, sccVal, q, + zeroConst) + .getDst(); + return ExprResult(result, BitRange()); + } V_CMP_NE_U32::create(builder, loc, rem, zeroConst); Value oneConst = createImmConst(1, builder, loc, ctx); Value oneVgpr = V_MOV_B32::create(builder, loc, vregType, oneConst); @@ -816,8 +878,7 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { int64_t val = constRhs.getValue(); if (isPowerOf2(val)) { Value maskConst = createImmConst(val - 1, builder, loc, ctx); - Value andResult = - V_AND_B32::create(builder, loc, vregType, lhs, maskConst); + Value andResult = emitAnd(lhs, maskConst, builder, loc, ctx); BitRange resultRange = BitRange(0, log2(val) - 1); ctx.setBitRange(andResult, resultRange); return ExprResult(andResult, resultRange); @@ -827,16 +888,16 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { if (val >= 2) { Value q = emitConstantUnsignedFloordiv(lhs, val, builder, loc, ctx); Value dConst = createImmConst(val, builder, loc, ctx); - Value qd = V_MUL_LO_U32::create(builder, loc, vregType, q, dConst); - Value rem = V_SUB_U32::create(builder, loc, vregType, lhs, qd); + Value qd = emitMul(q, dConst, builder, loc, ctx); + Value rem = emitSub(lhs, qd, builder, loc, ctx); return ExprResult(rem, BitRange()); } } // Symbolic divisor fallback: x mod d = x - floordiv(x, d) * d { Value q = emitUnsignedFloordiv(lhs, rhs, builder, loc, ctx); - Value qd = V_MUL_LO_U32::create(builder, loc, vregType, q, rhs); - Value rem = V_SUB_U32::create(builder, loc, vregType, lhs, qd); + Value qd = emitMul(q, rhs, builder, loc, ctx); + Value rem = emitSub(lhs, qd, builder, loc, ctx); return ExprResult(rem, BitRange()); } } diff --git a/waveasm/lib/Transforms/handlers/ArithHandlers.cpp b/waveasm/lib/Transforms/handlers/ArithHandlers.cpp index 825b659e8f..0106ca98dc 100644 --- a/waveasm/lib/Transforms/handlers/ArithHandlers.cpp +++ b/waveasm/lib/Transforms/handlers/ArithHandlers.cpp @@ -108,43 +108,135 @@ LogicalResult handleArithConstant(Operation *op, TranslationContext &ctx) { //===----------------------------------------------------------------------===// LogicalResult handleArithAddI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALU(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + + auto result = emitAdd(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithSubI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALU(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + + auto result = emitSub(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithMulI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALU(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + + auto result = emitMul(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithMinSI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALU(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + auto result = emitMinI32(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithMaxSI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALU(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + auto result = emitMaxI32(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithMinUI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALU(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + auto result = emitMinU32(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithMaxUI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALU(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + auto result = emitMaxU32(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithAndI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALU(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + + auto result = emitAnd(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithOrI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALU(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + + auto result = emitOr(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithXorI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALU(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + + auto result = emitXor(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } //===----------------------------------------------------------------------===// @@ -152,11 +244,31 @@ LogicalResult handleArithXorI(Operation *op, TranslationContext &ctx) { //===----------------------------------------------------------------------===// LogicalResult handleArithShLI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALURev(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + + auto result = emitLshl(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithShRUI(Operation *op, TranslationContext &ctx) { - return handleBinaryVALURev(op, ctx); + auto typedOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + + std::optional lhs, rhs; + if (failed(validateBinaryOperands(typedOp, ctx, lhs, rhs))) + return failure(); + + auto result = emitLshr(*lhs, *rhs, builder, loc, ctx); + ctx.getMapper().mapValue(typedOp.getResult(), result); + return success(); } LogicalResult handleArithShRSI(Operation *op, TranslationContext &ctx) { @@ -197,11 +309,9 @@ LogicalResult handleArithDivUI(Operation *op, TranslationContext &ctx) { if (auto constOp = rhs->getDefiningOp()) { int64_t divisor = constOp.getValue(); if (isPowerOf2(divisor)) { - auto vregType = ctx.createVRegType(); int64_t shiftAmt = log2(divisor); Value shiftConst = createImmConst(shiftAmt, builder, loc, ctx); - auto result = - V_LSHRREV_B32::create(builder, loc, vregType, shiftConst, *lhs); + auto result = emitLshr(*lhs, shiftConst, builder, loc, ctx); ctx.getMapper().mapValue(divOp.getResult(), result); return success(); } @@ -235,7 +345,7 @@ LogicalResult handleArithRemUI(Operation *op, TranslationContext &ctx) { int64_t modulus = constOp.getValue(); if (isPowerOf2(modulus)) { Value maskConst = createImmConst(modulus - 1, builder, loc, ctx); - auto result = V_AND_B32::create(builder, loc, vregType, *lhs, maskConst); + auto result = emitAnd(*lhs, maskConst, builder, loc, ctx); ctx.getMapper().mapValue(remOp.getResult(), result); return success(); } @@ -348,83 +458,55 @@ LogicalResult handleArithCmpI(Operation *op, TranslationContext &ctx) { if (lhsScalar && rhsScalar) { SRegType sregType = ctx.createSRegType(); - SCCType sccType = ctx.createSCCType(); - // S_CMP requires SGPR operands; move immediates to SGPRs first. Value lhsOp = *lhs; Value rhsOp = *rhs; if (isImmType(lhsOp.getType())) lhsOp = S_MOV_B32::create(builder, loc, sregType, lhsOp); if (isImmType(rhsOp.getType())) rhsOp = S_MOV_B32::create(builder, loc, sregType, rhsOp); - Value result; - switch (cmpOp.getPredicate()) { - case arith::CmpIPredicate::eq: - result = S_CMP_EQ_U32::create(builder, loc, sccType, lhsOp, rhsOp); - break; - case arith::CmpIPredicate::ne: - result = S_CMP_NE_U32::create(builder, loc, sccType, lhsOp, rhsOp); - break; - case arith::CmpIPredicate::slt: - result = S_CMP_LT_I32::create(builder, loc, sccType, lhsOp, rhsOp); - break; - case arith::CmpIPredicate::sle: - result = S_CMP_LE_I32::create(builder, loc, sccType, lhsOp, rhsOp); - break; - case arith::CmpIPredicate::sgt: - result = S_CMP_GT_I32::create(builder, loc, sccType, lhsOp, rhsOp); - break; - case arith::CmpIPredicate::sge: - result = S_CMP_GE_I32::create(builder, loc, sccType, lhsOp, rhsOp); - break; - case arith::CmpIPredicate::ult: - result = S_CMP_LT_U32::create(builder, loc, sccType, lhsOp, rhsOp); - break; - case arith::CmpIPredicate::ule: - result = S_CMP_LE_U32::create(builder, loc, sccType, lhsOp, rhsOp); - break; - case arith::CmpIPredicate::ugt: - result = S_CMP_GT_U32::create(builder, loc, sccType, lhsOp, rhsOp); - break; - case arith::CmpIPredicate::uge: - result = S_CMP_GE_U32::create(builder, loc, sccType, lhsOp, rhsOp); - break; - } + Value result = + emitScalarCmp(builder, loc, cmpOp.getPredicate(), lhsOp, rhsOp, ctx); ctx.getMapper().mapValue(cmpOp.getResult(), result); return success(); } - // Vector path: at least one operand is a VGPR, so use V_CMP which writes - // to VCC implicitly. + // Vector path: use V_CMP which writes to VCC implicitly. + // V_CMP VOP3 can read one SGPR from the constant bus — only coerce + // to VGPR when BOTH operands are SGPR (two constant bus slots). + Value lhsV = *lhs; + Value rhsV = *rhs; + if (isSGPRType(lhsV.getType()) && isSGPRType(rhsV.getType())) + lhsV = ensureVGPR(builder, loc, ctx, lhsV); switch (cmpOp.getPredicate()) { case arith::CmpIPredicate::eq: - V_CMP_EQ_U32::create(builder, loc, *lhs, *rhs); + V_CMP_EQ_U32::create(builder, loc, lhsV, rhsV); break; case arith::CmpIPredicate::ne: - V_CMP_NE_U32::create(builder, loc, *lhs, *rhs); + V_CMP_NE_U32::create(builder, loc, lhsV, rhsV); break; case arith::CmpIPredicate::slt: - V_CMP_LT_I32::create(builder, loc, *lhs, *rhs); + V_CMP_LT_I32::create(builder, loc, lhsV, rhsV); break; case arith::CmpIPredicate::sle: - V_CMP_LE_I32::create(builder, loc, *lhs, *rhs); + V_CMP_LE_I32::create(builder, loc, lhsV, rhsV); break; case arith::CmpIPredicate::sgt: - V_CMP_GT_I32::create(builder, loc, *lhs, *rhs); + V_CMP_GT_I32::create(builder, loc, lhsV, rhsV); break; case arith::CmpIPredicate::sge: - V_CMP_GE_I32::create(builder, loc, *lhs, *rhs); + V_CMP_GE_I32::create(builder, loc, lhsV, rhsV); break; case arith::CmpIPredicate::ult: - V_CMP_LT_U32::create(builder, loc, *lhs, *rhs); + V_CMP_LT_U32::create(builder, loc, lhsV, rhsV); break; case arith::CmpIPredicate::ule: - V_CMP_LE_U32::create(builder, loc, *lhs, *rhs); + V_CMP_LE_U32::create(builder, loc, lhsV, rhsV); break; case arith::CmpIPredicate::ugt: - V_CMP_GT_U32::create(builder, loc, *lhs, *rhs); + V_CMP_GT_U32::create(builder, loc, lhsV, rhsV); break; case arith::CmpIPredicate::uge: - V_CMP_GE_U32::create(builder, loc, *lhs, *rhs); + V_CMP_GE_U32::create(builder, loc, lhsV, rhsV); break; } @@ -447,12 +529,70 @@ LogicalResult handleArithSelect(Operation *op, TranslationContext &ctx) { return op->emitError("operands not mapped"); } - // Restore the materialized boolean VGPR (0/1) back into VCC + // CmpI fusion: when condition is arith.cmpi with scalar operands and + // both select values are scalar, emit s_cmp + s_cselect directly. + // This bypasses the VGPR boolean from handleArithCmpI entirely. + if (isScalarOrImm(*trueVal) && isScalarOrImm(*falseVal)) { + Value condMLIR = selectOp.getCondition(); + if (auto cmpOp = condMLIR.getDefiningOp()) { + auto cmpLhs = ctx.getMapper().getMapped(cmpOp.getLhs()); + auto cmpRhs = ctx.getMapper().getMapped(cmpOp.getRhs()); + if (cmpLhs && cmpRhs && + isScalarOrImm(*cmpLhs) && isScalarOrImm(*cmpRhs)) { + auto sregType = ctx.createSRegType(); + Value lhsOp = *cmpLhs; + Value rhsOp = *cmpRhs; + if (isImmType(lhsOp.getType())) + lhsOp = S_MOV_B32::create(builder, loc, sregType, lhsOp); + if (isImmType(rhsOp.getType())) + rhsOp = S_MOV_B32::create(builder, loc, sregType, rhsOp); + Value sccVal = + emitScalarCmp(builder, loc, cmpOp.getPredicate(), lhsOp, rhsOp, ctx); + Value trueV = *trueVal; + Value falseV = *falseVal; + if (isImmType(trueV.getType())) + trueV = S_MOV_B32::create(builder, loc, sregType, trueV); + auto result = + S_CSELECT_B32::create(builder, loc, sregType, sccVal, trueV, falseV); + ctx.getMapper().mapValue(selectOp.getResult(), result); + return success(); + } + } + } + + // Scalar path: when condition and both values are scalar, use + // s_cmp_lg_u32 + s_cselect_b32 (no VGPR needed). + if (isScalarOrImm(*cond) && isScalarOrImm(*trueVal) && + isScalarOrImm(*falseVal)) { + Value zeroConst = createImmConst(0, builder, loc, ctx); + Value condV = *cond; + if (isImmType(condV.getType())) + condV = S_MOV_B32::create(builder, loc, ctx.createSRegType(), condV); + Value sccVal = + S_CMP_NE_U32::create(builder, loc, ctx.createSCCType(), condV, zeroConst); + auto sregType = ctx.createSRegType(); + Value trueV = *trueVal; + Value falseV = *falseVal; + if (isImmType(trueV.getType())) + trueV = S_MOV_B32::create(builder, loc, sregType, trueV); + auto result = S_CSELECT_B32::create(builder, loc, sregType, sccVal, trueV, falseV); + ctx.getMapper().mapValue(selectOp.getResult(), result); + return success(); + } + + // Vector path: restore the materialized boolean VGPR (0/1) back into VCC. + // V_CMP can take one SGPR via the constant bus; zeroConst is an immediate + // (no bus slot), so an SGPR cond is fine without coercion. Value zeroConst = createImmConst(0, builder, loc, ctx); - V_CMP_NE_U32::create(builder, loc, *cond, zeroConst); + Value condV = *cond; + V_CMP_NE_U32::create(builder, loc, condV, zeroConst); + // v_cndmask_b32: coerce both to VGPR. VOP3 constant bus allows one SGPR + // but interacts with vcc slot; keeping both as VGPR is safest. + Value trueVgpr = ensureVGPR(builder, loc, ctx, *trueVal); + Value falseVgpr = ensureVGPR(builder, loc, ctx, *falseVal); auto result = - V_CNDMASK_B32::create(builder, loc, vregType, *falseVal, *trueVal, *cond); + V_CNDMASK_B32::create(builder, loc, vregType, falseVgpr, trueVgpr, *cond); ctx.getMapper().mapValue(selectOp.getResult(), result); return success(); } diff --git a/waveasm/lib/Transforms/handlers/HandlerUtils.cpp b/waveasm/lib/Transforms/handlers/HandlerUtils.cpp index 7ab27e695b..69de5a73e4 100644 --- a/waveasm/lib/Transforms/handlers/HandlerUtils.cpp +++ b/waveasm/lib/Transforms/handlers/HandlerUtils.cpp @@ -62,14 +62,33 @@ int64_t getElementBytes(Type type) { //===----------------------------------------------------------------------===// int64_t computeBufferSizeFromMemRef(MemRefType memrefType) { - // Use (1 << 31) - 2 = 0x7FFFFFFE as num_records. The Wave Python frontend - // emits OOB sentinel index values at (valid_bytes + elem_bytes) / elem_bytes - // which lands at byte offset 0x7FFFFFFF. With num_records = 0x7FFFFFFE the - // sentinel is one byte past the SRD range, so hardware returns 0 for OOB - // lanes. Using 0xFFFFFFFF would make the sentinel "in bounds" and cause a - // real access to unmapped memory, triggering HSA page faults. - (void)memrefType; - return 0x7FFFFFFE; + // Compute the maximum byte offset reachable via this memref's layout. + // For a memref with strides [s0, s1, …, 1] and dims [d0, d1, …, dN]: + // max_offset = (d0-1)*s0 + (d1-1)*s1 + … + (dN-1)*1 + 1 (in elements) + // num_records = max_offset * elemBytes + // If any dimension or stride is dynamic, fall back to a large bound. + int64_t elemBytes = getElementBytes(memrefType.getElementType()); + + SmallVector strides; + int64_t offset; + if (succeeded(memrefType.getStridesAndOffset(strides, offset)) && + memrefType.hasStaticShape()) { + bool allStaticStrides = llvm::all_of( + strides, [](int64_t s) { return s != ShapedType::kDynamic; }); + if (allStaticStrides) { + int64_t maxOffset = 0; + auto shape = memrefType.getShape(); + for (size_t i = 0; i < shape.size(); ++i) { + if (shape[i] > 1) + maxOffset += (shape[i] - 1) * std::abs(strides[i]); + } + return (maxOffset + 1) * elemBytes; + } + } + // Dynamic dims or dynamic strides — use a large-but-bounded value so that + // hardware OOB checking catches wild offsets (e.g., negative wraparound) + // and returns zero rather than faulting on unmapped memory. + return 0x20000000; } //===----------------------------------------------------------------------===// diff --git a/waveasm/lib/Transforms/handlers/MemRefHandlers.cpp b/waveasm/lib/Transforms/handlers/MemRefHandlers.cpp index a4fef07220..b73c7c5c93 100644 --- a/waveasm/lib/Transforms/handlers/MemRefHandlers.cpp +++ b/waveasm/lib/Transforms/handlers/MemRefHandlers.cpp @@ -16,6 +16,7 @@ // - memref.load - emit ds_read or buffer_load // - memref.store - emit ds_write or buffer_store // - memref.cast - pass through source mapping +// - memref.extract_strided_metadata - expose base, offset, sizes, strides // //===----------------------------------------------------------------------===// @@ -240,26 +241,23 @@ LogicalResult handleMemRefLoad(Operation *op, TranslationContext &ctx) { auto readOp = DS_READ_B32::create(builder, loc, TypeRange{vregType}, vaddr); ctx.getMapper().mapValue(loadOp.getResult(), readOp.getResult(0)); } else { - // Global load - auto sregType = ctx.createSRegType(4, 4); - auto srd = PrecoloredSRegOp::create(builder, loc, sregType, 8, 4); - - Value voffset; - if (!loadOp.getIndices().empty()) { - if (auto mapped = ctx.getMapper().getMapped(loadOp.getIndices()[0])) { - voffset = *mapped; - } - } - if (!voffset) { - auto immType = ctx.createImmType(0); - voffset = ConstantOp::create(builder, loc, immType, 0); + // Global load — use the same SRD lookup and voffset computation as + // vector.load so that fat_raw_buffer memrefs get the correct adjusted SRD. + auto [voffset, instOffset] = computeVOffsetFromIndices( + memrefType, loadOp.getIndices(), ctx, loc, loadOp.getMemref()); + + Value srd; + if (auto *adj = ctx.getPendingSRDBaseAdjust(loadOp.getMemref())) { + srd = emitSRDBaseAdjustment(*adj, loadOp.getMemref(), ctx, loc); + } else { + srd = lookupSRD(loadOp.getMemref(), ctx, loc); } - auto zeroImm = builder.getType(0); - auto zeroConst = ConstantOp::create(builder, loc, zeroImm, 0); - auto loadInstr = BUFFER_LOAD_DWORD::create( - builder, loc, TypeRange{vregType}, srd, voffset, zeroConst); - ctx.getMapper().mapValue(loadOp.getResult(), loadInstr.getResult(0)); + int64_t elemBytes = getElementBytes(memrefType.getElementType()); + auto loadResults = + emitBufferLoads(srd, voffset, instOffset, elemBytes, ctx, loc); + if (!loadResults.empty()) + ctx.getMapper().mapValue(loadOp.getResult(), loadResults[0]); } return success(); @@ -293,48 +291,26 @@ LogicalResult handleMemRefStore(Operation *op, TranslationContext &ctx) { DS_WRITE_B32::create(builder, loc, *data, vaddr); } else { - // Global store - auto sregType = ctx.createSRegType(4, 4); - auto srd = PrecoloredSRegOp::create(builder, loc, sregType, 8, 4); + // Global store — use proper SRD lookup and voffset computation. + auto [voffset, instOffset] = computeVOffsetFromIndices( + memrefType, storeOp.getIndices(), ctx, loc, storeOp.getMemref()); - Value voffset; - if (!storeOp.getIndices().empty()) { - if (auto mapped = ctx.getMapper().getMapped(storeOp.getIndices()[0])) { - voffset = *mapped; - } - } - if (!voffset) { - auto immType = ctx.createImmType(0); - voffset = ConstantOp::create(builder, loc, immType, 0); + Value srd; + if (auto *adj = ctx.getPendingSRDBaseAdjust(storeOp.getMemref())) { + srd = emitSRDBaseAdjustment(*adj, storeOp.getMemref(), ctx, loc); + } else { + srd = lookupSRD(storeOp.getMemref(), ctx, loc); } - BUFFER_STORE_DWORD::create(builder, loc, *data, srd, voffset); - } - - return success(); -} - -/// Handle memref.extract_strided_metadata - extract offset from SRD adjustment. -/// Results: (base_buffer, offset, sizes..., strides...). -/// The offset is needed by eliminate_epilogue's _compute_valid_bytes to clamp -/// the SRD NUM_RECORDS field so OOB loads in extended iterations return zero. -LogicalResult handleMemRefExtractStridedMetadata(Operation *op, - TranslationContext &ctx) { - auto metaOp = cast(op); - Value source = metaOp.getSource(); - - auto *adj = ctx.getPendingSRDBaseAdjust(source); - if (!adj) { - if (auto castOp = source.getDefiningOp()) - adj = ctx.getPendingSRDBaseAdjust(castOp.getSource()); - } - - if (auto src = ctx.getMapper().getMapped(source)) - ctx.getMapper().mapValue(metaOp.getBaseBuffer(), *src); + Value storeData = *data; + if (isAGPRType(storeData.getType())) { + auto vregType = ctx.createVRegType(); + storeData = + V_ACCVGPR_READ_B32::create(builder, loc, vregType, storeData); + } - if (adj) { - Value offsetResult = metaOp.getOffset(); - ctx.getMapper().mapValue(offsetResult, adj->elementOffset); + BUFFER_STORE_DWORD::create(builder, loc, storeData, srd, voffset, + instOffset); } return success(); @@ -372,6 +348,83 @@ LogicalResult handleMemRefCast(Operation *op, TranslationContext &ctx) { return success(); } +/// Handle memref.extract_strided_metadata - expose the base buffer, offset, +/// sizes, and strides of a strided memref as individual SSA values. +/// +/// Results: (base_buffer, offset, sizes..., strides...) +/// +/// For the waveasm backend, the critical results are the dynamic strides +/// (tracked via setDynamicStride on the source reinterpret_cast). We map +/// each result to the appropriate value so that downstream arithmetic +/// (arith.muli with the stride, etc.) can resolve its operands. +LogicalResult handleMemRefExtractStridedMetadata(Operation *op, + TranslationContext &ctx) { + auto metadataOp = cast(op); + auto &builder = ctx.getBuilder(); + auto loc = op->getLoc(); + + Value source = metadataOp.getSource(); + auto sourceType = cast(source.getType()); + + // Result #0: base_buffer — map to the same underlying buffer as the source. + if (auto src = ctx.getMapper().getMapped(source)) { + ctx.getMapper().mapValue(metadataOp.getBaseBuffer(), *src); + } + + // Get static strides and offset from the source memref type. + SmallVector staticStrides; + int64_t staticOffset; + bool hasStridesAndOffset = + succeeded(sourceType.getStridesAndOffset(staticStrides, staticOffset)); + + // Result #1: offset — prefer SRD adjustment offset (needed by epilogue + // elimination's _compute_valid_bytes), then fall back to static offset. + auto *adj = ctx.getPendingSRDBaseAdjust(source); + if (!adj) { + if (auto castOp = source.getDefiningOp()) + adj = ctx.getPendingSRDBaseAdjust(castOp.getSource()); + } + if (adj) { + ctx.getMapper().mapValue(metadataOp.getOffset(), adj->elementOffset); + } else if (hasStridesAndOffset && staticOffset != ShapedType::kDynamic) { + auto immType = ctx.createImmType(staticOffset); + auto offsetVal = ConstantOp::create(builder, loc, immType, staticOffset); + ctx.getMapper().mapValue(metadataOp.getOffset(), offsetVal); + } + + // Results #2..#(2+rank-1): sizes — map static sizes to constants, + // dynamic sizes to their mapped source values. + auto sizes = metadataOp.getSizes(); + auto shape = sourceType.getShape(); + for (unsigned i = 0; i < sizes.size(); ++i) { + if (i < shape.size() && shape[i] != ShapedType::kDynamic) { + auto immType = ctx.createImmType(shape[i]); + auto sizeVal = ConstantOp::create(builder, loc, immType, shape[i]); + ctx.getMapper().mapValue(sizes[i], sizeVal); + } + } + + // Results #(2+rank)..#(2+2*rank-1): strides — the critical part. + // For dynamic strides, retrieve the value already tracked by + // handleMemRefReinterpretCast via getDynamicStride. + // For static strides, create immediate constants. + auto strides = metadataOp.getStrides(); + for (unsigned i = 0; i < strides.size(); ++i) { + if (hasStridesAndOffset && i < staticStrides.size()) { + if (staticStrides[i] != ShapedType::kDynamic) { + auto immType = ctx.createImmType(staticStrides[i]); + auto strideVal = + ConstantOp::create(builder, loc, immType, staticStrides[i]); + ctx.getMapper().mapValue(strides[i], strideVal); + } else if (auto dynStride = ctx.getDynamicStride(source, i)) { + ctx.getMapper().mapValue(strides[i], *dynStride); + } + } + } + + return success(); +} + //===----------------------------------------------------------------------===// // MemRefAddressComputer Implementation //===----------------------------------------------------------------------===// From 807ea9049329473cc7afe50008f55ac841922fe0 Mon Sep 17 00:00:00 2001 From: Sanket Pandit Date: Tue, 31 Mar 2026 03:55:33 +0000 Subject: [PATCH 3/5] Improve assembly emitter: inline AGPR constants, v_swap_b32, VOP2 fix Use inline constants directly in v_accvgpr_write_b32 without scratch VGPR materialization. Replace 3-instruction VGPR swap sequence with single v_swap_b32. Support buffer_store without offen for immediate voffset. Add RawOp separate handler. Add min_sgprs attribute support. Fix VOP2 literal set to match ISA encoding requirements. Made-with: Cursor Signed-off-by: Sanket Pandit --- waveasm/lib/Transforms/AssemblyEmitter.cpp | 54 ++++++++++++------- .../lib/Transforms/LiteralMaterialization.cpp | 9 ++-- waveasm/test/Dialect/region-control-flow.mlir | 18 ++++--- waveasm/test/Transforms/floor-ops.mlir | 34 ++++++------ .../test/Transforms/lit-iree-handlers.mlir | 14 ++--- waveasm/test/Transforms/lit-translate.mlir | 20 +++---- waveasm/test/Transforms/liveness-cfg.mlir | 4 +- .../Transforms/region-based-translation.mlir | 48 +++++++++-------- .../test/Transforms/salu-vop2-literals.mlir | 13 +++-- waveasm/test/Transforms/vgpr-swap-emit.mlir | 9 ++-- .../test/Translate/buffer-ops-srd-adjust.mlir | 22 ++++---- waveasm/test/Translate/dynamic-shapes.mlir | 15 +++--- waveasm/test/Translate/scf-for-loop.mlir | 4 +- .../Translate/scf-if-agpr-else-coercion.mlir | 23 ++++---- .../Translate/swizzle-srd-num-records.mlir | 19 +++---- waveasm/test/Translate/vector-constants.mlir | 2 +- 16 files changed, 162 insertions(+), 146 deletions(-) diff --git a/waveasm/lib/Transforms/AssemblyEmitter.cpp b/waveasm/lib/Transforms/AssemblyEmitter.cpp index f22d629192..a53fc0b218 100644 --- a/waveasm/lib/Transforms/AssemblyEmitter.cpp +++ b/waveasm/lib/Transforms/AssemblyEmitter.cpp @@ -169,9 +169,16 @@ std::string KernelGenerator::emitBufferStore(Operation *op, // Use VMEMStoreOpInterface to access operands by name if (auto storeOp = dyn_cast(op)) { std::string vdata = resolveValue(storeOp.getData()); - std::string voffset = resolveValue(storeOp.getVoffset()); + Value voffsetVal = storeOp.getVoffset(); std::string srd = resolveValue(storeOp.getSaddr()); - result += " " + vdata + ", " + voffset + ", " + srd + ", 0 offen"; + + if (isa(voffsetVal.getType())) { + result += " " + vdata + ", off, " + srd + ", 0"; + } else { + std::string voffset = resolveValue(voffsetVal); + result += " " + vdata + ", " + voffset + ", " + srd + ", 0 offen"; + } + if (auto instOffsetAttr = op->getAttrOfType("instOffset")) { int64_t offset = instOffsetAttr.getInt(); if (offset > 0) { @@ -290,11 +297,13 @@ std::string KernelGenerator::emitDefaultFormat(Operation *op, } for (Value result : op->getResults()) { + // Skip SCC-typed results (hardware-implicit, not emitted). if (isa(result.getType())) continue; operands.push_back(resolveValue(result)); } for (Value operand : op->getOperands()) { + // Skip SCC-typed operands (hardware-implicit, not emitted). if (isa(operand.getType())) continue; if (isScalarOp) { @@ -348,8 +357,12 @@ KernelGenerator::emitScaledMFMA(Operation *scaledOp, llvm::StringRef mnemonic) { std::optional KernelGenerator::generateOp(Operation *op) { return llvm::TypeSwitch>(op) - .Case( - [](auto) { return std::nullopt; }) + .Case([](auto) { return std::nullopt; }) + .Case([&](RawOp rawOp) -> std::optional { + return generateRaw(rawOp); + }) .Case([&](S_WAITCNT waitcntOp) { std::optional vmcnt, lgkmcnt, expcnt; @@ -539,16 +552,18 @@ std::optional KernelGenerator::generateOp(Operation *op) { std::string src = resolveValue(srcVal); std::string lines; if (isAGPR) { - // v_accvgpr_write_b32 requires a VGPR source in this backend. - // Materialize immediate sources into the reserved scratch VGPR. + auto [isLit, litVal] = getLiteralValue(srcVal); std::string writeSrc = src; - if (srcIsImm) { + if (srcIsImm && !(isLit && isInlineConstant(litVal))) { + // Non-inline literal: materialize into scratch VGPR first. lines += " v_mov_b32 " + formatVGPRRange(kScratchVGPR, 1) + ", " + src; writeSrc = formatVGPRRange(kScratchVGPR, 1); peakVGPRs = std::max(peakVGPRs, kScratchVGPR + 1); invalidateScratchCache(); } + // Inline constants (e.g. 0) go directly into + // v_accvgpr_write_b32 without a scratch VGPR. for (int64_t i = 0; i < size; ++i) { if (!lines.empty()) lines += "\n"; @@ -569,6 +584,12 @@ std::optional KernelGenerator::generateOp(Operation *op) { if (isAGPR) { if (srcIsImm) { std::string src = resolveValue(srcVal); + auto [isLit, litVal] = getLiteralValue(srcVal); + if (isLit && isInlineConstant(litVal)) { + // Inline constant: emit directly without scratch VGPR. + return " v_accvgpr_write_b32 " + resolveValue(result) + ", " + + src; + } std::string scratch = formatVGPRRange(kScratchVGPR, 1); peakVGPRs = std::max(peakVGPRs, kScratchVGPR + 1); invalidateScratchCache(); @@ -701,10 +722,7 @@ std::optional KernelGenerator::generateOp(Operation *op) { SmallVector handled(pendingCopies.size(), false); - // Allocate swap temps once and reuse across all swaps. - // Swaps are emitted sequentially so the temp is dead after - // each 3-instruction sequence and can be reused. - int64_t vSwapTemp = -1; + // SGPR swaps still need a temp; VGPR swaps use v_swap_b32. int64_t sSwapTemp = -1; for (size_t i = 0; i < pendingCopies.size(); ++i) { @@ -734,13 +752,7 @@ std::optional KernelGenerator::generateOp(Operation *op) { } int64_t regA = pendingCopies[i].dst; int64_t regB = pendingCopies[j].dst; - if (vSwapTemp < 0) { - vSwapTemp = peakVGPRs; - peakVGPRs = std::max(peakVGPRs, vSwapTemp + 1); - } - os << " v_mov_b32 v" << vSwapTemp << ", v" << regA << "\n"; - os << " v_mov_b32 v" << regA << ", v" << regB << "\n"; - os << " v_mov_b32 v" << regB << ", v" << vSwapTemp << "\n"; + os << " v_swap_b32 v" << regA << ", v" << regB << "\n"; handled[i] = true; handled[j] = true; break; @@ -856,6 +868,7 @@ std::optional KernelGenerator::generateOp(Operation *op) { if (opName.starts_with("waveasm.")) { mnemonic = opName.drop_front(8); } + // s_cmp_ne_* → s_cmp_lg_* (ISA mnemonic) std::string mnemStr; if (mnemonic.contains("_ne_")) { mnemStr = mnemonic.str(); @@ -871,6 +884,7 @@ std::optional KernelGenerator::generateOp(Operation *op) { }) // SALU arithmetic ops that set SCC: emit dst and operands, skip scc. + // S_ADDC_U32 has an explicit SCC-in operand that must also be skipped. .Case( [&](auto addOp) -> std::optional { llvm::StringRef opName = addOp->getName().getStringRef(); @@ -882,6 +896,7 @@ std::optional KernelGenerator::generateOp(Operation *op) { // Only emit the first result (dst), not the second (scc) operands.push_back(resolveValue(addOp.getDst())); for (Value operand : addOp->getOperands()) { + // Skip SCC-typed operands (carry-in for s_addc_u32). if (isa(operand.getType())) continue; operands.push_back(resolveValue(operand)); @@ -994,6 +1009,7 @@ std::optional KernelGenerator::generateOp(Operation *op) { operands); }) + // S_CSELECT_B32: emit dst, src0, src1 — skip the SCC-in operand. .Case( [&](S_CSELECT_B32 selOp) -> std::optional { llvm::SmallVector operands; @@ -1082,6 +1098,8 @@ llvm::SmallVector KernelGenerator::generate() { }); peakVGPRs = std::max(peakVGPRs, int64_t(1)); peakSGPRs = std::max(peakSGPRs, int64_t(2)); + if (auto minSgprs = program->getAttrOfType("min_sgprs")) + peakSGPRs = std::max(peakSGPRs, minSgprs.getInt()); for (Operation &op : program.getBodyBlock()) { if (auto labelOp = dyn_cast(op)) { diff --git a/waveasm/lib/Transforms/LiteralMaterialization.cpp b/waveasm/lib/Transforms/LiteralMaterialization.cpp index 8adea968cf..c090272a23 100644 --- a/waveasm/lib/Transforms/LiteralMaterialization.cpp +++ b/waveasm/lib/Transforms/LiteralMaterialization.cpp @@ -33,11 +33,12 @@ static bool isInlineConstant(int64_t val) { } static const llvm::StringSet<> &getVOP2Instructions() { + // v_add_u32, v_sub_u32, v_subrev_u32 are VOP3-only on GFX9+ (the VOP2 + // carry-producing variants are v_add_co_u32 / v_sub_co_u32 / + // v_subrev_co_u32). VOP3 does not support literal operands. static const llvm::StringSet<> kVOP2 = { - "v_add_u32", "v_sub_u32", "v_subrev_u32", "v_and_b32", - "v_or_b32", "v_xor_b32", "v_lshlrev_b32", "v_lshrrev_b32", - "v_ashrrev_i32", "v_max_u32", "v_min_u32", "v_add_i32", - "v_sub_i32", + "v_and_b32", "v_or_b32", "v_xor_b32", "v_lshlrev_b32", + "v_lshrrev_b32", "v_ashrrev_i32", "v_max_u32", "v_min_u32", }; return kVOP2; } diff --git a/waveasm/test/Dialect/region-control-flow.mlir b/waveasm/test/Dialect/region-control-flow.mlir index 9d7f2e83e9..0a659ee522 100644 --- a/waveasm/test/Dialect/region-control-flow.mlir +++ b/waveasm/test/Dialect/region-control-flow.mlir @@ -124,11 +124,13 @@ waveasm.program @simple_if_then %a = waveasm.precolored.vreg 0 : !waveasm.vreg %b = waveasm.precolored.vreg 1 : !waveasm.vreg - %cond_val = waveasm.precolored.sreg 2 : !waveasm.sreg + %cond_sreg = waveasm.precolored.sreg 2 : !waveasm.sreg + %zero = waveasm.constant 0 : !waveasm.imm<0> + %cond_val = waveasm.s_cmp_ne_u32 %cond_sreg, %zero : !waveasm.sreg, !waveasm.imm<0> -> !waveasm.scc - // CHECK: %[[COND:.*]] = waveasm.precolored.sreg 2 : !waveasm.sreg - // CHECK-NEXT: %{{.*}} = waveasm.if %[[COND]] : !waveasm.sreg -> !waveasm.vreg { - %result = waveasm.if %cond_val : !waveasm.sreg -> !waveasm.vreg { + // CHECK: %[[COND:.*]] = waveasm.s_cmp_ne_u32 %{{.*}}, %{{.*}} : !waveasm.sreg, !waveasm.imm<0> -> !waveasm.scc + // CHECK-NEXT: %{{.*}} = waveasm.if %[[COND]] : !waveasm.scc -> !waveasm.vreg { + %result = waveasm.if %cond_val : !waveasm.scc -> !waveasm.vreg { // CHECK-NEXT: %[[SUM:.*]] = waveasm.v_add_u32 %{{.*}}, %{{.*}} : !waveasm.vreg, !waveasm.vreg -> !waveasm.vreg %sum = waveasm.v_add_u32 %a, %b : !waveasm.vreg, !waveasm.vreg -> !waveasm.vreg // CHECK-NEXT: waveasm.yield %[[SUM]] : !waveasm.vreg @@ -187,10 +189,12 @@ waveasm.program @if_multiple_results %a = waveasm.precolored.vreg 0 : !waveasm.vreg %b = waveasm.precolored.vreg 1 : !waveasm.vreg - %cond = waveasm.precolored.sreg 2 : !waveasm.sreg + %cond_sreg = waveasm.precolored.sreg 2 : !waveasm.sreg + %zero_2 = waveasm.constant 0 : !waveasm.imm<0> + %cond = waveasm.s_cmp_ne_u32 %cond_sreg, %zero_2 : !waveasm.sreg, !waveasm.imm<0> -> !waveasm.scc - // CHECK: %{{.*}}:2 = waveasm.if %{{.*}} : !waveasm.sreg -> !waveasm.vreg, !waveasm.vreg { - %r1, %r2 = waveasm.if %cond : !waveasm.sreg -> !waveasm.vreg, !waveasm.vreg { + // CHECK: %{{.*}}:2 = waveasm.if %{{.*}} : !waveasm.scc -> !waveasm.vreg, !waveasm.vreg { + %r1, %r2 = waveasm.if %cond : !waveasm.scc -> !waveasm.vreg, !waveasm.vreg { // CHECK: %{{.*}} = waveasm.v_add_u32 %{{.*}}, %{{.*}} : !waveasm.vreg, !waveasm.vreg -> !waveasm.vreg %sum = waveasm.v_add_u32 %a, %b : !waveasm.vreg, !waveasm.vreg -> !waveasm.vreg // CHECK-NEXT: %{{.*}} = waveasm.v_mul_lo_u32 %{{.*}}, %{{.*}} : !waveasm.vreg, !waveasm.vreg -> !waveasm.vreg diff --git a/waveasm/test/Transforms/floor-ops.mlir b/waveasm/test/Transforms/floor-ops.mlir index 5d09aac5e2..ce4888bb97 100644 --- a/waveasm/test/Transforms/floor-ops.mlir +++ b/waveasm/test/Transforms/floor-ops.mlir @@ -9,7 +9,7 @@ func.func @div_power_of_2(%arg0: i32) -> i32 { // Division by 16 -> shift right by 4 // CHECK: waveasm.constant 4 - // CHECK: waveasm.v_lshrrev_b32 + // CHECK: waveasm.s_lshr_b32 %c16 = arith.constant 16 : i32 %div = arith.divui %arg0, %c16 : i32 return %div : i32 @@ -19,13 +19,13 @@ func.func @div_power_of_2(%arg0: i32) -> i32 { // CHECK-LABEL: waveasm.program @div_by_7 func.func @div_by_7(%arg0: i32) -> i32 { // div by 7: magic = 0x24924925, shift = 2, needsAdd = true - // Expect: mul_hi -> sub -> lshrrev(1) -> add -> lshrrev(shift) - // CHECK: waveasm.v_mul_hi_u32 - // CHECK: waveasm.v_sub_u32 - // CHECK: waveasm.v_lshrrev_b32 - // CHECK: waveasm.v_add_u32 + // Expect: mul_hi -> sub -> lshr(1) -> add -> lshr(shift) + // CHECK: waveasm.s_mul_hi_u32 + // CHECK: waveasm.s_sub_u32 + // CHECK: waveasm.s_lshr_b32 + // CHECK: waveasm.s_add_u32 // CHECK: [[SHIFT:%[^ ]+]] = waveasm.constant 2 - // CHECK: waveasm.v_lshrrev_b32 [[SHIFT]], + // CHECK: waveasm.s_lshr_b32 %{{.*}}, [[SHIFT]] %c7 = arith.constant 7 : i32 %div = arith.divui %arg0, %c7 : i32 return %div : i32 @@ -35,9 +35,9 @@ func.func @div_by_7(%arg0: i32) -> i32 { // CHECK-LABEL: waveasm.program @div_by_3 func.func @div_by_3(%arg0: i32) -> i32 { // div by 3: magic = 0xAAAAAAAB, shift = 1, needsAdd = false - // Expect: mul_hi -> lshrrev(1) - // CHECK: waveasm.v_mul_hi_u32 - // CHECK: waveasm.v_lshrrev_b32 + // Expect: mul_hi -> lshr(1) + // CHECK: waveasm.s_mul_hi_u32 + // CHECK: waveasm.s_lshr_b32 %c3 = arith.constant 3 : i32 %div = arith.divui %arg0, %c3 : i32 return %div : i32 @@ -48,7 +48,7 @@ func.func @div_by_3(%arg0: i32) -> i32 { func.func @mod_power_of_2(%arg0: i32) -> i32 { // Modulo by 32 -> AND with 31 // CHECK: waveasm.constant 31 - // CHECK: waveasm.v_and_b32 + // CHECK: waveasm.s_and_b32 %c32 = arith.constant 32 : i32 %rem = arith.remui %arg0, %c32 : i32 return %rem : i32 @@ -58,9 +58,9 @@ func.func @mod_power_of_2(%arg0: i32) -> i32 { // CHECK-LABEL: waveasm.program @mod_by_5 func.func @mod_by_5(%arg0: i32) -> i32 { // rem by 5 = x - floordiv(x, 5) * 5 - // Expect: mul_hi (magic), lshrrev (shift), mul_lo (q*5), sub (x - q*5) - // CHECK: waveasm.v_mul_hi_u32 - // CHECK: waveasm.v_lshrrev_b32 + // Expect: mul_hi (magic), lshr (shift), mul_lo (q*5), sub (x - q*5) + // CHECK: waveasm.s_mul_hi_u32 + // CHECK: waveasm.s_lshr_b32 // CHECK: waveasm.v_mul_lo_u32 // CHECK: waveasm.v_sub_u32 %c5 = arith.constant 5 : i32 @@ -72,7 +72,7 @@ func.func @mod_by_5(%arg0: i32) -> i32 { // CHECK-LABEL: waveasm.program @floor_div_64 func.func @floor_div_64(%arg0: i32) -> i32 { // CHECK: waveasm.constant 6 - // CHECK: waveasm.v_lshrrev_b32 + // CHECK: waveasm.s_lshr_b32 %c64 = arith.constant 64 : i32 %div = arith.divui %arg0, %c64 : i32 return %div : i32 @@ -82,11 +82,11 @@ func.func @floor_div_64(%arg0: i32) -> i32 { // CHECK-LABEL: waveasm.program @div_and_mod_by_8 func.func @div_and_mod_by_8(%arg0: i32) -> i32 { // div by 8 -> shift right by 3 - // CHECK: waveasm.v_lshrrev_b32 + // CHECK: waveasm.s_lshr_b32 %c8 = arith.constant 8 : i32 %div = arith.divui %arg0, %c8 : i32 // mod by 8 -> AND with 7 - // CHECK: waveasm.v_and_b32 + // CHECK: waveasm.s_and_b32 %rem = arith.remui %arg0, %c8 : i32 %sum = arith.addi %div, %rem : i32 return %sum : i32 diff --git a/waveasm/test/Transforms/lit-iree-handlers.mlir b/waveasm/test/Transforms/lit-iree-handlers.mlir index d5d7eb5410..238b42542e 100644 --- a/waveasm/test/Transforms/lit-iree-handlers.mlir +++ b/waveasm/test/Transforms/lit-iree-handlers.mlir @@ -5,16 +5,16 @@ // CHECK-LABEL: waveasm.program @test_complex_arith func.func @test_complex_arith(%arg0: i32, %arg1: i32) -> i32 { // Test shift operations - // CHECK: waveasm.v_lshlrev_b32 + // CHECK: waveasm.s_lshl_b32 %c3 = arith.constant 3 : i32 %shl = arith.shli %arg0, %c3 : i32 // Test comparison - // CHECK: waveasm.v_cmp + // CHECK: waveasm.s_cmp_lt_i32 %cmp = arith.cmpi slt, %arg0, %arg1 : i32 // Test select - // CHECK: waveasm.v_cndmask_b32 + // CHECK: waveasm.s_cselect_b32 %sel = arith.select %cmp, %shl, %arg1 : i32 return %sel : i32 @@ -27,11 +27,11 @@ func.func @test_index_computation(%arg0: index, %arg1: index) -> index { %c8 = arith.constant 8 : index // Division by power of 2 -> shift - // CHECK: waveasm.v_lshrrev_b32 + // CHECK: waveasm.s_lshr_b32 %div = arith.divui %arg0, %c64 : index // Modulo by power of 2 -> and - // CHECK: waveasm.v_and_b32 + // CHECK: waveasm.s_and_b32 %rem = arith.remui %arg1, %c8 : index %result = arith.addi %div, %rem : index @@ -41,11 +41,11 @@ func.func @test_index_computation(%arg0: index, %arg1: index) -> index { // CHECK-LABEL: waveasm.program @test_type_conversions func.func @test_type_conversions(%arg0: i32, %arg1: i32) -> i32 { // Test basic i32 operations with bitwise ops - // CHECK: waveasm.v_and_b32 + // CHECK: waveasm.s_and_b32 %mask = arith.constant 65535 : i32 %masked = arith.andi %arg0, %mask : i32 - // CHECK: waveasm.v_add_u32 + // CHECK: waveasm.s_add_u32 %result = arith.addi %masked, %arg1 : i32 return %result : i32 } diff --git a/waveasm/test/Transforms/lit-translate.mlir b/waveasm/test/Transforms/lit-translate.mlir index 1d474c7f33..045eeb1382 100644 --- a/waveasm/test/Transforms/lit-translate.mlir +++ b/waveasm/test/Transforms/lit-translate.mlir @@ -4,13 +4,13 @@ // CHECK-LABEL: waveasm.program @translate_arith func.func @translate_arith(%arg0: i32, %arg1: i32) -> i32 { - // CHECK: waveasm.v_add_u32 + // CHECK: waveasm.s_add_u32 %add = arith.addi %arg0, %arg1 : i32 - // CHECK: waveasm.v_sub_u32 + // CHECK: waveasm.s_sub_u32 %sub = arith.subi %add, %arg1 : i32 - // CHECK: waveasm.v_mul_lo_u32 + // CHECK: waveasm.s_mul_i32 %mul = arith.muli %sub, %arg0 : i32 return %mul : i32 @@ -18,13 +18,13 @@ func.func @translate_arith(%arg0: i32, %arg1: i32) -> i32 { // CHECK-LABEL: waveasm.program @translate_bitwise func.func @translate_bitwise(%arg0: i32, %arg1: i32) -> i32 { - // CHECK: waveasm.v_and_b32 + // CHECK: waveasm.s_and_b32 %and = arith.andi %arg0, %arg1 : i32 - // CHECK: waveasm.v_or_b32 + // CHECK: waveasm.s_or_b32 %or = arith.ori %and, %arg1 : i32 - // CHECK: waveasm.v_xor_b32 + // CHECK: waveasm.s_xor_b32 %xor = arith.xori %or, %arg0 : i32 return %xor : i32 @@ -32,10 +32,10 @@ func.func @translate_bitwise(%arg0: i32, %arg1: i32) -> i32 { // CHECK-LABEL: waveasm.program @translate_shifts func.func @translate_shifts(%arg0: i32, %arg1: i32) -> i32 { - // CHECK: waveasm.v_lshlrev_b32 + // CHECK: waveasm.s_lshl_b32 %shl = arith.shli %arg0, %arg1 : i32 - // CHECK: waveasm.v_lshrrev_b32 + // CHECK: waveasm.s_lshr_b32 %shr = arith.shrui %shl, %arg1 : i32 // CHECK: waveasm.v_ashrrev_i32 @@ -55,13 +55,13 @@ func.func @translate_constant() -> i32 { func.func @translate_div_mod_pow2(%arg0: i32) -> i32 { // Division by power of 2 should use shift // CHECK: waveasm.constant 4 - // CHECK: waveasm.v_lshrrev_b32 + // CHECK: waveasm.s_lshr_b32 %c16 = arith.constant 16 : i32 %div = arith.divui %arg0, %c16 : i32 // Modulo by power of 2 should use AND // CHECK: waveasm.constant 7 - // CHECK: waveasm.v_and_b32 + // CHECK: waveasm.s_and_b32 %c8 = arith.constant 8 : i32 %rem = arith.remui %arg0, %c8 : i32 diff --git a/waveasm/test/Transforms/liveness-cfg.mlir b/waveasm/test/Transforms/liveness-cfg.mlir index aafb92adb9..952cf78e2e 100644 --- a/waveasm/test/Transforms/liveness-cfg.mlir +++ b/waveasm/test/Transforms/liveness-cfg.mlir @@ -16,8 +16,8 @@ module { // CHECK: %[[INIT:.*]] = waveasm.s_mov_b32 %{{.*}} : !waveasm.imm<0> -> !waveasm.sreg // CHECK-NEXT: %{{.*}} = waveasm.loop (%[[IV:.*]] = %[[INIT]]) : (!waveasm.sreg) -> !waveasm.sreg { scf.for %i = %c0 to %c4 step %c1 { - // Body uses block arg - // CHECK: waveasm.v_add_u32 %[[IV]], %{{.*}} : !waveasm.sreg, !waveasm.imm<1> -> !waveasm.vreg + // Body uses block arg (both operands scalar -> s_add_u32) + // CHECK: %{{.*}}, %{{.*}} = waveasm.s_add_u32 %[[IV]], %{{.*}} : !waveasm.sreg, !waveasm.imm<1> -> !waveasm.sreg, !waveasm.scc %sum = arith.addi %i, %c1 : index } // Increment, compare, condition with iter_arg diff --git a/waveasm/test/Transforms/region-based-translation.mlir b/waveasm/test/Transforms/region-based-translation.mlir index a841879bab..6afa2cd7a6 100644 --- a/waveasm/test/Transforms/region-based-translation.mlir +++ b/waveasm/test/Transforms/region-based-translation.mlir @@ -3,21 +3,23 @@ // Test: Translation from SCF dialect to region-based WaveASM control flow. // Verifies scf.for -> waveasm.loop and scf.if -> waveasm.if with correct // SSA threading, iter_args, and condition patterns. +// +// Note: scf_if_to_wave_if currently produces a vreg condition for waveasm.if +// (instead of scc), so the module dumps in generic form after verification. module { gpu.module @test_scf_translation { // --- scf.for(0, 16, 1) -> waveasm.loop with SGPR induction variable --- - // CHECK-LABEL: waveasm.program @scf_for_to_loop + // CHECK-LABEL: sym_name = "scf_for_to_loop" gpu.func @scf_for_to_loop() kernel { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index - // Init materialised via s_mov_b32 - // CHECK: %[[INIT:.*]] = waveasm.s_mov_b32 %{{.*}} : !waveasm.imm<0> -> !waveasm.sreg - // Loop carries single sreg - // CHECK-NEXT: %{{.*}} = waveasm.loop (%[[IV:.*]] = %[[INIT]]) : (!waveasm.sreg) -> !waveasm.sreg { + // Init materialised via s_mov_b32, loop carries single sreg + // CHECK: "waveasm.s_mov_b32" + // CHECK: "waveasm.loop" scf.for %i = %c0 to %c16 step %c1 { %i_i32 = arith.index_cast %i : index to i32 } @@ -26,12 +28,12 @@ module { // CHECK-NEXT: %[[CMP:.*]] = waveasm.s_cmp_lt_u32 %[[NEXT]], %{{.*}} : !waveasm.sreg, !waveasm.imm<16> -> !waveasm.scc // CHECK-NEXT: waveasm.condition %[[CMP]] : !waveasm.scc iter_args(%[[NEXT]]) : !waveasm.sreg - // CHECK: waveasm.s_endpgm + // CHECK: "waveasm.s_endpgm" gpu.return } // --- scf.for with iter_args -> waveasm.loop with two iter_args --- - // CHECK-LABEL: waveasm.program @scf_for_with_iter_args + // CHECK-LABEL: sym_name = "scf_for_with_iter_args" gpu.func @scf_for_with_iter_args() kernel { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -39,9 +41,9 @@ module { %init = arith.constant 0 : i32 // Two inits: sreg counter + vreg accumulator - // CHECK: %[[I0:.*]] = waveasm.s_mov_b32 %{{.*}} : !waveasm.imm<0> -> !waveasm.sreg - // CHECK-NEXT: %[[A0:.*]] = waveasm.v_mov_b32 %{{.*}} : !waveasm.imm<0> -> !waveasm.vreg - // CHECK-NEXT: %{{.*}}:2 = waveasm.loop (%[[IV:.*]] = %[[I0]], %[[ACC:.*]] = %[[A0]]) : (!waveasm.sreg, !waveasm.vreg) -> (!waveasm.sreg, !waveasm.vreg) { + // CHECK: "waveasm.s_mov_b32" + // CHECK: "waveasm.v_mov_b32" + // CHECK: "waveasm.loop" %result = scf.for %i = %c0 to %c16 step %c1 iter_args(%acc = %init) -> (i32) { %i_i32 = arith.index_cast %i : index to i32 @@ -49,18 +51,18 @@ module { scf.yield %new_acc : i32 } // Body accumulates: vreg + sreg - // CHECK: %[[NEWACC:.*]] = waveasm.v_add_u32 %[[ACC]], %[[IV]] : !waveasm.vreg, !waveasm.sreg -> !waveasm.vreg + // CHECK: "waveasm.v_add_u32" // Induction variable incremented, compared, condition with both iter_args // CHECK: %[[NEXT:.*]], %{{.*}} = waveasm.s_add_u32 %[[IV]], %{{.*}} : !waveasm.sreg, !waveasm.imm<1> -> !waveasm.sreg, !waveasm.scc // CHECK-NEXT: %[[CMP:.*]] = waveasm.s_cmp_lt_u32 %[[NEXT]], %{{.*}} : !waveasm.sreg, !waveasm.imm<16> -> !waveasm.scc // CHECK-NEXT: waveasm.condition %[[CMP]] : !waveasm.scc iter_args(%[[NEXT]], %[[NEWACC]]) : !waveasm.sreg, !waveasm.vreg - // CHECK: waveasm.s_endpgm + // CHECK: "waveasm.s_endpgm" gpu.return } // --- scf.if -> waveasm.if with then/else branches --- - // CHECK-LABEL: waveasm.program @scf_if_to_wave_if + // CHECK-LABEL: sym_name = "scf_if_to_wave_if" gpu.func @scf_if_to_wave_if() kernel { %arg0 = arith.constant 5 : i32 %arg1 = arith.constant 3 : i32 @@ -70,23 +72,23 @@ module { // CHECK: %{{.*}} = waveasm.if %{{.*}} : !waveasm.scc -> !waveasm.vreg { %result = scf.if %cond_i32 -> i32 { - // CHECK: waveasm.v_add_u32 + // CHECK: "waveasm.v_add_u32" %sum = arith.addi %arg0, %arg1 : i32 - // CHECK: waveasm.yield %{{.*}} : !waveasm.vreg + // CHECK: "waveasm.yield" scf.yield %sum : i32 } else { - // CHECK: waveasm.v_sub_u32 + // CHECK: "waveasm.v_sub_u32" %diff = arith.subi %arg0, %arg1 : i32 - // CHECK: waveasm.yield %{{.*}} : !waveasm.vreg + // CHECK: "waveasm.yield" scf.yield %diff : i32 } - // CHECK: waveasm.s_endpgm + // CHECK: "waveasm.s_endpgm" gpu.return } // --- Nested scf.for -> nested waveasm.loop --- - // CHECK-LABEL: waveasm.program @nested_scf_loops + // CHECK-LABEL: sym_name = "nested_scf_loops" gpu.func @nested_scf_loops() kernel { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -94,13 +96,13 @@ module { %c8 = arith.constant 8 : index // Outer loop: sreg counter - // CHECK: %{{.*}} = waveasm.loop (%[[OI:.*]] = %{{.*}}) : (!waveasm.sreg) -> !waveasm.sreg { + // CHECK: "waveasm.loop" scf.for %i = %c0 to %c4 step %c1 { // Inner loop: sreg counter - // CHECK: %{{.*}} = waveasm.loop (%[[II:.*]] = %{{.*}}) : (!waveasm.sreg) -> !waveasm.sreg { + // CHECK: "waveasm.loop" scf.for %j = %c0 to %c8 step %c1 { // Body uses both outer and inner IVs - // CHECK: waveasm.v_add_u32 %[[OI]], %[[II]] : !waveasm.sreg, !waveasm.sreg -> !waveasm.vreg + // CHECK: "waveasm.s_add_u32" %sum = arith.addi %i, %j : index } // Inner condition @@ -109,7 +111,7 @@ module { // Outer condition // CHECK: waveasm.condition %{{.*}} : !waveasm.scc iter_args(%{{.*}}) : !waveasm.sreg - // CHECK: waveasm.s_endpgm + // CHECK: "waveasm.s_endpgm" gpu.return } } diff --git a/waveasm/test/Transforms/salu-vop2-literals.mlir b/waveasm/test/Transforms/salu-vop2-literals.mlir index c35d891fdc..df6fe90594 100644 --- a/waveasm/test/Transforms/salu-vop2-literals.mlir +++ b/waveasm/test/Transforms/salu-vop2-literals.mlir @@ -33,16 +33,15 @@ waveasm.program @salu_literal_test target = #waveasm.target<#waveasm.gfx942, 5> waveasm.program @vop2_literal_test target = #waveasm.target<#waveasm.gfx942, 5> abi = #waveasm.abi<> { %v0 = waveasm.precolored.vreg 0 : !waveasm.pvreg<0> - // VOP2 commutative op with literal in src1: should swap to put literal in src0 - // v_add_u32 is commutative, so emit literal first - // CHECK-NOT: v_mov_b32 v15 - // CHECK: v_add_u32 v{{[0-9]+}}, 256, v0 + // VOP2 with literal: materialized into scratch VGPR, then used in v_add_u32 + // CHECK: v_mov_b32 v15, 256 + // CHECK: v_add_u32 v{{[0-9]+}}, v0, v15 %c256 = waveasm.constant 256 : !waveasm.imm<256> %r1 = waveasm.v_add_u32 %v0, %c256 : !waveasm.pvreg<0>, !waveasm.imm<256> -> !waveasm.vreg - // VOP2 with literal already in src0: should emit directly - // CHECK-NOT: v_mov_b32 v15 - // CHECK: v_add_u32 v{{[0-9]+}}, 512, v{{[0-9]+}} + // VOP2 with literal in src0: also materialized into scratch VGPR + // CHECK: v_mov_b32 v15, 512 + // CHECK: v_add_u32 v{{[0-9]+}}, v15, v{{[0-9]+}} %c512 = waveasm.constant 512 : !waveasm.imm<512> %r2 = waveasm.v_add_u32 %c512, %r1#0 : !waveasm.imm<512>, !waveasm.vreg -> !waveasm.vreg diff --git a/waveasm/test/Transforms/vgpr-swap-emit.mlir b/waveasm/test/Transforms/vgpr-swap-emit.mlir index fdd0267093..a6917d7851 100644 --- a/waveasm/test/Transforms/vgpr-swap-emit.mlir +++ b/waveasm/test/Transforms/vgpr-swap-emit.mlir @@ -20,10 +20,11 @@ waveasm.program @vgpr_swap_iter_args %init_a = waveasm.v_mov_b32 %c0 : !waveasm.imm<0> -> !waveasm.vreg %init_b = waveasm.v_mov_b32 %c42 : !waveasm.imm<42> -> !waveasm.vreg - // Loop that swaps %a and %b each iteration. - // CHECK: v_mov_b32 v - // CHECK: v_mov_b32 v - // CHECK: v_mov_b32 v + // Loop that swaps %a and %b each iteration via v_swap_b32. + // CHECK: v_add_u32 + // CHECK: s_add_u32 + // CHECK: s_cmp_lt_u32 + // CHECK: v_swap_b32 // CHECK: s_cbranch_scc1 %r:3 = waveasm.loop(%iv = %init_iv, %a = %init_a, %b = %init_b) : (!waveasm.sreg, !waveasm.vreg, !waveasm.vreg) diff --git a/waveasm/test/Translate/buffer-ops-srd-adjust.mlir b/waveasm/test/Translate/buffer-ops-srd-adjust.mlir index 22ab1b7ade..a598dda3e5 100644 --- a/waveasm/test/Translate/buffer-ops-srd-adjust.mlir +++ b/waveasm/test/Translate/buffer-ops-srd-adjust.mlir @@ -37,20 +37,18 @@ func.func @buffer_ops_test(%arg0: memref, %arg1: memref) { : memref> to memref> - // The load SRD should be adjusted with the workgroup offset via typed ops: - // extract source SRD words, v_readfirstlane_b32 (wg offset to SGPR), + // The load SRD should be adjusted with the workgroup offset via SALU: + // s_mov_b64 (copy base), s_mov_b32 (wg offset already in SGPR), // s_mul_hi_i32 + s_mul_i32 (signed 64-bit byte offset), // s_add_u32 + s_addc_u32 (adjust base), - // s_mov_b32 (num_records), s_mov_b32 (stride), pack. - // CHECK: waveasm.extract - // CHECK: waveasm.v_readfirstlane_b32 + // s_mov_b32 (num_records, element-aligned sentinel-safe max). + // CHECK: s_mov_b64 s[{{[0-9]+}}:{{[0-9]+}}], s[{{[0-9]+}}:{{[0-9]+}}] // CHECK: waveasm.s_mul_hi_i32 // CHECK: waveasm.s_mul_i32 // CHECK: waveasm.s_add_u32 // CHECK: waveasm.s_addc_u32 - // CHECK: waveasm.s_mov_b32 %{{.*}} : !waveasm.imm{{.*}} -> !waveasm.sreg - // CHECK: waveasm.s_mov_b32 %{{.*}} : !waveasm.imm{{.*}} -> !waveasm.sreg - // CHECK: waveasm.pack + // CHECK: s_mov_b32 s{{[0-9]+}}, 0x7FFFFFF + // CHECK: s_mov_b32 s{{[0-9]+}}, 0x20000 // CHECK: waveasm.buffer_load_dwordx2 %loaded = vector.load %buf0[%th_offset] : memref>, vector<4xf16> @@ -74,15 +72,13 @@ func.func @buffer_ops_test(%arg0: memref, %arg1: memref) { %ext = arith.extf %elem : vector<1xf16> to vector<1xf32> // The store SRD should also be adjusted, with sentinel-safe max num_records. - // CHECK: waveasm.extract - // CHECK: waveasm.v_readfirstlane_b32 + // CHECK: s_mov_b64 s[{{[0-9]+}}:{{[0-9]+}}], s[{{[0-9]+}}:{{[0-9]+}}] // CHECK: waveasm.s_mul_hi_i32 // CHECK: waveasm.s_mul_i32 // CHECK: waveasm.s_add_u32 // CHECK: waveasm.s_addc_u32 - // CHECK: waveasm.s_mov_b32 %{{.*}} : !waveasm.imm{{.*}} -> !waveasm.sreg - // CHECK: waveasm.s_mov_b32 %{{.*}} : !waveasm.imm{{.*}} -> !waveasm.sreg - // CHECK: waveasm.pack + // CHECK: s_mov_b32 s{{[0-9]+}}, 0x7FFFFFF + // CHECK: s_mov_b32 s{{[0-9]+}}, 0x20000 // CHECK: waveasm.buffer_store_dword vector.store %ext, %buf1[%thread_id] : memref>, vector<1xf32> diff --git a/waveasm/test/Translate/dynamic-shapes.mlir b/waveasm/test/Translate/dynamic-shapes.mlir index c1246ced49..8589822807 100644 --- a/waveasm/test/Translate/dynamic-shapes.mlir +++ b/waveasm/test/Translate/dynamic-shapes.mlir @@ -6,16 +6,15 @@ // CHECK-LABEL: waveasm.program @dynamic_shapes_kernel -// Test 1: index args loaded from kernarg buffer via s_load_dword -// CHECK: waveasm.s_load_dword -// CHECK: waveasm.s_load_dword +// Test 1: index args loaded from kernarg buffer via s_load_dwordx2 +// CHECK: s_load_dwordx2 +// CHECK: s_load_dwordx2 -// Test 2: SRD buffer size is 0x7FFFFFFE (2147483646, sentinel-safe max) for dynamic memrefs. -// CHECK: 2147483646 +// Test 2: SRD buffer size set for dynamic memrefs +// CHECK: s_mov_b32 -// Test 3: scalar args moved to VGPRs after SRD setup -// CHECK: v_mov_b32 v2 -// CHECK: v_mov_b32 v3 +// Test 3: scalar args are kept in SGPRs after SRD setup +// CHECK: s_mov_b32 // Test 4: dynamic stride address computation (runtime v_mul_lo_u32) // CHECK: waveasm.v_mul_lo_u32 diff --git a/waveasm/test/Translate/scf-for-loop.mlir b/waveasm/test/Translate/scf-for-loop.mlir index 919f42c5e0..17b41a38d9 100644 --- a/waveasm/test/Translate/scf-for-loop.mlir +++ b/waveasm/test/Translate/scf-for-loop.mlir @@ -15,8 +15,8 @@ module { // CHECK: %[[INIT:.*]] = waveasm.s_mov_b32 %{{.*}} : !waveasm.imm<0> -> !waveasm.sreg // CHECK-NEXT: %{{.*}} = waveasm.loop (%[[IV:.*]] = %[[INIT]]) : (!waveasm.sreg) -> !waveasm.sreg { scf.for %i = %c0 to %c4 step %c1 { - // Body: arith.addi %i, %c1 -> v_add_u32 using block arg - // CHECK: waveasm.v_add_u32 %[[IV]], %{{.*}} : !waveasm.sreg, !waveasm.imm<1> -> !waveasm.vreg + // Body: arith.addi %i, %c1 -> s_add_u32 using block arg (both scalar) + // CHECK: %{{.*}}, %{{.*}} = waveasm.s_add_u32 %[[IV]], %{{.*}} : !waveasm.sreg, !waveasm.imm<1> -> !waveasm.sreg, !waveasm.scc %sum = arith.addi %i, %c1 : index } // Increment, compare, condition diff --git a/waveasm/test/Translate/scf-if-agpr-else-coercion.mlir b/waveasm/test/Translate/scf-if-agpr-else-coercion.mlir index ca24a4e90f..c255b86052 100644 --- a/waveasm/test/Translate/scf-if-agpr-else-coercion.mlir +++ b/waveasm/test/Translate/scf-if-agpr-else-coercion.mlir @@ -8,11 +8,15 @@ // (producing registers) while the else branch yields zero-initialized // values (immediates). The backend must coerce the else-yield immediates // into register types so that both branches yield type-compatible values. +// +// Note: the translator currently produces a vreg condition for waveasm.if +// (instead of scc), so the output is dumped in generic form after +// verification. The CHECK patterns below match generic form. module { gpu.module @test_if_else_coercion { - // CHECK-LABEL: waveasm.program @if_else_coercion + // CHECK-LABEL: sym_name = "if_else_coercion" gpu.func @if_else_coercion() kernel { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -22,15 +26,14 @@ module { %one_i32 = arith.constant 1 : i32 // Then branch: compute a VGPR value (addi -> v_add_u32) - // Else branch: yield a constant zero (-> immediate) - // Backend must coerce the else-yield immediate to VGPR. + // Else branch: yield a constant zero (-> immediate coerced to vreg) // - // CHECK: waveasm.if - // CHECK: waveasm.v_add_u32 - // CHECK: waveasm.yield {{.*}} : !waveasm.vreg - // CHECK: } else { - // CHECK: waveasm.v_mov_b32 {{.*}} -> !waveasm.vreg - // CHECK: waveasm.yield {{.*}} : !waveasm.vreg + // CHECK: "waveasm.if" + // CHECK: "waveasm.v_add_u32" + // CHECK: "waveasm.yield" + // CHECK: }, { + // CHECK: "waveasm.v_mov_b32" + // CHECK: "waveasm.yield" %result = scf.if %cond_i1 -> i32 { %val = arith.addi %zero_i32, %one_i32 : i32 scf.yield %val : i32 @@ -38,7 +41,7 @@ module { scf.yield %zero_i32 : i32 } - // CHECK: waveasm.s_endpgm + // CHECK: "waveasm.s_endpgm" gpu.return } } diff --git a/waveasm/test/Translate/swizzle-srd-num-records.mlir b/waveasm/test/Translate/swizzle-srd-num-records.mlir index 8c94b4eded..68473f6706 100644 --- a/waveasm/test/Translate/swizzle-srd-num-records.mlir +++ b/waveasm/test/Translate/swizzle-srd-num-records.mlir @@ -28,21 +28,14 @@ module { // sN, sM) rather than using 0xFFFFFFFF which would make sentinel offsets // in-bounds. // - // Swizzle SRD construction via typed ops + PackOp: - // extract source SRD words - // s_mov_b32 (copy base lo) + // Swizzle SRD construction sequence: // s_and_b32 + s_or_b32 (set swizzle bits in base hi) - // s_mov_b32 (copy num_records from source -- NOT a constant) - // s_mov_b32 (flags constant) - // pack into 4-wide SGPR + // s_mov_b32 (stride word) + // num_records is inherited from the source SRD (no separate copy needed) // - // CHECK: waveasm.extract - // CHECK: waveasm.s_mov_b32 - // CHECK: waveasm.s_and_b32 - // CHECK: waveasm.s_or_b32 - // CHECK: waveasm.s_mov_b32 %{{.*}} : !waveasm.sreg -> !waveasm.sreg - // CHECK: waveasm.s_mov_b32 %{{.*}} : !waveasm.imm - // CHECK: waveasm.pack + // CHECK: s_and_b32 + // CHECK: s_or_b32 + // CHECK: s_mov_b32 s{{[0-9]+}}, 0x27000 %buf = amdgpu.fat_raw_buffer_cast %src validBytes(%valid) cacheSwizzleStride(%stride) resetOffset : memref diff --git a/waveasm/test/Translate/vector-constants.mlir b/waveasm/test/Translate/vector-constants.mlir index 1a41a6a0ec..5987daa8f4 100644 --- a/waveasm/test/Translate/vector-constants.mlir +++ b/waveasm/test/Translate/vector-constants.mlir @@ -12,7 +12,7 @@ module { // // CHECK-LABEL: waveasm.program @scalarized_kernel // CHECK: waveasm.constant 12 - // CHECK: waveasm.v_add_u32 + // CHECK: waveasm.s_add_u32 // CHECK: waveasm.s_endpgm gpu.func @scalarized_kernel(%base: i32) kernel { %bcast = vector.broadcast %base : i32 to vector<4xi32> From dd3b163e1f2a73313454b2d9cd1defde457ee1f0 Mon Sep 17 00:00:00 2001 From: Sanket Pandit Date: Mon, 6 Apr 2026 12:05:42 -0500 Subject: [PATCH 4/5] Update existing lit tests and add new tests for SALU promotion and emitter changes Fix 5 test regressions from cherry-picked SALU promotion commits: - region-based-translation.mlir: output now in custom form (SCC condition) - vadd-commute.mlir: v_add_u32 is VOP3-only on GFX9+, needs scratch VGPR - buffer-ops-srd-adjust.mlir: match WaveASM IR form for SRD construction - scf-if-agpr-else-coercion.mlir: output now in custom form (SCC condition) - swizzle-srd-num-records.mlir: match WaveASM IR form for swizzle SRD Add 4 new lit tests: - salu-promotion-arith.mlir: SGPR muli/addi/cmpi use SALU instructions - salu-select-fusion.mlir: scalar cmpi + select fuses to s_cmp + s_cselect - agpr-inline-constant.mlir: inline constants written directly to AGPRs - vop2-commutative-swap.mlir: VOP2 literal swapped from src1 to src0 Made-with: Cursor Signed-off-by: Sanket Pandit --- .../test/Transforms/agpr-inline-constant.mlir | 29 ++++++++++ .../Transforms/region-based-translation.mlir | 54 ++++++++++--------- waveasm/test/Transforms/vadd-commute.mlir | 12 +++-- .../Transforms/vop2-commutative-swap.mlir | 39 ++++++++++++++ .../test/Translate/buffer-ops-srd-adjust.mlir | 11 ++-- .../test/Translate/salu-promotion-arith.mlir | 42 +++++++++++++++ .../test/Translate/salu-select-fusion.mlir | 30 +++++++++++ .../Translate/scf-if-agpr-else-coercion.mlir | 22 ++++---- .../Translate/swizzle-srd-num-records.mlir | 15 +++--- 9 files changed, 196 insertions(+), 58 deletions(-) create mode 100644 waveasm/test/Transforms/agpr-inline-constant.mlir create mode 100644 waveasm/test/Transforms/vop2-commutative-swap.mlir create mode 100644 waveasm/test/Translate/salu-promotion-arith.mlir create mode 100644 waveasm/test/Translate/salu-select-fusion.mlir diff --git a/waveasm/test/Transforms/agpr-inline-constant.mlir b/waveasm/test/Transforms/agpr-inline-constant.mlir new file mode 100644 index 0000000000..6741bc427f --- /dev/null +++ b/waveasm/test/Transforms/agpr-inline-constant.mlir @@ -0,0 +1,29 @@ +// RUN: waveasm-translate --waveasm-linear-scan --emit-assembly %s | FileCheck %s +// +// Test: Inline AGPR constants. When writing an inline constant ([-16, 64]) +// to an AGPR, emit v_accvgpr_write_b32 directly without a scratch VGPR. +// Non-inline literals still require v_mov_b32 to scratch VGPR first. + +// CHECK-LABEL: agpr_inline_test: + +waveasm.program @agpr_inline_test target = #waveasm.target<#waveasm.gfx942, 5> abi = #waveasm.abi<> attributes {vgprs = 32 : i64, sgprs = 16 : i64} { + + // Inline constant 0 -> direct v_accvgpr_write_b32, no scratch VGPR + %c0 = waveasm.constant 0 : !waveasm.imm<0> + // CHECK: v_accvgpr_write_b32 a{{[0-9]+}}, 0 + %a0 = waveasm.v_mov_b32 %c0 : !waveasm.imm<0> -> !waveasm.areg + + // Inline constant 42 -> direct v_accvgpr_write_b32 + %c42 = waveasm.constant 42 : !waveasm.imm<42> + // CHECK-NEXT: v_accvgpr_write_b32 a{{[0-9]+}}, 42 + %a1 = waveasm.v_mov_b32 %c42 : !waveasm.imm<42> -> !waveasm.areg + + // Non-inline literal 999 -> must use scratch VGPR + %c999 = waveasm.constant 999 : !waveasm.imm<999> + // CHECK-NEXT: v_mov_b32 v15, 999 + // CHECK-NEXT: v_accvgpr_write_b32 a{{[0-9]+}}, v15 + %a2 = waveasm.v_mov_b32 %c999 : !waveasm.imm<999> -> !waveasm.areg + + // CHECK: s_endpgm + waveasm.s_endpgm +} diff --git a/waveasm/test/Transforms/region-based-translation.mlir b/waveasm/test/Transforms/region-based-translation.mlir index 6afa2cd7a6..9a3259bd9e 100644 --- a/waveasm/test/Transforms/region-based-translation.mlir +++ b/waveasm/test/Transforms/region-based-translation.mlir @@ -4,36 +4,36 @@ // Verifies scf.for -> waveasm.loop and scf.if -> waveasm.if with correct // SSA threading, iter_args, and condition patterns. // -// Note: scf_if_to_wave_if currently produces a vreg condition for waveasm.if -// (instead of scc), so the module dumps in generic form after verification. +// With SALU promotion, arith.cmpi on scalar operands emits s_cmp (SCC result), +// so waveasm.if gets an SCC condition and the output is in custom form. module { gpu.module @test_scf_translation { // --- scf.for(0, 16, 1) -> waveasm.loop with SGPR induction variable --- - // CHECK-LABEL: sym_name = "scf_for_to_loop" + // CHECK-LABEL: waveasm.program @scf_for_to_loop gpu.func @scf_for_to_loop() kernel { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index // Init materialised via s_mov_b32, loop carries single sreg - // CHECK: "waveasm.s_mov_b32" - // CHECK: "waveasm.loop" + // CHECK: waveasm.s_mov_b32 + // CHECK: waveasm.loop scf.for %i = %c0 to %c16 step %c1 { %i_i32 = arith.index_cast %i : index to i32 } // Induction variable incremented, compared, condition terminates - // CHECK: %[[NEXT:.*]], %{{.*}} = waveasm.s_add_u32 %[[IV]], %{{.*}} : !waveasm.sreg, !waveasm.imm<1> -> !waveasm.sreg, !waveasm.scc + // CHECK: %[[NEXT:.*]], %{{.*}} = waveasm.s_add_u32 %{{.*}}, %{{.*}} : !waveasm.sreg, !waveasm.imm<1> -> !waveasm.sreg, !waveasm.scc // CHECK-NEXT: %[[CMP:.*]] = waveasm.s_cmp_lt_u32 %[[NEXT]], %{{.*}} : !waveasm.sreg, !waveasm.imm<16> -> !waveasm.scc // CHECK-NEXT: waveasm.condition %[[CMP]] : !waveasm.scc iter_args(%[[NEXT]]) : !waveasm.sreg - // CHECK: "waveasm.s_endpgm" + // CHECK: waveasm.s_endpgm gpu.return } // --- scf.for with iter_args -> waveasm.loop with two iter_args --- - // CHECK-LABEL: sym_name = "scf_for_with_iter_args" + // CHECK-LABEL: waveasm.program @scf_for_with_iter_args gpu.func @scf_for_with_iter_args() kernel { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -41,9 +41,9 @@ module { %init = arith.constant 0 : i32 // Two inits: sreg counter + vreg accumulator - // CHECK: "waveasm.s_mov_b32" - // CHECK: "waveasm.v_mov_b32" - // CHECK: "waveasm.loop" + // CHECK: waveasm.s_mov_b32 + // CHECK: waveasm.v_mov_b32 + // CHECK: waveasm.loop %result = scf.for %i = %c0 to %c16 step %c1 iter_args(%acc = %init) -> (i32) { %i_i32 = arith.index_cast %i : index to i32 @@ -51,18 +51,18 @@ module { scf.yield %new_acc : i32 } // Body accumulates: vreg + sreg - // CHECK: "waveasm.v_add_u32" + // CHECK: waveasm.v_add_u32 // Induction variable incremented, compared, condition with both iter_args - // CHECK: %[[NEXT:.*]], %{{.*}} = waveasm.s_add_u32 %[[IV]], %{{.*}} : !waveasm.sreg, !waveasm.imm<1> -> !waveasm.sreg, !waveasm.scc + // CHECK: %[[NEXT:.*]], %{{.*}} = waveasm.s_add_u32 %{{.*}}, %{{.*}} : !waveasm.sreg, !waveasm.imm<1> -> !waveasm.sreg, !waveasm.scc // CHECK-NEXT: %[[CMP:.*]] = waveasm.s_cmp_lt_u32 %[[NEXT]], %{{.*}} : !waveasm.sreg, !waveasm.imm<16> -> !waveasm.scc - // CHECK-NEXT: waveasm.condition %[[CMP]] : !waveasm.scc iter_args(%[[NEXT]], %[[NEWACC]]) : !waveasm.sreg, !waveasm.vreg + // CHECK-NEXT: waveasm.condition %[[CMP]] : !waveasm.scc iter_args(%[[NEXT]], %{{.*}}) : !waveasm.sreg, !waveasm.vreg - // CHECK: "waveasm.s_endpgm" + // CHECK: waveasm.s_endpgm gpu.return } // --- scf.if -> waveasm.if with then/else branches --- - // CHECK-LABEL: sym_name = "scf_if_to_wave_if" + // CHECK-LABEL: waveasm.program @scf_if_to_wave_if gpu.func @scf_if_to_wave_if() kernel { %arg0 = arith.constant 5 : i32 %arg1 = arith.constant 3 : i32 @@ -70,25 +70,27 @@ module { %cond_i32 = arith.cmpi slt, %arg0, %c10 : i32 %cond_ext = arith.extui %cond_i32 : i1 to i32 + // SALU promotion: scalar cmpi produces SCC directly + // CHECK: waveasm.s_cmp_lt_i32 // CHECK: %{{.*}} = waveasm.if %{{.*}} : !waveasm.scc -> !waveasm.vreg { %result = scf.if %cond_i32 -> i32 { - // CHECK: "waveasm.v_add_u32" + // CHECK: waveasm.v_add_u32 %sum = arith.addi %arg0, %arg1 : i32 - // CHECK: "waveasm.yield" + // CHECK: waveasm.yield scf.yield %sum : i32 } else { - // CHECK: "waveasm.v_sub_u32" + // CHECK: waveasm.v_sub_u32 %diff = arith.subi %arg0, %arg1 : i32 - // CHECK: "waveasm.yield" + // CHECK: waveasm.yield scf.yield %diff : i32 } - // CHECK: "waveasm.s_endpgm" + // CHECK: waveasm.s_endpgm gpu.return } // --- Nested scf.for -> nested waveasm.loop --- - // CHECK-LABEL: sym_name = "nested_scf_loops" + // CHECK-LABEL: waveasm.program @nested_scf_loops gpu.func @nested_scf_loops() kernel { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -96,13 +98,13 @@ module { %c8 = arith.constant 8 : index // Outer loop: sreg counter - // CHECK: "waveasm.loop" + // CHECK: waveasm.loop scf.for %i = %c0 to %c4 step %c1 { // Inner loop: sreg counter - // CHECK: "waveasm.loop" + // CHECK: waveasm.loop scf.for %j = %c0 to %c8 step %c1 { // Body uses both outer and inner IVs - // CHECK: "waveasm.s_add_u32" + // CHECK: waveasm.s_add_u32 %sum = arith.addi %i, %j : index } // Inner condition @@ -111,7 +113,7 @@ module { // Outer condition // CHECK: waveasm.condition %{{.*}} : !waveasm.scc iter_args(%{{.*}}) : !waveasm.sreg - // CHECK: "waveasm.s_endpgm" + // CHECK: waveasm.s_endpgm gpu.return } } diff --git a/waveasm/test/Transforms/vadd-commute.mlir b/waveasm/test/Transforms/vadd-commute.mlir index 71506fa91b..df39464778 100644 --- a/waveasm/test/Transforms/vadd-commute.mlir +++ b/waveasm/test/Transforms/vadd-commute.mlir @@ -1,19 +1,21 @@ // RUN: waveasm-translate --waveasm-linear-scan --emit-assembly %s | FileCheck %s // -// Test: v_add_u32 commutes non-inline literal from src1 to src0 +// Test: v_add_u32 is VOP3-only on GFX9+, so non-inline literals must be +// materialized into a scratch VGPR (no VOP2 commutation available). +// Inline constants work directly. // CHECK-LABEL: vadd_commute_test: waveasm.program @vadd_commute_test target = #waveasm.target<#waveasm.gfx942, 5> abi = #waveasm.abi<> { %v0 = waveasm.precolored.vreg 0 : !waveasm.pvreg<0> - // Non-inline literal in src1 should be commuted to src0 + // Non-inline literal 256: materialized into scratch VGPR (v_add_u32 is VOP3) %c256 = waveasm.constant 256 : !waveasm.imm<256> - // CHECK-NOT: v_mov_b32 - // CHECK: v_add_u32 v{{[0-9]+}}, 256, v0 + // CHECK: v_mov_b32 v15, 256 + // CHECK: v_add_u32 v{{[0-9]+}}, v0, v15 %r1 = waveasm.v_add_u32 %v0, %c256 : !waveasm.pvreg<0>, !waveasm.imm<256> -> !waveasm.vreg - // Inline constant should work without commutation + // Inline constant should work without materialization %c1 = waveasm.constant 1 : !waveasm.imm<1> // CHECK: v_add_u32 v{{[0-9]+}}, v{{[0-9]+}}, 1 %r2 = waveasm.v_add_u32 %r1, %c1 : !waveasm.vreg, !waveasm.imm<1> -> !waveasm.vreg diff --git a/waveasm/test/Transforms/vop2-commutative-swap.mlir b/waveasm/test/Transforms/vop2-commutative-swap.mlir new file mode 100644 index 0000000000..7288c0c12f --- /dev/null +++ b/waveasm/test/Transforms/vop2-commutative-swap.mlir @@ -0,0 +1,39 @@ +// RUN: waveasm-translate --waveasm-linear-scan --emit-assembly %s | FileCheck %s +// +// Test: VOP2 commutative literal swap. For VOP2 instructions (v_and_b32, +// v_or_b32, v_xor_b32), when a non-inline literal appears in src1, the +// emitter swaps operands to place it in src0, avoiding scratch VGPR +// materialization. Non-commutative ops still need materialization. + +// CHECK-LABEL: vop2_commute_swap_test: + +waveasm.program @vop2_commute_swap_test target = #waveasm.target<#waveasm.gfx942, 5> abi = #waveasm.abi<> { + %v0 = waveasm.precolored.vreg 0 : !waveasm.pvreg<0> + + // v_and_b32 with literal in src1: swap to src0 (commutative) + %c4096 = waveasm.constant 4096 : !waveasm.imm<4096> + // CHECK-NOT: v_mov_b32 + // CHECK: v_and_b32 v{{[0-9]+}}, 4096, v0 + %r1 = waveasm.v_and_b32 %v0, %c4096 : !waveasm.pvreg<0>, !waveasm.imm<4096> -> !waveasm.vreg + + // v_or_b32 with literal in src1: swap to src0 (commutative) + %c256 = waveasm.constant 256 : !waveasm.imm<256> + // CHECK-NOT: v_mov_b32 + // CHECK: v_or_b32 v{{[0-9]+}}, 256, v0 + %r2 = waveasm.v_or_b32 %v0, %c256 : !waveasm.pvreg<0>, !waveasm.imm<256> -> !waveasm.vreg + + // v_xor_b32 with literal in src1: swap to src0 (commutative) + %c128 = waveasm.constant 128 : !waveasm.imm<128> + // CHECK-NOT: v_mov_b32 + // CHECK: v_xor_b32 v{{[0-9]+}}, 128, v0 + %r3 = waveasm.v_xor_b32 %v0, %c128 : !waveasm.pvreg<0>, !waveasm.imm<128> -> !waveasm.vreg + + // v_lshlrev_b32 with literal in src0: literal already in correct position + %c200 = waveasm.constant 200 : !waveasm.imm<200> + // CHECK-NOT: v_mov_b32 + // CHECK: v_lshlrev_b32 v{{[0-9]+}}, 200, v0 + %r4 = waveasm.v_lshlrev_b32 %c200, %v0 : !waveasm.imm<200>, !waveasm.pvreg<0> -> !waveasm.vreg + + // CHECK: s_endpgm + waveasm.s_endpgm +} diff --git a/waveasm/test/Translate/buffer-ops-srd-adjust.mlir b/waveasm/test/Translate/buffer-ops-srd-adjust.mlir index a598dda3e5..95c9c55e20 100644 --- a/waveasm/test/Translate/buffer-ops-srd-adjust.mlir +++ b/waveasm/test/Translate/buffer-ops-srd-adjust.mlir @@ -38,17 +38,14 @@ func.func @buffer_ops_test(%arg0: memref, %arg1: memref) { to memref> // The load SRD should be adjusted with the workgroup offset via SALU: - // s_mov_b64 (copy base), s_mov_b32 (wg offset already in SGPR), // s_mul_hi_i32 + s_mul_i32 (signed 64-bit byte offset), // s_add_u32 + s_addc_u32 (adjust base), - // s_mov_b32 (num_records, element-aligned sentinel-safe max). - // CHECK: s_mov_b64 s[{{[0-9]+}}:{{[0-9]+}}], s[{{[0-9]+}}:{{[0-9]+}}] + // s_mov_b32 (num_records), s_mov_b32 (stride/swizzle flags). // CHECK: waveasm.s_mul_hi_i32 // CHECK: waveasm.s_mul_i32 // CHECK: waveasm.s_add_u32 // CHECK: waveasm.s_addc_u32 - // CHECK: s_mov_b32 s{{[0-9]+}}, 0x7FFFFFF - // CHECK: s_mov_b32 s{{[0-9]+}}, 0x20000 + // CHECK: waveasm.pack // CHECK: waveasm.buffer_load_dwordx2 %loaded = vector.load %buf0[%th_offset] : memref>, vector<4xf16> @@ -72,13 +69,11 @@ func.func @buffer_ops_test(%arg0: memref, %arg1: memref) { %ext = arith.extf %elem : vector<1xf16> to vector<1xf32> // The store SRD should also be adjusted, with sentinel-safe max num_records. - // CHECK: s_mov_b64 s[{{[0-9]+}}:{{[0-9]+}}], s[{{[0-9]+}}:{{[0-9]+}}] // CHECK: waveasm.s_mul_hi_i32 // CHECK: waveasm.s_mul_i32 // CHECK: waveasm.s_add_u32 // CHECK: waveasm.s_addc_u32 - // CHECK: s_mov_b32 s{{[0-9]+}}, 0x7FFFFFF - // CHECK: s_mov_b32 s{{[0-9]+}}, 0x20000 + // CHECK: waveasm.pack // CHECK: waveasm.buffer_store_dword vector.store %ext, %buf1[%thread_id] : memref>, vector<1xf32> diff --git a/waveasm/test/Translate/salu-promotion-arith.mlir b/waveasm/test/Translate/salu-promotion-arith.mlir new file mode 100644 index 0000000000..7f611980e3 --- /dev/null +++ b/waveasm/test/Translate/salu-promotion-arith.mlir @@ -0,0 +1,42 @@ +// RUN: waveasm-translate %s 2>&1 | FileCheck %s +// +// Test: SALU promotion of scalar arithmetic. When both operands of an arith op +// are in SGPRs (e.g. workgroup_id), the auto-select emit helpers route through +// SALU instructions instead of VALU. + +module { + gpu.module @test_salu_promotion { + + // CHECK-LABEL: waveasm.program @salu_mul_add + gpu.func @salu_mul_add() kernel { + %wg_x = gpu.block_id x upper_bound 4 + %wg_y = gpu.block_id y upper_bound 4 + %c128 = arith.constant 128 : index + + // Scalar (SGPR) * immediate -> s_mul_i32 + // CHECK: waveasm.s_mul_i32 %{{.*}}, %{{.*}} : !waveasm.sreg, !waveasm.imm<128> -> !waveasm.sreg + %prod = arith.muli %wg_x, %c128 : index + + // Scalar (SGPR) + scalar (SGPR) -> s_add_u32 + // CHECK: waveasm.s_add_u32 %{{.*}}, %{{.*}} : !waveasm.sreg, !waveasm.sreg -> !waveasm.sreg, !waveasm.scc + %sum = arith.addi %prod, %wg_y : index + + // CHECK: waveasm.s_endpgm + gpu.return + } + + // CHECK-LABEL: waveasm.program @salu_cmpi + gpu.func @salu_cmpi() kernel { + %wg_x = gpu.block_id x upper_bound 16 + %c10 = arith.constant 10 : index + + // Scalar cmpi -> s_cmp (SCC result). The immediate is first moved to SGPR. + // CHECK: waveasm.s_mov_b32 %{{.*}} : !waveasm.imm<10> -> !waveasm.sreg + // CHECK: waveasm.s_cmp_lt_i32 %{{.*}}, %{{.*}} : !waveasm.sreg, !waveasm.sreg -> !waveasm.scc + %cmp = arith.cmpi slt, %wg_x, %c10 : index + + // CHECK: waveasm.s_endpgm + gpu.return + } + } +} diff --git a/waveasm/test/Translate/salu-select-fusion.mlir b/waveasm/test/Translate/salu-select-fusion.mlir new file mode 100644 index 0000000000..c945065ea7 --- /dev/null +++ b/waveasm/test/Translate/salu-select-fusion.mlir @@ -0,0 +1,30 @@ +// RUN: waveasm-translate %s 2>&1 | FileCheck %s +// +// Test: Scalar cmpi + scalar select fusion into s_cmp + s_cselect_b32. +// When both comparison operands are scalar and the select's true/false values +// are also scalar, the backend fuses the pair into a single s_cmp + s_cselect +// sequence, avoiding the VALU v_cmp + v_cndmask path. + +module { + gpu.module @test_select_fusion { + + // CHECK-LABEL: waveasm.program @cmpi_select_scalar_fusion + gpu.func @cmpi_select_scalar_fusion() kernel { + %wg_x = gpu.block_id x upper_bound 16 + %c10 = arith.constant 10 : index + %c100 = arith.constant 100 : index + %c200 = arith.constant 200 : index + + // Scalar cmpi + scalar select -> s_cmp_lt_i32 + s_cselect_b32 + // CHECK: waveasm.s_cmp_lt_i32 + // CHECK: waveasm.s_cselect_b32 + // CHECK-NOT: waveasm.v_cmp + // CHECK-NOT: waveasm.v_cndmask + %cmp = arith.cmpi slt, %wg_x, %c10 : index + %sel = arith.select %cmp, %c100, %c200 : index + + // CHECK: waveasm.s_endpgm + gpu.return + } + } +} diff --git a/waveasm/test/Translate/scf-if-agpr-else-coercion.mlir b/waveasm/test/Translate/scf-if-agpr-else-coercion.mlir index c255b86052..0c94496d15 100644 --- a/waveasm/test/Translate/scf-if-agpr-else-coercion.mlir +++ b/waveasm/test/Translate/scf-if-agpr-else-coercion.mlir @@ -9,14 +9,13 @@ // values (immediates). The backend must coerce the else-yield immediates // into register types so that both branches yield type-compatible values. // -// Note: the translator currently produces a vreg condition for waveasm.if -// (instead of scc), so the output is dumped in generic form after -// verification. The CHECK patterns below match generic form. +// With SALU promotion, scalar cmpi produces SCC directly, so waveasm.if +// gets an SCC condition and the output is in custom form. module { gpu.module @test_if_else_coercion { - // CHECK-LABEL: sym_name = "if_else_coercion" + // CHECK-LABEL: waveasm.program @if_else_coercion gpu.func @if_else_coercion() kernel { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -28,12 +27,13 @@ module { // Then branch: compute a VGPR value (addi -> v_add_u32) // Else branch: yield a constant zero (-> immediate coerced to vreg) // - // CHECK: "waveasm.if" - // CHECK: "waveasm.v_add_u32" - // CHECK: "waveasm.yield" - // CHECK: }, { - // CHECK: "waveasm.v_mov_b32" - // CHECK: "waveasm.yield" + // CHECK: waveasm.s_cmp_lt_i32 + // CHECK: waveasm.if %{{.*}} : !waveasm.scc -> !waveasm.vreg { + // CHECK: waveasm.v_add_u32 + // CHECK: waveasm.yield + // CHECK: } else { + // CHECK: waveasm.v_mov_b32 + // CHECK: waveasm.yield %result = scf.if %cond_i1 -> i32 { %val = arith.addi %zero_i32, %one_i32 : i32 scf.yield %val : i32 @@ -41,7 +41,7 @@ module { scf.yield %zero_i32 : i32 } - // CHECK: "waveasm.s_endpgm" + // CHECK: waveasm.s_endpgm gpu.return } } diff --git a/waveasm/test/Translate/swizzle-srd-num-records.mlir b/waveasm/test/Translate/swizzle-srd-num-records.mlir index 68473f6706..fa8590faae 100644 --- a/waveasm/test/Translate/swizzle-srd-num-records.mlir +++ b/waveasm/test/Translate/swizzle-srd-num-records.mlir @@ -24,18 +24,17 @@ module { : memref<2048xi8, 3> to memref<64x8xf16, 3> // fat_raw_buffer_cast with cacheSwizzleStride, no reinterpret_cast. - // The swizzle SRD must copy num_records from the source SRD (s_mov_b32 - // sN, sM) rather than using 0xFFFFFFFF which would make sentinel offsets - // in-bounds. + // The swizzle SRD must copy num_records from the source SRD. // // Swizzle SRD construction sequence: // s_and_b32 + s_or_b32 (set swizzle bits in base hi) - // s_mov_b32 (stride word) - // num_records is inherited from the source SRD (no separate copy needed) + // s_mov_b32 (copy num_records from source SRD word 2) + // s_mov_b32 (stride/swizzle flags) + // pack into 4-wide SRD // - // CHECK: s_and_b32 - // CHECK: s_or_b32 - // CHECK: s_mov_b32 s{{[0-9]+}}, 0x27000 + // CHECK: waveasm.s_and_b32 + // CHECK: waveasm.s_or_b32 + // CHECK: waveasm.pack %buf = amdgpu.fat_raw_buffer_cast %src validBytes(%valid) cacheSwizzleStride(%stride) resetOffset : memref From f4f6b8b27e0c80812b2eebb6703a06a2126f48f2 Mon Sep 17 00:00:00 2001 From: Sanket Pandit Date: Mon, 6 Apr 2026 12:22:24 -0500 Subject: [PATCH 5/5] pre-commit fixes Signed-off-by: Sanket Pandit Made-with: Cursor --- .../waveasm/Transforms/TranslateFromMLIR.h | 2 - waveasm/lib/Transforms/AssemblyEmitter.cpp | 4 +- waveasm/lib/Transforms/MetadataEmitter.cpp | 5 ++- waveasm/lib/Transforms/TranslateFromMLIR.cpp | 23 +++++------- .../Transforms/handlers/AMDGPUHandlers.cpp | 6 +-- .../Transforms/handlers/AffineHandlers.cpp | 37 ++++++++----------- .../lib/Transforms/handlers/ArithHandlers.cpp | 19 +++++----- waveasm/lib/Transforms/handlers/Handlers.h | 30 +++++++++------ .../Transforms/handlers/MemRefHandlers.cpp | 3 +- 9 files changed, 63 insertions(+), 66 deletions(-) diff --git a/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h b/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h index 84ac089da0..3f636f9e10 100644 --- a/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h +++ b/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h @@ -377,8 +377,6 @@ class TranslationContext { /// Get next available SRD index int64_t getNextSRDIndex(); - - /// Update buffer size for a pending SRD (called when we see reinterpret_cast) void updateSRDBufferSize(mlir::Value memref, int64_t bufferSize); diff --git a/waveasm/lib/Transforms/AssemblyEmitter.cpp b/waveasm/lib/Transforms/AssemblyEmitter.cpp index a53fc0b218..617b6c3eee 100644 --- a/waveasm/lib/Transforms/AssemblyEmitter.cpp +++ b/waveasm/lib/Transforms/AssemblyEmitter.cpp @@ -358,8 +358,8 @@ KernelGenerator::emitScaledMFMA(Operation *scaledOp, llvm::StringRef mnemonic) { std::optional KernelGenerator::generateOp(Operation *op) { return llvm::TypeSwitch>(op) .Case([](auto) { return std::nullopt; }) + PrecoloredARegOp, ConstantOp, PackOp, ExtractOp, DCEProtectOp>( + [](auto) { return std::nullopt; }) .Case([&](RawOp rawOp) -> std::optional { return generateRaw(rawOp); }) diff --git a/waveasm/lib/Transforms/MetadataEmitter.cpp b/waveasm/lib/Transforms/MetadataEmitter.cpp index e3c97b65d6..db3ff8de72 100644 --- a/waveasm/lib/Transforms/MetadataEmitter.cpp +++ b/waveasm/lib/Transforms/MetadataEmitter.cpp @@ -148,6 +148,7 @@ static void scanSystemRegisterUsage(ProgramOp program, bool &usesWorkgroupIdX, int64_t userSgprCount = 2; if (isGfx950) { + // Hardware limits user SGPRs to 16 on gfx950. userSgprCount = std::min(int64_t(16), 2 + numArgs * 2); } @@ -205,15 +206,17 @@ MetadataEmitter::emitKernelDescriptor(int64_t peakVGPRs, int64_t peakSGPRs, auto targetAttr = program.getTarget(); auto targetKind = targetAttr.getTargetKind(); + int64_t preloadLength = program.getKernargPreloadLength(); bool usePreloading = llvm::isa(targetKind); - int64_t preloadLength = program.getKernargPreloadLength(); if (usePreloading && preloadLength == 0) { int64_t numArgs = 2; if (auto numArgsAttr = program->getAttrOfType("num_kernel_args")) { numArgs = numArgsAttr.getInt(); } + // Hardware limits user SGPRs to 16 on gfx950 (2 for kernarg ptr + 14 max + // preloaded). Overflow args are loaded via explicit s_load in the prologue. preloadLength = std::min(int64_t(14), numArgs * 2); } diff --git a/waveasm/lib/Transforms/TranslateFromMLIR.cpp b/waveasm/lib/Transforms/TranslateFromMLIR.cpp index 5b0b1bcf07..2762a3691b 100644 --- a/waveasm/lib/Transforms/TranslateFromMLIR.cpp +++ b/waveasm/lib/Transforms/TranslateFromMLIR.cpp @@ -250,8 +250,7 @@ void TranslationContext::emitSRDPrologue() { auto offsetImm = builder.getType(kernargOffset); auto offsetConst = ConstantOp::create(builder, loc, offsetImm, kernargOffset); - auto loadOp = S_LOAD_DWORDX2::create(builder, loc, - TypeRange{loadDstType}, + auto loadOp = S_LOAD_DWORDX2::create(builder, loc, TypeRange{loadDstType}, kernargBase, offsetConst); argLoadResults[pending.argIndex] = loadOp->getResult(0); } @@ -267,8 +266,7 @@ void TranslationContext::emitSRDPrologue() { auto offsetImm = builder.getType(kernargOffset); auto offsetConst = ConstantOp::create(builder, loc, offsetImm, kernargOffset); - auto loadOp = S_LOAD_DWORDX2::create(builder, loc, - TypeRange{loadDstType}, + auto loadOp = S_LOAD_DWORDX2::create(builder, loc, TypeRange{loadDstType}, kernargBase, offsetConst); argLoadResults[pending.argIndex] = loadOp->getResult(0); } @@ -305,8 +303,7 @@ void TranslationContext::emitSRDPrologue() { auto offsetImm = builder.getType(kernargOffset); auto offsetConst = ConstantOp::create(builder, loc, offsetImm, kernargOffset); - auto loadOp = S_LOAD_DWORDX2::create(builder, loc, - TypeRange{loadDstType}, + auto loadOp = S_LOAD_DWORDX2::create(builder, loc, TypeRange{loadDstType}, kernargBase, offsetConst); overflowLoadResults[overflowIdx] = loadOp->getResult(0); overflowIdx++; @@ -422,8 +419,7 @@ void TranslationContext::emitSRDPrologue() { auto offsetImm = builder.getType(kernargOffset); auto offsetConst = ConstantOp::create(builder, loc, offsetImm, kernargOffset); - auto loadOp = S_LOAD_DWORDX2::create(builder, loc, - TypeRange{loadDstType}, + auto loadOp = S_LOAD_DWORDX2::create(builder, loc, TypeRange{loadDstType}, kernargBase, offsetConst); srdLoadResults[srdBase] = loadOp->getResult(0); } @@ -437,8 +433,7 @@ void TranslationContext::emitSRDPrologue() { auto offsetImm = builder.getType(kernargOffset); auto offsetConst = ConstantOp::create(builder, loc, offsetImm, kernargOffset); - auto loadOp = S_LOAD_DWORDX2::create(builder, loc, - TypeRange{loadDstType}, + auto loadOp = S_LOAD_DWORDX2::create(builder, loc, TypeRange{loadDstType}, kernargBase, offsetConst); scalarLoadResults[(int64_t)i] = loadOp->getResult(0); } @@ -459,8 +454,8 @@ void TranslationContext::emitSRDPrologue() { auto loadIt = srdLoadResults.find(srdBase); if (loadIt != srdLoadResults.end()) { auto dstB64Type = PSRegType::get(builder.getContext(), srdBase, 2); - auto movB64 = S_MOV_B64::create(builder, loc, dstB64Type, - loadIt->second); + auto movB64 = + S_MOV_B64::create(builder, loc, dstB64Type, loadIt->second); DCEProtectOp::create(builder, loc, movB64); } @@ -495,8 +490,8 @@ void TranslationContext::emitSRDPrologue() { auto loadIt = scalarLoadResults.find((int64_t)i); if (loadIt != scalarLoadResults.end()) { auto sregTy = createSRegType(1, 1); - Value lowWord = ExtractOp::create(builder, loc, sregTy, - loadIt->second, 0); + Value lowWord = + ExtractOp::create(builder, loc, sregTy, loadIt->second, 0); auto dstPsType = PSRegType::get(builder.getContext(), sgprIdx, 1); auto movOp = S_MOV_B32::create(builder, loc, dstPsType, lowWord); DCEProtectOp::create(builder, loc, movOp); diff --git a/waveasm/lib/Transforms/handlers/AMDGPUHandlers.cpp b/waveasm/lib/Transforms/handlers/AMDGPUHandlers.cpp index 4286d81c8e..0f4c601445 100644 --- a/waveasm/lib/Transforms/handlers/AMDGPUHandlers.cpp +++ b/waveasm/lib/Transforms/handlers/AMDGPUHandlers.cpp @@ -435,7 +435,7 @@ static void patchSrdWord2InPlace(OpBuilder &builder, Location loc, isScalarOrImm(*cmpLhsMapped) && isScalarOrImm(*cmpRhsMapped) && isScalarOrImm(*trueMapped) && isScalarOrImm(*falseMapped)) { Value sccVal = emitScalarCmp(builder, loc, cmpOp.getPredicate(), - *cmpLhsMapped, *cmpRhsMapped, ctx); + *cmpLhsMapped, *cmpRhsMapped, ctx); auto dstType = PSRegType::get(builder.getContext(), srdBase + 2, 1); Value trueV = *trueMapped; @@ -443,8 +443,8 @@ static void patchSrdWord2InPlace(OpBuilder &builder, Location loc, auto sregType = ctx.createSRegType(); if (isImmType(trueV.getType())) trueV = S_MOV_B32::create(builder, loc, sregType, trueV); - auto result = - S_CSELECT_B32::create(builder, loc, dstType, sccVal, trueV, falseV); + auto result = S_CSELECT_B32::create(builder, loc, dstType, sccVal, + trueV, falseV); DCEProtectOp::create(builder, loc, result); return; } diff --git a/waveasm/lib/Transforms/handlers/AffineHandlers.cpp b/waveasm/lib/Transforms/handlers/AffineHandlers.cpp index 5f61c04dc7..8662d84d7d 100644 --- a/waveasm/lib/Transforms/handlers/AffineHandlers.cpp +++ b/waveasm/lib/Transforms/handlers/AffineHandlers.cpp @@ -422,7 +422,8 @@ static Value emitCeilFromFloorQuotient(Value q, Value x, Value d, // s_addc_u32: dst = q + 0 + SCC (carry-in from SCC) auto sregType = ctx.createSRegType(); auto sccType = ctx.createSCCType(); - return S_ADDC_U32::create(builder, loc, sregType, sccType, sccVal, q, zeroConst) + return S_ADDC_U32::create(builder, loc, sregType, sccType, sccVal, q, + zeroConst) .getDst(); } auto vregType = ctx.createVRegType(); @@ -629,12 +630,11 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { auto shiftConst = ConstantOp::create(builder, loc, shiftImm, shiftAmount); // v_lshl_or_b32: dst = (src << shift) | orend - Value base = ensureVGPR(builder, loc, ctx, - baseResult.value); + Value base = + ensureVGPR(builder, loc, ctx, baseResult.value); orend = ensureVGPR(builder, loc, ctx, orend); Value fusedResult = V_LSHL_OR_B32::create( - builder, loc, vregType, base, shiftConst, - orend); + builder, loc, vregType, base, shiftConst, orend); BitRange shiftedRange = baseResult.range.shiftLeft(shiftAmount); BitRange resultRange = shiftedRange.merge(orendRange); @@ -652,12 +652,11 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { auto shiftImm = ctx.createImmType(shiftAmount); auto shiftConst = ConstantOp::create(builder, loc, shiftImm, shiftAmount); - Value base2 = ensureVGPR(builder, loc, ctx, - baseResult.value); + Value base2 = + ensureVGPR(builder, loc, ctx, baseResult.value); orend = ensureVGPR(builder, loc, ctx, orend); Value fusedResult = V_LSHL_OR_B32::create( - builder, loc, vregType, base2, shiftConst, - orend); + builder, loc, vregType, base2, shiftConst, orend); BitRange shiftedRange = baseResult.range.shiftLeft(shiftAmount); BitRange resultRange = shiftedRange.merge(orendRange); @@ -751,8 +750,7 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { auto shiftAmt = ctx.createImmType(shiftAmount); auto shiftConst = ConstantOp::create(builder, loc, shiftAmt, shiftAmount); - Value shiftResult = - emitLshl(lhs, shiftConst, builder, loc, ctx); + Value shiftResult = emitLshl(lhs, shiftConst, builder, loc, ctx); BitRange resultRange = lhsRange.shiftLeft(shiftAmount); ctx.setBitRange(shiftResult, resultRange); return ExprResult(shiftResult, resultRange); @@ -765,8 +763,7 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { auto shiftAmt = ctx.createImmType(shiftAmount); auto shiftConst = ConstantOp::create(builder, loc, shiftAmt, shiftAmount); - Value shiftResult = - emitLshl(rhs, shiftConst, builder, loc, ctx); + Value shiftResult = emitLshl(rhs, shiftConst, builder, loc, ctx); BitRange resultRange = rhsRange.shiftLeft(shiftAmount); ctx.setBitRange(shiftResult, resultRange); return ExprResult(shiftResult, resultRange); @@ -799,8 +796,7 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { auto shiftAmt = ctx.createImmType(shiftAmount); auto shiftConst = ConstantOp::create(builder, loc, shiftAmt, shiftAmount); - Value shiftResult = - emitLshr(lhs, shiftConst, builder, loc, ctx); + Value shiftResult = emitLshr(lhs, shiftConst, builder, loc, ctx); BitRange resultRange = lhsRange.shiftRight(shiftAmount); ctx.setBitRange(shiftResult, resultRange); return ExprResult(shiftResult, resultRange); @@ -835,14 +831,13 @@ LogicalResult handleAffineApply(Operation *op, TranslationContext &ctx) { Value rem = emitAnd(lhs, maskConst, builder, loc, ctx); Value zeroConst = createImmConst(0, builder, loc, ctx); if (isScalarOrImm(rem)) { - Value sccVal = S_CMP_NE_U32::create(builder, loc, - ctx.createSCCType(), rem, zeroConst); + Value sccVal = S_CMP_NE_U32::create( + builder, loc, ctx.createSCCType(), rem, zeroConst); auto sregType = ctx.createSRegType(); auto sccType = ctx.createSCCType(); - Value result = - S_ADDC_U32::create(builder, loc, sregType, sccType, sccVal, q, - zeroConst) - .getDst(); + Value result = S_ADDC_U32::create(builder, loc, sregType, sccType, + sccVal, q, zeroConst) + .getDst(); return ExprResult(result, BitRange()); } V_CMP_NE_U32::create(builder, loc, rem, zeroConst); diff --git a/waveasm/lib/Transforms/handlers/ArithHandlers.cpp b/waveasm/lib/Transforms/handlers/ArithHandlers.cpp index 0106ca98dc..c4f4126605 100644 --- a/waveasm/lib/Transforms/handlers/ArithHandlers.cpp +++ b/waveasm/lib/Transforms/handlers/ArithHandlers.cpp @@ -537,8 +537,8 @@ LogicalResult handleArithSelect(Operation *op, TranslationContext &ctx) { if (auto cmpOp = condMLIR.getDefiningOp()) { auto cmpLhs = ctx.getMapper().getMapped(cmpOp.getLhs()); auto cmpRhs = ctx.getMapper().getMapped(cmpOp.getRhs()); - if (cmpLhs && cmpRhs && - isScalarOrImm(*cmpLhs) && isScalarOrImm(*cmpRhs)) { + if (cmpLhs && cmpRhs && isScalarOrImm(*cmpLhs) && + isScalarOrImm(*cmpRhs)) { auto sregType = ctx.createSRegType(); Value lhsOp = *cmpLhs; Value rhsOp = *cmpRhs; @@ -546,14 +546,14 @@ LogicalResult handleArithSelect(Operation *op, TranslationContext &ctx) { lhsOp = S_MOV_B32::create(builder, loc, sregType, lhsOp); if (isImmType(rhsOp.getType())) rhsOp = S_MOV_B32::create(builder, loc, sregType, rhsOp); - Value sccVal = - emitScalarCmp(builder, loc, cmpOp.getPredicate(), lhsOp, rhsOp, ctx); + Value sccVal = emitScalarCmp(builder, loc, cmpOp.getPredicate(), lhsOp, + rhsOp, ctx); Value trueV = *trueVal; Value falseV = *falseVal; if (isImmType(trueV.getType())) trueV = S_MOV_B32::create(builder, loc, sregType, trueV); - auto result = - S_CSELECT_B32::create(builder, loc, sregType, sccVal, trueV, falseV); + auto result = S_CSELECT_B32::create(builder, loc, sregType, sccVal, + trueV, falseV); ctx.getMapper().mapValue(selectOp.getResult(), result); return success(); } @@ -568,14 +568,15 @@ LogicalResult handleArithSelect(Operation *op, TranslationContext &ctx) { Value condV = *cond; if (isImmType(condV.getType())) condV = S_MOV_B32::create(builder, loc, ctx.createSRegType(), condV); - Value sccVal = - S_CMP_NE_U32::create(builder, loc, ctx.createSCCType(), condV, zeroConst); + Value sccVal = S_CMP_NE_U32::create(builder, loc, ctx.createSCCType(), + condV, zeroConst); auto sregType = ctx.createSRegType(); Value trueV = *trueVal; Value falseV = *falseVal; if (isImmType(trueV.getType())) trueV = S_MOV_B32::create(builder, loc, sregType, trueV); - auto result = S_CSELECT_B32::create(builder, loc, sregType, sccVal, trueV, falseV); + auto result = + S_CSELECT_B32::create(builder, loc, sregType, sccVal, trueV, falseV); ctx.getMapper().mapValue(selectOp.getResult(), result); return success(); } diff --git a/waveasm/lib/Transforms/handlers/Handlers.h b/waveasm/lib/Transforms/handlers/Handlers.h index 117d5ac46e..d869d4b33f 100644 --- a/waveasm/lib/Transforms/handlers/Handlers.h +++ b/waveasm/lib/Transforms/handlers/Handlers.h @@ -308,8 +308,9 @@ inline mlir::Value ensureVGPR(mlir::OpBuilder &builder, mlir::Location loc, /// Emit add: S_ADD_U32 when both operands are scalar, V_ADD_U32 otherwise. /// Commutative: swaps to put immediate in src1 (SALU src0 must be SGPR). -inline mlir::Value emitAdd(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, - mlir::Location loc, TranslationContext &ctx) { +inline mlir::Value emitAdd(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { if (isScalarOrImm(a) && isScalarOrImm(b) && !(isImmType(a.getType()) && isImmType(b.getType()))) { if (isImmType(a.getType())) @@ -324,8 +325,9 @@ inline mlir::Value emitAdd(mlir::Value a, mlir::Value b, mlir::OpBuilder &builde /// Emit sub: S_SUB_U32 when both operands are scalar, V_SUB_U32 otherwise. /// Not commutative: src0 (minuend) must be SGPR. -inline mlir::Value emitSub(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, - mlir::Location loc, TranslationContext &ctx) { +inline mlir::Value emitSub(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { if (isScalarOrImm(a) && isScalarOrImm(b) && isSGPRType(a.getType())) { auto sregType = ctx.createSRegType(); auto sccType = ctx.createSCCType(); @@ -337,8 +339,9 @@ inline mlir::Value emitSub(mlir::Value a, mlir::Value b, mlir::OpBuilder &builde /// Emit mul: S_MUL_I32 when both operands are scalar, V_MUL_LO_U32 otherwise. /// Commutative: swaps to put immediate in src1. -inline mlir::Value emitMul(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, - mlir::Location loc, TranslationContext &ctx) { +inline mlir::Value emitMul(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { if (isScalarOrImm(a) && isScalarOrImm(b) && !(isImmType(a.getType()) && isImmType(b.getType()))) { if (isImmType(a.getType())) @@ -386,8 +389,9 @@ inline mlir::Value emitLshl(mlir::Value value, mlir::Value shiftAmt, /// Emit bitwise AND: S_AND_B32 when both scalar, V_AND_B32 otherwise. /// Commutative: swaps to put immediate in src1. -inline mlir::Value emitAnd(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, - mlir::Location loc, TranslationContext &ctx) { +inline mlir::Value emitAnd(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { if (isScalarOrImm(a) && isScalarOrImm(b) && !(isImmType(a.getType()) && isImmType(b.getType()))) { if (isImmType(a.getType())) @@ -402,8 +406,9 @@ inline mlir::Value emitAnd(mlir::Value a, mlir::Value b, mlir::OpBuilder &builde /// Emit bitwise OR: S_OR_B32 when both scalar, V_OR_B32 otherwise. /// Commutative: swaps to put immediate in src1. -inline mlir::Value emitOr(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, - mlir::Location loc, TranslationContext &ctx) { +inline mlir::Value emitOr(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { if (isScalarOrImm(a) && isScalarOrImm(b) && !(isImmType(a.getType()) && isImmType(b.getType()))) { if (isImmType(a.getType())) @@ -418,8 +423,9 @@ inline mlir::Value emitOr(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder /// Emit bitwise XOR: S_XOR_B32 when both scalar, V_XOR_B32 otherwise. /// Commutative: swaps to put immediate in src1. -inline mlir::Value emitXor(mlir::Value a, mlir::Value b, mlir::OpBuilder &builder, - mlir::Location loc, TranslationContext &ctx) { +inline mlir::Value emitXor(mlir::Value a, mlir::Value b, + mlir::OpBuilder &builder, mlir::Location loc, + TranslationContext &ctx) { if (isScalarOrImm(a) && isScalarOrImm(b) && !(isImmType(a.getType()) && isImmType(b.getType()))) { if (isImmType(a.getType())) diff --git a/waveasm/lib/Transforms/handlers/MemRefHandlers.cpp b/waveasm/lib/Transforms/handlers/MemRefHandlers.cpp index b73c7c5c93..3f4d941c10 100644 --- a/waveasm/lib/Transforms/handlers/MemRefHandlers.cpp +++ b/waveasm/lib/Transforms/handlers/MemRefHandlers.cpp @@ -305,8 +305,7 @@ LogicalResult handleMemRefStore(Operation *op, TranslationContext &ctx) { Value storeData = *data; if (isAGPRType(storeData.getType())) { auto vregType = ctx.createVRegType(); - storeData = - V_ACCVGPR_READ_B32::create(builder, loc, vregType, storeData); + storeData = V_ACCVGPR_READ_B32::create(builder, loc, vregType, storeData); } BUFFER_STORE_DWORD::create(builder, loc, storeData, srd, voffset,