From 40eacb8acb4cd721a59808b809d602d32dec32eb Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 15:05:24 +0200 Subject: [PATCH 01/19] Fix emit_host_func stride arg mismatch and add wave_get_stride runtime helper emit_host_func was building arg_types from linear_bindings only (buffers + scalars + symbols) while the kernel function also carried trailing stride index args added by emit_func when dynamic_strides is active. The count mismatch caused a segfault during canonicalization because old kernel arguments were left with dangling uses after erase. - Match stride arg count in emit_host_func so the gpu.func wrapper has the correct signature. - Add wave_get_stride runtime function (mirrors wave_get_dim but calls tensor.stride() instead of tensor.size()). - Pass stride values through gpu.launch_func kernel operands. - Add GEMM waveasm e2e test skeleton. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Ivan Butygin --- tests/kernel/e2e/test_gemm_waveasm.py | 64 +++++++++++++++++++ .../kernel/compiler/wave_codegen/emitter.py | 35 +++++++++- .../wave/execution_engine/buffer_utils.cpp | 35 ++++++++++ .../wave/execution_engine/buffer_utils.h | 14 ++++ .../wave/execution_engine/execution_engine.py | 1 + 5 files changed, 148 insertions(+), 1 deletion(-) create mode 100644 tests/kernel/e2e/test_gemm_waveasm.py diff --git a/tests/kernel/e2e/test_gemm_waveasm.py b/tests/kernel/e2e/test_gemm_waveasm.py new file mode 100644 index 0000000000..eedd97e06c --- /dev/null +++ b/tests/kernel/e2e/test_gemm_waveasm.py @@ -0,0 +1,64 @@ +# Copyright 2026 The Wave Authors +# +# Licensed under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +"""GEMM through the water+waveasm pipeline (LLVM dialect → WaveASM → binary).""" + +import pytest +import torch +from torch.testing import assert_close + +from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile +from wave_lang.kernel.wave.constraints import MMAType +from wave_lang.kernel.wave.templates.gemm import get_gemm_kernel +from wave_lang.kernel.wave.utils.run_utils import set_default_run_config +from wave_lang.kernel.wave.utils.torch_utils import device_randn, device_zeros + +from ..common.utils import require_cdna_3_or_4, require_e2e + + +@require_e2e +@require_cdna_3_or_4 +@pytest.mark.parametrize( + "shape,block_shape,waves_per_block", + [ + ((64, 64, 64), (32, 32, 16), (1, 1)), + ], + ids=["64x64x64"], +) +def test_gemm_water_waveasm( + shape: tuple[int, int, int], + block_shape: tuple[int, int, int], + waves_per_block: tuple[int, int], +) -> None: + """Test GEMM through the water+waveasm pipeline.""" + m, n, k = shape + + gemm, hyperparams, _ = get_gemm_kernel( + shape=shape, + dynamic_dims=False, + mfma_variant=MMAType.F32_16x16x16_F16, + block_shape=block_shape, + waves_per_block=waves_per_block, + ) + + options = WaveCompileOptions( + subs=hyperparams, + canonicalize=True, + use_water_backend=True, + use_buffer_ops=True, + backend="asm", + wave_runtime=True, + ) + options = set_default_run_config(options) + compiled = wave_compile(options, gemm) + + a = device_randn((m, k), dtype=torch.float16) + b = device_randn((n, k), dtype=torch.float16) + c = device_zeros((m, n), dtype=torch.float32) + compiled(a, b, c) + + expected = torch.matmul(a, b.T).float() + assert_close(c, expected, rtol=1e-3, atol=1e-3) diff --git a/wave_lang/kernel/compiler/wave_codegen/emitter.py b/wave_lang/kernel/compiler/wave_codegen/emitter.py index 36cb2ce912..9c9f39da6d 100644 --- a/wave_lang/kernel/compiler/wave_codegen/emitter.py +++ b/wave_lang/kernel/compiler/wave_codegen/emitter.py @@ -322,8 +322,18 @@ def abi_type(binding: BindingDesc): arg_types = [abi_type(b) for b in bindings] + # Match stride args added by emit_func when dynamic_strides is active. + stride_arg_count = 0 + if self.options.dynamic_strides: + stride_arg_count = sum( + max(0, len(b.kernel_buffer_type.symbolic_shape) - 1) + for b in self.root_sig.sig.kernel_buffer_bindings + ) + if stride_arg_count > 0: + arg_types += [IndexType.get()] * stride_arg_count + ftype = FunctionType.get(arg_types, []) - locs = [a.location for a in kernel_func.body.blocks[0].arguments] + locs = [Location.unknown() for _ in arg_types] gpu_module = gpu_d.module("gpu_module") gpu_module.parent.operation.attributes["gpu.container_module"] = UnitAttr.get() @@ -396,6 +406,11 @@ def abi_type(binding: BindingDesc): "wave_get_float64", [ptr], [f64], emit_c_interface=True ) + # Get tensor stride function from PyObject*. + get_stride_func, get_stride_func_symbol = self._declare_runtime_func( + "wave_get_stride", [ptr, i32], [i64], emit_c_interface=True + ) + # Declare host function # First argument is stream pointer # Rest are kernel arguments as PyObject* @@ -490,6 +505,24 @@ def abi_type(binding: BindingDesc): else: raise CodegenError(f"Unsupported binding type: {binding}") + # Append stride arguments matching the trailing index args + # added by emit_func. One stride per leading dimension (rank-1) + # for each kernel buffer. + if self.options.dynamic_strides: + for binding in self.root_sig.sig.kernel_buffer_bindings: + rank = len(binding.kernel_buffer_type.symbolic_shape) + leading_count = max(0, rank - 1) + arg = func_args[binding_map[id(binding)]] + for dim_idx in range(leading_count): + dim = arith_d.constant(i32, dim_idx) + stride = func_d.call( + get_stride_func.type.results, + get_stride_func_symbol, + [arg, dim], + ) + stride = arith_d.index_cast(IndexType.get(), stride) + launch_args.append(stride) + gpu_d.launch_func( kernel=[gpu_module.sym_name.value, self.kernel_name], grid_size=grid, diff --git a/wave_lang/kernel/wave/execution_engine/buffer_utils.cpp b/wave_lang/kernel/wave/execution_engine/buffer_utils.cpp index 7a82700f96..afa4d0c7c7 100644 --- a/wave_lang/kernel/wave/execution_engine/buffer_utils.cpp +++ b/wave_lang/kernel/wave/execution_engine/buffer_utils.cpp @@ -121,3 +121,38 @@ extern "C" int64_t _mlir_ciface_wave_get_dim(PyObject *obj_ptr, return dim_size; } + +extern "C" int64_t _mlir_ciface_wave_get_stride(PyObject *obj_ptr, + int32_t dim_idx) { + GILState gil_state; + + PyObjectPtr stride_method(PyObject_GetAttrString(obj_ptr, "stride")); + if (!stride_method) { + PyErr_Clear(); + throw std::runtime_error( + "wave_get_stride: object does not have 'stride' attribute"); + } + + PyObjectPtr dim_arg(PyLong_FromLong(dim_idx)); + if (!dim_arg) { + PyErr_Clear(); + throw std::runtime_error( + "wave_get_stride: failed to create dimension argument"); + } + + PyObjectPtr stride_result( + PyObject_CallOneArg(stride_method.get(), dim_arg.get())); + if (!stride_result) { + PyErr_Clear(); + throw std::runtime_error("wave_get_stride: failed to call stride()"); + } + + int64_t stride_val = PyLong_AsLongLong(stride_result.get()); + if (PyErr_Occurred()) { + PyErr_Clear(); + throw std::runtime_error( + "wave_get_stride: stride() returned invalid value"); + } + + return stride_val; +} diff --git a/wave_lang/kernel/wave/execution_engine/buffer_utils.h b/wave_lang/kernel/wave/execution_engine/buffer_utils.h index 25938c6d47..54c89dfb7e 100644 --- a/wave_lang/kernel/wave/execution_engine/buffer_utils.h +++ b/wave_lang/kernel/wave/execution_engine/buffer_utils.h @@ -60,4 +60,18 @@ double _mlir_ciface_wave_get_float64(PyObject *obj); /// std::runtime_error if the object doesn't have a size() method or /// if the dimension index is invalid int64_t _mlir_ciface_wave_get_dim(PyObject *obj, int32_t dim_idx); + +/// Extract the stride of a specific dimension from a PyObject (PyTorch tensor). +/// +/// Args: +/// obj: PyObject* pointing to a PyTorch tensor +/// dim_idx: Dimension index to query (0-based) +/// +/// Returns: +/// Stride of the specified dimension in elements as int64_t. +/// +/// Throws: +/// std::runtime_error if the object doesn't have a stride() method or +/// if the dimension index is invalid. +int64_t _mlir_ciface_wave_get_stride(PyObject *obj, int32_t dim_idx); } diff --git a/wave_lang/kernel/wave/execution_engine/execution_engine.py b/wave_lang/kernel/wave/execution_engine/execution_engine.py index ff2c5dd8ba..93ca05bd3e 100644 --- a/wave_lang/kernel/wave/execution_engine/execution_engine.py +++ b/wave_lang/kernel/wave/execution_engine/execution_engine.py @@ -104,6 +104,7 @@ def _load_runtime_helpers(): _add_symbol(symbol_map, lib, "_mlir_ciface_wave_get_int64") _add_symbol(symbol_map, lib, "_mlir_ciface_wave_get_float64") _add_symbol(symbol_map, lib, "_mlir_ciface_wave_get_dim") + _add_symbol(symbol_map, lib, "_mlir_ciface_wave_get_stride") return symbol_map From 80add6d471550b18fda5731fe9a6a9bf2a9f2af2 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 16:38:02 +0200 Subject: [PATCH 02/19] Add GEMM op handlers to LLVM->WaveASM translation Extend the waveasm-translate-from-llvm pass to handle all LLVM dialect ops generated by the GEMM kernel, including: - LDS: llvm.mlir.addressof for globals, ptr<3> GEPs, ds_read/ds_write - Arithmetic: llvm.sdiv/srem (power-of-2 -> shift/mask) - Barriers: llvm.fence (no-op), rocdl.s.barrier -> s_barrier - MFMA: rocdl.mfma.f32.16x16x16f16 -> v_mfma_f32_16x16x16_f16 - Vector: llvm.shufflevector (single-element extract) - Constants: dense vector splat (MFMA accumulator init) - Control flow: scf.for -> waveasm.loop with do-while semantics Also update the water+waveasm lowering pipeline to preserve scf.for (remove convert-scf-to-cf) and use alloca-to-global for LDS, matching what the translation pass expects. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Ivan Butygin --- wave_lang/kernel/wave/water.py | 28 +- .../Transforms/TranslateFromLLVMDialect.cpp | 473 +++++++++++++++++- ...slate-from-llvm-error-multi-index-gep.mlir | 2 +- 3 files changed, 481 insertions(+), 22 deletions(-) diff --git a/wave_lang/kernel/wave/water.py b/wave_lang/kernel/wave/water.py index 2af389e294..759033e1b6 100644 --- a/wave_lang/kernel/wave/water.py +++ b/wave_lang/kernel/wave/water.py @@ -390,6 +390,27 @@ def add_opt(pipeline): return [pipeline] return [] + def add_transform(transform: str, entry_point: str) -> tuple[str, dict[str, Any]]: + nonlocal mlir_asm + # Erase the last '}' closing the module, append the transform, re-close. + last_close = mlir_asm.rfind("}") + if last_close != -1: + mlir_asm = mlir_asm[:last_close] + mlir_asm += transform + mlir_asm += "}\n" + return ("transform-interpreter", {"entry-point": entry_point}) + + alloca_to_global = """ + transform.named_sequence @__transform_alloca_to_global(%arg0: !transform.any_op {transform.readonly}) { + %alloca = transform.structured.match ops{["memref.alloca"]} in %arg0 + : (!transform.any_op) -> !transform.op<"memref.alloca"> + %get_global, %global = transform.memref.alloca_to_global %alloca + : (!transform.op<"memref.alloca">) + -> (!transform.any_op, !transform.any_op) + transform.yield + } +""" + canonicalize_cse = "composite-fixed-point-pass", { "name": "canonicalize_cse", "pipeline": "any(canonicalize,cse)", @@ -401,13 +422,16 @@ def add_opt(pipeline): } # Step 1: water-opt lowers host + device to LLVM dialect. + # Note: convert-scf-to-cf is NOT used here -- waveasm-translate handles + # scf.for directly via waveasm.loop structured ops. lowering_pipeline = [ "water-memref-decomposition", *add_opt(canonicalize_cse), "lower-affine", *add_opt(int_range_optimizations), *add_opt("loop-invariant-code-motion"), - "convert-scf-to-cf", + ("water-alloc-to-alloca", {}, "gpu.module"), + add_transform(alloca_to_global, "__transform_alloca_to_global"), ("convert-amdgpu-to-rocdl", {"chipset": target_chip}), ( "convert-gpu-to-rocdl", @@ -417,6 +441,8 @@ def add_opt(pipeline): ("gpu-to-llvm", {"use-bare-pointers-for-kernels": "1"}), "convert-vector-to-llvm", "reconcile-unrealized-casts", + "water-drop-transform-ops", + "symbol-dce", *add_opt(canonicalize_cse), ] diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 409cba963b..4904b48122 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -21,6 +21,7 @@ #include "mlir/Dialect/GPU/IR/GPUDialect.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/LLVMIR/ROCDLDialect.h" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Utils/StaticValueUtils.h" #include "mlir/IR/Builders.h" #include "llvm/ADT/StringExtras.h" @@ -86,10 +87,15 @@ class LLVMTranslationState { void setBaseOffset(Value ptr, Value offset) { baseOffsets[ptr] = offset; } Value lookupBaseOffset(Value ptr) const { return baseOffsets.lookup(ptr); } + /// Track LDS base values (from addressof / ptr<3> GEPs). + void setLDSBase(Value v) { ldsBaseValues.insert(v); } + bool isLDSBase(Value v) const { return ldsBaseValues.contains(v); } + private: DenseMap rsrcToSRD; DenseMap gepMap; DenseMap baseOffsets; + DenseSet ldsBaseValues; }; //===----------------------------------------------------------------------===// @@ -195,6 +201,64 @@ static Type inferResultType(ValueRange operands, TranslationContext &ctx) { return ctx.createSRegType(); } +/// Try to extract a constant integer from an LLVM SSA value. +static std::optional getConstantInt(Value v) { + if (auto constOp = v.getDefiningOp()) + if (auto intAttr = dyn_cast(constOp.getValue())) + return intAttr.getValue().getSExtValue(); + return std::nullopt; +} + +/// Compute the byte offset for a GEP as a single WaveASM Value. +/// For all-zero constant indices, returns std::nullopt (offset is 0). +static std::optional computeGEPOffset(LLVM::GEPOp op, + TranslationContext &ctx) { + auto &builder = ctx.getBuilder(); + auto loc = op.getLoc(); + auto indices = op.getIndices(); + + // Check if all indices are constant zero. + bool allZero = true; + for (auto idx : indices) { + if (auto val = idx.dyn_cast()) { + allZero = false; + break; + } + if (cast(idx).getInt() != 0) { + allZero = false; + break; + } + } + if (allZero) + return std::nullopt; + + // Single dynamic index -- common case. + if (indices.size() == 1) { + auto idx = indices[0]; + if (auto val = idx.dyn_cast()) { + FailureOr resolved = resolve(val, ctx); + if (failed(resolved)) + return std::nullopt; + return *resolved; + } + int64_t constVal = cast(idx).getInt(); + auto immTy = ctx.createImmType(constVal); + return ConstantOp::create(builder, loc, immTy, constVal)->getResult(0); + } + + // Multi-index: accumulate constant indices. + int64_t totalOffset = 0; + for (auto idx : indices) { + if (idx.dyn_cast()) + return std::nullopt; + totalOffset += cast(idx).getInt(); + } + if (totalOffset == 0) + return std::nullopt; + auto immTy = ctx.createImmType(totalOffset); + return ConstantOp::create(builder, loc, immTy, totalOffset)->getResult(0); +} + //===----------------------------------------------------------------------===// // Op handlers //===----------------------------------------------------------------------===// @@ -241,6 +305,22 @@ static LogicalResult handleConstant(LLVM::ConstantOp op, auto loc = op.getLoc(); auto valAttr = op.getValue(); + + // Dense vector constant (e.g. MFMA accumulator init). + if (auto denseAttr = dyn_cast(valAttr)) { + if (!denseAttr.isSplat()) + return op->emitOpError("non-splat dense constant not yet supported"); + int64_t numElems = denseAttr.getNumElements(); + APFloat splatVal = denseAttr.getSplatValue(); + int64_t rawBits = splatVal.bitcastToAPInt().getZExtValue(); + auto immTy = ctx.createImmType(rawBits); + auto immOp = ConstantOp::create(builder, loc, immTy, rawBits); + auto vregTy = ctx.createVRegType(numElems, numElems); + Value mov = V_MOV_B32::create(builder, loc, vregTy, immOp); + ctx.getMapper().mapValue(op.getResult(), mov); + return success(); + } + int64_t intVal = 0; if (auto intAttr = dyn_cast(valAttr)) intVal = intAttr.getValue().getSExtValue(); @@ -492,36 +572,54 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { auto loc = op.getLoc(); Value base = op.getBase(); - // GEP index is a dynamic Value (not a constant attr). - auto indices = op.getIndices(); - if (indices.size() != 1) - return op->emitOpError("GEP with multiple indices not yet supported"); - auto idx = indices[0].dyn_cast(); - if (!idx) - return op->emitOpError("GEP with constant index attr not yet supported"); - - FailureOr resolved = resolve(idx, ctx); - if (failed(resolved)) - return op->emitOpError("unmapped GEP index"); - Value newOffset = *resolved; - - // Bare-pointer GEP (!llvm.ptr, not <7>): 64-bit pointer arithmetic before - // make.buffer.rsrc. Propagate the mapper entry and accumulate the byte - // offset so it can be folded into the SRD base later. auto baseTy = op.getBase().getType(); unsigned addrSpace = 0; if (auto ptrTy = dyn_cast(baseTy)) addrSpace = ptrTy.getAddressSpace(); + // LDS GEP (ptr<3>): compute byte offset for ds_read/ds_write. + // DS instructions only accept a 32-bit vaddr, so truncate wide offsets. + if (addrSpace == 3) { + FailureOr baseOff = resolve(base, ctx); + if (failed(baseOff)) + return op->emitOpError("unmapped LDS GEP base"); + auto vregTy = ctx.createVRegType(); + if (getRegSize(baseOff->getType()) > 1) + *baseOff = ArithTruncOp::create(builder, loc, vregTy, *baseOff); + auto maybeOffset = computeGEPOffset(op, ctx); + if (!maybeOffset) { + ctx.getMapper().mapValue(op.getResult(), *baseOff); + } else { + Value off = *maybeOffset; + if (getRegSize(off.getType()) > 1) + off = ArithTruncOp::create(builder, loc, vregTy, off); + Value sum = ArithAddOp::create(builder, loc, vregTy, *baseOff, off); + ctx.getMapper().mapValue(op.getResult(), sum); + } + st.setLDSBase(op.getResult()); + return success(); + } + if (addrSpace != 0 && addrSpace != 7) return op->emitOpError("unsupported address space ") << addrSpace; + // Bare-pointer GEP (!llvm.ptr, not <7>): 64-bit pointer arithmetic before + // make.buffer.rsrc. Propagate the mapper entry and accumulate the byte + // offset so it can be folded into the SRD base later. if (addrSpace == 0) { + auto maybeOffset = computeGEPOffset(op, ctx); // Forward mapper entry so make.buffer.rsrc can find the SRD. if (auto mapped = ctx.getMapper().getMapped(base)) ctx.getMapper().mapValue(op.getResult(), *mapped); - // Accumulate base offset with 64-bit add (pointer arithmetic). + if (!maybeOffset) { + Value prevOffset = st.lookupBaseOffset(base); + if (prevOffset) + st.setBaseOffset(op.getResult(), prevOffset); + return success(); + } + + Value newOffset = *maybeOffset; Value prevOffset = st.lookupBaseOffset(base); if (prevOffset) { Type resTy = inferResultType({prevOffset, newOffset}, ctx); @@ -532,10 +630,19 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { return success(); } - // Buffer voffsets are 32-bit. Truncate i64 GEP indices. - newOffset = truncToI32(newOffset, idx.getType(), builder, loc, ctx); + // Buffer GEP (ptr<7>): single dynamic index. + auto indices = op.getIndices(); + if (indices.size() != 1) + return op->emitOpError("buffer GEP must have a single index"); + auto idx = indices[0].dyn_cast(); + if (!idx) + return op->emitOpError("buffer GEP with constant index not yet supported"); + + FailureOr resolved = resolve(idx, ctx); + if (failed(resolved)) + return op->emitOpError("unmapped GEP index"); + Value newOffset = truncToI32(*resolved, idx.getType(), builder, loc, ctx); - // Buffer GEP (ptr<7>): decompose into (SRD, voffset). // Check gepMap first -- covers both chained buffer GEPs and // make.buffer.rsrc entries seeded with a bare-pointer base offset. if (std::optional baseGEP = st.lookupGEP(base)) { @@ -566,11 +673,21 @@ static int64_t getBufferAccessBytes(Type ty) { return 0; } +/// Forward declaration for LDS handlers. +static LogicalResult handleLDSLoad(LLVM::LoadOp op, Value addr, + LLVMTranslationState &st); +static LogicalResult handleLDSStore(LLVM::StoreOp op, Value addr, + LLVMTranslationState &st); + static LogicalResult handleLoad(LLVM::LoadOp op, LLVMTranslationState &st) { auto &ctx = st.ctx; auto &builder = ctx.getBuilder(); auto loc = op.getLoc(); + // LDS load (ptr<3>). + if (st.isLDSBase(op.getAddr())) + return handleLDSLoad(op, op.getAddr(), st); + std::optional ptr = st.lookupGEP(op.getAddr()); if (!ptr) return op->emitOpError("load address not from a tracked GEP"); @@ -617,6 +734,10 @@ static LogicalResult handleStore(LLVM::StoreOp op, LLVMTranslationState &st) { auto &builder = ctx.getBuilder(); auto loc = op.getLoc(); + // LDS store (ptr<3>). + if (st.isLDSBase(op.getAddr())) + return handleLDSStore(op, op.getAddr(), st); + std::optional ptr = st.lookupGEP(op.getAddr()); if (!ptr) return op->emitOpError("store address not from a tracked GEP"); @@ -647,6 +768,306 @@ static LogicalResult handleStore(LLVM::StoreOp op, LLVMTranslationState &st) { return success(); } +//===----------------------------------------------------------------------===// +// LDS global / addressof handlers +//===----------------------------------------------------------------------===// + +static LogicalResult handleAddressOf(LLVM::AddressOfOp op, + LLVMTranslationState &st) { + auto &ctx = st.ctx; + // Look up the global to determine LDS size. + auto global = SymbolTable::lookupNearestSymbolFrom( + op, op.getGlobalNameAttr()); + if (global) { + auto globalTy = global.getType(); + int64_t sizeBytes = 0; + if (auto arrTy = dyn_cast(globalTy)) + sizeBytes = arrTy.getNumElements(); + if (sizeBytes > 0) + ctx.addLDSSize(sizeBytes); + } + + // Map to a constant 0 offset in a VGPR -- LDS addressing is relative. + auto &builder = ctx.getBuilder(); + auto immTy = ctx.createImmType(0); + auto zero = ConstantOp::create(builder, op.getLoc(), immTy, int64_t{0}); + auto vregTy = ctx.createVRegType(); + Value mov = V_MOV_B32::create(builder, op.getLoc(), vregTy, zero); + ctx.getMapper().mapValue(op.getResult(), mov); + st.setLDSBase(op.getResult()); + return success(); +} + +//===----------------------------------------------------------------------===// +// LDS load/store handlers +//===----------------------------------------------------------------------===// + +static LogicalResult handleLDSLoad(LLVM::LoadOp op, Value addr, + LLVMTranslationState &st) { + auto &ctx = st.ctx; + auto &builder = ctx.getBuilder(); + auto loc = op.getLoc(); + + int64_t numBytes = getBufferAccessBytes(op.getResult().getType()); + FailureOr offset = resolve(addr, ctx); + if (failed(offset)) + return op->emitOpError("unmapped LDS address"); + // DS instructions use a 32-bit vaddr. Truncate wide offsets. + auto vregTy = ctx.createVRegType(); + if (getRegSize(offset->getType()) > 1) + *offset = ArithTruncOp::create(builder, loc, vregTy, *offset); + + Operation *loadOp = nullptr; + if (numBytes == 2) + loadOp = DS_READ_U16::create(builder, loc, TypeRange{vregTy}, *offset); + else if (numBytes == 4) + loadOp = DS_READ_B32::create(builder, loc, TypeRange{vregTy}, *offset); + else if (numBytes == 8) { + auto wideTy = ctx.createVRegType(2, 2); + loadOp = DS_READ_B64::create(builder, loc, TypeRange{wideTy}, *offset); + } else if (numBytes == 16) { + auto wideTy = ctx.createVRegType(4, 4); + loadOp = DS_READ_B128::create(builder, loc, TypeRange{wideTy}, *offset); + } else + return op->emitOpError("unsupported LDS load size: ") + << numBytes << " bytes"; + + ctx.getMapper().mapValue(op.getResult(), loadOp->getResult(0)); + return success(); +} + +static LogicalResult handleLDSStore(LLVM::StoreOp op, Value addr, + LLVMTranslationState &st) { + auto &ctx = st.ctx; + auto &builder = ctx.getBuilder(); + auto loc = op.getLoc(); + + FailureOr data = resolve(op.getValue(), ctx); + FailureOr offset = resolve(addr, ctx); + if (failed(data) || failed(offset)) + return op->emitOpError("unmapped LDS operand"); + // DS instructions use a 32-bit vaddr. Truncate wide offsets. + auto vregTy = ctx.createVRegType(); + if (getRegSize(offset->getType()) > 1) + *offset = ArithTruncOp::create(builder, loc, vregTy, *offset); + int64_t numBytes = getBufferAccessBytes(op.getValue().getType()); + + if (numBytes == 2) + DS_WRITE_B16::create(builder, loc, *data, *offset); + else if (numBytes == 4) + DS_WRITE_B32::create(builder, loc, *data, *offset); + else if (numBytes == 8) + DS_WRITE_B64::create(builder, loc, *data, *offset); + else if (numBytes == 16) + DS_WRITE_B128::create(builder, loc, *data, *offset); + else + return op->emitOpError("unsupported LDS store size: ") + << numBytes << " bytes"; + + return success(); +} + +//===----------------------------------------------------------------------===// +// Signed div/rem handlers +//===----------------------------------------------------------------------===// + +static LogicalResult handleSDiv(LLVM::SDivOp op, LLVMTranslationState &st) { + auto &ctx = st.ctx; + auto &builder = ctx.getBuilder(); + FailureOr lhs = resolve(op.getLhs(), ctx); + if (failed(lhs)) + return op->emitOpError("unmapped operand in sdiv"); + + // Power-of-2 constant divisor -> arithmetic shift right. + if (auto constVal = getConstantInt(op.getRhs())) { + int64_t divisor = *constVal; + if (divisor > 0 && (divisor & (divisor - 1)) == 0) { + int64_t shiftAmt = llvm::Log2_64(divisor); + auto immTy = ctx.createImmType(shiftAmt); + auto shiftConst = + ConstantOp::create(builder, op.getLoc(), immTy, shiftAmt); + auto vregTy = ctx.createVRegType(); + Value result = + V_ASHRREV_I32::create(builder, op.getLoc(), vregTy, shiftConst, *lhs); + ctx.getMapper().mapValue(op.getResult(), result); + return success(); + } + } + return op->emitOpError("signed division by non-power-of-2 not yet supported"); +} + +static LogicalResult handleSRem(LLVM::SRemOp op, LLVMTranslationState &st) { + auto &ctx = st.ctx; + auto &builder = ctx.getBuilder(); + FailureOr lhs = resolve(op.getLhs(), ctx); + if (failed(lhs)) + return op->emitOpError("unmapped operand in srem"); + + // Power-of-2 constant divisor -> v_and_b32 with (divisor - 1). + // Correct for non-negative dividends (thread IDs, indices). + if (auto constVal = getConstantInt(op.getRhs())) { + int64_t divisor = *constVal; + if (divisor > 0 && (divisor & (divisor - 1)) == 0) { + int64_t mask = divisor - 1; + auto immTy = ctx.createImmType(mask); + auto maskConst = ConstantOp::create(builder, op.getLoc(), immTy, mask); + auto vregTy = ctx.createVRegType(); + Value result = + V_AND_B32::create(builder, op.getLoc(), vregTy, *lhs, maskConst); + ctx.getMapper().mapValue(op.getResult(), result); + return success(); + } + } + return op->emitOpError( + "signed remainder by non-power-of-2 not yet supported"); +} + +//===----------------------------------------------------------------------===// +// Memory fence / barrier handlers +//===----------------------------------------------------------------------===// + +static LogicalResult handleFence(LLVM::FenceOp, LLVMTranslationState &) { + // Memory fences are handled implicitly by s_barrier and waitcnt insertion. + return success(); +} + +template +static LogicalResult handleBarrier(OpTy op, LLVMTranslationState &st) { + S_BARRIER::create(st.ctx.getBuilder(), op.getLoc()); + return success(); +} + +//===----------------------------------------------------------------------===// +// Vector shuffle handler +//===----------------------------------------------------------------------===// + +static LogicalResult handleShuffleVector(LLVM::ShuffleVectorOp op, + LLVMTranslationState &st) { + auto &ctx = st.ctx; + auto &builder = ctx.getBuilder(); + FailureOr src = resolve(op.getV1(), ctx); + if (failed(src)) + return op->emitOpError("unmapped operand in shufflevector"); + + // shufflevector with a single index extracts one element. + auto mask = op.getMask(); + if (mask.size() == 1) { + int64_t idx = mask[0]; + auto vregTy = ctx.createVRegType(); + Value extract = ExtractOp::create(builder, op.getLoc(), vregTy, *src, idx); + ctx.getMapper().mapValue(op.getResult(), extract); + return success(); + } + + return op->emitOpError("multi-element shufflevector not yet supported"); +} + +//===----------------------------------------------------------------------===// +// MFMA handler +//===----------------------------------------------------------------------===// + +static LogicalResult handleMFMA_F32_16x16x16_F16(ROCDL::mfma_f32_16x16x16f16 op, + LLVMTranslationState &st) { + auto &ctx = st.ctx; + auto &builder = ctx.getBuilder(); + auto loc = op.getLoc(); + + FailureOr a = resolve(op.getA(), ctx); + FailureOr b = resolve(op.getB(), ctx); + FailureOr c = resolve(op.getC(), ctx); + if (failed(a) || failed(b) || failed(c)) + return op->emitOpError("unmapped operand in MFMA"); + + // Result and accumulator are vector<4xf32> -> 4 VGPRs. + auto accTy = ctx.createVRegType(4, 4); + Value mfma = V_MFMA_F32_16X16X16_F16::create(builder, loc, accTy, *a, *b, *c); + ctx.getMapper().mapValue(op.getResult(), mfma); + return success(); +} + +//===----------------------------------------------------------------------===// +// SCF for/yield handler +//===----------------------------------------------------------------------===// + +/// Forward declaration for recursive op translation. +static LogicalResult translateOp(Operation *op, LLVMTranslationState &st); + +static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { + auto &ctx = st.ctx; + auto &builder = ctx.getBuilder(); + auto loc = op.getLoc(); + + FailureOr lb = resolve(op.getLowerBound(), ctx); + FailureOr ub = resolve(op.getUpperBound(), ctx); + FailureOr step = resolve(op.getStep(), ctx); + if (failed(lb) || failed(ub) || failed(step)) + return op->emitOpError("unmapped operand in scf.for"); + + // Build init args: [lower_bound, iter_args...]. + SmallVector initArgs; + initArgs.push_back(*lb); + for (Value arg : op.getInitArgs()) { + FailureOr resolved = resolve(arg, ctx); + if (failed(resolved)) + return op->emitOpError("unmapped init arg in scf.for"); + initArgs.push_back(*resolved); + } + + // Create the waveasm.loop (do-while semantics). + auto loopOp = LoopOp::create(builder, loc, initArgs); + Block &bodyBlock = loopOp.getBodyBlock(); + + // Map the induction variable (block arg 0). + ctx.getMapper().mapValue(op.getInductionVar(), bodyBlock.getArgument(0)); + + // Map iter_args (block args 1..N). + for (unsigned i = 0; i < op.getInitArgs().size(); ++i) + ctx.getMapper().mapValue(op.getRegionIterArgs()[i], + bodyBlock.getArgument(i + 1)); + + // Translate the loop body. + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(&bodyBlock); + for (Operation &bodyOp : op.getBody()->without_terminator()) + if (failed(translateOp(&bodyOp, st))) + return failure(); + + // Build loop increment and condition. + Value inductionVar = bodyBlock.getArgument(0); + auto ivTy = inductionVar.getType(); + Value nextIV = ArithAddOp::create(builder, loc, ivTy, inductionVar, *step); + // Condition must be SGPR/SCC for waveasm.condition. + auto sregTy = ctx.createSRegType(); + auto sccTy = ctx.createSCCType(); + Value scalarIV = nextIV; + Value scalarUB = *ub; + if (isVGPRType(nextIV.getType())) + scalarIV = V_READFIRSTLANE_B32::create(builder, loc, sregTy, nextIV); + if (isVGPRType(ub->getType())) + scalarUB = V_READFIRSTLANE_B32::create(builder, loc, sregTy, *ub); + Value cond = S_CMP_LT_U32::create(builder, loc, sccTy, scalarIV, scalarUB); + + // Collect iter args from yield. + auto yieldOp = cast(op.getBody()->getTerminator()); + SmallVector condIterArgs; + condIterArgs.push_back(nextIV); + for (Value v : yieldOp.getOperands()) { + FailureOr resolved = resolve(v, ctx); + if (failed(resolved)) + return op->emitOpError("unmapped yield operand in scf.for"); + condIterArgs.push_back(*resolved); + } + + ConditionOp::create(builder, loc, cond, condIterArgs); + + // Map loop results. scf.for results are iter_args only (no IV), + // but waveasm.loop results include the IV at index 0. + for (unsigned i = 0; i < op.getNumResults(); ++i) + ctx.getMapper().mapValue(op.getResult(i), loopOp.getResult(i + 1)); + + return success(); +} + //===----------------------------------------------------------------------===// // Op dispatch //===----------------------------------------------------------------------===// @@ -688,6 +1109,18 @@ static LogicalResult translateOp(Operation *op, LLVMTranslationState &st) { .Case([&](LLVM::GEPOp o) { return handleGEP(o, st); }) .Case([&](LLVM::LoadOp o) { return handleLoad(o, st); }) .Case([&](LLVM::StoreOp o) { return handleStore(o, st); }) + .Case([&](LLVM::AddressOfOp o) { return handleAddressOf(o, st); }) + .Case([&](LLVM::SDivOp o) { return handleSDiv(o, st); }) + .Case([&](LLVM::SRemOp o) { return handleSRem(o, st); }) + .Case([&](LLVM::FenceOp o) { return handleFence(o, st); }) + .Case([&](LLVM::ShuffleVectorOp o) { return handleShuffleVector(o, st); }) + .Case([&](ROCDL::BarrierOp o) { return handleBarrier(o, st); }) + .Case([&](ROCDL::SBarrierOp o) { return handleBarrier(o, st); }) + .Case([&](ROCDL::mfma_f32_16x16x16f16 o) { + return handleMFMA_F32_16x16x16_F16(o, st); + }) + .Case([&](scf::ForOp o) { return handleSCFFor(o, st); }) + .Case([&](scf::YieldOp) { return success(); }) .Default([](Operation *op) { return op->emitOpError("unhandled op in LLVM->WaveASM translation"); }); diff --git a/waveasm/test/Transforms/translate-from-llvm-error-multi-index-gep.mlir b/waveasm/test/Transforms/translate-from-llvm-error-multi-index-gep.mlir index aa926f03e7..42e61d33e7 100644 --- a/waveasm/test/Transforms/translate-from-llvm-error-multi-index-gep.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-error-multi-index-gep.mlir @@ -1,7 +1,7 @@ // RUN: not waveasm-translate %s --waveasm-translate-from-llvm 2>&1 | FileCheck %s // Verify that multi-index GEPs produce a diagnostic instead of crashing. -// CHECK: GEP with multiple indices not yet supported +// CHECK: buffer GEP must have a single index gpu.module @gpu_module { llvm.func @test(%arg0: !llvm.ptr) attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { From 69f565aad6f79d6ff842641f81dcb2eebb5e3f7a Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 16:52:04 +0200 Subject: [PATCH 03/19] Add lit tests for GEMM-related LLVM->WaveASM handlers - translate-from-llvm-mfma.mlir: MFMA, dense vector constant, shufflevector extract - translate-from-llvm-scf-for.mlir: scf.for -> waveasm.loop with condition - translate-from-llvm-lds.mlir: addressof, ptr<3> GEP, ds_read/ds_write - translate-from-llvm-barrier.mlir: rocdl.barrier, rocdl.s.barrier, llvm.fence - translate-from-llvm-sdiv-srem.mlir: power-of-2 sdiv/srem -> shift/mask Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Ivan Butygin --- .../translate-from-llvm-barrier.mlir | 19 ++++++++++ .../Transforms/translate-from-llvm-lds.mlir | 37 +++++++++++++++++++ .../Transforms/translate-from-llvm-mfma.mlir | 31 ++++++++++++++++ .../translate-from-llvm-scf-for.mlir | 29 +++++++++++++++ .../translate-from-llvm-sdiv-srem.mlir | 20 ++++++++++ 5 files changed, 136 insertions(+) create mode 100644 waveasm/test/Transforms/translate-from-llvm-barrier.mlir create mode 100644 waveasm/test/Transforms/translate-from-llvm-lds.mlir create mode 100644 waveasm/test/Transforms/translate-from-llvm-mfma.mlir create mode 100644 waveasm/test/Transforms/translate-from-llvm-scf-for.mlir create mode 100644 waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir diff --git a/waveasm/test/Transforms/translate-from-llvm-barrier.mlir b/waveasm/test/Transforms/translate-from-llvm-barrier.mlir new file mode 100644 index 0000000000..723419fa1a --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-barrier.mlir @@ -0,0 +1,19 @@ +// RUN: waveasm-translate %s --waveasm-translate-from-llvm | FileCheck %s +// Verify barrier and fence translation. +// rocdl.barrier -> s_barrier; rocdl.s.barrier -> s_barrier; llvm.fence -> no-op. + +// CHECK: waveasm.program @test__waveasm +// CHECK: waveasm.s_barrier +// CHECK: waveasm.s_barrier +// fence produces no output. +// CHECK-NOT: fence +// CHECK: waveasm.s_endpgm + +gpu.module @gpu_module attributes {llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"} { + llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + rocdl.barrier + rocdl.s.barrier + llvm.fence acquire + llvm.return + } +} diff --git a/waveasm/test/Transforms/translate-from-llvm-lds.mlir b/waveasm/test/Transforms/translate-from-llvm-lds.mlir new file mode 100644 index 0000000000..e1ab19d60c --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-lds.mlir @@ -0,0 +1,37 @@ +// RUN: waveasm-translate %s --waveasm-translate-from-llvm | FileCheck %s +// Verify LDS addressof, GEP, load, and store translation. +// addressof maps to VGPR zero; LDS GEPs produce arith.add offsets; +// loads/stores dispatch to ds_read/ds_write by width. + +// CHECK: waveasm.program @test__waveasm +// addressof -> v_mov_b32 of zero. +// CHECK: waveasm.v_mov_b32 +// GEP [0,0] on the array type passes through (all-zero indices). +// GEP with constant offset 512 produces an arith.add. +// CHECK: waveasm.arith.add +// 4-byte LDS load -> ds_read_b32. +// CHECK: waveasm.ds_read_b32 +// 4-byte LDS store -> ds_write_b32. +// CHECK: waveasm.ds_write_b32 +// CHECK: waveasm.s_endpgm + +gpu.module @gpu_module attributes {llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"} { + llvm.mlir.global private @alloca() {addr_space = 3 : i32} : !llvm.array<1024 x i8> + llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %0 = llvm.mlir.addressof @alloca : !llvm.ptr<3> + %1 = llvm.mlir.constant(42 : i32) : i32 + // Multi-index GEP with all-zero indices -> passthrough. + %2 = llvm.getelementptr %0[0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.array<1024 x i8> + // Constant-offset GEP. + %3 = llvm.getelementptr %2[512] : (!llvm.ptr<3>) -> !llvm.ptr<3>, i8 + // Dynamic-offset GEP. + %tid = rocdl.workitem.id.x range : i32 + %tidext = llvm.sext %tid : i32 to i64 + %4 = llvm.getelementptr nusw %2[%tidext] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i8 + // LDS load (4 bytes). + %5 = llvm.load %4 : !llvm.ptr<3> -> i32 + // LDS store (4 bytes). + llvm.store %1, %3 : i32, !llvm.ptr<3> + llvm.return + } +} diff --git a/waveasm/test/Transforms/translate-from-llvm-mfma.mlir b/waveasm/test/Transforms/translate-from-llvm-mfma.mlir new file mode 100644 index 0000000000..610de8ff93 --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-mfma.mlir @@ -0,0 +1,31 @@ +// RUN: waveasm-translate %s --waveasm-translate-from-llvm | FileCheck %s +// Verify MFMA and shufflevector (element extraction) translation. +// Dense vector constant -> wide v_mov_b32; mfma -> v_mfma; shuffle -> extract. + +// CHECK: waveasm.program @test__waveasm +// Zero-init accumulator (dense<0.0> : vector<4xf32>). +// CHECK: waveasm.v_mov_b32 {{.*}} -> !waveasm.vreg<4, 4> +// MFMA instruction. +// CHECK: waveasm.v_mfma_f32_16x16x16_f16 +// Extract elements from MFMA result. +// CHECK: waveasm.extract {{.*}}[0] : !waveasm.vreg<4, 4> -> !waveasm.vreg +// CHECK: waveasm.extract {{.*}}[1] : !waveasm.vreg<4, 4> -> !waveasm.vreg +// CHECK: waveasm.extract {{.*}}[2] : !waveasm.vreg<4, 4> -> !waveasm.vreg +// CHECK: waveasm.extract {{.*}}[3] : !waveasm.vreg<4, 4> -> !waveasm.vreg +// CHECK: waveasm.s_endpgm + +gpu.module @gpu_module attributes {llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"} { + llvm.func @test(%arg0: !llvm.ptr) attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %c0 = llvm.mlir.constant(dense<0.000000e+00> : vector<4xf32>) : vector<4xf32> + // Use zero for A and B inputs (vector<4xf16>). + %a = llvm.mlir.constant(dense<0.000000e+00> : vector<4xf16>) : vector<4xf16> + %b = llvm.mlir.constant(dense<0.000000e+00> : vector<4xf16>) : vector<4xf16> + %mfma = rocdl.mfma.f32.16x16x16f16 %a, %b, %c0, 0, 0, 0 : (vector<4xf16>, vector<4xf16>, vector<4xf32>) -> vector<4xf32> + // Extract individual elements. + %e0 = llvm.shufflevector %mfma, %mfma [0] : vector<4xf32> + %e1 = llvm.shufflevector %mfma, %mfma [1] : vector<4xf32> + %e2 = llvm.shufflevector %mfma, %mfma [2] : vector<4xf32> + %e3 = llvm.shufflevector %mfma, %mfma [3] : vector<4xf32> + llvm.return + } +} diff --git a/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir b/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir new file mode 100644 index 0000000000..bdacf1be35 --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir @@ -0,0 +1,29 @@ +// RUN: waveasm-translate %s --waveasm-translate-from-llvm | FileCheck %s +// Verify scf.for translation to waveasm.loop with condition terminator. +// The loop body contains an add; iter_args carry the accumulator. + +// CHECK: waveasm.program @test__waveasm +// Loop with one iter_arg (the accumulator). +// CHECK: waveasm.loop +// Loop body: arith.add for the accumulation. +// CHECK: waveasm.arith.add +// IV increment. +// CHECK: waveasm.arith.add +// Condition: s_cmp_lt_u32 for the loop back-edge. +// CHECK: waveasm.s_cmp_lt_u32 +// CHECK: waveasm.condition +// CHECK: waveasm.s_endpgm + +gpu.module @gpu_module attributes {llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"} { + llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %lb = llvm.mlir.constant(0 : i32) : i32 + %ub = llvm.mlir.constant(64 : i32) : i32 + %step = llvm.mlir.constant(16 : i32) : i32 + %init = llvm.mlir.constant(0 : i32) : i32 + %result = scf.for %iv = %lb to %ub step %step iter_args(%acc = %init) -> (i32) : i32 { + %sum = llvm.add %acc, %iv : i32 + scf.yield %sum : i32 + } + llvm.return + } +} diff --git a/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir b/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir new file mode 100644 index 0000000000..ab509fd906 --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir @@ -0,0 +1,20 @@ +// RUN: waveasm-translate %s --waveasm-translate-from-llvm | FileCheck %s +// Verify sdiv and srem with power-of-2 constants. +// sdiv -> v_ashrrev_i32; srem -> v_and_b32. + +// CHECK: waveasm.program @test__waveasm +// sdiv by 16 -> shift right by 4. +// CHECK: waveasm.v_ashrrev_i32 +// srem by 16 -> and with 15. +// CHECK: waveasm.v_and_b32 +// CHECK: waveasm.s_endpgm + +gpu.module @gpu_module attributes {llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"} { + llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %tid = rocdl.workitem.id.x range : i32 + %16 = llvm.mlir.constant(16 : i32) : i32 + %div = llvm.sdiv %tid, %16 : i32 + %rem = llvm.srem %tid, %16 : i32 + llvm.return + } +} From 7b864b2cd011e0a1d4f37b8d3c08b9d105c444df Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 17:24:29 +0200 Subject: [PATCH 04/19] Propagate correct register widths in LLVM->WaveASM translation inferResultType and handleCastOp always produced 1-wide register types regardless of the LLVM operation's actual bit width. This made arith.trunc a no-op for i64->i32 truncations (both sides were vreg), so legalization removed it as a passthrough before the sext it guarded was expanded to its true 2-wide form -- leaking vreg<2,2> values into LDS address chains and producing invalid ds_read/ds_write assembly. Fix inferResultType to propagate max operand width and handleCastOp to derive width from the LLVM result type. Now sext i32->i64 produces vreg<2,2>, trunc i64->i32 produces vreg, and legalization handles both correctly without the trunc being silently eliminated. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Ivan Butygin --- .../Transforms/TranslateFromLLVMDialect.cpp | 21 +++++++++++++------ .../translate-from-llvm-barrier.mlir | 2 +- .../Transforms/translate-from-llvm-lds.mlir | 2 +- .../Transforms/translate-from-llvm-mfma.mlir | 2 +- .../translate-from-llvm-scf-for.mlir | 2 +- .../translate-from-llvm-sdiv-srem.mlir | 2 +- 6 files changed, 20 insertions(+), 11 deletions(-) diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 4904b48122..4408cc5257 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -193,12 +193,16 @@ static Value truncToI32(Value v, Type llvmType, OpBuilder &builder, } /// Infer the pseudo-op result type from operand types. -/// If any operand is VGPR -> VReg; otherwise SReg. +/// Register file: VGPR if any operand is VGPR, else SGPR. +/// Width: max operand width (in 32-bit dwords). static Type inferResultType(ValueRange operands, TranslationContext &ctx) { + int64_t width = 1; + for (Value v : operands) + width = std::max(width, getRegSize(v.getType())); for (Value v : operands) if (isVGPRType(v.getType())) - return ctx.createVRegType(); - return ctx.createSRegType(); + return ctx.createVRegType(width, width); + return ctx.createSRegType(width, width); } /// Try to extract a constant integer from an LLVM SSA value. @@ -396,8 +400,9 @@ static LogicalResult handleWorkgroupId(OpTy op, LLVMTranslationState &st, return success(); } -// Emit arith pseudo-ops for i32/i64 casts -- legalization pass handles width. /// Translate an LLVM cast op to a WaveASM arithmetic pseudo-op. +/// Width is derived from the LLVM result type so that sext i32->i64 +/// produces a 2-wide register and trunc i64->i32 produces a 1-wide one. template static LogicalResult handleCastOp(LLVMOp op, LLVMTranslationState &st) { auto &ctx = st.ctx; @@ -405,8 +410,12 @@ static LogicalResult handleCastOp(LLVMOp op, LLVMTranslationState &st) { FailureOr src = resolve(op.getOperand(), ctx); if (failed(src)) return op->emitOpError("unmapped operand in cast"); - Type resTy = isVGPRType(src->getType()) ? (Type)ctx.createVRegType() - : (Type)ctx.createSRegType(); + int64_t width = 1; + if (auto intTy = dyn_cast(op.getResult().getType())) + width = (intTy.getWidth() + 31) / 32; + Type resTy = isVGPRType(src->getType()) + ? (Type)ctx.createVRegType(width, width) + : (Type)ctx.createSRegType(width, width); Value pseudo = WaveASMOp::create(builder, op.getLoc(), resTy, *src); ctx.getMapper().mapValue(op.getResult(), pseudo); return success(); diff --git a/waveasm/test/Transforms/translate-from-llvm-barrier.mlir b/waveasm/test/Transforms/translate-from-llvm-barrier.mlir index 723419fa1a..eef8ddbc94 100644 --- a/waveasm/test/Transforms/translate-from-llvm-barrier.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-barrier.mlir @@ -9,7 +9,7 @@ // CHECK-NOT: fence // CHECK: waveasm.s_endpgm -gpu.module @gpu_module attributes {llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"} { +gpu.module @gpu_module { llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { rocdl.barrier rocdl.s.barrier diff --git a/waveasm/test/Transforms/translate-from-llvm-lds.mlir b/waveasm/test/Transforms/translate-from-llvm-lds.mlir index e1ab19d60c..b2e6bb497f 100644 --- a/waveasm/test/Transforms/translate-from-llvm-lds.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-lds.mlir @@ -15,7 +15,7 @@ // CHECK: waveasm.ds_write_b32 // CHECK: waveasm.s_endpgm -gpu.module @gpu_module attributes {llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"} { +gpu.module @gpu_module { llvm.mlir.global private @alloca() {addr_space = 3 : i32} : !llvm.array<1024 x i8> llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { %0 = llvm.mlir.addressof @alloca : !llvm.ptr<3> diff --git a/waveasm/test/Transforms/translate-from-llvm-mfma.mlir b/waveasm/test/Transforms/translate-from-llvm-mfma.mlir index 610de8ff93..1cee2889da 100644 --- a/waveasm/test/Transforms/translate-from-llvm-mfma.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-mfma.mlir @@ -14,7 +14,7 @@ // CHECK: waveasm.extract {{.*}}[3] : !waveasm.vreg<4, 4> -> !waveasm.vreg // CHECK: waveasm.s_endpgm -gpu.module @gpu_module attributes {llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"} { +gpu.module @gpu_module { llvm.func @test(%arg0: !llvm.ptr) attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { %c0 = llvm.mlir.constant(dense<0.000000e+00> : vector<4xf32>) : vector<4xf32> // Use zero for A and B inputs (vector<4xf16>). diff --git a/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir b/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir index bdacf1be35..fb9fc4cb44 100644 --- a/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir @@ -14,7 +14,7 @@ // CHECK: waveasm.condition // CHECK: waveasm.s_endpgm -gpu.module @gpu_module attributes {llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"} { +gpu.module @gpu_module { llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { %lb = llvm.mlir.constant(0 : i32) : i32 %ub = llvm.mlir.constant(64 : i32) : i32 diff --git a/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir b/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir index ab509fd906..1eb95e1add 100644 --- a/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir @@ -9,7 +9,7 @@ // CHECK: waveasm.v_and_b32 // CHECK: waveasm.s_endpgm -gpu.module @gpu_module attributes {llvm.data_layout = "e-p:64:64-p1:64:64-p2:32:32-p3:32:32-p4:64:64-p5:32:32-p6:32:32-p7:160:256:256:32-p8:128:128:128:48-p9:192:256:256:32-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-v2048:2048-n32:64-S32-A5-G1-ni:7:8:9"} { +gpu.module @gpu_module { llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { %tid = rocdl.workitem.id.x range : i32 %16 = llvm.mlir.constant(16 : i32) : i32 From 78b11756fed8f3716bb021000c16a697040529e7 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 20:48:53 +0200 Subject: [PATCH 05/19] Restrict GEMM waveasm test to CDNA4 The waveasm backend only supports gfx950 for now. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Ivan Butygin --- tests/kernel/e2e/test_gemm_waveasm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/kernel/e2e/test_gemm_waveasm.py b/tests/kernel/e2e/test_gemm_waveasm.py index eedd97e06c..4f49ae9249 100644 --- a/tests/kernel/e2e/test_gemm_waveasm.py +++ b/tests/kernel/e2e/test_gemm_waveasm.py @@ -16,11 +16,11 @@ from wave_lang.kernel.wave.utils.run_utils import set_default_run_config from wave_lang.kernel.wave.utils.torch_utils import device_randn, device_zeros -from ..common.utils import require_cdna_3_or_4, require_e2e +from ..common.utils import require_cdna4, require_e2e @require_e2e -@require_cdna_3_or_4 +@require_cdna4 @pytest.mark.parametrize( "shape,block_shape,waves_per_block", [ From 05835d4d69fdf4a9c1b6d9290b13219d192ef299 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 21:07:04 +0200 Subject: [PATCH 06/19] Drop isLDSBase tracking, check LLVM pointer address space directly The LDS base value set was redundant -- the LLVM load/store address operand already carries the pointer type with the address space. Check getLLVMAddrSpace(op.getAddr()) == 3 instead. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Ivan Butygin --- .../Transforms/TranslateFromLLVMDialect.cpp | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 4408cc5257..6af65fe85d 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -87,15 +87,10 @@ class LLVMTranslationState { void setBaseOffset(Value ptr, Value offset) { baseOffsets[ptr] = offset; } Value lookupBaseOffset(Value ptr) const { return baseOffsets.lookup(ptr); } - /// Track LDS base values (from addressof / ptr<3> GEPs). - void setLDSBase(Value v) { ldsBaseValues.insert(v); } - bool isLDSBase(Value v) const { return ldsBaseValues.contains(v); } - private: DenseMap rsrcToSRD; DenseMap gepMap; DenseMap baseOffsets; - DenseSet ldsBaseValues; }; //===----------------------------------------------------------------------===// @@ -205,6 +200,13 @@ static Type inferResultType(ValueRange operands, TranslationContext &ctx) { return ctx.createSRegType(width, width); } +/// Return the LLVM pointer address space, or 0 for non-pointer types. +static unsigned getLLVMAddrSpace(Value v) { + if (auto ptrTy = dyn_cast(v.getType())) + return ptrTy.getAddressSpace(); + return 0; +} + /// Try to extract a constant integer from an LLVM SSA value. static std::optional getConstantInt(Value v) { if (auto constOp = v.getDefiningOp()) @@ -605,7 +607,6 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { Value sum = ArithAddOp::create(builder, loc, vregTy, *baseOff, off); ctx.getMapper().mapValue(op.getResult(), sum); } - st.setLDSBase(op.getResult()); return success(); } @@ -694,7 +695,7 @@ static LogicalResult handleLoad(LLVM::LoadOp op, LLVMTranslationState &st) { auto loc = op.getLoc(); // LDS load (ptr<3>). - if (st.isLDSBase(op.getAddr())) + if (getLLVMAddrSpace(op.getAddr()) == 3) return handleLDSLoad(op, op.getAddr(), st); std::optional ptr = st.lookupGEP(op.getAddr()); @@ -744,7 +745,7 @@ static LogicalResult handleStore(LLVM::StoreOp op, LLVMTranslationState &st) { auto loc = op.getLoc(); // LDS store (ptr<3>). - if (st.isLDSBase(op.getAddr())) + if (getLLVMAddrSpace(op.getAddr()) == 3) return handleLDSStore(op, op.getAddr(), st); std::optional ptr = st.lookupGEP(op.getAddr()); @@ -803,7 +804,6 @@ static LogicalResult handleAddressOf(LLVM::AddressOfOp op, auto vregTy = ctx.createVRegType(); Value mov = V_MOV_B32::create(builder, op.getLoc(), vregTy, zero); ctx.getMapper().mapValue(op.getResult(), mov); - st.setLDSBase(op.getResult()); return success(); } From 71f2692cf61db42a24141ad0ac0ed251a13be14e Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 21:10:51 +0200 Subject: [PATCH 07/19] Remove incorrect multi-index GEP offset accumulation Summing GEP indices is wrong -- each index operates at a different level of the type hierarchy and must be multiplied by the element size at that level. The all-zero multi-index case (the only one we hit in practice) was already handled correctly. Remove the bogus fallback and return nullopt for unhandled multi-index GEPs. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Ivan Butygin --- .../Transforms/TranslateFromLLVMDialect.cpp | 34 ++++++------------- 1 file changed, 11 insertions(+), 23 deletions(-) diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 6af65fe85d..6186d635c8 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -238,31 +238,19 @@ static std::optional computeGEPOffset(LLVM::GEPOp op, if (allZero) return std::nullopt; - // Single dynamic index -- common case. - if (indices.size() == 1) { - auto idx = indices[0]; - if (auto val = idx.dyn_cast()) { - FailureOr resolved = resolve(val, ctx); - if (failed(resolved)) - return std::nullopt; - return *resolved; - } - int64_t constVal = cast(idx).getInt(); - auto immTy = ctx.createImmType(constVal); - return ConstantOp::create(builder, loc, immTy, constVal)->getResult(0); - } - - // Multi-index: accumulate constant indices. - int64_t totalOffset = 0; - for (auto idx : indices) { - if (idx.dyn_cast()) + // Single index -- the only case we handle (element type gives byte stride). + if (indices.size() != 1) + return std::nullopt; + auto idx = indices[0]; + if (auto val = idx.dyn_cast()) { + FailureOr resolved = resolve(val, ctx); + if (failed(resolved)) return std::nullopt; - totalOffset += cast(idx).getInt(); + return *resolved; } - if (totalOffset == 0) - return std::nullopt; - auto immTy = ctx.createImmType(totalOffset); - return ConstantOp::create(builder, loc, immTy, totalOffset)->getResult(0); + int64_t constVal = cast(idx).getInt(); + auto immTy = ctx.createImmType(constVal); + return ConstantOp::create(builder, loc, immTy, constVal)->getResult(0); } //===----------------------------------------------------------------------===// From c36e56c4586de61444c1e2fe906a0c15ab2d22b7 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 21:27:53 +0200 Subject: [PATCH 08/19] Fix computeGEPOffset: account for element type size, use proper types The GEP offset computation was treating all indices as byte offsets regardless of element type. Rename to computeGEPByteOffset and multiply dynamic indices by getGEPElementBytes(). Replace local getConstantInt with getConstantIntValue from StaticValueUtils. Use explicit types instead of auto where the type is not obvious. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Ivan Butygin --- .../Transforms/TranslateFromLLVMDialect.cpp | 91 ++++++++++--------- 1 file changed, 47 insertions(+), 44 deletions(-) diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 6186d635c8..1169496baf 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -208,49 +208,55 @@ static unsigned getLLVMAddrSpace(Value v) { } /// Try to extract a constant integer from an LLVM SSA value. -static std::optional getConstantInt(Value v) { - if (auto constOp = v.getDefiningOp()) - if (auto intAttr = dyn_cast(constOp.getValue())) - return intAttr.getValue().getSExtValue(); - return std::nullopt; +/// Return the byte stride for a single-index GEP element type. +/// Only handles sized scalars, vectors, and arrays of i8 (byte addressing). +/// Returns 0 for unsupported types. +static int64_t getGEPElementBytes(Type elemTy) { + if (elemTy.isIntOrFloat()) + return elemTy.getIntOrFloatBitWidth() / 8; + if (auto vecTy = dyn_cast(elemTy)) + return vecTy.getNumElements() * + vecTy.getElementType().getIntOrFloatBitWidth() / 8; + if (auto arrTy = dyn_cast(elemTy)) + return getGEPElementBytes(arrTy.getElementType()) * arrTy.getNumElements(); + return 0; } -/// Compute the byte offset for a GEP as a single WaveASM Value. -/// For all-zero constant indices, returns std::nullopt (offset is 0). -static std::optional computeGEPOffset(LLVM::GEPOp op, - TranslationContext &ctx) { - auto &builder = ctx.getBuilder(); - auto loc = op.getLoc(); +/// Compute the byte offset for a single-index GEP. +/// Returns std::nullopt when the offset is zero or unsupported. +static std::optional computeGEPByteOffset(LLVM::GEPOp op, + TranslationContext &ctx) { + OpBuilder &builder = ctx.getBuilder(); + Location loc = op.getLoc(); auto indices = op.getIndices(); - // Check if all indices are constant zero. - bool allZero = true; - for (auto idx : indices) { - if (auto val = idx.dyn_cast()) { - allZero = false; - break; - } - if (cast(idx).getInt() != 0) { - allZero = false; - break; - } - } - if (allZero) - return std::nullopt; - - // Single index -- the only case we handle (element type gives byte stride). + // Only handle single-index GEPs. Multi-index GEPs walk nested types + // and require full DataLayout support. if (indices.size() != 1) return std::nullopt; + + int64_t elemBytes = getGEPElementBytes(op.getElemType()); + auto idx = indices[0]; - if (auto val = idx.dyn_cast()) { - FailureOr resolved = resolve(val, ctx); + if (Value dynIdx = idx.dyn_cast()) { + FailureOr resolved = resolve(dynIdx, ctx); if (failed(resolved)) return std::nullopt; - return *resolved; + if (elemBytes <= 1) + return *resolved; + // Scale by element size: offset = index * elemBytes. + ImmType scaleTy = ctx.createImmType(elemBytes); + Value scale = ConstantOp::create(builder, loc, scaleTy, elemBytes); + Type resTy = inferResultType({*resolved, scale}, ctx); + return ArithMulOp::create(builder, loc, resTy, *resolved, scale); } - int64_t constVal = cast(idx).getInt(); - auto immTy = ctx.createImmType(constVal); - return ConstantOp::create(builder, loc, immTy, constVal)->getResult(0); + + int64_t constIdx = cast(idx).getInt(); + if (constIdx == 0) + return std::nullopt; + int64_t byteOffset = constIdx * std::max(elemBytes, int64_t{1}); + ImmType immTy = ctx.createImmType(byteOffset); + return ConstantOp::create(builder, loc, immTy, byteOffset); } //===----------------------------------------------------------------------===// @@ -571,10 +577,7 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { auto loc = op.getLoc(); Value base = op.getBase(); - auto baseTy = op.getBase().getType(); - unsigned addrSpace = 0; - if (auto ptrTy = dyn_cast(baseTy)) - addrSpace = ptrTy.getAddressSpace(); + unsigned addrSpace = getLLVMAddrSpace(op.getBase()); // LDS GEP (ptr<3>): compute byte offset for ds_read/ds_write. // DS instructions only accept a 32-bit vaddr, so truncate wide offsets. @@ -582,10 +585,10 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { FailureOr baseOff = resolve(base, ctx); if (failed(baseOff)) return op->emitOpError("unmapped LDS GEP base"); - auto vregTy = ctx.createVRegType(); + VRegType vregTy = ctx.createVRegType(); if (getRegSize(baseOff->getType()) > 1) *baseOff = ArithTruncOp::create(builder, loc, vregTy, *baseOff); - auto maybeOffset = computeGEPOffset(op, ctx); + std::optional maybeOffset = computeGEPByteOffset(op, ctx); if (!maybeOffset) { ctx.getMapper().mapValue(op.getResult(), *baseOff); } else { @@ -605,9 +608,9 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { // make.buffer.rsrc. Propagate the mapper entry and accumulate the byte // offset so it can be folded into the SRD base later. if (addrSpace == 0) { - auto maybeOffset = computeGEPOffset(op, ctx); + std::optional maybeOffset = computeGEPByteOffset(op, ctx); // Forward mapper entry so make.buffer.rsrc can find the SRD. - if (auto mapped = ctx.getMapper().getMapped(base)) + if (std::optional mapped = ctx.getMapper().getMapped(base)) ctx.getMapper().mapValue(op.getResult(), *mapped); if (!maybeOffset) { @@ -632,7 +635,7 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { auto indices = op.getIndices(); if (indices.size() != 1) return op->emitOpError("buffer GEP must have a single index"); - auto idx = indices[0].dyn_cast(); + Value idx = indices[0].template dyn_cast(); if (!idx) return op->emitOpError("buffer GEP with constant index not yet supported"); @@ -876,7 +879,7 @@ static LogicalResult handleSDiv(LLVM::SDivOp op, LLVMTranslationState &st) { return op->emitOpError("unmapped operand in sdiv"); // Power-of-2 constant divisor -> arithmetic shift right. - if (auto constVal = getConstantInt(op.getRhs())) { + if (auto constVal = getConstantIntValue(op.getRhs())) { int64_t divisor = *constVal; if (divisor > 0 && (divisor & (divisor - 1)) == 0) { int64_t shiftAmt = llvm::Log2_64(divisor); @@ -902,7 +905,7 @@ static LogicalResult handleSRem(LLVM::SRemOp op, LLVMTranslationState &st) { // Power-of-2 constant divisor -> v_and_b32 with (divisor - 1). // Correct for non-negative dividends (thread IDs, indices). - if (auto constVal = getConstantInt(op.getRhs())) { + if (auto constVal = getConstantIntValue(op.getRhs())) { int64_t divisor = *constVal; if (divisor > 0 && (divisor & (divisor - 1)) == 0) { int64_t mask = divisor - 1; From 5145ebcdf18a1286daf94f2a04c3ff81a7f5bc85 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 21:32:19 +0200 Subject: [PATCH 09/19] Add lit tests for unsupported GEP configurations - Buffer GEP with constant attr index (not a dynamic Value). - GEP on unsupported address space (e.g. ptr<5>). Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Ivan Butygin --- .../translate-from-llvm-error-gep-addrspace.mlir | 12 ++++++++++++ ...ranslate-from-llvm-error-gep-const-index.mlir | 16 ++++++++++++++++ 2 files changed, 28 insertions(+) create mode 100644 waveasm/test/Transforms/translate-from-llvm-error-gep-addrspace.mlir create mode 100644 waveasm/test/Transforms/translate-from-llvm-error-gep-const-index.mlir diff --git a/waveasm/test/Transforms/translate-from-llvm-error-gep-addrspace.mlir b/waveasm/test/Transforms/translate-from-llvm-error-gep-addrspace.mlir new file mode 100644 index 0000000000..c0c2fd6472 --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-error-gep-addrspace.mlir @@ -0,0 +1,12 @@ +// RUN: not waveasm-translate %s --waveasm-translate-from-llvm 2>&1 | FileCheck %s +// Verify that GEPs on unsupported address spaces produce a diagnostic. + +// CHECK: unsupported address space 5 + +gpu.module @gpu_module { + llvm.func @test(%arg0: !llvm.ptr<5>) attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %0 = llvm.mlir.constant(0 : i32) : i32 + %1 = llvm.getelementptr %arg0[%0] : (!llvm.ptr<5>, i32) -> !llvm.ptr<5>, i8 + llvm.return + } +} diff --git a/waveasm/test/Transforms/translate-from-llvm-error-gep-const-index.mlir b/waveasm/test/Transforms/translate-from-llvm-error-gep-const-index.mlir new file mode 100644 index 0000000000..f107d750d1 --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-error-gep-const-index.mlir @@ -0,0 +1,16 @@ +// RUN: not waveasm-translate %s --waveasm-translate-from-llvm 2>&1 | FileCheck %s +// Verify that buffer GEPs with constant attr indices (not dynamic Values) +// produce a diagnostic. + +// CHECK: buffer GEP with constant index not yet supported + +gpu.module @gpu_module { + llvm.func @test(%arg0: !llvm.ptr) attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %0 = llvm.mlir.constant(0 : i16) : i16 + %1 = llvm.mlir.constant(2147483645 : i64) : i64 + %2 = llvm.mlir.constant(822243328 : i32) : i32 + %3 = rocdl.make.buffer.rsrc %arg0, %0, %1, %2 : !llvm.ptr to <7> + %4 = llvm.getelementptr %3[42] : (!llvm.ptr<7>) -> !llvm.ptr<7>, i8 + llvm.return + } +} From 593790f544d39017de42f8bda1ddeb2bc8389096 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 21:55:00 +0200 Subject: [PATCH 10/19] Fix LLVM->WaveASM edge-case lowering semantics Reject unsupported non-zero structural GEPs and preserve signed power-of-2 div/rem semantics for negative inputs. Keep scf.for loop control scalar and add regression tests for the guarded cases. Made-with: Cursor Signed-off-by: Ivan Butygin --- .../Transforms/TranslateFromLLVMDialect.cpp | 170 ++++++++++++++---- ...m-llvm-error-bare-ptr-multi-index-gep.mlir | 12 ++ ...translate-from-llvm-error-scf-for-i64.mlir | 17 ++ .../translate-from-llvm-scf-for.mlir | 4 +- .../translate-from-llvm-sdiv-srem.mlir | 19 +- 5 files changed, 183 insertions(+), 39 deletions(-) create mode 100644 waveasm/test/Transforms/translate-from-llvm-error-bare-ptr-multi-index-gep.mlir create mode 100644 waveasm/test/Transforms/translate-from-llvm-error-scf-for-i64.mlir diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 1169496baf..d56abfc96d 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -222,6 +222,22 @@ static int64_t getGEPElementBytes(Type elemTy) { return 0; } +/// Return true iff a GEP index is statically known to be zero. +static bool isZeroGEPIndex(llvm::PointerUnion idx) { + if (isa(idx)) + return isConstantIntValue(cast(idx), 0); + return isConstantIntValue(cast(idx), 0); +} + +/// Structural GEPs like [0, 0] are a no-op and can be forwarded even though we +/// do not model nested aggregate layouts yet. +static bool isAllZeroIndexGEP(LLVM::GEPOp op) { + for (auto idx : op.getIndices()) + if (!isZeroGEPIndex(idx)) + return false; + return true; +} + /// Compute the byte offset for a single-index GEP. /// Returns std::nullopt when the offset is zero or unsupported. static std::optional computeGEPByteOffset(LLVM::GEPOp op, @@ -576,19 +592,26 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { auto &builder = ctx.getBuilder(); auto loc = op.getLoc(); Value base = op.getBase(); + auto indices = op.getIndices(); + bool isSingleIndex = indices.size() == 1; + bool isAllZeroGEP = isAllZeroIndexGEP(op); unsigned addrSpace = getLLVMAddrSpace(op.getBase()); // LDS GEP (ptr<3>): compute byte offset for ds_read/ds_write. // DS instructions only accept a 32-bit vaddr, so truncate wide offsets. if (addrSpace == 3) { + if (!isSingleIndex && !isAllZeroGEP) + return op->emitOpError( + "LDS GEP must have a single index or all-zero structural indices"); FailureOr baseOff = resolve(base, ctx); if (failed(baseOff)) return op->emitOpError("unmapped LDS GEP base"); VRegType vregTy = ctx.createVRegType(); if (getRegSize(baseOff->getType()) > 1) *baseOff = ArithTruncOp::create(builder, loc, vregTy, *baseOff); - std::optional maybeOffset = computeGEPByteOffset(op, ctx); + std::optional maybeOffset = + isSingleIndex ? computeGEPByteOffset(op, ctx) : std::nullopt; if (!maybeOffset) { ctx.getMapper().mapValue(op.getResult(), *baseOff); } else { @@ -608,7 +631,11 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { // make.buffer.rsrc. Propagate the mapper entry and accumulate the byte // offset so it can be folded into the SRD base later. if (addrSpace == 0) { - std::optional maybeOffset = computeGEPByteOffset(op, ctx); + if (!isSingleIndex && !isAllZeroGEP) + return op->emitOpError("bare-pointer GEP must have a single index or " + "all-zero structural indices"); + std::optional maybeOffset = + isSingleIndex ? computeGEPByteOffset(op, ctx) : std::nullopt; // Forward mapper entry so make.buffer.rsrc can find the SRD. if (std::optional mapped = ctx.getMapper().getMapped(base)) ctx.getMapper().mapValue(op.getResult(), *mapped); @@ -632,7 +659,6 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { } // Buffer GEP (ptr<7>): single dynamic index. - auto indices = op.getIndices(); if (indices.size() != 1) return op->emitOpError("buffer GEP must have a single index"); Value idx = indices[0].template dyn_cast(); @@ -874,52 +900,116 @@ static LogicalResult handleLDSStore(LLVM::StoreOp op, Value addr, static LogicalResult handleSDiv(LLVM::SDivOp op, LLVMTranslationState &st) { auto &ctx = st.ctx; auto &builder = ctx.getBuilder(); + auto loc = op.getLoc(); FailureOr lhs = resolve(op.getLhs(), ctx); if (failed(lhs)) return op->emitOpError("unmapped operand in sdiv"); + if (getRegSize(lhs->getType()) != 1) + return op->emitOpError( + "signed division currently supports only i32 operands"); - // Power-of-2 constant divisor -> arithmetic shift right. + // Signed division by a positive power-of-2 constant: + // q = (x + (x < 0 ? divisor - 1 : 0)) >> log2(divisor) + // This preserves LLVM's trunc-toward-zero semantics for negative dividends. if (auto constVal = getConstantIntValue(op.getRhs())) { int64_t divisor = *constVal; if (divisor > 0 && (divisor & (divisor - 1)) == 0) { int64_t shiftAmt = llvm::Log2_64(divisor); - auto immTy = ctx.createImmType(shiftAmt); - auto shiftConst = - ConstantOp::create(builder, op.getLoc(), immTy, shiftAmt); - auto vregTy = ctx.createVRegType(); - Value result = - V_ASHRREV_I32::create(builder, op.getLoc(), vregTy, shiftConst, *lhs); + auto zero = ConstantOp::create(builder, loc, ctx.createImmType(0), 0); + auto biasImm = ConstantOp::create( + builder, loc, ctx.createImmType(divisor - 1), divisor - 1); + auto shiftConst = ConstantOp::create( + builder, loc, ctx.createImmType(shiftAmt), shiftAmt); + Value result; + if (isSGPRType(lhs->getType())) { + auto sregTy = ctx.createSRegType(); + auto sccTy = ctx.createSCCType(); + Value isNegative = + S_CMP_LT_I32::create(builder, loc, sccTy, *lhs, zero); + Value bias = S_CSELECT_B32::create(builder, loc, sregTy, isNegative, + biasImm, zero); + Value biased = ArithAddOp::create(builder, loc, sregTy, *lhs, bias); + result = + S_ASHR_I32::create(builder, loc, sregTy, sccTy, biased, shiftConst) + .getDst(); + } else { + auto vregTy = ctx.createVRegType(); + Value isNegative = ArithCmpOp::create(builder, loc, vregTy, + CmpPredicate::slt, *lhs, zero); + Value bias = ArithSelectOp::create(builder, loc, vregTy, zero, biasImm, + isNegative); + Value biased = ArithAddOp::create(builder, loc, vregTy, *lhs, bias); + result = + V_ASHRREV_I32::create(builder, loc, vregTy, shiftConst, biased); + } ctx.getMapper().mapValue(op.getResult(), result); return success(); } } - return op->emitOpError("signed division by non-power-of-2 not yet supported"); + return op->emitOpError( + "signed division currently supports only positive power-of-2 constants"); } static LogicalResult handleSRem(LLVM::SRemOp op, LLVMTranslationState &st) { auto &ctx = st.ctx; auto &builder = ctx.getBuilder(); + auto loc = op.getLoc(); FailureOr lhs = resolve(op.getLhs(), ctx); if (failed(lhs)) return op->emitOpError("unmapped operand in srem"); - - // Power-of-2 constant divisor -> v_and_b32 with (divisor - 1). - // Correct for non-negative dividends (thread IDs, indices). + if (getRegSize(lhs->getType()) != 1) + return op->emitOpError( + "signed remainder currently supports only i32 operands"); + + // Signed remainder by a positive power-of-2 constant: + // r = x & (divisor - 1) + // if (x < 0 && r != 0) r -= divisor + // This keeps the remainder's sign consistent with LLVM srem. if (auto constVal = getConstantIntValue(op.getRhs())) { int64_t divisor = *constVal; if (divisor > 0 && (divisor & (divisor - 1)) == 0) { - int64_t mask = divisor - 1; - auto immTy = ctx.createImmType(mask); - auto maskConst = ConstantOp::create(builder, op.getLoc(), immTy, mask); - auto vregTy = ctx.createVRegType(); - Value result = - V_AND_B32::create(builder, op.getLoc(), vregTy, *lhs, maskConst); + auto zero = ConstantOp::create(builder, loc, ctx.createImmType(0), 0); + auto maskConst = ConstantOp::create( + builder, loc, ctx.createImmType(divisor - 1), divisor - 1); + auto negDivisor = ConstantOp::create( + builder, loc, ctx.createImmType(-divisor), -divisor); + Value result; + if (isSGPRType(lhs->getType())) { + auto sregTy = ctx.createSRegType(); + auto sccTy = ctx.createSCCType(); + Value rawRem = + ArithAndOp::create(builder, loc, sregTy, *lhs, maskConst); + Value isNegative = + S_CMP_LT_I32::create(builder, loc, sccTy, *lhs, zero); + Value isNonZero = + S_CMP_NE_I32::create(builder, loc, sccTy, rawRem, zero); + Value adjusted = + ArithAddOp::create(builder, loc, sregTy, rawRem, negDivisor); + Value maybeAdjusted = S_CSELECT_B32::create( + builder, loc, sregTy, isNonZero, adjusted, rawRem); + result = S_CSELECT_B32::create(builder, loc, sregTy, isNegative, + maybeAdjusted, rawRem); + } else { + auto vregTy = ctx.createVRegType(); + Value rawRem = + ArithAndOp::create(builder, loc, vregTy, *lhs, maskConst); + Value isNegative = ArithCmpOp::create(builder, loc, vregTy, + CmpPredicate::slt, *lhs, zero); + Value isNonZero = ArithCmpOp::create(builder, loc, vregTy, + CmpPredicate::ne, rawRem, zero); + Value adjusted = + ArithAddOp::create(builder, loc, vregTy, rawRem, negDivisor); + Value maybeAdjusted = ArithSelectOp::create( + builder, loc, vregTy, rawRem, adjusted, isNonZero); + result = ArithSelectOp::create(builder, loc, vregTy, rawRem, + maybeAdjusted, isNegative); + } ctx.getMapper().mapValue(op.getResult(), result); return success(); } } return op->emitOpError( - "signed remainder by non-power-of-2 not yet supported"); + "signed remainder currently supports only positive power-of-2 constants"); } //===----------------------------------------------------------------------===// @@ -1003,9 +1093,31 @@ static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { if (failed(lb) || failed(ub) || failed(step)) return op->emitOpError("unmapped operand in scf.for"); + auto toI32SReg = [&](Value v, StringRef name) -> FailureOr { + if (getRegSize(v.getType()) != 1) { + op->emitOpError(name) << " must lower to i32"; + return failure(); + } + auto sregTy = ctx.createSRegType(); + if (isSGPRType(v.getType())) + return v; + if (isImmType(v.getType())) + return S_MOV_B32::create(builder, loc, sregTy, v).getResult(); + if (isVGPRType(v.getType())) + return V_READFIRSTLANE_B32::create(builder, loc, sregTy, v).getResult(); + op->emitOpError(name) << " must lower to an SGPR or VGPR i32"; + return failure(); + }; + + FailureOr lbScalar = toI32SReg(*lb, "scf.for lower bound"); + FailureOr ubScalar = toI32SReg(*ub, "scf.for upper bound"); + FailureOr stepScalar = toI32SReg(*step, "scf.for step"); + if (failed(lbScalar) || failed(ubScalar) || failed(stepScalar)) + return failure(); + // Build init args: [lower_bound, iter_args...]. SmallVector initArgs; - initArgs.push_back(*lb); + initArgs.push_back(*lbScalar); for (Value arg : op.getInitArgs()) { FailureOr resolved = resolve(arg, ctx); if (failed(resolved)) @@ -1034,18 +1146,12 @@ static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { // Build loop increment and condition. Value inductionVar = bodyBlock.getArgument(0); - auto ivTy = inductionVar.getType(); - Value nextIV = ArithAddOp::create(builder, loc, ivTy, inductionVar, *step); - // Condition must be SGPR/SCC for waveasm.condition. auto sregTy = ctx.createSRegType(); auto sccTy = ctx.createSCCType(); - Value scalarIV = nextIV; - Value scalarUB = *ub; - if (isVGPRType(nextIV.getType())) - scalarIV = V_READFIRSTLANE_B32::create(builder, loc, sregTy, nextIV); - if (isVGPRType(ub->getType())) - scalarUB = V_READFIRSTLANE_B32::create(builder, loc, sregTy, *ub); - Value cond = S_CMP_LT_U32::create(builder, loc, sccTy, scalarIV, scalarUB); + Value nextIV = + S_ADD_U32::create(builder, loc, sregTy, sccTy, inductionVar, *stepScalar) + .getDst(); + Value cond = S_CMP_LT_U32::create(builder, loc, sccTy, nextIV, *ubScalar); // Collect iter args from yield. auto yieldOp = cast(op.getBody()->getTerminator()); diff --git a/waveasm/test/Transforms/translate-from-llvm-error-bare-ptr-multi-index-gep.mlir b/waveasm/test/Transforms/translate-from-llvm-error-bare-ptr-multi-index-gep.mlir new file mode 100644 index 0000000000..06562490ba --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-error-bare-ptr-multi-index-gep.mlir @@ -0,0 +1,12 @@ +// RUN: not waveasm-translate %s --waveasm-translate-from-llvm 2>&1 | FileCheck %s +// Verify that non-zero multi-index bare-pointer GEPs fail instead of silently +// dropping the aggregate offset. + +// CHECK: bare-pointer GEP must have a single index or all-zero structural indices + +gpu.module @gpu_module { + llvm.func @test(%arg0: !llvm.ptr) attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %0 = llvm.getelementptr %arg0[0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<4 x i32> + llvm.return + } +} diff --git a/waveasm/test/Transforms/translate-from-llvm-error-scf-for-i64.mlir b/waveasm/test/Transforms/translate-from-llvm-error-scf-for-i64.mlir new file mode 100644 index 0000000000..33bef58425 --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-error-scf-for-i64.mlir @@ -0,0 +1,17 @@ +// RUN: not waveasm-translate %s --waveasm-translate-from-llvm 2>&1 | FileCheck %s +// Verify that scf.for loop control stays explicitly limited to i32 lowering. + +// CHECK: scf.for lower bound must lower to i32 + +gpu.module @gpu_module { + llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %lb = llvm.mlir.constant(0 : i64) : i64 + %ub = llvm.mlir.constant(64 : i64) : i64 + %step = llvm.mlir.constant(16 : i64) : i64 + %init = llvm.mlir.constant(0 : i32) : i32 + %result = scf.for %iv = %lb to %ub step %step iter_args(%acc = %init) -> (i32) : i64 { + scf.yield %acc : i32 + } + llvm.return + } +} diff --git a/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir b/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir index fb9fc4cb44..e7ae148c02 100644 --- a/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-scf-for.mlir @@ -7,8 +7,8 @@ // CHECK: waveasm.loop // Loop body: arith.add for the accumulation. // CHECK: waveasm.arith.add -// IV increment. -// CHECK: waveasm.arith.add +// IV increment stays scalar. +// CHECK: waveasm.s_add_u32 // Condition: s_cmp_lt_u32 for the loop back-edge. // CHECK: waveasm.s_cmp_lt_u32 // CHECK: waveasm.condition diff --git a/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir b/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir index 1eb95e1add..8fbf93a8ff 100644 --- a/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir @@ -1,20 +1,29 @@ // RUN: waveasm-translate %s --waveasm-translate-from-llvm | FileCheck %s -// Verify sdiv and srem with power-of-2 constants. -// sdiv -> v_ashrrev_i32; srem -> v_and_b32. +// Verify sdiv and srem with positive power-of-2 constants. +// Negative dividends require bias/correction so we should see compare/select +// scaffolding in addition to the final shift/mask-shaped arithmetic. // CHECK: waveasm.program @test__waveasm -// sdiv by 16 -> shift right by 4. +// sdiv lowers through signed-bias correction before the arithmetic shift. +// CHECK: waveasm.arith.cmp slt +// CHECK: waveasm.arith.select // CHECK: waveasm.v_ashrrev_i32 -// srem by 16 -> and with 15. -// CHECK: waveasm.v_and_b32 +// srem keeps LLVM sign semantics via mask + conditional correction. +// CHECK: waveasm.arith.and +// CHECK: waveasm.arith.add +// CHECK: waveasm.arith.select // CHECK: waveasm.s_endpgm gpu.module @gpu_module { llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { %tid = rocdl.workitem.id.x range : i32 %16 = llvm.mlir.constant(16 : i32) : i32 + %neg5 = llvm.mlir.constant(-5 : i32) : i32 + %4 = llvm.mlir.constant(4 : i32) : i32 %div = llvm.sdiv %tid, %16 : i32 %rem = llvm.srem %tid, %16 : i32 + %negdiv = llvm.sdiv %neg5, %4 : i32 + %negrem = llvm.srem %neg5, %4 : i32 llvm.return } } From 000a3eadc6d6b8a2c02d9750b73ec70a936e8a13 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 22:18:15 +0200 Subject: [PATCH 11/19] Tighten LLVM->WaveASM GEP byte-size handling Compute LDS/global sizes in bytes and reject unsupported GEP element types so the translator fails instead of silently misaddressing memory. Restore source locations for original wrapper args, add lit coverage for the new behavior, and leave a TODO for the remaining scf.for zero-trip semantic gap. Made-with: Cursor Signed-off-by: Ivan Butygin --- .../kernel/compiler/wave_codegen/emitter.py | 5 +- .../Transforms/TranslateFromLLVMDialect.cpp | 77 ++++++++++--------- ...anslate-from-llvm-error-gep-elem-type.mlir | 13 ++++ .../Transforms/translate-from-llvm-lds.mlir | 15 ++-- 4 files changed, 65 insertions(+), 45 deletions(-) create mode 100644 waveasm/test/Transforms/translate-from-llvm-error-gep-elem-type.mlir diff --git a/wave_lang/kernel/compiler/wave_codegen/emitter.py b/wave_lang/kernel/compiler/wave_codegen/emitter.py index 9c9f39da6d..a72886b596 100644 --- a/wave_lang/kernel/compiler/wave_codegen/emitter.py +++ b/wave_lang/kernel/compiler/wave_codegen/emitter.py @@ -333,7 +333,10 @@ def abi_type(binding: BindingDesc): arg_types += [IndexType.get()] * stride_arg_count ftype = FunctionType.get(arg_types, []) - locs = [Location.unknown() for _ in arg_types] + # Preserve source locations for original args; synthesized stride args + # do not have an originating kernel SSA value. + locs = [a.location for a in kernel_func.body.blocks[0].arguments] + locs += [Location.unknown()] * (len(arg_types) - len(locs)) gpu_module = gpu_d.module("gpu_module") gpu_module.parent.operation.attributes["gpu.container_module"] = UnitAttr.get() diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index d56abfc96d..71577c8eb5 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -207,18 +207,16 @@ static unsigned getLLVMAddrSpace(Value v) { return 0; } -/// Try to extract a constant integer from an LLVM SSA value. -/// Return the byte stride for a single-index GEP element type. -/// Only handles sized scalars, vectors, and arrays of i8 (byte addressing). +/// Return the byte size for a sized LLVM scalar, vector, or array type. /// Returns 0 for unsupported types. -static int64_t getGEPElementBytes(Type elemTy) { - if (elemTy.isIntOrFloat()) - return elemTy.getIntOrFloatBitWidth() / 8; - if (auto vecTy = dyn_cast(elemTy)) +static int64_t getLLVMTypeBytes(Type ty) { + if (ty.isIntOrFloat()) + return ty.getIntOrFloatBitWidth() / 8; + if (auto vecTy = dyn_cast(ty)) return vecTy.getNumElements() * vecTy.getElementType().getIntOrFloatBitWidth() / 8; - if (auto arrTy = dyn_cast(elemTy)) - return getGEPElementBytes(arrTy.getElementType()) * arrTy.getNumElements(); + if (auto arrTy = dyn_cast(ty)) + return getLLVMTypeBytes(arrTy.getElementType()) * arrTy.getNumElements(); return 0; } @@ -238,10 +236,11 @@ static bool isAllZeroIndexGEP(LLVM::GEPOp op) { return true; } -/// Compute the byte offset for a single-index GEP. -/// Returns std::nullopt when the offset is zero or unsupported. -static std::optional computeGEPByteOffset(LLVM::GEPOp op, - TranslationContext &ctx) { +/// Compute the non-zero byte offset for a supported single-index GEP. +/// Emits a diagnostic and returns failure for unsupported element types or +/// unmapped dynamic indices. +static FailureOr computeGEPByteOffset(LLVM::GEPOp op, + TranslationContext &ctx) { OpBuilder &builder = ctx.getBuilder(); Location loc = op.getLoc(); auto indices = op.getIndices(); @@ -249,30 +248,32 @@ static std::optional computeGEPByteOffset(LLVM::GEPOp op, // Only handle single-index GEPs. Multi-index GEPs walk nested types // and require full DataLayout support. if (indices.size() != 1) - return std::nullopt; + return op->emitOpError("GEP byte offset requires a single index"); - int64_t elemBytes = getGEPElementBytes(op.getElemType()); + int64_t elemBytes = getLLVMTypeBytes(op.getElemType()); + if (elemBytes == 0) + return op->emitOpError( + "unsupported GEP element type for byte offset computation"); auto idx = indices[0]; if (Value dynIdx = idx.dyn_cast()) { FailureOr resolved = resolve(dynIdx, ctx); if (failed(resolved)) - return std::nullopt; - if (elemBytes <= 1) + return op->emitOpError("unmapped GEP index"); + if (elemBytes == 1) return *resolved; // Scale by element size: offset = index * elemBytes. ImmType scaleTy = ctx.createImmType(elemBytes); Value scale = ConstantOp::create(builder, loc, scaleTy, elemBytes); Type resTy = inferResultType({*resolved, scale}, ctx); - return ArithMulOp::create(builder, loc, resTy, *resolved, scale); + return ArithMulOp::create(builder, loc, resTy, *resolved, scale) + .getResult(); } int64_t constIdx = cast(idx).getInt(); - if (constIdx == 0) - return std::nullopt; - int64_t byteOffset = constIdx * std::max(elemBytes, int64_t{1}); + int64_t byteOffset = constIdx * elemBytes; ImmType immTy = ctx.createImmType(byteOffset); - return ConstantOp::create(builder, loc, immTy, byteOffset); + return ConstantOp::create(builder, loc, immTy, byteOffset).getResult(); } //===----------------------------------------------------------------------===// @@ -610,17 +611,18 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { VRegType vregTy = ctx.createVRegType(); if (getRegSize(baseOff->getType()) > 1) *baseOff = ArithTruncOp::create(builder, loc, vregTy, *baseOff); - std::optional maybeOffset = - isSingleIndex ? computeGEPByteOffset(op, ctx) : std::nullopt; - if (!maybeOffset) { + if (isAllZeroGEP) { ctx.getMapper().mapValue(op.getResult(), *baseOff); - } else { - Value off = *maybeOffset; - if (getRegSize(off.getType()) > 1) - off = ArithTruncOp::create(builder, loc, vregTy, off); - Value sum = ArithAddOp::create(builder, loc, vregTy, *baseOff, off); - ctx.getMapper().mapValue(op.getResult(), sum); + return success(); } + FailureOr maybeOffset = computeGEPByteOffset(op, ctx); + if (failed(maybeOffset)) + return failure(); + Value off = *maybeOffset; + if (getRegSize(off.getType()) > 1) + off = ArithTruncOp::create(builder, loc, vregTy, off); + Value sum = ArithAddOp::create(builder, loc, vregTy, *baseOff, off); + ctx.getMapper().mapValue(op.getResult(), sum); return success(); } @@ -634,19 +636,20 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { if (!isSingleIndex && !isAllZeroGEP) return op->emitOpError("bare-pointer GEP must have a single index or " "all-zero structural indices"); - std::optional maybeOffset = - isSingleIndex ? computeGEPByteOffset(op, ctx) : std::nullopt; // Forward mapper entry so make.buffer.rsrc can find the SRD. if (std::optional mapped = ctx.getMapper().getMapped(base)) ctx.getMapper().mapValue(op.getResult(), *mapped); - if (!maybeOffset) { + if (isAllZeroGEP) { Value prevOffset = st.lookupBaseOffset(base); if (prevOffset) st.setBaseOffset(op.getResult(), prevOffset); return success(); } + FailureOr maybeOffset = computeGEPByteOffset(op, ctx); + if (failed(maybeOffset)) + return failure(); Value newOffset = *maybeOffset; Value prevOffset = st.lookupBaseOffset(base); if (prevOffset) { @@ -807,9 +810,7 @@ static LogicalResult handleAddressOf(LLVM::AddressOfOp op, op, op.getGlobalNameAttr()); if (global) { auto globalTy = global.getType(); - int64_t sizeBytes = 0; - if (auto arrTy = dyn_cast(globalTy)) - sizeBytes = arrTy.getNumElements(); + int64_t sizeBytes = getLLVMTypeBytes(globalTy); if (sizeBytes > 0) ctx.addLDSSize(sizeBytes); } @@ -1093,6 +1094,8 @@ static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { if (failed(lb) || failed(ub) || failed(step)) return op->emitOpError("unmapped operand in scf.for"); + // TODO: Preserve full scf.for zero-trip semantics. waveasm.loop is do-while + // and this lowering currently assumes the trip count is known positive. auto toI32SReg = [&](Value v, StringRef name) -> FailureOr { if (getRegSize(v.getType()) != 1) { op->emitOpError(name) << " must lower to i32"; diff --git a/waveasm/test/Transforms/translate-from-llvm-error-gep-elem-type.mlir b/waveasm/test/Transforms/translate-from-llvm-error-gep-elem-type.mlir new file mode 100644 index 0000000000..4556186ab7 --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-error-gep-elem-type.mlir @@ -0,0 +1,13 @@ +// RUN: not waveasm-translate %s --waveasm-translate-from-llvm 2>&1 | FileCheck %s +// Verify that single-index GEPs with unsupported element types fail instead of +// silently treating the index as a byte offset. + +// CHECK: unsupported GEP element type for byte offset computation + +gpu.module @gpu_module { + llvm.func @test(%arg0: !llvm.ptr) attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %idx = llvm.mlir.constant(1 : i64) : i64 + %gep = llvm.getelementptr %arg0[%idx] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.struct<(i32, i32)> + llvm.return + } +} diff --git a/waveasm/test/Transforms/translate-from-llvm-lds.mlir b/waveasm/test/Transforms/translate-from-llvm-lds.mlir index b2e6bb497f..d9647301bd 100644 --- a/waveasm/test/Transforms/translate-from-llvm-lds.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-lds.mlir @@ -1,9 +1,10 @@ // RUN: waveasm-translate %s --waveasm-translate-from-llvm | FileCheck %s // Verify LDS addressof, GEP, load, and store translation. -// addressof maps to VGPR zero; LDS GEPs produce arith.add offsets; -// loads/stores dispatch to ds_read/ds_write by width. +// addressof maps to VGPR zero; LDS GEPs produce byte offsets; loads/stores +// dispatch to ds_read/ds_write by width; and LDS size is counted in bytes. // CHECK: waveasm.program @test__waveasm +// CHECK: lds_size = 1024 : i64 // addressof -> v_mov_b32 of zero. // CHECK: waveasm.v_mov_b32 // GEP [0,0] on the array type passes through (all-zero indices). @@ -16,18 +17,18 @@ // CHECK: waveasm.s_endpgm gpu.module @gpu_module { - llvm.mlir.global private @alloca() {addr_space = 3 : i32} : !llvm.array<1024 x i8> + llvm.mlir.global private @alloca() {addr_space = 3 : i32} : !llvm.array<256 x i32> llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { %0 = llvm.mlir.addressof @alloca : !llvm.ptr<3> %1 = llvm.mlir.constant(42 : i32) : i32 // Multi-index GEP with all-zero indices -> passthrough. - %2 = llvm.getelementptr %0[0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.array<1024 x i8> - // Constant-offset GEP. - %3 = llvm.getelementptr %2[512] : (!llvm.ptr<3>) -> !llvm.ptr<3>, i8 + %2 = llvm.getelementptr %0[0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.array<256 x i32> + // Constant-offset GEP. 128 * sizeof(i32) = 512 bytes. + %3 = llvm.getelementptr %2[128] : (!llvm.ptr<3>) -> !llvm.ptr<3>, i32 // Dynamic-offset GEP. %tid = rocdl.workitem.id.x range : i32 %tidext = llvm.sext %tid : i32 to i64 - %4 = llvm.getelementptr nusw %2[%tidext] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i8 + %4 = llvm.getelementptr nusw %2[%tidext] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i32 // LDS load (4 bytes). %5 = llvm.load %4 : !llvm.ptr<3> -> i32 // LDS store (4 bytes). From ea6948aaa21d5def940d0a56c72d6ab9666b2125 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 22:42:10 +0200 Subject: [PATCH 12/19] Make LLVM->WaveASM translator types explicit Use concrete local types in the new lowering paths so the translator is easier to audit and maintain without changing behavior. Made-with: Cursor Signed-off-by: Ivan Butygin --- .../Transforms/TranslateFromLLVMDialect.cpp | 197 +++++++++--------- 1 file changed, 103 insertions(+), 94 deletions(-) diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 71577c8eb5..0b6900b089 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -202,7 +202,8 @@ static Type inferResultType(ValueRange operands, TranslationContext &ctx) { /// Return the LLVM pointer address space, or 0 for non-pointer types. static unsigned getLLVMAddrSpace(Value v) { - if (auto ptrTy = dyn_cast(v.getType())) + if (LLVM::LLVMPointerType ptrTy = + dyn_cast(v.getType())) return ptrTy.getAddressSpace(); return 0; } @@ -212,10 +213,10 @@ static unsigned getLLVMAddrSpace(Value v) { static int64_t getLLVMTypeBytes(Type ty) { if (ty.isIntOrFloat()) return ty.getIntOrFloatBitWidth() / 8; - if (auto vecTy = dyn_cast(ty)) + if (VectorType vecTy = dyn_cast(ty)) return vecTy.getNumElements() * vecTy.getElementType().getIntOrFloatBitWidth() / 8; - if (auto arrTy = dyn_cast(ty)) + if (LLVM::LLVMArrayType arrTy = dyn_cast(ty)) return getLLVMTypeBytes(arrTy.getElementType()) * arrTy.getNumElements(); return 0; } @@ -230,7 +231,7 @@ static bool isZeroGEPIndex(llvm::PointerUnion idx) { /// Structural GEPs like [0, 0] are a no-op and can be forwarded even though we /// do not model nested aggregate layouts yet. static bool isAllZeroIndexGEP(LLVM::GEPOp op) { - for (auto idx : op.getIndices()) + for (llvm::PointerUnion idx : op.getIndices()) if (!isZeroGEPIndex(idx)) return false; return true; @@ -243,7 +244,7 @@ static FailureOr computeGEPByteOffset(LLVM::GEPOp op, TranslationContext &ctx) { OpBuilder &builder = ctx.getBuilder(); Location loc = op.getLoc(); - auto indices = op.getIndices(); + LLVM::GEPIndicesAdaptor indices = op.getIndices(); // Only handle single-index GEPs. Multi-index GEPs walk nested types // and require full DataLayout support. @@ -255,7 +256,7 @@ static FailureOr computeGEPByteOffset(LLVM::GEPOp op, return op->emitOpError( "unsupported GEP element type for byte offset computation"); - auto idx = indices[0]; + llvm::PointerUnion idx = indices[0]; if (Value dynIdx = idx.dyn_cast()) { FailureOr resolved = resolve(dynIdx, ctx); if (failed(resolved)) @@ -317,22 +318,22 @@ static LogicalResult handlePoison(LLVM::PoisonOp op, LLVMTranslationState &st) { static LogicalResult handleConstant(LLVM::ConstantOp op, LLVMTranslationState &st) { - auto &ctx = st.ctx; - auto &builder = ctx.getBuilder(); - auto loc = op.getLoc(); + TranslationContext &ctx = st.ctx; + OpBuilder &builder = ctx.getBuilder(); + Location loc = op.getLoc(); - auto valAttr = op.getValue(); + Attribute valAttr = op.getValue(); // Dense vector constant (e.g. MFMA accumulator init). - if (auto denseAttr = dyn_cast(valAttr)) { + if (DenseElementsAttr denseAttr = dyn_cast(valAttr)) { if (!denseAttr.isSplat()) return op->emitOpError("non-splat dense constant not yet supported"); int64_t numElems = denseAttr.getNumElements(); APFloat splatVal = denseAttr.getSplatValue(); int64_t rawBits = splatVal.bitcastToAPInt().getZExtValue(); - auto immTy = ctx.createImmType(rawBits); - auto immOp = ConstantOp::create(builder, loc, immTy, rawBits); - auto vregTy = ctx.createVRegType(numElems, numElems); + ImmType immTy = ctx.createImmType(rawBits); + Value immOp = ConstantOp::create(builder, loc, immTy, rawBits); + VRegType vregTy = ctx.createVRegType(numElems, numElems); Value mov = V_MOV_B32::create(builder, loc, vregTy, immOp); ctx.getMapper().mapValue(op.getResult(), mov); return success(); @@ -344,7 +345,7 @@ static LogicalResult handleConstant(LLVM::ConstantOp op, else return op->emitOpError("unsupported constant type"); - auto intType = dyn_cast(op.getResult().getType()); + IntegerType intType = dyn_cast(op.getResult().getType()); if (!intType) return op->emitOpError("expected integer constant"); @@ -424,7 +425,7 @@ static LogicalResult handleCastOp(LLVMOp op, LLVMTranslationState &st) { if (failed(src)) return op->emitOpError("unmapped operand in cast"); int64_t width = 1; - if (auto intTy = dyn_cast(op.getResult().getType())) + if (IntegerType intTy = dyn_cast(op.getResult().getType())) width = (intTy.getWidth() + 31) / 32; Type resTy = isVGPRType(src->getType()) ? (Type)ctx.createVRegType(width, width) @@ -535,17 +536,18 @@ static LogicalResult handleMakeBufferRsrc(ROCDL::MakeBufferRsrcOp op, Value word3 = ExtractOp::create(builder, loc, sregTy, *srdVal, 3); // Clear stride/swizzle bits in SRD word 1 (keep only base_addr[47:32]). - auto maskImm = ctx.createImmType(kSRDWord1BaseMask); - auto maskVal = ConstantOp::create(builder, loc, maskImm, kSRDWord1BaseMask); + ImmType maskImm = ctx.createImmType(kSRDWord1BaseMask); + Value maskVal = + ConstantOp::create(builder, loc, maskImm, kSRDWord1BaseMask); word1 = S_AND_B32::create(builder, loc, sregTy, ctx.createSCCType(), word1, maskVal) .getDst(); // Patch SRD[3] with the actual flags from make.buffer.rsrc. - auto flags = getConstantIntValue(op.getFlags()); + std::optional flags = getConstantIntValue(op.getFlags()); if (flags && *flags != kSRDDefaultFlags) { - auto flagsImm = ctx.createImmType(*flags); - auto flagsVal = ConstantOp::create(builder, loc, flagsImm, *flags); + ImmType flagsImm = ctx.createImmType(*flags); + Value flagsVal = ConstantOp::create(builder, loc, flagsImm, *flags); word3 = S_MOV_B32::create(builder, loc, sregTy, flagsVal); } @@ -569,7 +571,8 @@ static LogicalResult handleMakeBufferRsrc(ROCDL::MakeBufferRsrcOp op, // Adjust base: S_ADD_U32 sets SCC, S_ADDC_U32 reads it. SCCType sccTy = ctx.createSCCType(); - auto addLo = S_ADD_U32::create(builder, loc, sregTy, sccTy, word0, offLo); + S_ADD_U32 addLo = + S_ADD_U32::create(builder, loc, sregTy, sccTy, word0, offLo); word0 = addLo.getDst(); word1 = S_ADDC_U32::create(builder, loc, sregTy, sccTy, addLo.getScc(), word1, offHi) @@ -577,7 +580,7 @@ static LogicalResult handleMakeBufferRsrc(ROCDL::MakeBufferRsrcOp op, } // Pack into a 4-wide SGPR SRD. - auto srdType = ctx.createSRegType(4, 4); + SRegType srdType = ctx.createSRegType(4, 4); Value newSrd = PackOp::create(builder, loc, srdType, ValueRange{word0, word1, word2, word3}); st.mapBufferRsrc(op.getResult(), newSrd); @@ -804,22 +807,22 @@ static LogicalResult handleStore(LLVM::StoreOp op, LLVMTranslationState &st) { static LogicalResult handleAddressOf(LLVM::AddressOfOp op, LLVMTranslationState &st) { - auto &ctx = st.ctx; + TranslationContext &ctx = st.ctx; // Look up the global to determine LDS size. - auto global = SymbolTable::lookupNearestSymbolFrom( + LLVM::GlobalOp global = SymbolTable::lookupNearestSymbolFrom( op, op.getGlobalNameAttr()); if (global) { - auto globalTy = global.getType(); + Type globalTy = global.getType(); int64_t sizeBytes = getLLVMTypeBytes(globalTy); if (sizeBytes > 0) ctx.addLDSSize(sizeBytes); } // Map to a constant 0 offset in a VGPR -- LDS addressing is relative. - auto &builder = ctx.getBuilder(); - auto immTy = ctx.createImmType(0); - auto zero = ConstantOp::create(builder, op.getLoc(), immTy, int64_t{0}); - auto vregTy = ctx.createVRegType(); + OpBuilder &builder = ctx.getBuilder(); + ImmType immTy = ctx.createImmType(0); + Value zero = ConstantOp::create(builder, op.getLoc(), immTy, int64_t{0}); + VRegType vregTy = ctx.createVRegType(); Value mov = V_MOV_B32::create(builder, op.getLoc(), vregTy, zero); ctx.getMapper().mapValue(op.getResult(), mov); return success(); @@ -831,16 +834,16 @@ static LogicalResult handleAddressOf(LLVM::AddressOfOp op, static LogicalResult handleLDSLoad(LLVM::LoadOp op, Value addr, LLVMTranslationState &st) { - auto &ctx = st.ctx; - auto &builder = ctx.getBuilder(); - auto loc = op.getLoc(); + TranslationContext &ctx = st.ctx; + OpBuilder &builder = ctx.getBuilder(); + Location loc = op.getLoc(); int64_t numBytes = getBufferAccessBytes(op.getResult().getType()); FailureOr offset = resolve(addr, ctx); if (failed(offset)) return op->emitOpError("unmapped LDS address"); // DS instructions use a 32-bit vaddr. Truncate wide offsets. - auto vregTy = ctx.createVRegType(); + VRegType vregTy = ctx.createVRegType(); if (getRegSize(offset->getType()) > 1) *offset = ArithTruncOp::create(builder, loc, vregTy, *offset); @@ -850,10 +853,10 @@ static LogicalResult handleLDSLoad(LLVM::LoadOp op, Value addr, else if (numBytes == 4) loadOp = DS_READ_B32::create(builder, loc, TypeRange{vregTy}, *offset); else if (numBytes == 8) { - auto wideTy = ctx.createVRegType(2, 2); + VRegType wideTy = ctx.createVRegType(2, 2); loadOp = DS_READ_B64::create(builder, loc, TypeRange{wideTy}, *offset); } else if (numBytes == 16) { - auto wideTy = ctx.createVRegType(4, 4); + VRegType wideTy = ctx.createVRegType(4, 4); loadOp = DS_READ_B128::create(builder, loc, TypeRange{wideTy}, *offset); } else return op->emitOpError("unsupported LDS load size: ") @@ -865,16 +868,16 @@ static LogicalResult handleLDSLoad(LLVM::LoadOp op, Value addr, static LogicalResult handleLDSStore(LLVM::StoreOp op, Value addr, LLVMTranslationState &st) { - auto &ctx = st.ctx; - auto &builder = ctx.getBuilder(); - auto loc = op.getLoc(); + TranslationContext &ctx = st.ctx; + OpBuilder &builder = ctx.getBuilder(); + Location loc = op.getLoc(); FailureOr data = resolve(op.getValue(), ctx); FailureOr offset = resolve(addr, ctx); if (failed(data) || failed(offset)) return op->emitOpError("unmapped LDS operand"); // DS instructions use a 32-bit vaddr. Truncate wide offsets. - auto vregTy = ctx.createVRegType(); + VRegType vregTy = ctx.createVRegType(); if (getRegSize(offset->getType()) > 1) *offset = ArithTruncOp::create(builder, loc, vregTy, *offset); int64_t numBytes = getBufferAccessBytes(op.getValue().getType()); @@ -899,9 +902,9 @@ static LogicalResult handleLDSStore(LLVM::StoreOp op, Value addr, //===----------------------------------------------------------------------===// static LogicalResult handleSDiv(LLVM::SDivOp op, LLVMTranslationState &st) { - auto &ctx = st.ctx; - auto &builder = ctx.getBuilder(); - auto loc = op.getLoc(); + TranslationContext &ctx = st.ctx; + OpBuilder &builder = ctx.getBuilder(); + Location loc = op.getLoc(); FailureOr lhs = resolve(op.getLhs(), ctx); if (failed(lhs)) return op->emitOpError("unmapped operand in sdiv"); @@ -912,19 +915,19 @@ static LogicalResult handleSDiv(LLVM::SDivOp op, LLVMTranslationState &st) { // Signed division by a positive power-of-2 constant: // q = (x + (x < 0 ? divisor - 1 : 0)) >> log2(divisor) // This preserves LLVM's trunc-toward-zero semantics for negative dividends. - if (auto constVal = getConstantIntValue(op.getRhs())) { + if (std::optional constVal = getConstantIntValue(op.getRhs())) { int64_t divisor = *constVal; if (divisor > 0 && (divisor & (divisor - 1)) == 0) { int64_t shiftAmt = llvm::Log2_64(divisor); - auto zero = ConstantOp::create(builder, loc, ctx.createImmType(0), 0); - auto biasImm = ConstantOp::create( + Value zero = ConstantOp::create(builder, loc, ctx.createImmType(0), 0); + Value biasImm = ConstantOp::create( builder, loc, ctx.createImmType(divisor - 1), divisor - 1); - auto shiftConst = ConstantOp::create( + Value shiftConst = ConstantOp::create( builder, loc, ctx.createImmType(shiftAmt), shiftAmt); Value result; if (isSGPRType(lhs->getType())) { - auto sregTy = ctx.createSRegType(); - auto sccTy = ctx.createSCCType(); + SRegType sregTy = ctx.createSRegType(); + SCCType sccTy = ctx.createSCCType(); Value isNegative = S_CMP_LT_I32::create(builder, loc, sccTy, *lhs, zero); Value bias = S_CSELECT_B32::create(builder, loc, sregTy, isNegative, @@ -934,7 +937,7 @@ static LogicalResult handleSDiv(LLVM::SDivOp op, LLVMTranslationState &st) { S_ASHR_I32::create(builder, loc, sregTy, sccTy, biased, shiftConst) .getDst(); } else { - auto vregTy = ctx.createVRegType(); + VRegType vregTy = ctx.createVRegType(); Value isNegative = ArithCmpOp::create(builder, loc, vregTy, CmpPredicate::slt, *lhs, zero); Value bias = ArithSelectOp::create(builder, loc, vregTy, zero, biasImm, @@ -952,9 +955,9 @@ static LogicalResult handleSDiv(LLVM::SDivOp op, LLVMTranslationState &st) { } static LogicalResult handleSRem(LLVM::SRemOp op, LLVMTranslationState &st) { - auto &ctx = st.ctx; - auto &builder = ctx.getBuilder(); - auto loc = op.getLoc(); + TranslationContext &ctx = st.ctx; + OpBuilder &builder = ctx.getBuilder(); + Location loc = op.getLoc(); FailureOr lhs = resolve(op.getLhs(), ctx); if (failed(lhs)) return op->emitOpError("unmapped operand in srem"); @@ -966,18 +969,18 @@ static LogicalResult handleSRem(LLVM::SRemOp op, LLVMTranslationState &st) { // r = x & (divisor - 1) // if (x < 0 && r != 0) r -= divisor // This keeps the remainder's sign consistent with LLVM srem. - if (auto constVal = getConstantIntValue(op.getRhs())) { + if (std::optional constVal = getConstantIntValue(op.getRhs())) { int64_t divisor = *constVal; if (divisor > 0 && (divisor & (divisor - 1)) == 0) { - auto zero = ConstantOp::create(builder, loc, ctx.createImmType(0), 0); - auto maskConst = ConstantOp::create( + Value zero = ConstantOp::create(builder, loc, ctx.createImmType(0), 0); + Value maskConst = ConstantOp::create( builder, loc, ctx.createImmType(divisor - 1), divisor - 1); - auto negDivisor = ConstantOp::create( + Value negDivisor = ConstantOp::create( builder, loc, ctx.createImmType(-divisor), -divisor); Value result; if (isSGPRType(lhs->getType())) { - auto sregTy = ctx.createSRegType(); - auto sccTy = ctx.createSCCType(); + SRegType sregTy = ctx.createSRegType(); + SCCType sccTy = ctx.createSCCType(); Value rawRem = ArithAndOp::create(builder, loc, sregTy, *lhs, maskConst); Value isNegative = @@ -991,7 +994,7 @@ static LogicalResult handleSRem(LLVM::SRemOp op, LLVMTranslationState &st) { result = S_CSELECT_B32::create(builder, loc, sregTy, isNegative, maybeAdjusted, rawRem); } else { - auto vregTy = ctx.createVRegType(); + VRegType vregTy = ctx.createVRegType(); Value rawRem = ArithAndOp::create(builder, loc, vregTy, *lhs, maskConst); Value isNegative = ArithCmpOp::create(builder, loc, vregTy, @@ -1034,17 +1037,17 @@ static LogicalResult handleBarrier(OpTy op, LLVMTranslationState &st) { static LogicalResult handleShuffleVector(LLVM::ShuffleVectorOp op, LLVMTranslationState &st) { - auto &ctx = st.ctx; - auto &builder = ctx.getBuilder(); + TranslationContext &ctx = st.ctx; + OpBuilder &builder = ctx.getBuilder(); FailureOr src = resolve(op.getV1(), ctx); if (failed(src)) return op->emitOpError("unmapped operand in shufflevector"); // shufflevector with a single index extracts one element. - auto mask = op.getMask(); + llvm::ArrayRef mask = op.getMask(); if (mask.size() == 1) { int64_t idx = mask[0]; - auto vregTy = ctx.createVRegType(); + VRegType vregTy = ctx.createVRegType(); Value extract = ExtractOp::create(builder, op.getLoc(), vregTy, *src, idx); ctx.getMapper().mapValue(op.getResult(), extract); return success(); @@ -1059,9 +1062,9 @@ static LogicalResult handleShuffleVector(LLVM::ShuffleVectorOp op, static LogicalResult handleMFMA_F32_16x16x16_F16(ROCDL::mfma_f32_16x16x16f16 op, LLVMTranslationState &st) { - auto &ctx = st.ctx; - auto &builder = ctx.getBuilder(); - auto loc = op.getLoc(); + TranslationContext &ctx = st.ctx; + OpBuilder &builder = ctx.getBuilder(); + Location loc = op.getLoc(); FailureOr a = resolve(op.getA(), ctx); FailureOr b = resolve(op.getB(), ctx); @@ -1070,7 +1073,7 @@ static LogicalResult handleMFMA_F32_16x16x16_F16(ROCDL::mfma_f32_16x16x16f16 op, return op->emitOpError("unmapped operand in MFMA"); // Result and accumulator are vector<4xf32> -> 4 VGPRs. - auto accTy = ctx.createVRegType(4, 4); + VRegType accTy = ctx.createVRegType(4, 4); Value mfma = V_MFMA_F32_16X16X16_F16::create(builder, loc, accTy, *a, *b, *c); ctx.getMapper().mapValue(op.getResult(), mfma); return success(); @@ -1083,10 +1086,29 @@ static LogicalResult handleMFMA_F32_16x16x16_F16(ROCDL::mfma_f32_16x16x16f16 op, /// Forward declaration for recursive op translation. static LogicalResult translateOp(Operation *op, LLVMTranslationState &st); +static FailureOr materializeI32SReg(Value v, StringRef name, + scf::ForOp op, + TranslationContext &ctx, + OpBuilder &builder, Location loc) { + if (getRegSize(v.getType()) != 1) { + op->emitOpError(name) << " must lower to i32"; + return failure(); + } + SRegType sregTy = ctx.createSRegType(); + if (isSGPRType(v.getType())) + return v; + if (isImmType(v.getType())) + return S_MOV_B32::create(builder, loc, sregTy, v).getResult(); + if (isVGPRType(v.getType())) + return V_READFIRSTLANE_B32::create(builder, loc, sregTy, v).getResult(); + op->emitOpError(name) << " must lower to an SGPR or VGPR i32"; + return failure(); +} + static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { - auto &ctx = st.ctx; - auto &builder = ctx.getBuilder(); - auto loc = op.getLoc(); + TranslationContext &ctx = st.ctx; + OpBuilder &builder = ctx.getBuilder(); + Location loc = op.getLoc(); FailureOr lb = resolve(op.getLowerBound(), ctx); FailureOr ub = resolve(op.getUpperBound(), ctx); @@ -1096,25 +1118,12 @@ static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { // TODO: Preserve full scf.for zero-trip semantics. waveasm.loop is do-while // and this lowering currently assumes the trip count is known positive. - auto toI32SReg = [&](Value v, StringRef name) -> FailureOr { - if (getRegSize(v.getType()) != 1) { - op->emitOpError(name) << " must lower to i32"; - return failure(); - } - auto sregTy = ctx.createSRegType(); - if (isSGPRType(v.getType())) - return v; - if (isImmType(v.getType())) - return S_MOV_B32::create(builder, loc, sregTy, v).getResult(); - if (isVGPRType(v.getType())) - return V_READFIRSTLANE_B32::create(builder, loc, sregTy, v).getResult(); - op->emitOpError(name) << " must lower to an SGPR or VGPR i32"; - return failure(); - }; - - FailureOr lbScalar = toI32SReg(*lb, "scf.for lower bound"); - FailureOr ubScalar = toI32SReg(*ub, "scf.for upper bound"); - FailureOr stepScalar = toI32SReg(*step, "scf.for step"); + FailureOr lbScalar = + materializeI32SReg(*lb, "scf.for lower bound", op, ctx, builder, loc); + FailureOr ubScalar = + materializeI32SReg(*ub, "scf.for upper bound", op, ctx, builder, loc); + FailureOr stepScalar = + materializeI32SReg(*step, "scf.for step", op, ctx, builder, loc); if (failed(lbScalar) || failed(ubScalar) || failed(stepScalar)) return failure(); @@ -1129,7 +1138,7 @@ static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { } // Create the waveasm.loop (do-while semantics). - auto loopOp = LoopOp::create(builder, loc, initArgs); + LoopOp loopOp = LoopOp::create(builder, loc, initArgs); Block &bodyBlock = loopOp.getBodyBlock(); // Map the induction variable (block arg 0). @@ -1149,15 +1158,15 @@ static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { // Build loop increment and condition. Value inductionVar = bodyBlock.getArgument(0); - auto sregTy = ctx.createSRegType(); - auto sccTy = ctx.createSCCType(); + SRegType sregTy = ctx.createSRegType(); + SCCType sccTy = ctx.createSCCType(); Value nextIV = S_ADD_U32::create(builder, loc, sregTy, sccTy, inductionVar, *stepScalar) .getDst(); Value cond = S_CMP_LT_U32::create(builder, loc, sccTy, nextIV, *ubScalar); // Collect iter args from yield. - auto yieldOp = cast(op.getBody()->getTerminator()); + scf::YieldOp yieldOp = cast(op.getBody()->getTerminator()); SmallVector condIterArgs; condIterArgs.push_back(nextIV); for (Value v : yieldOp.getOperands()) { From 86ff7adc42e36f63313afe2ac29fec1955462f24 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 9 Apr 2026 22:48:44 +0200 Subject: [PATCH 13/19] unicode Signed-off-by: Ivan Butygin --- tests/kernel/e2e/test_gemm_waveasm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernel/e2e/test_gemm_waveasm.py b/tests/kernel/e2e/test_gemm_waveasm.py index 4f49ae9249..8b9fa86634 100644 --- a/tests/kernel/e2e/test_gemm_waveasm.py +++ b/tests/kernel/e2e/test_gemm_waveasm.py @@ -4,7 +4,7 @@ # See https://llvm.org/LICENSE.txt for license information. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -"""GEMM through the water+waveasm pipeline (LLVM dialect → WaveASM → binary).""" +"""GEMM through the water+waveasm pipeline (LLVM dialect -> WaveASM -> binary).""" import pytest import torch From 4fbfae1d464f2d7e747095d968b4584761806486 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Fri, 10 Apr 2026 19:44:21 +0200 Subject: [PATCH 14/19] Harden WaveASM lowering and ABI handling Move signed div/rem strength reduction into a legalization pass, keep LDS globals distinct, and fix dynamic stride ordering so translation and runtime dispatch agree on kernel ABI. Made-with: Cursor Signed-off-by: Ivan Butygin --- lit_tests/kernel/wave/water_host_wrapper.py | 51 +++ .../kernel/compiler/wave_codegen/emitter.py | 125 +++--- wave_lang/kernel/wave/water.py | 109 +++-- waveasm/include/waveasm/Transforms/Passes.td | 16 + waveasm/lib/Transforms/CMakeLists.txt | 1 + .../Transforms/LLVMSDivSRemLegalization.cpp | 139 ++++++ .../Transforms/TranslateFromLLVMDialect.cpp | 421 ++++++++++-------- .../llvm-sdiv-srem-legalization.mlir | 31 ++ ...ate-from-llvm-error-addressof-non-lds.mlir | 12 + .../Transforms/translate-from-llvm-lds.mlir | 16 +- .../translate-from-llvm-sdiv-srem.mlir | 8 +- 11 files changed, 619 insertions(+), 310 deletions(-) create mode 100644 waveasm/lib/Transforms/LLVMSDivSRemLegalization.cpp create mode 100644 waveasm/test/Transforms/llvm-sdiv-srem-legalization.mlir create mode 100644 waveasm/test/Transforms/translate-from-llvm-error-addressof-non-lds.mlir diff --git a/lit_tests/kernel/wave/water_host_wrapper.py b/lit_tests/kernel/wave/water_host_wrapper.py index bc047f7a10..5e6e7a552b 100644 --- a/lit_tests/kernel/wave/water_host_wrapper.py +++ b/lit_tests/kernel/wave/water_host_wrapper.py @@ -31,6 +31,7 @@ def get_wave_compile_options( canonicalize: bool = False, dynamic_symbols=[], additional_symbols={}, + wave_runtime: bool = False, location_capture_config=LocationCaptureConfig( level=LocationCaptureLevel.FILE_LINE_COL ), @@ -61,6 +62,7 @@ def get_wave_compile_options( location_capture_config=location_capture_config, drop_debug_info_before_mlir=drop_debug_info_before_mlir, use_water_backend=True, + wave_runtime=wave_runtime, ) @@ -227,3 +229,52 @@ def read_write( # CHECK-SAME: blocks in (%[[C1]], %[[C1]], %[[C1]]) # CHECK-SAME: threads in (%[[C64]], %[[C1]], %[[C1]]) # CHECK: return + + +@run_test +def test_dynamic_strides_output_placeholder_first(): + constraints = get_constraints() + + @tkw.wave(constraints) + def output_then_input( + out: tkl.Memory[M, N, ADDRESS_SPACE, tkl.f16], + inp: tkl.Memory[M, N, ADDRESS_SPACE, tkl.f16], + ): + res = tkw.read(inp) + tkw.write(res, out) + + output_then_input = wave_compile( + get_wave_compile_options( + canonicalize=True, + wave_runtime=True, + drop_debug_info_before_mlir=True, + ), + output_then_input, + ) + print(output_then_input.asm) + + # Even though the Python placeholder order is (out, inp), kernel ABI order + # is linearized to input first, output second. + # CHECK-LABEL: test_dynamic_strides_output_placeholder_first + # CHECK: gpu.func @output_then_input + # CHECK-SAME: (%[[IN:.*]]: memref {llvm.inreg}, %[[OUT:.*]]: memref {llvm.inreg}, %[[IN_STRIDE:.*]]: index {llvm.inreg}, %[[OUT_STRIDE:.*]]: index {llvm.inreg}) + # CHECK: %[[IN_VIEW:.*]] = memref.reinterpret_cast %[[IN]] to offset: [0], sizes: [16, 16], strides: [%[[IN_STRIDE]], 1] + # CHECK: %[[OUT_VIEW:.*]] = memref.reinterpret_cast %[[OUT]] to offset: [0], sizes: [16, 16], strides: [%[[OUT_STRIDE]], 1] + + # Host wrapper must mirror that same ABI order for both buffers and strides. + # CHECK-LABEL: func.func @isolated_benchmark + # CHECK-SAME: (%[[STREAM:.*]]: !llvm.ptr, %[[OUT_ARG:.*]]: !llvm.ptr, %[[IN_ARG:.*]]: !llvm.ptr) + # CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index + # CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index + # CHECK-DAG: %[[C64:.*]] = arith.constant 64 : index + # CHECK-DAG: %[[C0_I32:.*]] = arith.constant 0 : i32 + # CHECK: %[[IN_BUF:.*]] = call @wave_get_buffer(%[[IN_ARG]]) : (!llvm.ptr) -> memref + # CHECK: %[[IN_VIEW0:.*]] = memref.view %[[IN_BUF]][%[[C0]]][] : memref to memref + # CHECK: %[[OUT_BUF:.*]] = call @wave_get_buffer(%[[OUT_ARG]]) : (!llvm.ptr) -> memref + # CHECK: %[[OUT_VIEW0:.*]] = memref.view %[[OUT_BUF]][%[[C0]]][] : memref to memref + # CHECK: %[[IN_STRIDE_I64:.*]] = call @wave_get_stride(%[[IN_ARG]], %[[C0_I32]]) : (!llvm.ptr, i32) -> i64 + # CHECK: %[[IN_STRIDE_IDX:.*]] = arith.index_cast %[[IN_STRIDE_I64]] : i64 to index + # CHECK: %[[OUT_STRIDE_I64:.*]] = call @wave_get_stride(%[[OUT_ARG]], %[[C0_I32]]) : (!llvm.ptr, i32) -> i64 + # CHECK: %[[OUT_STRIDE_IDX:.*]] = arith.index_cast %[[OUT_STRIDE_I64]] : i64 to index + # CHECK: gpu.launch_func @gpu_module::@output_then_input blocks in (%[[C1]], %[[C1]], %[[C1]]) threads in (%[[C64]], %[[C1]], %[[C1]]) args(%[[IN_VIEW0]] : memref, %[[OUT_VIEW0]] : memref, %[[IN_STRIDE_IDX]] : index, %[[OUT_STRIDE_IDX]] : index) + # CHECK: return diff --git a/wave_lang/kernel/compiler/wave_codegen/emitter.py b/wave_lang/kernel/compiler/wave_codegen/emitter.py index a72886b596..b17bc3adab 100644 --- a/wave_lang/kernel/compiler/wave_codegen/emitter.py +++ b/wave_lang/kernel/compiler/wave_codegen/emitter.py @@ -71,6 +71,7 @@ BindingDesc, BoundKernelSignature, create_argument_locations, + get_dynamic_stride_arg_count, ) from ...lang.wave_types import IndexSymbol from ...wave.compile_options import WaveCompileOptions @@ -137,40 +138,70 @@ def emit_program_invariants(self): ), ] - def emit_func(self) -> Operation: - bindings = self.root_sig.sig.linear_bindings + def _abi_type(self, binding: BindingDesc) -> IrType: + if binding.binding_type == BindingType.KERNEL_BUFFER: + # Buffer passed to kernel as 0D memrefs to simplify ABI. + element_type = IrType.parse(binding.kernel_buffer_type.dtype.ir_type_asm()) + return MemRefType.get([], element_type=element_type) - def abi_type(binding: BindingDesc): - if binding.binding_type == BindingType.KERNEL_BUFFER: - # Buffer passed to kernel as 0D memrefs to simplify ABI. - element_type = IrType.parse( - binding.kernel_buffer_type.dtype.ir_type_asm() - ) - return MemRefType.get([], element_type=element_type) + # Scalars are passed as is. + return binding.as_mlir_type() - # Scalars are passed as is. - return binding.as_mlir_type() + def _create_stride_argument_locations( + self, bindings: list[BindingDesc], base_locations: list[Location] + ) -> list[Location]: + if not self.options.dynamic_strides: + return [] - arg_types = [abi_type(b) for b in bindings] + stride_locations = [] + for binding, base_location in zip(bindings, base_locations): + if binding.binding_type != BindingType.KERNEL_BUFFER: + continue + leading_count = max(0, len(binding.kernel_buffer_type.symbolic_shape) - 1) + argument_name = binding.name or "argument" + for dim_idx in range(leading_count): + stride_locations.append( + Location.name( + f"{argument_name}.stride.{dim_idx}", + base_location, + ) + ) + return stride_locations + + def _ordered_kernel_buffer_bindings( + self, bindings: list[BindingDesc] + ) -> list[BindingDesc]: + return [ + binding + for binding in bindings + if binding.binding_type == BindingType.KERNEL_BUFFER + ] - # Dynamic strides only with Wave runtime and LLVM backend (not ASM). - # Stride args only for leading dimensions (innermost always has unit stride). - stride_arg_count = 0 - if self.options.dynamic_strides: - stride_arg_count = sum( - max(0, len(b.kernel_buffer_type.symbolic_shape) - 1) - for b in self.root_sig.sig.kernel_buffer_bindings - ) - if stride_arg_count > 0: - arg_types += [IndexType.get()] * stride_arg_count + def _kernel_argument_abi( + self, bindings: list[BindingDesc] + ) -> tuple[list[IrType], list[Location], int]: + arg_types = [self._abi_type(binding) for binding in bindings] + arg_locations = create_argument_locations(bindings) + kernel_buffer_bindings = self._ordered_kernel_buffer_bindings(bindings) + stride_arg_count = get_dynamic_stride_arg_count( + self.options.dynamic_strides, + kernel_buffer_bindings, + ) + stride_locations = self._create_stride_argument_locations( + bindings, arg_locations + ) + assert len(stride_locations) == stride_arg_count + if stride_arg_count > 0: + arg_types += [IndexType.get()] * stride_arg_count + arg_locations += stride_locations + return arg_types, arg_locations, stride_arg_count + def emit_func(self) -> Operation: + bindings = self.root_sig.sig.linear_bindings + arg_types, arg_locations, stride_arg_count = self._kernel_argument_abi(bindings) ftype = FunctionType.get(arg_types, []) func_op = func_d.FuncOp(self.kernel_name, ftype, visibility="private") - - locs = create_argument_locations(bindings) - if stride_arg_count > 0: - locs += [Location.unknown()] * stride_arg_count - entry_block = func_op.add_entry_block(locs) + entry_block = func_op.add_entry_block(arg_locations) # Map dynamic symbols to buffer argument indices and dimensions. for bind, arg in zip(bindings, entry_block.arguments): @@ -285,8 +316,9 @@ def _declare_runtime_func( return func_op, symbol def emit_host_func(self, kernel_func: Operation) -> Operation: - # TODO: kernel bindings order may not be the same as the kernel function - # arguments order, so map kernel order to host function arguments order. + # Host placeholders follow the original Python signature, while kernel ABI + # arguments use linear_bindings order. Track both so launch args can be + # materialized in kernel ABI order from host-visible inputs. binding_map = {} symbol_map = {} @@ -308,35 +340,10 @@ def emit_host_func(self, kernel_func: Operation) -> Operation: bindings = self.root_sig.sig.linear_bindings ptr = llvm_d.PointerType.get() - - def abi_type(binding: BindingDesc): - if binding.binding_type == BindingType.KERNEL_BUFFER: - # Buffer passed to kernel as 0D memrefs to simplify ABI. - element_type = IrType.parse( - binding.kernel_buffer_type.dtype.ir_type_asm() - ) - return MemRefType.get([], element_type=element_type) - - # Scalars are passed as is. - return binding.as_mlir_type() - - arg_types = [abi_type(b) for b in bindings] - - # Match stride args added by emit_func when dynamic_strides is active. - stride_arg_count = 0 - if self.options.dynamic_strides: - stride_arg_count = sum( - max(0, len(b.kernel_buffer_type.symbolic_shape) - 1) - for b in self.root_sig.sig.kernel_buffer_bindings - ) - if stride_arg_count > 0: - arg_types += [IndexType.get()] * stride_arg_count - + arg_types, kernel_arg_locations, stride_arg_count = self._kernel_argument_abi( + bindings + ) ftype = FunctionType.get(arg_types, []) - # Preserve source locations for original args; synthesized stride args - # do not have an originating kernel SSA value. - locs = [a.location for a in kernel_func.body.blocks[0].arguments] - locs += [Location.unknown()] * (len(arg_types) - len(locs)) gpu_module = gpu_d.module("gpu_module") gpu_module.parent.operation.attributes["gpu.container_module"] = UnitAttr.get() @@ -358,7 +365,7 @@ def abi_type(binding: BindingDesc): new_kernel_entry_block = kernel_func_wrapper.body.blocks.append( *arg_types, - arg_locs=locs, + arg_locs=kernel_arg_locations, ) # Inline the kernel function into the gpu module function body and erase the original function @@ -512,7 +519,7 @@ def abi_type(binding: BindingDesc): # added by emit_func. One stride per leading dimension (rank-1) # for each kernel buffer. if self.options.dynamic_strides: - for binding in self.root_sig.sig.kernel_buffer_bindings: + for binding in self._ordered_kernel_buffer_bindings(bindings): rank = len(binding.kernel_buffer_type.symbolic_shape) leading_count = max(0, rank - 1) arg = func_args[binding_map[id(binding)]] diff --git a/wave_lang/kernel/wave/water.py b/wave_lang/kernel/wave/water.py index 759033e1b6..b1e242c764 100644 --- a/wave_lang/kernel/wave/water.py +++ b/wave_lang/kernel/wave/water.py @@ -14,6 +14,13 @@ import math from typing import Any, Sequence +from iree.compiler.dialects import ( + _structured_transform_ops_gen as structured_transform_ops, +) +from iree.compiler.dialects.transform import ( + any_op_t, +) + from wave_lang.kernel.wave.compile_options import WaveCompileOptions from wave_lang.support.detect_water import get_water_mlir_pkg_path, get_water_opt from wave_lang.support.detect_waveasm import get_waveasm_translate @@ -21,17 +28,21 @@ Attribute, BlockArgument, FunctionType, + IrType, InsertionPoint, IntegerType, + Location, MemRefType, Module, Operation, TypeAttr, + UnitAttr, WalkResult, gpu_d, llvm_d, memref_d, stream_d, + transform_d, ) @@ -216,6 +227,41 @@ def make_pass_arguments( ) +_ALLOCA_TO_GLOBAL_ENTRYPOINT = "__transform_alloca_to_global" +_TRANSFORM_MEMREF_ALLOCA_TYPE = '!transform.op<"memref.alloca">' + + +def _module_asm_with_alloca_to_global_transform(module: Module) -> str: + """Return a cloned module with a named alloca-to-global transform attached.""" + + with module.context, Location.unknown(): + transformed_module = Module.parse( + module.operation.get_asm(), context=module.context + ) + transformed_module.operation.attributes["transform.with_named_sequence"] = ( + UnitAttr.get() + ) + with InsertionPoint(transformed_module.body): + named_sequence = transform_d.NamedSequenceOp( + _ALLOCA_TO_GLOBAL_ENTRYPOINT, [any_op_t()], [] + ) + with InsertionPoint(named_sequence.body): + target = named_sequence.body.arguments[0] + alloca_handle_type = IrType.parse(_TRANSFORM_MEMREF_ALLOCA_TYPE) + alloca = structured_transform_ops.structured_match( + alloca_handle_type, + target, + ops=["memref.alloca"], + ) + Operation.create( + "transform.memref.alloca_to_global", + results=[any_op_t(), any_op_t()], + operands=[alloca], + ) + transform_d.YieldOp([]) + return transformed_module.operation.get_asm() + + def water_leak_in_bounds_check(module: Module, override_ir: str = ""): binary = get_water_opt() generic_mlir = _deiree(module) if override_ir == "" else override_ir @@ -381,7 +427,7 @@ def water_waveasm_lowering_pipeline( Step 3 (water-opt): host runtime wrapping (gpu.binary -> runtime calls). """ water_opt = get_water_opt() - mlir_asm = module.operation.get_asm() + mlir_asm = _module_asm_with_alloca_to_global_transform(module) target_chip = options.target lld_path = get_water_mlir_pkg_path() / "llvm" / "bin" / "ld.lld" @@ -390,27 +436,6 @@ def add_opt(pipeline): return [pipeline] return [] - def add_transform(transform: str, entry_point: str) -> tuple[str, dict[str, Any]]: - nonlocal mlir_asm - # Erase the last '}' closing the module, append the transform, re-close. - last_close = mlir_asm.rfind("}") - if last_close != -1: - mlir_asm = mlir_asm[:last_close] - mlir_asm += transform - mlir_asm += "}\n" - return ("transform-interpreter", {"entry-point": entry_point}) - - alloca_to_global = """ - transform.named_sequence @__transform_alloca_to_global(%arg0: !transform.any_op {transform.readonly}) { - %alloca = transform.structured.match ops{["memref.alloca"]} in %arg0 - : (!transform.any_op) -> !transform.op<"memref.alloca"> - %get_global, %global = transform.memref.alloca_to_global %alloca - : (!transform.op<"memref.alloca">) - -> (!transform.any_op, !transform.any_op) - transform.yield - } -""" - canonicalize_cse = "composite-fixed-point-pass", { "name": "canonicalize_cse", "pipeline": "any(canonicalize,cse)", @@ -431,7 +456,7 @@ def add_transform(transform: str, entry_point: str) -> tuple[str, dict[str, Any] *add_opt(int_range_optimizations), *add_opt("loop-invariant-code-motion"), ("water-alloc-to-alloca", {}, "gpu.module"), - add_transform(alloca_to_global, "__transform_alloca_to_global"), + ("transform-interpreter", {"entry-point": _ALLOCA_TO_GLOBAL_ENTRYPOINT}), ("convert-amdgpu-to-rocdl", {"chipset": target_chip}), ( "convert-gpu-to-rocdl", @@ -479,6 +504,7 @@ def run_subprocess(args, input_text, tool_name): waveasm_translate = get_waveasm_translate() waveasm_args = [ waveasm_translate, + "--waveasm-llvm-sdiv-srem-legalization", f"--waveasm-translate-from-llvm=target={target_chip}", "--waveasm-arith-legalization", "--waveasm-scoped-cse", @@ -520,7 +546,7 @@ def run_subprocess(args, input_text, tool_name): def water_lowering_pipeline(module: Module, options: WaveCompileOptions) -> Module: binary = get_water_opt() - mlir_asm = module.operation.get_asm() + mlir_asm = _module_asm_with_alloca_to_global_transform(module) target_chip = options.target def add_opt(pipeline): @@ -529,38 +555,6 @@ def add_opt(pipeline): return [] - def add_transform(transform: str, entry_point: str) -> tuple[str, dict[str, Any]]: - nonlocal mlir_asm - # Erase the last occurrence of '}' from mlir_asm which closes the module operation - last_close = mlir_asm.rfind("}") - if last_close != -1: - mlir_asm = mlir_asm[:last_close] - mlir_asm += transform - mlir_asm += "}\n" - return ("transform-interpreter", {"entry-point": entry_point}) - - # TODO: this transform refuses to work. - alloc_to_alloca = """ - transform.named_sequence @__transform_alloc_to_alloca(%arg0: !transform.any_op {transform.readonly}) { - %0 = transform.structured.match ops{["gpu.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op - transform.apply_patterns to %0 { - transform.apply_patterns.memref.alloc_to_alloca - } : !transform.any_op - transform.yield - } -""" - - alloca_to_global = """ - transform.named_sequence @__transform_alloca_to_global(%arg0: !transform.any_op {transform.readonly}) { - %alloca = transform.structured.match ops{["memref.alloca"]} in %arg0 - : (!transform.any_op) -> !transform.op<"memref.alloca"> - %get_global, %global = transform.memref.alloca_to_global %alloca - : (!transform.op<"memref.alloca">) - -> (!transform.any_op, !transform.any_op) - transform.yield - } -""" - canonicalize_cse = "composite-fixed-point-pass", { "name": "canonicalize_cse", "pipeline": "any(canonicalize,cse)", @@ -582,8 +576,7 @@ def add_transform(transform: str, entry_point: str) -> tuple[str, dict[str, Any] *add_opt(int_range_optimizations), *add_opt("loop-invariant-code-motion"), ("water-alloc-to-alloca", {}, "gpu.module"), - # add_transform(alloc_to_alloca, "__transform_alloc_to_alloca"), - add_transform(alloca_to_global, "__transform_alloca_to_global"), + ("transform-interpreter", {"entry-point": _ALLOCA_TO_GLOBAL_ENTRYPOINT}), "convert-scf-to-cf", ("convert-amdgpu-to-rocdl", {"chipset": target_chip}), ("convert-gpu-to-rocdl", {"use-bare-ptr-memref-call-conv": "1"}, "gpu.module"), diff --git a/waveasm/include/waveasm/Transforms/Passes.td b/waveasm/include/waveasm/Transforms/Passes.td index 8178213319..d3bc948418 100644 --- a/waveasm/include/waveasm/Transforms/Passes.td +++ b/waveasm/include/waveasm/Transforms/Passes.td @@ -373,6 +373,22 @@ def WAVEASMArithLegalization : Pass<"waveasm-arith-legalization"> { let dependentDialects = ["::waveasm::WaveASMDialect"]; } +def WAVEASMLLVMSDivSRemLegalization + : Pass<"waveasm-llvm-sdiv-srem-legalization"> { + let summary = "Rewrite LLVM signed div/rem by positive power-of-two constants"; + let description = [{ + Rewrites `llvm.sdiv` and `llvm.srem` on i32 values with positive + power-of-two constant divisors into equivalent LLVM dialect compare/select/ + add/and/ashr sequences. + + This keeps the strength reduction in LLVM dialect, before LLVM->WaveASM + translation changes the abstraction level, and makes the rewrite + independently testable. + }]; + + let dependentDialects = ["::mlir::LLVM::LLVMDialect"]; +} + def WAVEASMTranslateFromLLVM : Pass<"waveasm-translate-from-llvm"> { let summary = "Translate LLVM dialect kernels to WaveASM IR"; let description = [{ diff --git a/waveasm/lib/Transforms/CMakeLists.txt b/waveasm/lib/Transforms/CMakeLists.txt index b226c536a8..314c57e1c7 100644 --- a/waveasm/lib/Transforms/CMakeLists.txt +++ b/waveasm/lib/Transforms/CMakeLists.txt @@ -24,6 +24,7 @@ add_mlir_dialect_library(MLIRWaveASMTransforms LinearScanRegAlloc.cpp LiteralMaterialization.cpp Liveness.cpp + LLVMSDivSRemLegalization.cpp LoopAddressPromotion.cpp M0RedundancyElimination.cpp MemoryOffsetOptimization.cpp diff --git a/waveasm/lib/Transforms/LLVMSDivSRemLegalization.cpp b/waveasm/lib/Transforms/LLVMSDivSRemLegalization.cpp new file mode 100644 index 0000000000..82e06ac5e1 --- /dev/null +++ b/waveasm/lib/Transforms/LLVMSDivSRemLegalization.cpp @@ -0,0 +1,139 @@ +// Copyright 2026 The Wave Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//===----------------------------------------------------------------------===// +// LLVM sdiv/srem legalization +// +// Rewrites signed division and remainder by positive power-of-two constants +// into equivalent LLVM dialect compare/select/add/and/ashr sequences before +// LLVM->WaveASM translation changes the abstraction level. +//===----------------------------------------------------------------------===// + +#include "waveasm/Transforms/Passes.h" + +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "llvm/Support/MathExtras.h" + +namespace waveasm { +#define GEN_PASS_DEF_WAVEASMLLVMSDIVSREMLEGALIZATION +#include "waveasm/Transforms/Passes.h.inc" +} // namespace waveasm + +using namespace mlir; + +namespace { + +static std::optional getConstantI32Value(Value value) { + LLVM::ConstantOp constOp = value.getDefiningOp(); + if (!constOp) + return std::nullopt; + IntegerAttr intAttr = dyn_cast(constOp.getValue()); + if (!intAttr) + return std::nullopt; + return intAttr.getInt(); +} + +static std::optional> +matchPositivePowerOfTwoI32Divisor(Value rhs) { + std::optional constVal = getConstantI32Value(rhs); + if (!constVal || *constVal <= 0 || !llvm::isPowerOf2_64(*constVal)) + return std::nullopt; + int64_t divisor = *constVal; + int64_t shiftAmt = llvm::Log2_64(static_cast(divisor)); + return std::pair{divisor, shiftAmt}; +} + +static Value createI32Constant(PatternRewriter &rewriter, Location loc, + int64_t value) { + Type i32 = rewriter.getI32Type(); + return LLVM::ConstantOp::create(rewriter, loc, i32, + rewriter.getIntegerAttr(i32, value)); +} + +struct LegalizePowerOfTwoSDivPattern : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(LLVM::SDivOp op, + PatternRewriter &rewriter) const override { + if (!op.getType().isSignlessInteger(32)) + return failure(); + std::optional> divisorAndShift = + matchPositivePowerOfTwoI32Divisor(op.getRhs()); + if (!divisorAndShift) + return failure(); + + int64_t divisor = divisorAndShift->first; + int64_t shiftAmt = divisorAndShift->second; + Location loc = op.getLoc(); + + Value zero = createI32Constant(rewriter, loc, 0); + Value biasImm = createI32Constant(rewriter, loc, divisor - 1); + Value shiftConst = createI32Constant(rewriter, loc, shiftAmt); + Value isNegative = + LLVM::ICmpOp::create(rewriter, loc, LLVM::ICmpPredicate::slt, + op.getLhs(), zero); + Value bias = LLVM::SelectOp::create(rewriter, loc, isNegative, biasImm, + zero, LLVM::FastmathFlags::none); + Value biased = LLVM::AddOp::create( + rewriter, loc, op.getLhs(), bias, LLVM::IntegerOverflowFlags::none); + Value result = LLVM::AShrOp::create(rewriter, loc, biased, shiftConst, + /*isExact=*/false); + rewriter.replaceOp(op, result); + return success(); + } +}; + +struct LegalizePowerOfTwoSRemPattern : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(LLVM::SRemOp op, + PatternRewriter &rewriter) const override { + if (!op.getType().isSignlessInteger(32)) + return failure(); + std::optional> divisorAndShift = + matchPositivePowerOfTwoI32Divisor(op.getRhs()); + if (!divisorAndShift) + return failure(); + + int64_t divisor = divisorAndShift->first; + Location loc = op.getLoc(); + + Value zero = createI32Constant(rewriter, loc, 0); + Value maskConst = createI32Constant(rewriter, loc, divisor - 1); + Value negDivisor = createI32Constant(rewriter, loc, -divisor); + Value rawRem = LLVM::AndOp::create(rewriter, loc, op.getLhs(), maskConst); + Value isNegative = + LLVM::ICmpOp::create(rewriter, loc, LLVM::ICmpPredicate::slt, + op.getLhs(), zero); + Value isNonZero = LLVM::ICmpOp::create(rewriter, loc, LLVM::ICmpPredicate::ne, + rawRem, zero); + Value needsAdjust = LLVM::AndOp::create(rewriter, loc, isNegative, isNonZero); + Value adjust = LLVM::SelectOp::create(rewriter, loc, needsAdjust, negDivisor, + zero, LLVM::FastmathFlags::none); + Value result = LLVM::AddOp::create( + rewriter, loc, rawRem, adjust, LLVM::IntegerOverflowFlags::none); + rewriter.replaceOp(op, result); + return success(); + } +}; + +struct LLVMSDivSRemLegalizationPass + : public waveasm::impl::WAVEASMLLVMSDivSRemLegalizationBase< + LLVMSDivSRemLegalizationPass> { + void runOnOperation() override { + RewritePatternSet patterns(&getContext()); + patterns.add( + &getContext()); + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) + signalPassFailure(); + } +}; + +} // namespace diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 0b6900b089..95336f930b 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -24,9 +24,13 @@ #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Utils/StaticValueUtils.h" #include "mlir/IR/Builders.h" +#include "mlir/IR/SymbolTable.h" #include "llvm/ADT/StringExtras.h" +#include "llvm/ADT/StringMap.h" #include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/Debug.h" +#include "llvm/Support/MathExtras.h" +#include #define DEBUG_TYPE "waveasm-translate-llvm" @@ -64,7 +68,14 @@ struct BufferPtrInfo { /// State for LLVM->WaveASM translation, layered on top of TranslationContext. class LLVMTranslationState { public: - explicit LLVMTranslationState(TranslationContext &ctx) : ctx(ctx) {} + explicit LLVMTranslationState(TranslationContext &ctx, + Operation *symbolTableOp) + : ctx(ctx) { + assert(symbolTableOp && "expected symbol table operation"); + symbolTableOp->walk([&](LLVM::GlobalOp global) { + globalsByName[global.getSymName()] = global; + }); + } TranslationContext &ctx; @@ -87,10 +98,30 @@ class LLVMTranslationState { void setBaseOffset(Value ptr, Value offset) { baseOffsets[ptr] = offset; } Value lookupBaseOffset(Value ptr) const { return baseOffsets.lookup(ptr); } + /// Track the assigned LDS byte offset for each referenced llvm.mlir.global. + void setLDSGlobalOffset(LLVM::GlobalOp global, int64_t offset) { + ldsGlobalOffsets[global.getOperation()] = offset; + } + std::optional lookupLDSGlobalOffset(LLVM::GlobalOp global) const { + auto it = ldsGlobalOffsets.find(global.getOperation()); + if (it != ldsGlobalOffsets.end()) + return it->second; + return std::nullopt; + } + + LLVM::GlobalOp lookupGlobal(StringRef name) const { + auto it = globalsByName.find(name); + if (it != globalsByName.end()) + return it->second; + return {}; + } + private: DenseMap rsrcToSRD; DenseMap gepMap; DenseMap baseOffsets; + DenseMap ldsGlobalOffsets; + llvm::StringMap globalsByName; }; //===----------------------------------------------------------------------===// @@ -165,13 +196,17 @@ static ProgramOp createProgramFromLLVMFunc(LLVM::LLVMFuncOp func, } /// Look up the WaveASM value that an LLVM value was translated to. -/// Returns failure if the value was never mapped -- silently returning -/// the original (soon-to-be-erased) LLVM value is a use-after-free bug. +/// Returns failure if the value was never mapped. Never fall back to the +/// original LLVM SSA value: translation erases the LLVM ops after building +/// WaveASM, so reusing the source SSA value here would leave a dangling +/// reference. static FailureOr resolve(Value v, TranslationContext &ctx) { if (auto mapped = ctx.getMapper().getMapped(v)) return *mapped; - // Block arguments (func params) are mapped during prologue setup. - // If we get here, an LLVM op was skipped or handled incorrectly. + // Block arguments (func params) are mapped during prologue setup. Reaching + // this point means either malformed IR referenced an unmapped SSA value or + // translation skipped a producer; surface that as a hard failure in the + // caller instead of guessing through it. return failure(); } @@ -208,33 +243,51 @@ static unsigned getLLVMAddrSpace(Value v) { return 0; } -/// Return the byte size for a sized LLVM scalar, vector, or array type. -/// Returns 0 for unsupported types. -static int64_t getLLVMTypeBytes(Type ty) { +/// WaveASM register sizes are tracked in 32-bit lanes (dwords). +static int64_t getWaveASMLaneCount(int64_t bitWidth) { + return llvm::divideCeil(bitWidth, int64_t{32}); +} + +/// Return the byte size for an LLVM type whose layout is obvious without a +/// DataLayout query. +/// +/// We intentionally support only scalars, fixed vectors, and arrays thereof. +/// Nested aggregates such as structs require real DataLayout-based layout +/// reasoning; approximating them here would make GEP and LDS offsets +/// target-dependent in subtle ways. +static FailureOr getLLVMTypeBytesWithoutDataLayout(Type ty) { if (ty.isIntOrFloat()) return ty.getIntOrFloatBitWidth() / 8; - if (VectorType vecTy = dyn_cast(ty)) - return vecTy.getNumElements() * - vecTy.getElementType().getIntOrFloatBitWidth() / 8; - if (LLVM::LLVMArrayType arrTy = dyn_cast(ty)) - return getLLVMTypeBytes(arrTy.getElementType()) * arrTy.getNumElements(); - return 0; + if (VectorType vecTy = dyn_cast(ty)) { + FailureOr elemBytes = + getLLVMTypeBytesWithoutDataLayout(vecTy.getElementType()); + if (failed(elemBytes)) + return failure(); + return *elemBytes * vecTy.getNumElements(); + } + if (LLVM::LLVMArrayType arrTy = dyn_cast(ty)) { + FailureOr elemBytes = + getLLVMTypeBytesWithoutDataLayout(arrTy.getElementType()); + if (failed(elemBytes)) + return failure(); + return *elemBytes * arrTy.getNumElements(); + } + return failure(); } /// Return true iff a GEP index is statically known to be zero. static bool isZeroGEPIndex(llvm::PointerUnion idx) { - if (isa(idx)) - return isConstantIntValue(cast(idx), 0); - return isConstantIntValue(cast(idx), 0); + if (Value value = dyn_cast(idx)) + return isConstantIntValue(value, 0); + IntegerAttr intAttr = dyn_cast(idx); + assert(intAttr && "GEP index must be a Value or IntegerAttr"); + return isConstantIntValue(intAttr, 0); } /// Structural GEPs like [0, 0] are a no-op and can be forwarded even though we /// do not model nested aggregate layouts yet. static bool isAllZeroIndexGEP(LLVM::GEPOp op) { - for (llvm::PointerUnion idx : op.getIndices()) - if (!isZeroGEPIndex(idx)) - return false; - return true; + return llvm::all_of(op.getIndices(), isZeroGEPIndex); } /// Compute the non-zero byte offset for a supported single-index GEP. @@ -251,32 +304,49 @@ static FailureOr computeGEPByteOffset(LLVM::GEPOp op, if (indices.size() != 1) return op->emitOpError("GEP byte offset requires a single index"); - int64_t elemBytes = getLLVMTypeBytes(op.getElemType()); - if (elemBytes == 0) + FailureOr elemBytes = + getLLVMTypeBytesWithoutDataLayout(op.getElemType()); + if (failed(elemBytes)) return op->emitOpError( - "unsupported GEP element type for byte offset computation"); + "unsupported GEP element type for byte offset computation " + "without DataLayout: ") + << op.getElemType(); llvm::PointerUnion idx = indices[0]; - if (Value dynIdx = idx.dyn_cast()) { + if (Value dynIdx = dyn_cast(idx)) { FailureOr resolved = resolve(dynIdx, ctx); if (failed(resolved)) return op->emitOpError("unmapped GEP index"); - if (elemBytes == 1) + if (*elemBytes == 1) return *resolved; // Scale by element size: offset = index * elemBytes. - ImmType scaleTy = ctx.createImmType(elemBytes); - Value scale = ConstantOp::create(builder, loc, scaleTy, elemBytes); + ImmType scaleTy = ctx.createImmType(*elemBytes); + Value scale = ConstantOp::create(builder, loc, scaleTy, *elemBytes); Type resTy = inferResultType({*resolved, scale}, ctx); return ArithMulOp::create(builder, loc, resTy, *resolved, scale) .getResult(); } - int64_t constIdx = cast(idx).getInt(); - int64_t byteOffset = constIdx * elemBytes; + IntegerAttr constIdxAttr = dyn_cast(idx); + assert(constIdxAttr && "GEP index must be a Value or IntegerAttr"); + int64_t constIdx = constIdxAttr.getInt(); + int64_t byteOffset = constIdx * *elemBytes; ImmType immTy = ctx.createImmType(byteOffset); return ConstantOp::create(builder, loc, immTy, byteOffset).getResult(); } +/// DS instructions address LDS by a 32-bit byte offset, not a generic pointer. +/// Valid in-range LDS accesses are bounded by the kernel's per-workgroup LDS +/// allocation, so well-formed LLVM IR may represent the offset as i64 but still +/// fits in the hardware vaddr field after truncation. +static Value materializeLDSVAddr(Value offset, OpBuilder &builder, Location loc, + TranslationContext &ctx) { + if (getRegSize(offset.getType()) <= 1) + return offset; + VRegType vregTy = ctx.createVRegType(); + return ArithTruncOp::create(builder, loc, vregTy, offset); +} + //===----------------------------------------------------------------------===// // Op handlers //===----------------------------------------------------------------------===// @@ -426,7 +496,7 @@ static LogicalResult handleCastOp(LLVMOp op, LLVMTranslationState &st) { return op->emitOpError("unmapped operand in cast"); int64_t width = 1; if (IntegerType intTy = dyn_cast(op.getResult().getType())) - width = (intTy.getWidth() + 31) / 32; + width = getWaveASMLaneCount(intTy.getWidth()); Type resTy = isVGPRType(src->getType()) ? (Type)ctx.createVRegType(width, width) : (Type)ctx.createSRegType(width, width); @@ -509,6 +579,36 @@ static LogicalResult handleBinaryOp(LLVMOp op, LLVMTranslationState &st) { return success(); } +static LogicalResult handleAShr(LLVM::AShrOp op, LLVMTranslationState &st) { + TranslationContext &ctx = st.ctx; + OpBuilder &builder = ctx.getBuilder(); + Location loc = op.getLoc(); + FailureOr lhs = resolve(op.getLhs(), ctx); + FailureOr rhs = resolve(op.getRhs(), ctx); + if (failed(lhs) || failed(rhs)) + return op->emitOpError("unmapped operand in ashr"); + if (getRegSize(lhs->getType()) != 1 || getRegSize(rhs->getType()) != 1) + return op->emitOpError( + "arithmetic shift currently supports only i32 operands"); + + Value result; + if (isSGPRType(lhs->getType())) { + if (isVGPRType(rhs->getType())) + return op->emitOpError( + "SGPR arithmetic shift requires an SGPR or immediate shift amount"); + SRegType sregTy = ctx.createSRegType(); + SCCType sccTy = ctx.createSCCType(); + result = + S_ASHR_I32::create(builder, loc, sregTy, sccTy, *lhs, *rhs).getDst(); + } else { + VRegType vregTy = ctx.createVRegType(); + result = V_ASHRREV_I32::create(builder, loc, vregTy, *rhs, *lhs); + } + + ctx.getMapper().mapValue(op.getResult(), result); + return success(); +} + static LogicalResult handleMakeBufferRsrc(ROCDL::MakeBufferRsrcOp op, LLVMTranslationState &st) { auto &ctx = st.ctx; @@ -521,10 +621,10 @@ static LogicalResult handleMakeBufferRsrc(ROCDL::MakeBufferRsrcOp op, if (!srdVal) return op->emitOpError("SRD not found for base pointer"); - // The prologue used s_mov_b64 to copy the 64-bit pointer into SRD[0:1]. - // This corrupts SRD word 1 bits [31:16] (stride/swizzle) with pointer bits. - // Also, the prologue hardcodes SRD[3]=0x20000 but make.buffer.rsrc may - // want different flags. Rebuild the SRD with corrected words via PackOp. + // The prologue seeds SRD[0:1] from the raw 64-bit pointer. That leaves SRD + // word 1 bits [31:16] (stride/swizzle) populated with pointer bits, and + // initializes SRD[3] to the default flag value. Rebuild the descriptor here + // so make.buffer.rsrc controls both fields explicitly. bool needsPatch = (srdVal->getDefiningOp() != nullptr); if (needsPatch) { SRegType sregTy = ctx.createSRegType(); @@ -543,7 +643,7 @@ static LogicalResult handleMakeBufferRsrc(ROCDL::MakeBufferRsrcOp op, maskVal) .getDst(); - // Patch SRD[3] with the actual flags from make.buffer.rsrc. + // Patch SRD[3] with the requested descriptor flags. std::optional flags = getConstantIntValue(op.getFlags()); if (flags && *flags != kSRDDefaultFlags) { ImmType flagsImm = ctx.createImmType(*flags); @@ -611,9 +711,7 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { FailureOr baseOff = resolve(base, ctx); if (failed(baseOff)) return op->emitOpError("unmapped LDS GEP base"); - VRegType vregTy = ctx.createVRegType(); - if (getRegSize(baseOff->getType()) > 1) - *baseOff = ArithTruncOp::create(builder, loc, vregTy, *baseOff); + *baseOff = materializeLDSVAddr(*baseOff, builder, loc, ctx); if (isAllZeroGEP) { ctx.getMapper().mapValue(op.getResult(), *baseOff); return success(); @@ -621,9 +719,8 @@ static LogicalResult handleGEP(LLVM::GEPOp op, LLVMTranslationState &st) { FailureOr maybeOffset = computeGEPByteOffset(op, ctx); if (failed(maybeOffset)) return failure(); - Value off = *maybeOffset; - if (getRegSize(off.getType()) > 1) - off = ArithTruncOp::create(builder, loc, vregTy, off); + Value off = materializeLDSVAddr(*maybeOffset, builder, loc, ctx); + VRegType vregTy = ctx.createVRegType(); Value sum = ArithAddOp::create(builder, loc, vregTy, *baseOff, off); ctx.getMapper().mapValue(op.getResult(), sum); return success(); @@ -808,22 +905,45 @@ static LogicalResult handleStore(LLVM::StoreOp op, LLVMTranslationState &st) { static LogicalResult handleAddressOf(LLVM::AddressOfOp op, LLVMTranslationState &st) { TranslationContext &ctx = st.ctx; - // Look up the global to determine LDS size. - LLVM::GlobalOp global = SymbolTable::lookupNearestSymbolFrom( - op, op.getGlobalNameAttr()); - if (global) { - Type globalTy = global.getType(); - int64_t sizeBytes = getLLVMTypeBytes(globalTy); - if (sizeBytes > 0) - ctx.addLDSSize(sizeBytes); + LLVM::GlobalOp global = st.lookupGlobal(op.getGlobalName()); + if (!global) + return op->emitOpError("referenced global not found"); + + IntegerAttr addrSpaceAttr = global->getAttrOfType("addr_space"); + if (!addrSpaceAttr || addrSpaceAttr.getInt() != 3) + return op->emitOpError( + "llvm.mlir.addressof currently supports only LDS globals in " + "addrspace(3)"); + + FailureOr sizeBytes = + getLLVMTypeBytesWithoutDataLayout(global.getType()); + if (failed(sizeBytes)) + return op->emitOpError( + "unsupported LDS global type for byte size computation " + "without DataLayout: ") + << global.getType(); + + int64_t baseOffset = 0; + if (std::optional existingOffset = + st.lookupLDSGlobalOffset(global)) { + baseOffset = *existingOffset; + } else { + baseOffset = ctx.getLDSAllocOffset(); + assert(baseOffset >= 0 && + static_cast(baseOffset) <= + std::numeric_limits::max() && + "LDS base offset must fit in the 32-bit DS vaddr field"); + st.setLDSGlobalOffset(global, baseOffset); + ctx.addLDSSize(*sizeBytes); + ctx.advanceLDSAllocOffset(*sizeBytes); } - // Map to a constant 0 offset in a VGPR -- LDS addressing is relative. + // Map to the byte offset assigned to this LDS global. OpBuilder &builder = ctx.getBuilder(); - ImmType immTy = ctx.createImmType(0); - Value zero = ConstantOp::create(builder, op.getLoc(), immTy, int64_t{0}); + ImmType immTy = ctx.createImmType(baseOffset); + Value baseImm = ConstantOp::create(builder, op.getLoc(), immTy, baseOffset); VRegType vregTy = ctx.createVRegType(); - Value mov = V_MOV_B32::create(builder, op.getLoc(), vregTy, zero); + Value mov = V_MOV_B32::create(builder, op.getLoc(), vregTy, baseImm); ctx.getMapper().mapValue(op.getResult(), mov); return success(); } @@ -842,10 +962,8 @@ static LogicalResult handleLDSLoad(LLVM::LoadOp op, Value addr, FailureOr offset = resolve(addr, ctx); if (failed(offset)) return op->emitOpError("unmapped LDS address"); - // DS instructions use a 32-bit vaddr. Truncate wide offsets. + *offset = materializeLDSVAddr(*offset, builder, loc, ctx); VRegType vregTy = ctx.createVRegType(); - if (getRegSize(offset->getType()) > 1) - *offset = ArithTruncOp::create(builder, loc, vregTy, *offset); Operation *loadOp = nullptr; if (numBytes == 2) @@ -876,10 +994,7 @@ static LogicalResult handleLDSStore(LLVM::StoreOp op, Value addr, FailureOr offset = resolve(addr, ctx); if (failed(data) || failed(offset)) return op->emitOpError("unmapped LDS operand"); - // DS instructions use a 32-bit vaddr. Truncate wide offsets. - VRegType vregTy = ctx.createVRegType(); - if (getRegSize(offset->getType()) > 1) - *offset = ArithTruncOp::create(builder, loc, vregTy, *offset); + *offset = materializeLDSVAddr(*offset, builder, loc, ctx); int64_t numBytes = getBufferAccessBytes(op.getValue().getType()); if (numBytes == 2) @@ -902,118 +1017,13 @@ static LogicalResult handleLDSStore(LLVM::StoreOp op, Value addr, //===----------------------------------------------------------------------===// static LogicalResult handleSDiv(LLVM::SDivOp op, LLVMTranslationState &st) { - TranslationContext &ctx = st.ctx; - OpBuilder &builder = ctx.getBuilder(); - Location loc = op.getLoc(); - FailureOr lhs = resolve(op.getLhs(), ctx); - if (failed(lhs)) - return op->emitOpError("unmapped operand in sdiv"); - if (getRegSize(lhs->getType()) != 1) - return op->emitOpError( - "signed division currently supports only i32 operands"); - - // Signed division by a positive power-of-2 constant: - // q = (x + (x < 0 ? divisor - 1 : 0)) >> log2(divisor) - // This preserves LLVM's trunc-toward-zero semantics for negative dividends. - if (std::optional constVal = getConstantIntValue(op.getRhs())) { - int64_t divisor = *constVal; - if (divisor > 0 && (divisor & (divisor - 1)) == 0) { - int64_t shiftAmt = llvm::Log2_64(divisor); - Value zero = ConstantOp::create(builder, loc, ctx.createImmType(0), 0); - Value biasImm = ConstantOp::create( - builder, loc, ctx.createImmType(divisor - 1), divisor - 1); - Value shiftConst = ConstantOp::create( - builder, loc, ctx.createImmType(shiftAmt), shiftAmt); - Value result; - if (isSGPRType(lhs->getType())) { - SRegType sregTy = ctx.createSRegType(); - SCCType sccTy = ctx.createSCCType(); - Value isNegative = - S_CMP_LT_I32::create(builder, loc, sccTy, *lhs, zero); - Value bias = S_CSELECT_B32::create(builder, loc, sregTy, isNegative, - biasImm, zero); - Value biased = ArithAddOp::create(builder, loc, sregTy, *lhs, bias); - result = - S_ASHR_I32::create(builder, loc, sregTy, sccTy, biased, shiftConst) - .getDst(); - } else { - VRegType vregTy = ctx.createVRegType(); - Value isNegative = ArithCmpOp::create(builder, loc, vregTy, - CmpPredicate::slt, *lhs, zero); - Value bias = ArithSelectOp::create(builder, loc, vregTy, zero, biasImm, - isNegative); - Value biased = ArithAddOp::create(builder, loc, vregTy, *lhs, bias); - result = - V_ASHRREV_I32::create(builder, loc, vregTy, shiftConst, biased); - } - ctx.getMapper().mapValue(op.getResult(), result); - return success(); - } - } return op->emitOpError( - "signed division currently supports only positive power-of-2 constants"); + "llvm.sdiv must be legalized before LLVM->WaveASM translation"); } static LogicalResult handleSRem(LLVM::SRemOp op, LLVMTranslationState &st) { - TranslationContext &ctx = st.ctx; - OpBuilder &builder = ctx.getBuilder(); - Location loc = op.getLoc(); - FailureOr lhs = resolve(op.getLhs(), ctx); - if (failed(lhs)) - return op->emitOpError("unmapped operand in srem"); - if (getRegSize(lhs->getType()) != 1) - return op->emitOpError( - "signed remainder currently supports only i32 operands"); - - // Signed remainder by a positive power-of-2 constant: - // r = x & (divisor - 1) - // if (x < 0 && r != 0) r -= divisor - // This keeps the remainder's sign consistent with LLVM srem. - if (std::optional constVal = getConstantIntValue(op.getRhs())) { - int64_t divisor = *constVal; - if (divisor > 0 && (divisor & (divisor - 1)) == 0) { - Value zero = ConstantOp::create(builder, loc, ctx.createImmType(0), 0); - Value maskConst = ConstantOp::create( - builder, loc, ctx.createImmType(divisor - 1), divisor - 1); - Value negDivisor = ConstantOp::create( - builder, loc, ctx.createImmType(-divisor), -divisor); - Value result; - if (isSGPRType(lhs->getType())) { - SRegType sregTy = ctx.createSRegType(); - SCCType sccTy = ctx.createSCCType(); - Value rawRem = - ArithAndOp::create(builder, loc, sregTy, *lhs, maskConst); - Value isNegative = - S_CMP_LT_I32::create(builder, loc, sccTy, *lhs, zero); - Value isNonZero = - S_CMP_NE_I32::create(builder, loc, sccTy, rawRem, zero); - Value adjusted = - ArithAddOp::create(builder, loc, sregTy, rawRem, negDivisor); - Value maybeAdjusted = S_CSELECT_B32::create( - builder, loc, sregTy, isNonZero, adjusted, rawRem); - result = S_CSELECT_B32::create(builder, loc, sregTy, isNegative, - maybeAdjusted, rawRem); - } else { - VRegType vregTy = ctx.createVRegType(); - Value rawRem = - ArithAndOp::create(builder, loc, vregTy, *lhs, maskConst); - Value isNegative = ArithCmpOp::create(builder, loc, vregTy, - CmpPredicate::slt, *lhs, zero); - Value isNonZero = ArithCmpOp::create(builder, loc, vregTy, - CmpPredicate::ne, rawRem, zero); - Value adjusted = - ArithAddOp::create(builder, loc, vregTy, rawRem, negDivisor); - Value maybeAdjusted = ArithSelectOp::create( - builder, loc, vregTy, rawRem, adjusted, isNonZero); - result = ArithSelectOp::create(builder, loc, vregTy, rawRem, - maybeAdjusted, isNegative); - } - ctx.getMapper().mapValue(op.getResult(), result); - return success(); - } - } return op->emitOpError( - "signed remainder currently supports only positive power-of-2 constants"); + "llvm.srem must be legalized before LLVM->WaveASM translation"); } //===----------------------------------------------------------------------===// @@ -1105,31 +1115,61 @@ static FailureOr materializeI32SReg(Value v, StringRef name, return failure(); } -static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { +struct SCFForLoopBounds { + Value initialIV; + Value upperBound; + Value step; +}; + +/// Resolve and scalarize scf.for loop control in the preheader. +/// +/// scf.for evaluates lb/ub/step once before the first iteration. Capture those +/// values outside the lowered waveasm.loop body so the upper bound is not +/// implicitly rematerialized on the backedge. +static FailureOr +resolveSCFForLoopBounds(scf::ForOp op, LLVMTranslationState &st) { TranslationContext &ctx = st.ctx; OpBuilder &builder = ctx.getBuilder(); Location loc = op.getLoc(); - FailureOr lb = resolve(op.getLowerBound(), ctx); - FailureOr ub = resolve(op.getUpperBound(), ctx); + FailureOr lowerBound = resolve(op.getLowerBound(), ctx); + FailureOr upperBound = resolve(op.getUpperBound(), ctx); FailureOr step = resolve(op.getStep(), ctx); - if (failed(lb) || failed(ub) || failed(step)) - return op->emitOpError("unmapped operand in scf.for"); + if (failed(lowerBound) || failed(upperBound) || failed(step)) { + op->emitOpError("unmapped operand in scf.for"); + return failure(); + } // TODO: Preserve full scf.for zero-trip semantics. waveasm.loop is do-while // and this lowering currently assumes the trip count is known positive. - FailureOr lbScalar = - materializeI32SReg(*lb, "scf.for lower bound", op, ctx, builder, loc); - FailureOr ubScalar = - materializeI32SReg(*ub, "scf.for upper bound", op, ctx, builder, loc); - FailureOr stepScalar = + FailureOr initialIV = materializeI32SReg( + *lowerBound, "scf.for lower bound", op, ctx, builder, loc); + FailureOr loopUpperBound = materializeI32SReg( + *upperBound, "scf.for upper bound", op, ctx, builder, loc); + FailureOr loopStep = materializeI32SReg(*step, "scf.for step", op, ctx, builder, loc); - if (failed(lbScalar) || failed(ubScalar) || failed(stepScalar)) + if (failed(initialIV) || failed(loopUpperBound) || failed(loopStep)) + return failure(); + + return SCFForLoopBounds{ + /*initialIV=*/*initialIV, + /*upperBound=*/*loopUpperBound, + /*step=*/*loopStep, + }; +} + +static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { + TranslationContext &ctx = st.ctx; + OpBuilder &builder = ctx.getBuilder(); + Location loc = op.getLoc(); + + FailureOr loopBounds = resolveSCFForLoopBounds(op, st); + if (failed(loopBounds)) return failure(); // Build init args: [lower_bound, iter_args...]. SmallVector initArgs; - initArgs.push_back(*lbScalar); + initArgs.push_back(loopBounds->initialIV); for (Value arg : op.getInitArgs()) { FailureOr resolved = resolve(arg, ctx); if (failed(resolved)) @@ -1160,10 +1200,11 @@ static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { Value inductionVar = bodyBlock.getArgument(0); SRegType sregTy = ctx.createSRegType(); SCCType sccTy = ctx.createSCCType(); - Value nextIV = - S_ADD_U32::create(builder, loc, sregTy, sccTy, inductionVar, *stepScalar) - .getDst(); - Value cond = S_CMP_LT_U32::create(builder, loc, sccTy, nextIV, *ubScalar); + Value nextIV = S_ADD_U32::create(builder, loc, sregTy, sccTy, inductionVar, + loopBounds->step) + .getDst(); + Value cond = + S_CMP_LT_U32::create(builder, loc, sccTy, nextIV, loopBounds->upperBound); // Collect iter args from yield. scf::YieldOp yieldOp = cast(op.getBody()->getTerminator()); @@ -1215,6 +1256,7 @@ static LogicalResult translateOp(Operation *op, LLVMTranslationState &st) { .Case([&](LLVM::AddOp o) { return handleBinaryOp(o, st); }) + .Case([&](LLVM::AShrOp o) { return handleAShr(o, st); }) .Case([&](LLVM::OrOp o) { return handleBinaryOp(o, st); }) @@ -1277,7 +1319,9 @@ static LogicalResult translateLLVMModule(Operation *rootOp, return failure(); builder.setInsertionPointToStart(&program.getBodyBlock()); TranslationContext ctx(builder, program, target); - LLVMTranslationState st(ctx); + Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(func); + assert(symbolTableOp && "expected nearest symbol table for llvm.func"); + LLVMTranslationState st(ctx, symbolTableOp); // Map llvm.func arguments: pointers get SRD setup, scalars get mapped // to their preloaded SGPR positions directly. @@ -1306,9 +1350,10 @@ static LogicalResult translateLLVMModule(Operation *rootOp, ctx.getMapper().mapValue(arg, sreg); } - // Enable all workgroup IDs so the SGPR layout is predictable. - // Note: LLVM enables them selectively via amdgpu-no-workgroup-id-{y,z} - // attributes. We enable all three unconditionally for simplicity. + // Keep all workgroup IDs enabled so the SGPR layout is stable across + // translated kernels. LLVM can disable y/z selectively via + // amdgpu-no-workgroup-id-{y,z}, but the fixed WaveASM prologue layout here + // assumes all three slots exist. ctx.enableAllWorkgroupIds(); for (Operation &op : func.getBody().front()) { diff --git a/waveasm/test/Transforms/llvm-sdiv-srem-legalization.mlir b/waveasm/test/Transforms/llvm-sdiv-srem-legalization.mlir new file mode 100644 index 0000000000..638f94d283 --- /dev/null +++ b/waveasm/test/Transforms/llvm-sdiv-srem-legalization.mlir @@ -0,0 +1,31 @@ +// RUN: waveasm-translate %s --waveasm-llvm-sdiv-srem-legalization | FileCheck %s +// Verify that the LLVM pre-pass rewrites signed div/rem by positive power-of-2 +// constants before LLVM->WaveASM translation. + +// CHECK-LABEL: llvm.func @test +// CHECK: llvm.icmp "slt" +// CHECK: llvm.select +// CHECK: llvm.ashr +// CHECK: llvm.and +// CHECK: llvm.icmp "ne" +// CHECK: llvm.select + +gpu.module @gpu_module { + llvm.mlir.global private @scratch() {addr_space = 3 : i32} : i32 + llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %scratch = llvm.mlir.addressof @scratch : !llvm.ptr<3> + %tid = rocdl.workitem.id.x range : i32 + %16 = llvm.mlir.constant(16 : i32) : i32 + %neg5 = llvm.mlir.constant(-5 : i32) : i32 + %4 = llvm.mlir.constant(4 : i32) : i32 + %div = llvm.sdiv %tid, %16 : i32 + %rem = llvm.srem %tid, %16 : i32 + %negdiv = llvm.sdiv %neg5, %4 : i32 + %negrem = llvm.srem %neg5, %4 : i32 + %sum0 = llvm.add %div, %rem : i32 + %sum1 = llvm.add %sum0, %negdiv : i32 + %sum2 = llvm.add %sum1, %negrem : i32 + llvm.store %sum2, %scratch : i32, !llvm.ptr<3> + llvm.return + } +} diff --git a/waveasm/test/Transforms/translate-from-llvm-error-addressof-non-lds.mlir b/waveasm/test/Transforms/translate-from-llvm-error-addressof-non-lds.mlir new file mode 100644 index 0000000000..16aa7cc6a8 --- /dev/null +++ b/waveasm/test/Transforms/translate-from-llvm-error-addressof-non-lds.mlir @@ -0,0 +1,12 @@ +// RUN: not waveasm-translate %s --waveasm-translate-from-llvm 2>&1 | FileCheck %s +// Verify that llvm.mlir.addressof is currently limited to LDS globals. + +// CHECK: llvm.mlir.addressof currently supports only LDS globals in addrspace(3) + +gpu.module @gpu_module { + llvm.mlir.global private @global() : i32 + llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %0 = llvm.mlir.addressof @global : !llvm.ptr + llvm.return + } +} diff --git a/waveasm/test/Transforms/translate-from-llvm-lds.mlir b/waveasm/test/Transforms/translate-from-llvm-lds.mlir index d9647301bd..dcbd579008 100644 --- a/waveasm/test/Transforms/translate-from-llvm-lds.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-lds.mlir @@ -1,12 +1,16 @@ // RUN: waveasm-translate %s --waveasm-translate-from-llvm | FileCheck %s // Verify LDS addressof, GEP, load, and store translation. -// addressof maps to VGPR zero; LDS GEPs produce byte offsets; loads/stores -// dispatch to ds_read/ds_write by width; and LDS size is counted in bytes. +// addressof assigns per-global LDS byte offsets; LDS GEPs produce byte +// offsets; loads/stores dispatch to ds_read/ds_write by width; and LDS size is +// counted in bytes. // CHECK: waveasm.program @test__waveasm -// CHECK: lds_size = 1024 : i64 -// addressof -> v_mov_b32 of zero. +// CHECK: lds_size = 1040 : i64 +// First addressof uses byte offset 0. +// CHECK: waveasm.constant 0 // CHECK: waveasm.v_mov_b32 +// Second addressof uses the next available LDS byte offset. +// CHECK: waveasm.constant 1024 // GEP [0,0] on the array type passes through (all-zero indices). // GEP with constant offset 512 produces an arith.add. // CHECK: waveasm.arith.add @@ -18,8 +22,10 @@ gpu.module @gpu_module { llvm.mlir.global private @alloca() {addr_space = 3 : i32} : !llvm.array<256 x i32> + llvm.mlir.global private @scratch() {addr_space = 3 : i32} : !llvm.array<4 x i32> llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { %0 = llvm.mlir.addressof @alloca : !llvm.ptr<3> + %scratch = llvm.mlir.addressof @scratch : !llvm.ptr<3> %1 = llvm.mlir.constant(42 : i32) : i32 // Multi-index GEP with all-zero indices -> passthrough. %2 = llvm.getelementptr %0[0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.array<256 x i32> @@ -33,6 +39,8 @@ gpu.module @gpu_module { %5 = llvm.load %4 : !llvm.ptr<3> -> i32 // LDS store (4 bytes). llvm.store %1, %3 : i32, !llvm.ptr<3> + // Distinct LDS global uses a distinct non-zero base offset. + llvm.store %1, %scratch : i32, !llvm.ptr<3> llvm.return } } diff --git a/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir b/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir index 8fbf93a8ff..b1bcd96f99 100644 --- a/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir +++ b/waveasm/test/Transforms/translate-from-llvm-sdiv-srem.mlir @@ -1,4 +1,4 @@ -// RUN: waveasm-translate %s --waveasm-translate-from-llvm | FileCheck %s +// RUN: waveasm-translate %s --waveasm-llvm-sdiv-srem-legalization --waveasm-translate-from-llvm | FileCheck %s // Verify sdiv and srem with positive power-of-2 constants. // Negative dividends require bias/correction so we should see compare/select // scaffolding in addition to the final shift/mask-shaped arithmetic. @@ -15,7 +15,9 @@ // CHECK: waveasm.s_endpgm gpu.module @gpu_module { + llvm.mlir.global private @scratch() {addr_space = 3 : i32} : i32 llvm.func @test() attributes {gpu.kernel, gpu.known_block_size = array, rocdl.kernel, rocdl.reqd_work_group_size = array} { + %scratch = llvm.mlir.addressof @scratch : !llvm.ptr<3> %tid = rocdl.workitem.id.x range : i32 %16 = llvm.mlir.constant(16 : i32) : i32 %neg5 = llvm.mlir.constant(-5 : i32) : i32 @@ -24,6 +26,10 @@ gpu.module @gpu_module { %rem = llvm.srem %tid, %16 : i32 %negdiv = llvm.sdiv %neg5, %4 : i32 %negrem = llvm.srem %neg5, %4 : i32 + %sum0 = llvm.add %div, %rem : i32 + %sum1 = llvm.add %sum0, %negdiv : i32 + %sum2 = llvm.add %sum1, %negrem : i32 + llvm.store %sum2, %scratch : i32, !llvm.ptr<3> llvm.return } } From 4d11a22f551862fdccc7717337e2bc885c6e85c9 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Fri, 10 Apr 2026 20:10:54 +0200 Subject: [PATCH 15/19] cleanups Signed-off-by: Ivan Butygin --- waveasm/include/waveasm/Transforms/Passes.td | 4 --- .../Transforms/LLVMSDivSRemLegalization.cpp | 34 +++++++++---------- .../Transforms/TranslateFromLLVMDialect.cpp | 33 ++++++++---------- 3 files changed, 31 insertions(+), 40 deletions(-) diff --git a/waveasm/include/waveasm/Transforms/Passes.td b/waveasm/include/waveasm/Transforms/Passes.td index d3bc948418..d4dd32df97 100644 --- a/waveasm/include/waveasm/Transforms/Passes.td +++ b/waveasm/include/waveasm/Transforms/Passes.td @@ -385,8 +385,6 @@ def WAVEASMLLVMSDivSRemLegalization translation changes the abstraction level, and makes the rewrite independently testable. }]; - - let dependentDialects = ["::mlir::LLVM::LLVMDialect"]; } def WAVEASMTranslateFromLLVM : Pass<"waveasm-translate-from-llvm"> { @@ -400,8 +398,6 @@ def WAVEASMTranslateFromLLVM : Pass<"waveasm-translate-from-llvm"> { Option<"targetArch", "target", "std::string", "\"gfx950\"", "Target GPU architecture"> ]; - - let dependentDialects = ["::waveasm::WaveASMDialect"]; } #endif // WaveASM_TRANSFORMS_PASSES diff --git a/waveasm/lib/Transforms/LLVMSDivSRemLegalization.cpp b/waveasm/lib/Transforms/LLVMSDivSRemLegalization.cpp index 82e06ac5e1..a669e277de 100644 --- a/waveasm/lib/Transforms/LLVMSDivSRemLegalization.cpp +++ b/waveasm/lib/Transforms/LLVMSDivSRemLegalization.cpp @@ -58,7 +58,7 @@ static Value createI32Constant(PatternRewriter &rewriter, Location loc, } struct LegalizePowerOfTwoSDivPattern : OpRewritePattern { - using OpRewritePattern::OpRewritePattern; + using Base::Base; LogicalResult matchAndRewrite(LLVM::SDivOp op, PatternRewriter &rewriter) const override { @@ -76,13 +76,12 @@ struct LegalizePowerOfTwoSDivPattern : OpRewritePattern { Value zero = createI32Constant(rewriter, loc, 0); Value biasImm = createI32Constant(rewriter, loc, divisor - 1); Value shiftConst = createI32Constant(rewriter, loc, shiftAmt); - Value isNegative = - LLVM::ICmpOp::create(rewriter, loc, LLVM::ICmpPredicate::slt, - op.getLhs(), zero); + Value isNegative = LLVM::ICmpOp::create( + rewriter, loc, LLVM::ICmpPredicate::slt, op.getLhs(), zero); Value bias = LLVM::SelectOp::create(rewriter, loc, isNegative, biasImm, zero, LLVM::FastmathFlags::none); - Value biased = LLVM::AddOp::create( - rewriter, loc, op.getLhs(), bias, LLVM::IntegerOverflowFlags::none); + Value biased = LLVM::AddOp::create(rewriter, loc, op.getLhs(), bias, + LLVM::IntegerOverflowFlags::none); Value result = LLVM::AShrOp::create(rewriter, loc, biased, shiftConst, /*isExact=*/false); rewriter.replaceOp(op, result); @@ -91,7 +90,7 @@ struct LegalizePowerOfTwoSDivPattern : OpRewritePattern { }; struct LegalizePowerOfTwoSRemPattern : OpRewritePattern { - using OpRewritePattern::OpRewritePattern; + using Base::Base; LogicalResult matchAndRewrite(LLVM::SRemOp op, PatternRewriter &rewriter) const override { @@ -109,16 +108,17 @@ struct LegalizePowerOfTwoSRemPattern : OpRewritePattern { Value maskConst = createI32Constant(rewriter, loc, divisor - 1); Value negDivisor = createI32Constant(rewriter, loc, -divisor); Value rawRem = LLVM::AndOp::create(rewriter, loc, op.getLhs(), maskConst); - Value isNegative = - LLVM::ICmpOp::create(rewriter, loc, LLVM::ICmpPredicate::slt, - op.getLhs(), zero); - Value isNonZero = LLVM::ICmpOp::create(rewriter, loc, LLVM::ICmpPredicate::ne, - rawRem, zero); - Value needsAdjust = LLVM::AndOp::create(rewriter, loc, isNegative, isNonZero); - Value adjust = LLVM::SelectOp::create(rewriter, loc, needsAdjust, negDivisor, - zero, LLVM::FastmathFlags::none); - Value result = LLVM::AddOp::create( - rewriter, loc, rawRem, adjust, LLVM::IntegerOverflowFlags::none); + Value isNegative = LLVM::ICmpOp::create( + rewriter, loc, LLVM::ICmpPredicate::slt, op.getLhs(), zero); + Value isNonZero = LLVM::ICmpOp::create( + rewriter, loc, LLVM::ICmpPredicate::ne, rawRem, zero); + Value needsAdjust = + LLVM::AndOp::create(rewriter, loc, isNegative, isNonZero); + Value adjust = + LLVM::SelectOp::create(rewriter, loc, needsAdjust, negDivisor, zero, + LLVM::FastmathFlags::none); + Value result = LLVM::AddOp::create(rewriter, loc, rawRem, adjust, + LLVM::IntegerOverflowFlags::none); rewriter.replaceOp(op, result); return success(); } diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 95336f930b..1981c0684a 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -243,8 +243,8 @@ static unsigned getLLVMAddrSpace(Value v) { return 0; } -/// WaveASM register sizes are tracked in 32-bit lanes (dwords). -static int64_t getWaveASMLaneCount(int64_t bitWidth) { +/// WaveASM register sizes are tracked in 32-bit dwords. +static int64_t getWaveASMDwordCount(int64_t bitWidth) { return llvm::divideCeil(bitWidth, int64_t{32}); } @@ -253,21 +253,18 @@ static int64_t getWaveASMLaneCount(int64_t bitWidth) { /// /// We intentionally support only scalars, fixed vectors, and arrays thereof. /// Nested aggregates such as structs require real DataLayout-based layout -/// reasoning; approximating them here would make GEP and LDS offsets -/// target-dependent in subtle ways. -static FailureOr getLLVMTypeBytesWithoutDataLayout(Type ty) { +/// reasoning. +static FailureOr getLLVMTypeBytes(Type ty) { if (ty.isIntOrFloat()) return ty.getIntOrFloatBitWidth() / 8; if (VectorType vecTy = dyn_cast(ty)) { - FailureOr elemBytes = - getLLVMTypeBytesWithoutDataLayout(vecTy.getElementType()); + FailureOr elemBytes = getLLVMTypeBytes(vecTy.getElementType()); if (failed(elemBytes)) return failure(); return *elemBytes * vecTy.getNumElements(); } if (LLVM::LLVMArrayType arrTy = dyn_cast(ty)) { - FailureOr elemBytes = - getLLVMTypeBytesWithoutDataLayout(arrTy.getElementType()); + FailureOr elemBytes = getLLVMTypeBytes(arrTy.getElementType()); if (failed(elemBytes)) return failure(); return *elemBytes * arrTy.getNumElements(); @@ -279,9 +276,8 @@ static FailureOr getLLVMTypeBytesWithoutDataLayout(Type ty) { static bool isZeroGEPIndex(llvm::PointerUnion idx) { if (Value value = dyn_cast(idx)) return isConstantIntValue(value, 0); - IntegerAttr intAttr = dyn_cast(idx); - assert(intAttr && "GEP index must be a Value or IntegerAttr"); - return isConstantIntValue(intAttr, 0); + + return isConstantIntValue(cast(idx), 0); } /// Structural GEPs like [0, 0] are a no-op and can be forwarded even though we @@ -304,8 +300,7 @@ static FailureOr computeGEPByteOffset(LLVM::GEPOp op, if (indices.size() != 1) return op->emitOpError("GEP byte offset requires a single index"); - FailureOr elemBytes = - getLLVMTypeBytesWithoutDataLayout(op.getElemType()); + FailureOr elemBytes = getLLVMTypeBytes(op.getElemType()); if (failed(elemBytes)) return op->emitOpError( "unsupported GEP element type for byte offset computation " @@ -485,8 +480,9 @@ static LogicalResult handleWorkgroupId(OpTy op, LLVMTranslationState &st, } /// Translate an LLVM cast op to a WaveASM arithmetic pseudo-op. -/// Width is derived from the LLVM result type so that sext i32->i64 -/// produces a 2-wide register and trunc i64->i32 produces a 1-wide one. +/// Result width is derived from the LLVM type in 32-bit dwords, so sext +/// i32->i64 produces a 2-wide register and trunc i64->i32 produces a 1-wide +/// one. template static LogicalResult handleCastOp(LLVMOp op, LLVMTranslationState &st) { auto &ctx = st.ctx; @@ -496,7 +492,7 @@ static LogicalResult handleCastOp(LLVMOp op, LLVMTranslationState &st) { return op->emitOpError("unmapped operand in cast"); int64_t width = 1; if (IntegerType intTy = dyn_cast(op.getResult().getType())) - width = getWaveASMLaneCount(intTy.getWidth()); + width = getWaveASMDwordCount(intTy.getWidth()); Type resTy = isVGPRType(src->getType()) ? (Type)ctx.createVRegType(width, width) : (Type)ctx.createSRegType(width, width); @@ -915,8 +911,7 @@ static LogicalResult handleAddressOf(LLVM::AddressOfOp op, "llvm.mlir.addressof currently supports only LDS globals in " "addrspace(3)"); - FailureOr sizeBytes = - getLLVMTypeBytesWithoutDataLayout(global.getType()); + FailureOr sizeBytes = getLLVMTypeBytes(global.getType()); if (failed(sizeBytes)) return op->emitOpError( "unsupported LDS global type for byte size computation " From 6657aa1b3c4ced16d5e56fa1045f4033dfb828be Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Fri, 10 Apr 2026 20:17:37 +0200 Subject: [PATCH 16/19] cleanup Signed-off-by: Ivan Butygin --- waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 1981c0684a..7cb427e296 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -110,10 +110,7 @@ class LLVMTranslationState { } LLVM::GlobalOp lookupGlobal(StringRef name) const { - auto it = globalsByName.find(name); - if (it != globalsByName.end()) - return it->second; - return {}; + return globalsByName.lookup(name); } private: From 7bcc50f52950722f0c1a95397999524ef60f095f Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Fri, 10 Apr 2026 20:19:34 +0200 Subject: [PATCH 17/19] llvm::seq Signed-off-by: Ivan Butygin --- waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index 7cb427e296..fb687c165d 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -1177,7 +1177,7 @@ static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { ctx.getMapper().mapValue(op.getInductionVar(), bodyBlock.getArgument(0)); // Map iter_args (block args 1..N). - for (unsigned i = 0; i < op.getInitArgs().size(); ++i) + for (auto i : llvm::seq(op.getInitArgs().size())) ctx.getMapper().mapValue(op.getRegionIterArgs()[i], bodyBlock.getArgument(i + 1)); @@ -1213,7 +1213,7 @@ static LogicalResult handleSCFFor(scf::ForOp op, LLVMTranslationState &st) { // Map loop results. scf.for results are iter_args only (no IV), // but waveasm.loop results include the IV at index 0. - for (unsigned i = 0; i < op.getNumResults(); ++i) + for (auto i : llvm::seq(op.getNumResults())) ctx.getMapper().mapValue(op.getResult(i), loopOp.getResult(i + 1)); return success(); From 89a110b5f35738f31cd5d8c27c1d381976bc23d9 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Fri, 10 Apr 2026 20:21:50 +0200 Subject: [PATCH 18/19] cleanup Signed-off-by: Ivan Butygin --- waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp index fb687c165d..3655337cc9 100644 --- a/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp +++ b/waveasm/lib/Transforms/TranslateFromLLVMDialect.cpp @@ -968,9 +968,10 @@ static LogicalResult handleLDSLoad(LLVM::LoadOp op, Value addr, } else if (numBytes == 16) { VRegType wideTy = ctx.createVRegType(4, 4); loadOp = DS_READ_B128::create(builder, loc, TypeRange{wideTy}, *offset); - } else + } else { return op->emitOpError("unsupported LDS load size: ") << numBytes << " bytes"; + } ctx.getMapper().mapValue(op.getResult(), loadOp->getResult(0)); return success(); @@ -1085,7 +1086,6 @@ static LogicalResult handleMFMA_F32_16x16x16_F16(ROCDL::mfma_f32_16x16x16f16 op, // SCF for/yield handler //===----------------------------------------------------------------------===// -/// Forward declaration for recursive op translation. static LogicalResult translateOp(Operation *op, LLVMTranslationState &st); static FailureOr materializeI32SReg(Value v, StringRef name, From aec69034d60f505c64842581b6c32942e35d5f9a Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Sat, 11 Apr 2026 00:14:25 +0200 Subject: [PATCH 19/19] Fix SALU bitwise immediate operand ordering Normalize commutative bitwise legalization so immediate-first operands still satisfy the concrete SALU operand constraints, and keep the case covered by a focused regression test. Made-with: Cursor Signed-off-by: Ivan Butygin --- waveasm/lib/Transforms/ArithLegalization.cpp | 4 ++++ .../test/Transforms/arith-legalization-bitwise.mlir | 10 ++++++++++ 2 files changed, 14 insertions(+) diff --git a/waveasm/lib/Transforms/ArithLegalization.cpp b/waveasm/lib/Transforms/ArithLegalization.cpp index 2ca093f076..75ef74be90 100644 --- a/waveasm/lib/Transforms/ArithLegalization.cpp +++ b/waveasm/lib/Transforms/ArithLegalization.cpp @@ -549,6 +549,8 @@ static LogicalResult legalizeBitwiseOp(ArithOp op, OpBuilder &builder) { auto vregTy = VRegType::get(builder.getContext(), 2); result = VALUOp64::create(builder, loc, vregTy, lhs, rhs); } else { + if (isa(lhs.getType()) && isSGPRType(rhs.getType())) + std::swap(lhs, rhs); auto sregTy = SRegType::get(builder.getContext(), 2, 2); auto sccTy = SCCType::get(builder.getContext()); result = SALUOp64::create(builder, loc, sregTy, sccTy, lhs, rhs).getDst(); @@ -559,6 +561,8 @@ static LogicalResult legalizeBitwiseOp(ArithOp op, OpBuilder &builder) { auto vregTy = VRegType::get(builder.getContext()); result = VALUOp32::create(builder, loc, vregTy, lhs, rhs); } else { + if (isa(lhs.getType()) && isSGPRType(rhs.getType())) + std::swap(lhs, rhs); auto sregTy = SRegType::get(builder.getContext(), 1, 1); auto sccTy = SCCType::get(builder.getContext()); result = SALUOp32::create(builder, loc, sregTy, sccTy, lhs, rhs).getDst(); diff --git a/waveasm/test/Transforms/arith-legalization-bitwise.mlir b/waveasm/test/Transforms/arith-legalization-bitwise.mlir index da3f9fda7b..dc7ffd1f04 100644 --- a/waveasm/test/Transforms/arith-legalization-bitwise.mlir +++ b/waveasm/test/Transforms/arith-legalization-bitwise.mlir @@ -14,6 +14,7 @@ waveasm.program @test_or_i32 %v0 = waveasm.precolored.vreg 0 : !waveasm.vreg %s0 = waveasm.precolored.sreg 0 : !waveasm.sreg %s1 = waveasm.precolored.sreg 1 : !waveasm.sreg + %c1 = waveasm.constant 1 : !waveasm.imm<1> %c42 = waveasm.constant 42 : !waveasm.imm<42> // VGPR | VGPR -> v_or_b32. @@ -33,6 +34,10 @@ waveasm.program @test_or_i32 // CHECK: waveasm.s_or_b32 %{{.*}}, %{{.*}} : !waveasm.sreg, !waveasm.imm<42> %or_si = waveasm.arith.or %s0, %c42 : (!waveasm.sreg, !waveasm.imm<42>) -> !waveasm.sreg + // imm | SGPR -> s_or_b32 with the SGPR normalized to operand 0. + // CHECK: waveasm.s_or_b32 %{{.*}}, %{{.*}} : !waveasm.sreg, !waveasm.imm<1> + %or_is = waveasm.arith.or %c1, %s0 : (!waveasm.imm<1>, !waveasm.sreg) -> !waveasm.sreg + // CHECK-NOT: waveasm.arith. waveasm.s_endpgm } @@ -50,6 +55,7 @@ waveasm.program @test_and_i32 %v0 = waveasm.precolored.vreg 0 : !waveasm.vreg %s0 = waveasm.precolored.sreg 0 : !waveasm.sreg %s1 = waveasm.precolored.sreg 1 : !waveasm.sreg + %c1 = waveasm.constant 1 : !waveasm.imm<1> // VGPR & VGPR -> v_and_b32. // CHECK: waveasm.v_and_b32 %{{.*}}, %{{.*}} : !waveasm.vreg, !waveasm.vreg @@ -64,6 +70,10 @@ waveasm.program @test_and_i32 // CHECK: waveasm.v_and_b32 %and_sv = waveasm.arith.and %s0, %v0 : (!waveasm.sreg, !waveasm.vreg) -> !waveasm.vreg + // imm & SGPR -> s_and_b32 with the SGPR normalized to operand 0. + // CHECK: waveasm.s_and_b32 %{{.*}}, %{{.*}} : !waveasm.sreg, !waveasm.imm<1> + %and_is = waveasm.arith.and %c1, %s0 : (!waveasm.imm<1>, !waveasm.sreg) -> !waveasm.sreg + // CHECK-NOT: waveasm.arith. waveasm.s_endpgm }