diff --git a/.github/workflows/ci-waveasm-mi2xx.yaml b/.github/workflows/ci-waveasm-mi2xx.yaml new file mode 100644 index 0000000000..84612a4824 --- /dev/null +++ b/.github/workflows/ci-waveasm-mi2xx.yaml @@ -0,0 +1,119 @@ +name: "WaveASM MI2xx CI" + +on: + workflow_dispatch: + pull_request: + types: [opened, synchronize, ready, ready_for_review, converted_to_draft] + push: + branches: + - main + +concurrency: + group: ${{ github.workflow }}-${{ github.event.number || github.sha }} + cancel-in-progress: true + +env: + LLVM_SHA_FILE: llvm-sha.txt + LLVM_CACHE_NUMBER: 2 + +jobs: + test_linux_waveasm_mi2xx: + name: "MI2xx/gfx90a :: SHARED_LIBS ${{ matrix.shared_libs }} :: Run waveasm tests" + strategy: + fail-fast: false + matrix: + version: [3.11] + shared_libs: ["ON", "OFF"] + runs-on: nodai-amdgpu-mi250-x86-64 + timeout-minutes: 240 + if: github.event_name != 'pull_request' || github.event.pull_request.draft == false + + steps: + - name: Checkout repo + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + fetch-depth: 0 + + - name: Setup Cache Vars + run: | + echo "LLVM_SHA=$(cat $GITHUB_WORKSPACE/water/$LLVM_SHA_FILE)" >> $GITHUB_ENV + + - name: Cache LLVM-MLIR + id: cache-llvm-mlir + uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5.0.4 + with: + path: llvm-mlir/_mlir_install/** + key: ${{ runner.os }}-mi2xx-build-llvm-${{ env.LLVM_CACHE_NUMBER }}-${{ env.LLVM_SHA }} + + - name: Setup env + run: | + sudo apt-get update + sudo apt-get install -y ninja-build cmake clang lld dwarfdump + + - name: "Setting up Python" + id: setup_python + uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: ${{ matrix.version }} + pip-install: -r water/requirements-dev.txt + + - name: Checkout LLVM + if: steps.cache-llvm-mlir.outputs.cache-hit != 'true' + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + repository: llvm/llvm-project + ref: ${{ env.LLVM_SHA }} + path: llvm-mlir/llvm-project + + - name: Build LLVM-MLIR + if: steps.cache-llvm-mlir.outputs.cache-hit != 'true' + run: | + pushd ${GITHUB_WORKSPACE}/llvm-mlir + echo "INFO: Need to rebuild LLVM-MLIR. Previous installation for MLIR not found" + np=`nproc` + echo "INFO: nproc $np" + mkdir _build + cd _build + export CC=clang + export CXX=clang++ + cmake ../llvm-project/llvm \ + -GNinja \ + -DCMAKE_BUILD_TYPE=Release \ + -DLLVM_ENABLE_PROJECTS="mlir;llvm;lld;clang" \ + -DLLVM_ENABLE_ASSERTIONS=ON \ + -DLLVM_INSTALL_UTILS=ON \ + -DLLVM_TARGETS_TO_BUILD="X86;AMDGPU" \ + -DLLVM_ENABLE_BINDINGS=OFF \ + -DLLVM_ENABLE_ZSTD=OFF \ + -DMLIR_INCLUDE_TESTS=OFF \ + -DLLVM_USE_LINKER=lld \ + -DLLVM_DISTRIBUTION_COMPONENTS="llvm-headers;llvm-libraries;cmake-exports;FileCheck;count;not;mlir-headers;mlir-libraries;mlir-cmake-exports;mlir-tblgen;mlir-python-sources;lld;clang;clang-resource-headers" \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ + -DCMAKE_INSTALL_PREFIX=${GITHUB_WORKSPACE}/llvm-mlir/_mlir_install + echo "INFO: working around a missing dependency on stubgen" + ninja MLIRPythonModules.extension._mlir.dso._mlir.type_stubs + ninja install-distribution-stripped + popd + + - name: Build waveasm + run: | + export EXTERNAL_LIT=${GITHUB_WORKSPACE}/water/scripts/runlit.py + export LLVM_DIR=${GITHUB_WORKSPACE}/llvm-mlir/_mlir_install + mkdir -p cmake_build_waveasm + cd cmake_build_waveasm + export CC=clang + export CXX=clang++ + cmake ${GITHUB_WORKSPACE}/waveasm \ + -GNinja \ + -DCMAKE_BUILD_TYPE=Release \ + -DLLVM_DIR=${LLVM_DIR}/lib/cmake/llvm \ + -DMLIR_DIR=${LLVM_DIR}/lib/cmake/mlir \ + -DBUILD_SHARED_LIBS=${{ matrix.shared_libs }} \ + -DLLVM_EXTERNAL_LIT=${EXTERNAL_LIT} + cmake --build . + + - name: Test waveasm + if: ${{ matrix.shared_libs == 'OFF' }} + run: | + cd cmake_build_waveasm + cmake --build . --target check-waveasm diff --git a/wave_lang/support/detect_waveasm.py b/wave_lang/support/detect_waveasm.py index 5de9fa6ebb..09d6167386 100644 --- a/wave_lang/support/detect_waveasm.py +++ b/wave_lang/support/detect_waveasm.py @@ -21,11 +21,23 @@ def get_waveasm_pkg_path() -> Path: def find_binary(name: str) -> str | None: """Returns the path to the waveasm binary with the given name.""" + waveasm_dir = os.getenv("WAVE_WAVEASM_DIR") + if waveasm_dir: + tool_path = Path(waveasm_dir) / "bin" / name + if tool_path.is_file() and os.access(tool_path, os.X_OK): + return str(tool_path) + tool_path = get_waveasm_pkg_path() / "bin" / name - if not tool_path.is_file() or not os.access(tool_path, os.X_OK): - return None + if tool_path.is_file() and os.access(tool_path, os.X_OK): + return str(tool_path) + + repo_tool_path = ( + Path(__file__).parent.parent.parent / "waveasm" / "build" / "bin" / name + ) + if repo_tool_path.is_file() and os.access(repo_tool_path, os.X_OK): + return str(repo_tool_path) - return str(tool_path) + return None @lru_cache diff --git a/waveasm/include/waveasm/Dialect/WaveASMAttrs.h b/waveasm/include/waveasm/Dialect/WaveASMAttrs.h index 68698010df..513ce49c05 100644 --- a/waveasm/include/waveasm/Dialect/WaveASMAttrs.h +++ b/waveasm/include/waveasm/Dialect/WaveASMAttrs.h @@ -19,17 +19,18 @@ namespace waveasm { enum class TargetFeature : uint32_t { None = 0, - HasMFMA = 1 << 0, // Matrix fused multiply-add - HasFP8 = 1 << 1, // FP8 support - HasPackedFP32 = 1 << 2, // Packed FP32 operations - HasWave32 = 1 << 3, // Wave32 mode support - HasWave64 = 1 << 4, // Wave64 mode support - HasXF32 = 1 << 5, // Extended FP32 (TF32) - HasScaledMFMA = 1 << 6, // Scaled MFMA instructions - HasAtomicFAdd = 1 << 7, // Atomic float add - HasGlobalLoadLDS = 1 << 8, // Global load to LDS - HasFlatScratch = 1 << 9, // Flat scratch support - HasAGPRs = 1 << 10, // Accumulator GPRs + HasMFMA = 1 << 0, // Matrix fused multiply-add + HasFP8 = 1 << 1, // FP8 support + HasPackedFP32 = 1 << 2, // Packed FP32 operations + HasWave32 = 1 << 3, // Wave32 mode support + HasWave64 = 1 << 4, // Wave64 mode support + HasXF32 = 1 << 5, // Extended FP32 (TF32) + HasScaledMFMA = 1 << 6, // Scaled MFMA instructions + HasAtomicFAdd = 1 << 7, // Atomic float add + HasGlobalLoadLDS = 1 << 8, // Global load to LDS + HasFlatScratch = 1 << 9, // Flat scratch support + HasAGPRs = 1 << 10, // Accumulator GPRs + HasKernargPreload = 1 << 11, // Kernel argument preload SGPRs }; inline TargetFeature operator|(TargetFeature a, TargetFeature b) { diff --git a/waveasm/include/waveasm/Dialect/WaveASMAttrs.td b/waveasm/include/waveasm/Dialect/WaveASMAttrs.td index 64c0685349..c0e90b1303 100644 --- a/waveasm/include/waveasm/Dialect/WaveASMAttrs.td +++ b/waveasm/include/waveasm/Dialect/WaveASMAttrs.td @@ -80,9 +80,11 @@ def WaveASM_TargetKind : I32EnumAttr<"TargetKind", "Supported GPU targets", [I32EnumAttrCase<"GFX942", 0, "gfx942">, I32EnumAttrCase<"GFX950", 1, "gfx950">, - I32EnumAttrCase<"GFX1250", 2, "gfx1250">]> { + I32EnumAttrCase<"GFX1250", 2, "gfx1250">, + I32EnumAttrCase<"GFX90A", 3, "gfx90a">]> { let summary = "Supported GPU targets"; let description = [{ + - gfx90a: AMD CDNA2 (MI200 series) - gfx942: AMD CDNA3 (MI300 series) - gfx950: AMD CDNA3+ (future MI series) - gfx1250: AMD RDNA4 @@ -195,6 +197,83 @@ class WaveASM_TargetKindAttr { + let archGeneration = "GFX9"; + let computeArch = "CDNA2"; + let maxVGPRs = 256; + let maxSGPRs = 106; + let maxAGPRs = 256; + let defaultWaveSize = 64; + let supportedWaveSizes = [64]; + let maxLDSSize = 65536; + let LDSBankCount = 32; + let LDSBankWidth = 4; + let globalLoadLatency = 100; + let LDSLoadLatency = 20; + let maxVmcnt = 63; + let maxLgkmcnt = 15; + let maxExpcnt = 7; + let defaultCodeObjectVersion = 5; + let supportedCodeObjectVersions = [4, 5]; + let targetDirective = ".amdgcn_target \\\"amdgcn-amd-amdhsa--gfx90a\\\""; + let ABIVersion = "amdhsa"; + + let extraClassDeclaration = [{ + TargetFeature getFeatures() const { + return TargetFeature::HasMFMA | TargetFeature::HasWave64 | + TargetFeature::HasAtomicFAdd | TargetFeature::HasFlatScratch | + TargetFeature::HasAGPRs; + } + + int64_t getMFMALatency(::llvm::StringRef instrName) const { + // Approximate CDNA2 MFMA latencies. + if (instrName.contains("f32_32x32")) + return 64; + if (instrName.contains("f32_16x16")) + return 32; + if (instrName.contains("f16_32x32")) + return 64; + if (instrName.contains("f16_16x16")) + return 32; + if (instrName.contains("bf16")) + return 32; + return 16; + } + + bool supportsInstruction(::llvm::StringRef instrName) const { + // CDNA2 does not support FP8/BF8 MFMA, scaled MFMA/MXFP, or XF32. + if (instrName.contains("fp8") || instrName.contains("bf8") || + instrName.contains("f8") || instrName.contains("mxfp") || + instrName.contains("scale") || instrName.contains("xf32")) + return false; + // The 16x16x32 F16/BF16 variants are gfx950+. + if (instrName.contains("16x16x32")) + return false; + // Global-to-LDS gather instructions are not available on gfx90a. + if (instrName.contains("buffer_load") && instrName.contains("_lds")) + return false; + return true; + } + + std::optional getTargetInstructionName(::llvm::StringRef genericName) const { + if (genericName.starts_with("v_mfma_")) { + // gfx90a's assembler spells the element type as part of the final + // shape token, e.g. v_mfma_f32_16x16x16f16. + for (::llvm::StringRef suffix : + {"_f16", "_bf16", "_i8", "_f32", "_f64"}) { + if (!genericName.ends_with(suffix)) + continue; + std::string targetName = genericName.str(); + targetName.erase(targetName.size() - suffix.size(), 1); + return targetName; + } + } + return std::nullopt; + } + }]#commonClassDeclaration; +} + def WaveASM_TargetKindAttr_GFX942 : WaveASM_TargetKindAttr<"GFX942Target", "gfx942"> { let archGeneration = "GFX9"; @@ -284,7 +363,8 @@ def WaveASM_TargetKindAttr_GFX950 return TargetFeature::HasMFMA | TargetFeature::HasFP8 | TargetFeature::HasWave64 | TargetFeature::HasAtomicFAdd | TargetFeature::HasFlatScratch | TargetFeature::HasAGPRs | - TargetFeature::HasScaledMFMA | TargetFeature::HasXF32; + TargetFeature::HasScaledMFMA | TargetFeature::HasXF32 | + TargetFeature::HasKernargPreload; } int64_t getMFMALatency(llvm::StringRef instrName) const { diff --git a/waveasm/include/waveasm/Dialect/WaveASMOps.td b/waveasm/include/waveasm/Dialect/WaveASMOps.td index c8bac3147b..2dfa0b6186 100644 --- a/waveasm/include/waveasm/Dialect/WaveASMOps.td +++ b/waveasm/include/waveasm/Dialect/WaveASMOps.td @@ -291,6 +291,7 @@ def WaveASM_ProgramOp : WAVEASMOp<"program", [ ); let regions = (region SizedRegion<1>:$body); + let hasVerifier = 1; let assemblyFormat = [{ $sym_name diff --git a/waveasm/include/waveasm/Transforms/AssemblyEmitter.h b/waveasm/include/waveasm/Transforms/AssemblyEmitter.h index 0d21830cc3..7c4f32108e 100644 --- a/waveasm/include/waveasm/Transforms/AssemblyEmitter.h +++ b/waveasm/include/waveasm/Transforms/AssemblyEmitter.h @@ -168,6 +168,9 @@ class KernelGenerator { /// Generate code for a raw op std::string generateRaw(RawOp rawOp); + /// Return the target-specific assembly mnemonic for a WaveASM op mnemonic. + std::string getTargetMnemonic(llvm::StringRef mnemonic); + //===--------------------------------------------------------------------===// // Helper methods for TypeSwitch-based code generation //===--------------------------------------------------------------------===// diff --git a/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h b/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h index cbfe19176e..3c3fd6d147 100644 --- a/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h +++ b/waveasm/include/waveasm/Transforms/TranslateFromMLIR.h @@ -202,7 +202,7 @@ class OpHandlerRegistry { /// Options for MLIR to waveasm translation struct TranslationOptions { - /// Target architecture (gfx942, gfx950, gfx1250) + /// Target architecture (gfx90a, gfx942, gfx950, gfx1250) std::string targetId = "gfx942"; /// Workgroup size (x, y, z). If any dimension is 0, use defaults. @@ -614,7 +614,7 @@ class TranslationContext { // Base: 2 SGPRs for kernarg_segment_ptr int64_t count = 2; // On gfx950+ with kernarg preloading, add preloaded args - if (llvm::isa(target)) { + if (target.hasFeature(TargetFeature::HasKernargPreload)) { // Each kernel arg uses 2 SGPRs, capped at 14 (hardware max 16 total). count += std::min(size_t(14), getNumKernelArgs() * 2); } diff --git a/waveasm/lib/Dialect/WaveASMAttrs.cpp b/waveasm/lib/Dialect/WaveASMAttrs.cpp index b3d4fd50c6..912f897f79 100644 --- a/waveasm/lib/Dialect/WaveASMAttrs.cpp +++ b/waveasm/lib/Dialect/WaveASMAttrs.cpp @@ -15,6 +15,8 @@ using namespace waveasm; TargetAttrInterface waveasm::getTargetKindAttr(mlir::MLIRContext *ctx, TargetKind targetKind) { switch (targetKind) { + case TargetKind::GFX90A: + return GFX90ATargetAttr::get(ctx); case TargetKind::GFX942: return GFX942TargetAttr::get(ctx); case TargetKind::GFX950: diff --git a/waveasm/lib/Dialect/WaveASMOps.cpp b/waveasm/lib/Dialect/WaveASMOps.cpp index 764ef0f58c..fffab7a499 100644 --- a/waveasm/lib/Dialect/WaveASMOps.cpp +++ b/waveasm/lib/Dialect/WaveASMOps.cpp @@ -25,6 +25,29 @@ using namespace waveasm; // Verification is handled by TableGen-generated code for basic structure. // Custom verification can be added here if needed. +LogicalResult ProgramOp::verify() { + TargetAttrInterface targetKind = getTarget().getTargetKind(); + WalkResult result = walk([&](Operation *op) -> WalkResult { + if (op == getOperation()) + return WalkResult::advance(); + + llvm::StringRef opName = op->getName().getStringRef(); + if (!opName.starts_with("waveasm.")) + return WalkResult::advance(); + + if (!targetKind.supportsInstruction(opName)) { + op->emitOpError() << "is not supported on target " + << targetKind.getComputeArch() << " (" + << targetKind.getTargetDirective() << ")"; + return WalkResult::interrupt(); + } + + return WalkResult::advance(); + }); + + return failure(result.wasInterrupted()); +} + //===----------------------------------------------------------------------===// // MFMA Operation Verifiers //===----------------------------------------------------------------------===// diff --git a/waveasm/lib/Transforms/AssemblyEmitter.cpp b/waveasm/lib/Transforms/AssemblyEmitter.cpp index d5ab9497fe..c9354f6cb5 100644 --- a/waveasm/lib/Transforms/AssemblyEmitter.cpp +++ b/waveasm/lib/Transforms/AssemblyEmitter.cpp @@ -40,6 +40,14 @@ KernelGenerator::KernelGenerator(ProgramOp program, TargetAttrInterface target) : program(program), mapping(mapping), target(target) {} +std::string KernelGenerator::getTargetMnemonic(llvm::StringRef mnemonic) { + if (std::optional targetName = + target.getTargetInstructionName(mnemonic)) { + return *targetName; + } + return mnemonic.str(); +} + std::string KernelGenerator::resolveValue(Value value) { Type ty = value.getType(); @@ -1113,13 +1121,18 @@ std::optional KernelGenerator::generateOp(Operation *op) { }) .Default([&](Operation *defaultOp) -> std::optional { + // Register alias and bookkeeping ops covered by NonEmittingOp trait. + if (defaultOp->hasTrait()) + return std::nullopt; + llvm::StringRef opName = defaultOp->getName().getStringRef(); llvm::StringRef mnemonic = opName; if (opName.starts_with("waveasm.")) mnemonic = opName.drop_front(8); + std::string targetMnemonic = getTargetMnemonic(mnemonic); - if (mnemonic.starts_with("v_cmp_")) { - std::string mnem64 = (mnemonic + "_e64").str(); + if (llvm::StringRef(targetMnemonic).starts_with("v_cmp_")) { + std::string mnem64 = targetMnemonic + "_e64"; llvm::SmallVector operands; operands.push_back("vcc"); for (Value operand : defaultOp->getOperands()) @@ -1127,7 +1140,7 @@ std::optional KernelGenerator::generateOp(Operation *op) { return formatter.format(mnem64, operands); } - return emitDefaultFormat(defaultOp, mnemonic); + return emitDefaultFormat(defaultOp, targetMnemonic); }); } diff --git a/waveasm/lib/Transforms/HazardMitigation.cpp b/waveasm/lib/Transforms/HazardMitigation.cpp index b91d9365fb..b4bdffe7a4 100644 --- a/waveasm/lib/Transforms/HazardMitigation.cpp +++ b/waveasm/lib/Transforms/HazardMitigation.cpp @@ -8,8 +8,8 @@ // Hazard Mitigation Pass - Insert s_nop instructions for hardware hazards // // This pass handles hardware-specific hazards that require NOP insertion: -// - VALU → v_readfirstlane hazard (gfx940+) -// - Trans → non-Trans VALU forwarding hazard (gfx940+) +// - VALU → v_readfirstlane hazard (gfx90a/gfx940+) +// - Trans → non-Trans VALU forwarding hazard (gfx90a/gfx940+) //===----------------------------------------------------------------------===// #include "waveasm/Dialect/WaveASMAttrs.h" @@ -136,8 +136,9 @@ bool hasVGPRConflict(Operation *producer, Operation *consumer) { /// Check if target requires VALU → readfirstlane hazard mitigation static bool needsVALUReadFirstLaneHazard(TargetAttrInterface target) { - // gfx940+ (CDNA3/4) architectures need this hazard mitigation - return isa(target); + // gfx90a and gfx940+ architectures need this hazard mitigation. + return isa(target); } //===----------------------------------------------------------------------===// @@ -154,8 +155,9 @@ struct HazardMitigationPass // Parse target arch from option. std::optional parsed = symbolizeTargetKind(targetArch); if (!parsed) { - module->emitError() << "Invalid target architecture: '" << targetArch - << "'. Supported targets: gfx942, gfx950, gfx1250"; + module->emitError() + << "Invalid target architecture: '" << targetArch + << "'. Supported targets: gfx90a, gfx942, gfx950, gfx1250"; return signalPassFailure(); } targetKindEnum = *parsed; diff --git a/waveasm/lib/Transforms/LiteralMaterialization.cpp b/waveasm/lib/Transforms/LiteralMaterialization.cpp index 8adea968cf..6e9407068c 100644 --- a/waveasm/lib/Transforms/LiteralMaterialization.cpp +++ b/waveasm/lib/Transforms/LiteralMaterialization.cpp @@ -68,6 +68,9 @@ KernelGenerator::generateOpWithLiteralHandling(Operation *op) { mnemonic = opName.drop_front(8); } + std::string targetMnemonic = getTargetMnemonic(mnemonic); + llvm::StringRef asmMnemonic(targetMnemonic); + bool hasNonInlineLiteral = false; int64_t literalValue = 0; int literalOperandIdx = -1; @@ -90,7 +93,7 @@ KernelGenerator::generateOpWithLiteralHandling(Operation *op) { } // SALU instructions support 32-bit literals natively - if (mnemonic.starts_with("s_")) { + if (asmMnemonic.starts_with("s_")) { if (auto line = generateOp(op)) { lines.push_back(*line); } @@ -106,15 +109,15 @@ KernelGenerator::generateOpWithLiteralHandling(Operation *op) { } // VOP3+ instructions need literal materialization into scratch VGPR. - if (needsLiteralMaterialization(mnemonic)) { - emitMaterializedLiteral(lines, op, mnemonic, literalOperandIdx, + if (needsLiteralMaterialization(asmMnemonic)) { + emitMaterializedLiteral(lines, op, asmMnemonic, literalOperandIdx, literalValue); return lines; } // v_cndmask_b32 has a dedicated Case handler in generateOp that // materializes literals AND drops the implicit VCC condition operand. - if (mnemonic == "v_cndmask_b32") { + if (asmMnemonic == "v_cndmask_b32") { if (auto line = generateOp(op)) { lines.push_back(*line); } @@ -139,11 +142,12 @@ KernelGenerator::generateOpWithLiteralHandling(Operation *op) { } operands.push_back(std::to_string(literalValue)); operands.push_back(resolveValue(op->getOperand(0))); - lines.push_back(formatter.format(mnemonic, operands)); + lines.push_back(formatter.format(asmMnemonic, operands)); return lines; } - emitMaterializedLiteral(lines, op, mnemonic, literalOperandIdx, literalValue); + emitMaterializedLiteral(lines, op, asmMnemonic, literalOperandIdx, + literalValue); return lines; } diff --git a/waveasm/lib/Transforms/MetadataEmitter.cpp b/waveasm/lib/Transforms/MetadataEmitter.cpp index db3ff8de72..b3b5a5e19a 100644 --- a/waveasm/lib/Transforms/MetadataEmitter.cpp +++ b/waveasm/lib/Transforms/MetadataEmitter.cpp @@ -20,6 +20,20 @@ using namespace mlir; namespace waveasm { +static bool supportsKernargPreload(TargetAttrInterface target) { + return target.hasFeature(TargetFeature::HasKernargPreload); +} + +static bool usesCDNAAccumOffset(TargetAttrInterface target) { + return llvm::isa( + target); +} + +static bool +supportsArchitectedFlatScratchDirectives(TargetAttrInterface target) { + return !llvm::isa(target); +} + //===----------------------------------------------------------------------===// // Instruction Formatter Implementation //===----------------------------------------------------------------------===// @@ -138,7 +152,7 @@ static void scanSystemRegisterUsage(ProgramOp program, bool &usesWorkgroupIdX, auto targetAttr = program.getTarget(); auto targetKind = targetAttr.getTargetKind(); - bool isGfx950 = llvm::isa(targetKind); + bool usePreloading = supportsKernargPreload(targetKind); int64_t numArgs = 2; if (auto numArgsAttr = @@ -147,8 +161,8 @@ static void scanSystemRegisterUsage(ProgramOp program, bool &usesWorkgroupIdX, } int64_t userSgprCount = 2; - if (isGfx950) { - // Hardware limits user SGPRs to 16 on gfx950. + if (usePreloading) { + // Hardware limits user SGPRs to 16 on preload targets. userSgprCount = std::min(int64_t(16), 2 + numArgs * 2); } @@ -207,7 +221,7 @@ 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); + bool usePreloading = supportsKernargPreload(targetKind); if (usePreloading && preloadLength == 0) { int64_t numArgs = 2; @@ -215,8 +229,9 @@ MetadataEmitter::emitKernelDescriptor(int64_t peakVGPRs, int64_t peakSGPRs, 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. + // Hardware limits user SGPRs to 16 on preload targets (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); } @@ -239,10 +254,11 @@ MetadataEmitter::emitKernelDescriptor(int64_t peakVGPRs, int64_t peakSGPRs, lines.push_back(" .amdhsa_user_sgpr_private_segment_size 0"); lines.push_back(" .amdhsa_uses_dynamic_stack 0"); - lines.push_back(" .amdhsa_enable_private_segment 0"); + if (supportsArchitectedFlatScratchDirectives(targetKind)) + lines.push_back(" .amdhsa_enable_private_segment 0"); int64_t vgprGranularity = 4; - if (llvm::isa(targetKind)) { + if (usesCDNAAccumOffset(targetKind)) { vgprGranularity = 8; } int64_t nextFreeVGPR = @@ -253,11 +269,11 @@ MetadataEmitter::emitKernelDescriptor(int64_t peakVGPRs, int64_t peakSGPRs, int64_t sgprGranularity = 8; int64_t nextFreeSGPR = ((peakSGPRs + sgprGranularity - 1) / sgprGranularity) * sgprGranularity; - if (llvm::isa(targetKind)) { + if (usesCDNAAccumOffset(targetKind)) { nextFreeSGPR = std::min(nextFreeSGPR, int64_t(102)); } - if (llvm::isa(targetKind)) { + if (usesCDNAAccumOffset(targetKind)) { int64_t accumOffset = std::max(int64_t(4), ((nextFreeVGPR + 3) / 4) * 4); accumOffset = std::min(accumOffset, int64_t(256)); lines.push_back(" .amdhsa_accum_offset " + std::to_string(accumOffset)); diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 409cba963b..49b650a0d9 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -128,7 +128,8 @@ static ProgramOp createProgramFromLLVMFunc(LLVM::LLVMFuncOp func, auto *mlirCtx = builder.getContext(); auto loc = func.getLoc(); - // Code object version 5: supports kernel argument preloading. + // Use code object version 5 consistently; target features decide whether + // kernarg preload metadata/prologue is emitted. auto targetAttr = TargetAttr::get(mlirCtx, getTargetKindAttr(mlirCtx, targetId), /*code_object_version=*/5); diff --git a/waveasm/lib/Transforms/TranslateFromMLIR.cpp b/waveasm/lib/Transforms/TranslateFromMLIR.cpp index a23eaeded3..63cde011ab 100644 --- a/waveasm/lib/Transforms/TranslateFromMLIR.cpp +++ b/waveasm/lib/Transforms/TranslateFromMLIR.cpp @@ -27,7 +27,6 @@ #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/Verifier.h" -#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/Debug.h" @@ -162,12 +161,21 @@ void TranslationContext::emitSRDPrologue() { srdPrologueEmitted = true; auto loc = builder.getUnknownLoc(); - bool isGFX95 = llvm::isa(target); - bool usePreloading = isGFX95; + // Targets with kernel argument preload require the preload prologue pattern + // with branch+alignment. gfx90a intentionally uses the non-preload path. + // Keep the gfx950 type check as a fallback while target feature plumbing is + // being migrated across commits. + bool supportsKernargPreload = + target.hasFeature(TargetFeature::HasKernargPreload) || + llvm::isa(target); // Recompute SRD base indices now that we know the total number of args. // SRDs must start after: user SGPRs + system SGPRs (workgroup IDs). - int64_t userSgprCount = getUserSgprCount(); + size_t numPreloadedArgs = getNumKernelArgs(); + int64_t userSgprCount = 2; // kernarg ptr + if (supportsKernargPreload) { + userSgprCount += std::min(int64_t(14), (int64_t)getNumKernelArgs() * 2); + } int64_t systemSgprCount = 3; // workgroup_id_x, y, z int64_t srdStartIndex = (userSgprCount + systemSgprCount + 3) & ~3; // Align to 4 @@ -199,72 +207,55 @@ void TranslationContext::emitSRDPrologue() { PrecoloredSRegOp::create(builder, loc, kernargPtrType, /*index=*/0, /*size=*/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 (supportsKernargPreload) { + // On targets with kernarg preload, 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. 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 (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; - if (reservedPreloadBases.insert(preloadBase).second) { - auto preloadType = createSRegType(2, 2); - PrecoloredSRegOp::create(builder, loc, preloadType, preloadBase, - /*size=*/2); - } + auto preloadType = createSRegType(2, 2); + PrecoloredSRegOp::create(builder, loc, preloadType, preloadBase, + /*size=*/2); } } - 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. + if (supportsKernargPreload) { + // Kernarg preload path: Use preload pattern with intermediate locations + // and s_mov_b64 copies. This matches the Python backend behavior for + // gfx950. + // + // Step 1: Load all kernel args into preload locations s[2:3], s[4:5], + // etc. 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. auto kernargSRegType = createSRegType(2, 2); auto kernargBase = PrecoloredSRegOp::create(builder, loc, kernargSRegType, 0, 2); - for (const auto &pending : pendingSRDs) { - 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); - S_LOAD_DWORDX2::create(builder, loc, TypeRange{loadDstType}, kernargBase, - offsetConst); - } - - // Also load scalar args that fit in the preload window. - for (const auto &pending : pendingScalarArgs) { - int64_t loadBase = 2 + pending.argIndex * 2; + llvm::DenseMap argLoadResults; + for (size_t i = 0; i < numPreloadedArgs; ++i) { + int64_t loadBase = 2 + i * 2; if (loadBase >= 16) - continue; - int64_t kernargOffset = pending.argIndex * 8; + continue; // Overflow arg: loaded via s_load_dword path below. + int64_t kernargOffset = i * 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); + argLoadResults[i] = loadOp->getResult(0); } - // Step 2: Branch to aligned entry point (gfx95* requirement). + // Step 2: Branch to aligned entry point (kernarg preload requirement). + // Keep any high-SGPR overflow loads after the aligned entry; LLVM does the + // same, and loading them before the branch leaves the overflow arg stale + // on hardware with kernarg preload. std::string kernelName = getKernelName(program).str(); std::string mainLabel = ".L_" + kernelName + "_main"; @@ -303,7 +294,10 @@ void TranslationContext::emitSRDPrologue() { /*expcnt=*/IntegerAttr{}); // Step 5: Copy from preload locations to SRD positions and fill - // size/stride. Use typed ops with DCEProtectOp to prevent elimination. + // size/stride. Use typed ops targeting precolored registers with + // DCEProtectOp to prevent elimination of Pure ops. + // When we have a captured S_LOAD_DWORDX2 result for this arg, use it + // directly as the S_MOV_B64 source to maintain the SSA def-use chain. for (size_t i = 0; i < pendingSRDs.size(); ++i) { const auto &pending = pendingSRDs[i]; int64_t srdBase = pending.srdBaseIndex; @@ -312,9 +306,15 @@ void TranslationContext::emitSRDPrologue() { auto srdType = createSRegType(4, 4); auto srdReg = PrecoloredSRegOp::create(builder, loc, srdType, srdBase, 4); - auto preloadType = createSRegType(2, preloadBase); - auto preloadSrc = - PrecoloredSRegOp::create(builder, loc, preloadType, preloadBase, 2); + Value preloadSrc; + auto loadIt = argLoadResults.find(pending.argIndex); + if (preloadBase < 16 && loadIt != argLoadResults.end()) { + preloadSrc = loadIt->second; + } else { + auto preloadType = createSRegType(2, preloadBase); + preloadSrc = + PrecoloredSRegOp::create(builder, loc, preloadType, preloadBase, 2); + } auto dstB64Type = PSRegType::get(builder.getContext(), srdBase, 2); auto movB64 = S_MOV_B64::create(builder, loc, dstB64Type, preloadSrc); DCEProtectOp::create(builder, loc, movB64); @@ -359,33 +359,42 @@ void TranslationContext::emitSRDPrologue() { mapper.mapValue(pending.blockArg, vreg); } } else { - // Direct-load path (non-GFX950 targets): - // Load base addresses directly into SRD[0:1] positions, then fill SRD[2:3]. + // Non-preload path (e.g., gfx90a/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. - PrecoloredSRegOp::create(builder, loc, createSRegType(2, 2), 0, 2); + auto kernargSRegType = createSRegType(2, 2); + auto kernargBase = + PrecoloredSRegOp::create(builder, loc, kernargSRegType, 0, 2); for (const auto &pending : pendingSRDs) { int64_t srdBase = pending.srdBaseIndex; int64_t kernargOffset = pending.argIndex * 8; - auto pairType = PSRegType::get(builder.getContext(), srdBase, 2); - PrecoloredSRegOp::create(builder, loc, pairType, srdBase, 2); - RawOp::create(builder, loc, - "s_load_dwordx2 s[" + std::to_string(srdBase) + ":" + - std::to_string(srdBase + 1) + "], s[0:1], " + - std::to_string(kernargOffset)); + // Load directly into SRD base: s[srdBase:srdBase+1]. + auto loadDstType = PSRegType::get(builder.getContext(), srdBase, 2); + auto offsetImm = builder.getType(kernargOffset); + auto offsetConst = + ConstantOp::create(builder, loc, offsetImm, kernargOffset); + S_LOAD_DWORDX2::create(builder, loc, TypeRange{loadDstType}, kernargBase, + offsetConst); } + // 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; for (size_t i = 0; i < pendingScalarArgs.size(); ++i) { const auto &pending = pendingScalarArgs[i]; - int64_t pairBase = overflowSgprBase + (int64_t)i * 2; + int64_t sgprIdx = scalarSgprBase + (int64_t)i; int64_t kernargOffset = pending.argIndex * 8; - auto pairType = PSRegType::get(builder.getContext(), pairBase, 2); - PrecoloredSRegOp::create(builder, loc, pairType, pairBase, 2); + PrecoloredSRegOp::create(builder, loc, createSRegType(1, 1), sgprIdx, 1); RawOp::create(builder, loc, - "s_load_dwordx2 s[" + std::to_string(pairBase) + ":" + - std::to_string(pairBase + 1) + "], s[0:1], " + + "s_load_dword s" + std::to_string(sgprIdx) + ", s[0:1], " + std::to_string(kernargOffset)); } @@ -420,16 +429,18 @@ void TranslationContext::emitSRDPrologue() { for (size_t i = 0; i < pendingScalarArgs.size(); ++i) { const auto &pending = pendingScalarArgs[i]; - int64_t srcSgpr = overflowSgprBase + (int64_t)i * 2; + int64_t sgprIdx = scalarSgprBase + (int64_t)i; 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(srcSgpr)); + std::to_string(sgprIdx)); mapper.mapValue(pending.blockArg, vreg); } + + overflowSgprBase = scalarSgprBase; } // For the direct-load path, record the pinned SGPR reservations as @@ -437,8 +448,8 @@ void TranslationContext::emitSRDPrologue() { // (they're Pure with no SSA users since the RawOp references registers // by string). The preloading path uses typed S_LOAD_DWORDX2 ops which // the allocator sees directly, so no extra attributes are needed. - if (!usePreloading && !pendingScalarArgs.empty()) { - int64_t scalarCount = (int64_t)pendingScalarArgs.size() * 2; + if (!supportsKernargPreload && !pendingScalarArgs.empty()) { + int64_t scalarCount = (int64_t)pendingScalarArgs.size(); int64_t scalarEnd = overflowSgprBase + scalarCount; program->setAttr("min_sgprs", builder.getI64IntegerAttr(scalarEnd)); program->setAttr("scalar_sgpr_base", diff --git a/waveasm/test/Integration/compile-and-assemble.mlir b/waveasm/test/Integration/compile-and-assemble.mlir index ec75418504..256f2bb18d 100644 --- a/waveasm/test/Integration/compile-and-assemble.mlir +++ b/waveasm/test/Integration/compile-and-assemble.mlir @@ -16,7 +16,7 @@ // CHECK: s_endpgm // CHECK: .amdhsa_kernel simple_add -// CHECK-HSACO: ELF 64-bit LSB shared object, AMD GPU +// CHECK-HSACO: ELF 64-bit LSB shared object waveasm.program @simple_add target = #waveasm.target<#waveasm.gfx942, 5> diff --git a/waveasm/test/Translate/dynamic-shapes.mlir b/waveasm/test/Translate/dynamic-shapes.mlir index c75884d1f0..c1246ced49 100644 --- a/waveasm/test/Translate/dynamic-shapes.mlir +++ b/waveasm/test/Translate/dynamic-shapes.mlir @@ -6,16 +6,16 @@ // CHECK-LABEL: waveasm.program @dynamic_shapes_kernel -// Test 1: index args loaded from kernarg buffer via s_load_dwordx2 -// CHECK: waveasm.raw "s_load_dwordx2 -// CHECK: waveasm.raw "s_load_dwordx2 +// Test 1: index args loaded from kernarg buffer via s_load_dword +// CHECK: waveasm.s_load_dword +// CHECK: waveasm.s_load_dword // Test 2: SRD buffer size is 0x7FFFFFFE (2147483646, sentinel-safe max) for dynamic memrefs. // CHECK: 2147483646 // Test 3: scalar args moved to VGPRs after SRD setup -// CHECK: waveasm.raw "v_mov_b32 -// CHECK: waveasm.raw "v_mov_b32 +// CHECK: v_mov_b32 v2 +// CHECK: v_mov_b32 v3 // Test 4: dynamic stride address computation (runtime v_mul_lo_u32) // CHECK: waveasm.v_mul_lo_u32 diff --git a/waveasm/tools/waveasm-translate/waveasm-translate.cpp b/waveasm/tools/waveasm-translate/waveasm-translate.cpp index 0daa6a64d5..db168d6c1f 100644 --- a/waveasm/tools/waveasm-translate/waveasm-translate.cpp +++ b/waveasm/tools/waveasm-translate/waveasm-translate.cpp @@ -64,7 +64,7 @@ static llvm::cl::opt static llvm::cl::opt targetId("target", llvm::cl::desc("Target GPU architecture"), - llvm::cl::value_desc("gfx942|gfx950|gfx1250"), + llvm::cl::value_desc("gfx90a|gfx942|gfx950|gfx1250"), llvm::cl::init("gfx942")); static llvm::cl::opt disablePassVerifier(