diff --git a/include/flydsl/Dialect/Fly/Utils/TypeUtils.h b/include/flydsl/Dialect/Fly/Utils/TypeUtils.h new file mode 100644 index 000000000..018d63e20 --- /dev/null +++ b/include/flydsl/Dialect/Fly/Utils/TypeUtils.h @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#ifndef FLYDSL_DIALECT_FLY_UTILS_TYPEUTILS_H +#define FLYDSL_DIALECT_FLY_UTILS_TYPEUTILS_H + +#include "mlir/IR/Types.h" + +namespace mlir::fly { + +/// Return the SSA value type that loading the storage element type \p elemTy produces. +/// +/// FlyDSL keeps the *logical* signedness of an integer dtype in the element type of a +/// storage descriptor (`!fly.memref`, `!fly.ptr`), because a signless +/// `i8` cannot say whether the bytes behind it are `uint8` or `int8`. SSA values, on the +/// other hand, are always signless: the `arith` dialect only accepts signless integers and +/// the signedness of an operation is carried by the opcode (`divui` vs `divsi`). This maps +/// a scalar storage element type to the value type that loading it produces; the Python +/// side additionally unwraps vectors. +/// +/// Signedness is spelled on storage only where the signless type would be ambiguous, so +/// the DSL emits `uiN` and plain `iN`, never `siN`. Accordingly only `uiN` is mapped here; +/// everything else, `siN` included, is returned unchanged so that an unexpected spelling +/// fails verification instead of being silently accepted. The Python +/// `_ssa_value_ir_type` in `expr/numeric.py` is the same function on the DSL side. +Type toSSAValueType(Type elemTy); + +} // namespace mlir::fly + +#endif // FLYDSL_DIALECT_FLY_UTILS_TYPEUTILS_H diff --git a/lib/Bindings/Python/DLTensorAdaptor.h b/lib/Bindings/Python/DLTensorAdaptor.h index e77ef93eb..b3bf6a55e 100644 --- a/lib/Bindings/Python/DLTensorAdaptor.h +++ b/lib/Bindings/Python/DLTensorAdaptor.h @@ -154,7 +154,9 @@ class DLTensorAdaptor { case kDLInt: return IntegerType::get(ctx, dtype.bits); case kDLUInt: - return IntegerType::get(ctx, dtype.bits); + // Unsigned storage keeps its signedness in the element type so that the DSL can + // reconstruct the logical dtype; SSA values loaded from it are signless. + return IntegerType::get(ctx, dtype.bits, IntegerType::Unsigned); case kDLBfloat: return BFloat16Type::get(ctx); case kDLBool: diff --git a/lib/Dialect/Fly/CMakeLists.txt b/lib/Dialect/Fly/CMakeLists.txt index e33ebaaf7..bf798a9b2 100644 --- a/lib/Dialect/Fly/CMakeLists.txt +++ b/lib/Dialect/Fly/CMakeLists.txt @@ -9,6 +9,7 @@ add_mlir_dialect_library(MLIRFlyDialect Utils/IntTupleUtils.cpp Utils/NormalForm.cpp Utils/PointerUtils.cpp + Utils/TypeUtils.cpp Transforms/LayoutLowering.cpp Transforms/Canonicalize.cpp Transforms/RewriteFuncSignature.cpp diff --git a/lib/Dialect/Fly/IR/FlyOps.cpp b/lib/Dialect/Fly/IR/FlyOps.cpp index 41b02b1b9..68eda8b11 100644 --- a/lib/Dialect/Fly/IR/FlyOps.cpp +++ b/lib/Dialect/Fly/IR/FlyOps.cpp @@ -10,6 +10,7 @@ #include "flydsl/Dialect/Fly/Utils/IntTupleUtils.h" #include "flydsl/Dialect/Fly/Utils/LayoutUtils.h" #include "flydsl/Dialect/Fly/Utils/TiledOpUtils.h" +#include "flydsl/Dialect/Fly/Utils/TypeUtils.h" #include #include @@ -1913,7 +1914,8 @@ FLY_INFER_RETURN_TYPES(DecompositionOp) { FLY_INFER_RETURN_TYPES(MemRefLoadOp) { if (auto memrefTy = dyn_cast(operands[0].getType())) { - inferredReturnTypes.push_back(memrefTy.getElemTy()); + // Signedness lives in the storage element type; the loaded SSA value is signless. + inferredReturnTypes.push_back(toSSAValueType(memrefTy.getElemTy())); return success(); } if (auto coordTensorTy = dyn_cast(operands[0].getType())) { @@ -1989,7 +1991,7 @@ FLY_INFER_RETURN_TYPES(MemRefLoadVecOp) { location, "MemRefLoadVecOp: layout size must be static and leaf int, got ", size); inferredReturnTypes.push_back( - VectorType::get({size.getLeafAsInt().getValue()}, memrefTy.getElemTy())); + VectorType::get({size.getLeafAsInt().getValue()}, toSSAValueType(memrefTy.getElemTy()))); return success(); } diff --git a/lib/Dialect/Fly/Utils/PointerUtils.cpp b/lib/Dialect/Fly/Utils/PointerUtils.cpp index 215e584b5..9860f6ca8 100644 --- a/lib/Dialect/Fly/Utils/PointerUtils.cpp +++ b/lib/Dialect/Fly/Utils/PointerUtils.cpp @@ -6,6 +6,7 @@ #include "flydsl/Dialect/Fly/Utils/LayoutUtils.h" #include "flydsl/Dialect/Fly/Utils/PointerUtils.h" +#include "flydsl/Dialect/Fly/Utils/TypeUtils.h" namespace mlir::fly { @@ -34,7 +35,9 @@ Type projectToLLVMCompatibleElemTy(Type elemTy) { if (width < 16) return IntegerType::get(elemTy.getContext(), width); } - return elemTy; + // A storage element type may carry logical signedness (`ui8`); LLVM and `arith` only + // know signless integers, so drop it once we are past the DSL type system. + return toSSAValueType(elemTy); } Type RegMem2SSAType(fly::MemRefType memRefTy, bool llvmCompatibleType) { diff --git a/lib/Dialect/Fly/Utils/TypeUtils.cpp b/lib/Dialect/Fly/Utils/TypeUtils.cpp new file mode 100644 index 000000000..6b3a55f70 --- /dev/null +++ b/lib/Dialect/Fly/Utils/TypeUtils.cpp @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 FlyDSL Project Contributors + +#include "flydsl/Dialect/Fly/Utils/TypeUtils.h" + +#include "mlir/IR/BuiltinTypes.h" + +namespace mlir::fly { + +Type toSSAValueType(Type elemTy) { + auto intTy = dyn_cast(elemTy); + if (!intTy || !intTy.isUnsigned()) + return elemTy; + return IntegerType::get(elemTy.getContext(), intTy.getWidth()); +} + +} // namespace mlir::fly diff --git a/python/flydsl/compiler/jit_argument.py b/python/flydsl/compiler/jit_argument.py index 5d5c1d0bd..0a14f6335 100644 --- a/python/flydsl/compiler/jit_argument.py +++ b/python/flydsl/compiler/jit_argument.py @@ -516,7 +516,9 @@ def ptr_fill(a, s, _open=_open, _shared=shared): torch.float32: T.f32, torch.float64: T.f64, torch.bool: lambda: ir.IntegerType.get_signless(1), - torch.uint8: lambda: ir.IntegerType.get_signless(8), + # Unsigned tensors keep their signedness in the memref element type so the kernel sees + # ``Uint*`` rather than ``Int*``; signed integers stay signless (their canonical form). + torch.uint8: lambda: ir.IntegerType.get_unsigned(8), torch.int8: lambda: ir.IntegerType.get_signless(8), torch.int16: lambda: ir.IntegerType.get_signless(16), torch.int32: lambda: ir.IntegerType.get_signless(32), @@ -617,7 +619,7 @@ def _trivial_alignment_bytes(element_type) -> int: def __get_ir_types__(self): ir_type = self.element_type if isinstance(ir_type, type) and issubclass(ir_type, Numeric): - ir_type = self.element_type.ir_type + ir_type = self.element_type.storage_ir_type return [PointerType.get(ir_type, self.address_space, self.alignment)] def __cache_signature__(self): diff --git a/python/flydsl/expr/derived.py b/python/flydsl/expr/derived.py index 9e81b3891..d42564fde 100644 --- a/python/flydsl/expr/derived.py +++ b/python/flydsl/expr/derived.py @@ -103,7 +103,7 @@ def make_rmem_tensor(shape_or_layout, dtype): else: layout = shape_or_layout - tensorTy = fly.MemRefType.get(dtype.ir_type, layout.type, fly.AddressSpace.Register) + tensorTy = fly.MemRefType.get(dtype.storage_ir_type, layout.type, fly.AddressSpace.Register) return memref_alloca(tensorTy, layout=layout) diff --git a/python/flydsl/expr/numeric.py b/python/flydsl/expr/numeric.py index c26af754d..393d2dfbc 100644 --- a/python/flydsl/expr/numeric.py +++ b/python/flydsl/expr/numeric.py @@ -24,6 +24,28 @@ ) +def _ssa_value_ir_type(ir_type): + """Value type produced by loading the unsigned storage element type *ir_type*. + + An unsigned integer storage type spells its logical signedness (``ui8``), but the SSA + value read out of it is signless, since ``arith`` accepts only signless integers. This + is the inverse of :attr:`NumericMeta.storage_ir_type`. Vectors are mapped element-wise, + preserving scalable dimensions; everything else is returned unchanged. + + ``siN`` is deliberately left alone: no DSL dtype spells storage that way, so mapping it + here would silently accept an input the type system never produces. The C++ + ``mlir::fly::toSSAValueType`` applies the same rule on the dialect side. + """ + if isinstance(ir_type, ir.VectorType): + elem = _ssa_value_ir_type(ir_type.element_type) + if elem == ir_type.element_type: + return ir_type + return ir.VectorType.get(list(ir_type.shape), elem, scalable=list(ir_type.scalable_dims)) + if isinstance(ir_type, ir.IntegerType) and ir_type.is_unsigned: + return ir.IntegerType.get_signless(ir_type.width) + return ir_type + + def _infer_np_dtype(width, signed, name): if signed is not None: if width == 1: @@ -286,6 +308,28 @@ def ir_type(cls): return cls._ir_type() return None + @property + def storage_ir_type(cls): + """MLIR type to use when this dtype is a *storage* element type. + + A signless ``i8`` cannot say whether the bytes behind it are ``uint8`` or + ``int8``, so unsigned integers carry their signedness in the element type of + the storage descriptor (``!fly.memref``, ``!fly.ptr``). + That is what lets :attr:`Tensor.element_type` reconstruct ``Uint8`` rather + than ``Int8`` after a round trip through IR. + + SSA values stay signless (:attr:`ir_type`): ``arith`` only accepts signless + integers and encodes signedness in the opcode (``divui`` vs ``divsi``). + Signed integers keep the signless spelling as their canonical storage form, + so existing IR is unchanged. Dtypes whose ``ir_type`` is not a builtin integer + keep it as well -- notably :class:`Index`, which is flagged unsigned but is + ``index``, a type that has no signed and unsigned spellings. + """ + ir_type = cls.ir_type + if cls.signed is False and isinstance(ir_type, ir.IntegerType): + return ir.IntegerType.get_unsigned(cls.width) + return ir_type + @property def is_integer(cls) -> bool: return cls.signed is not None diff --git a/python/flydsl/expr/primitive.py b/python/flydsl/expr/primitive.py index ab8af056f..7197e2f81 100644 --- a/python/flydsl/expr/primitive.py +++ b/python/flydsl/expr/primitive.py @@ -539,7 +539,11 @@ def make_fragment_layout_like(tensor): @dsl_loc_tracing def make_fragment_like(tensor, dtype=None): - if hasattr(dtype, "ir_type"): + # A Numeric dtype becomes the fragment's *storage* element type, which keeps unsigned + # signedness; other dtype-like objects keep the previous ``ir_type`` contract. + if hasattr(dtype, "storage_ir_type"): + dtype = dtype.storage_ir_type + elif hasattr(dtype, "ir_type"): dtype = dtype.ir_type return fly.make_fragment_like(tensor, dtype=dtype) @@ -1203,18 +1207,39 @@ def apply_swizzle(ptr, swizzle): @dsl_loc_tracing -@dsl_wrap_result def ptr_load(ptr, result_type=None): """Load one value (scalar or vector) from *ptr*; dtype defaults to ptr's element type. Examples: v = ptr_load(ptr) """ + from .numeric import Numeric, _ssa_value_ir_type + from .typing import Vector, as_dsl_value + if result_type is None: result_type = ptr.element_type + dtype = None if not isinstance(result_type, ir.Type): + # The loaded SSA value is signless, so wrap it as the dtype we asked for instead of + # re-deriving one from the result type -- otherwise ``Uint8`` comes back as ``Int8``. + # Other dtype-like objects only promise ``ir_type``, so they keep dispatching on it. + if isinstance(result_type, type) and issubclass(result_type, Numeric): + dtype = result_type result_type = result_type.ir_type - return fly.ptr_load(result_type, ptr) + else: + # A raw MLIR type may still spell its signedness (``ui8``, ``vector<4xui8>``), because + # that is how a storage element type is written. Take the dtype it names, but load it + # as the signless SSA form -- ``arith`` accepts nothing else. + signless = _ssa_value_ir_type(result_type) + if signless != result_type: + elem = result_type.element_type if isinstance(result_type, ir.VectorType) else result_type + dtype = Numeric.from_ir_type(elem) + result_type = signless + + value = fly.ptr_load(result_type, ptr) + if dtype is not None and isinstance(result_type, ir.VectorType): + return Vector(value, dtype=dtype) + return as_dsl_value(value, dtype) @dsl_loc_tracing @@ -1244,7 +1269,7 @@ def recast_iter(result_type, src): if isinstance(result_type, type): if issubclass(result_type, Numeric): - result_type = result_type.ir_type + result_type = result_type.storage_ir_type else: raise TypeError( f"result_type must be a Numeric subclass or a fly Pointer, got unsupported class {result_type!r}" @@ -1273,16 +1298,21 @@ def memref_store_vec(vector, memref): @dsl_loc_tracing -@dsl_wrap_result def memref_load(memref, indices): + from .typing import as_dsl_value + + # As in :func:`ptr_load`: the loaded value is signless, so the memref's own element type + # is what says whether it is unsigned. A CoordTensor has none, so it keeps dispatching + # on the result type. + dtype = None if isinstance(memref.type, CoordTensorType) else memref.dtype if isinstance(indices, ir.Value): if not _is_int_tuple_value(indices): indices = make_int_tuple(indices) - return fly.memref_load(memref, indices) + return as_dsl_value(fly.memref_load(memref, indices), dtype) indices = make_int_tuple(indices) _check_profile(is_profile_weakly_congruent, indices, memref) - return fly.memref_load(memref, indices) + return as_dsl_value(fly.memref_load(memref, indices), dtype) @dsl_loc_tracing diff --git a/python/flydsl/expr/utils/arith.py b/python/flydsl/expr/utils/arith.py index 3a06d8c0c..24794a9db 100644 --- a/python/flydsl/expr/utils/arith.py +++ b/python/flydsl/expr/utils/arith.py @@ -400,6 +400,8 @@ def __init__(self, v, signed=None): super().__init__(v) elem_ty = element_type(self.type) self.is_float = not is_integer_like_type(elem_ty) + # ``i1`` has no sign bit to spare, so reading it as signed makes ``true`` compare + # as -1 and inverts every ordering. Width-1 integers are always unsigned {0, 1}. self.signed = signed and elem_ty.width > 1 def with_signedness(self, signed): diff --git a/tests/language/test_unsigned_semantics.py b/tests/language/test_unsigned_semantics.py new file mode 100644 index 000000000..6fb5066ca --- /dev/null +++ b/tests/language/test_unsigned_semantics.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""Logical signedness of integer dtypes (ROCm/FlyDSL#701). + +MLIR integers are signless, so the *storage* descriptor of an unsigned dtype spells its +element type ``uiN``. That is what lets a DSL value reconstructed from IR come back as +``Uint*`` instead of ``Int*``, which in turn selects the unsigned form of every operation +whose signed and unsigned forms differ. +""" + +import pytest +from lang_utils import source_ir + +import flydsl.expr as fx +from flydsl._mlir import ir +from flydsl.expr.numeric import Numeric, _ssa_value_ir_type + +pytestmark = pytest.mark.l1a_compile_no_target_dialect + + +UNSIGNED_SIGNED = [(fx.Uint8, fx.Int8), (fx.Uint16, fx.Int16), (fx.Uint32, fx.Int32), (fx.Uint64, fx.Int64)] + + +@pytest.mark.parametrize("udtype,sdtype", UNSIGNED_SIGNED) +def test_storage_type_carries_signedness(udtype, sdtype): + """Storage keeps the signedness that SSA values cannot express; values stay signless.""" + with ir.Context(): + assert str(udtype.storage_ir_type) == f"ui{udtype.width}" + assert str(udtype.ir_type) == f"i{udtype.width}" + # Signed integers keep the signless spelling as their canonical storage form. + assert sdtype.storage_ir_type == sdtype.ir_type + + +def test_non_integer_storage_type_is_unchanged(): + with ir.Context(): + assert fx.Float32.storage_ir_type == fx.Float32.ir_type + assert fx.Boolean.storage_ir_type == fx.Boolean.ir_type + # ``Index`` is flagged unsigned, but ``index`` has no signed and unsigned spellings. + assert fx.Index.storage_ir_type == fx.Index.ir_type + + +@pytest.mark.parametrize( + "storage, value", + [ + ("ui8", "i8"), + ("ui64", "i64"), + ("si16", "si16"), # no DSL dtype maps to a signed spelling; leave it alone + ("i32", "i32"), + ("vector<4xui8>", "vector<4xi8>"), + ("vector<4xi8>", "vector<4xi8>"), + ("vector<[4]xui8>", "vector<[4]xi8>"), # scalable dimensions survive the rewrite + ("index", "index"), + ("f32", "f32"), + ], +) +def test_ssa_value_ir_type_undoes_storage_signedness(storage, value): + """The inverse of ``storage_ir_type``: what loading a storage element type produces.""" + with ir.Context(), ir.Location.unknown(): + assert str(_ssa_value_ir_type(ir.Type.parse(storage))) == value + + +def test_register_tensor_round_trips_its_dtype(): + """A fragment declared unsigned reads back unsigned; the loaded value is signless.""" + seen = {} + + def body(): + for dtype in (fx.Uint8, fx.Int8, fx.Uint32, fx.Float32): + frag = fx.make_rmem_tensor(fx.make_layout(4, 1), dtype) + seen[dtype] = frag.element_type + + text = source_ir(body) + assert seen == {fx.Uint8: fx.Uint8, fx.Int8: fx.Int8, fx.Uint32: fx.Uint32, fx.Float32: fx.Float32} + assert "!fly.memref> dtype(2) + ptr = fx.get_iter(frag) + ptr[3] = ptr[2] >> dtype(2) + + text = source_ir(body) + other = "arith.shrsi" if shift_op == "arith.shrui" else "arith.shrui" + assert text.count(shift_op) == 2 and other not in text + + +def _widen(v, dtype): + return v.to(fx.Int64) + + +def _to_float(v, dtype): + return v.to(fx.Float32) + + +def _div(v, dtype): + return v // dtype(3) + + +def _rem(v, dtype): + return v % dtype(3) + + +def _shr(v, dtype): + return v >> dtype(2) + + +def _gt(v, dtype): + return v > dtype(3) + + +def _gt_select(v, dtype): + # Ordering, spelled out of the pieces the DSL already has. + return (v > dtype(3)).select(v, dtype(3)) + + +def _apply_ir(apply, dtype): + """Trace ``apply`` over a run-time value of *dtype*; return the emitted IR.""" + + def body(v: dtype): + apply(v, dtype) + + return source_ir(body, 1) + + +@pytest.mark.parametrize( + "apply,unsigned_op,signed_op", + [ + (_widen, "arith.extui", "arith.extsi"), + (_to_float, "arith.uitofp", "arith.sitofp"), + (_div, "arith.divui", "arith.floordivsi"), + (_rem, "arith.remui", "arith.remsi"), + (_shr, "arith.shrui", "arith.shrsi"), + (_gt, "ugt", "sgt"), + (_gt_select, "ugt", "sgt"), + ], +) +def test_operation_form_follows_dtype_signedness(apply, unsigned_op, signed_op): + unsigned_ir = _apply_ir(apply, fx.Uint8) + signed_ir = _apply_ir(apply, fx.Int8) + assert unsigned_op in unsigned_ir and signed_op not in unsigned_ir + assert signed_op in signed_ir and unsigned_op not in signed_ir + + +def test_mixed_signedness_promotes_before_choosing_the_predicate(): + """A ``Uint8`` against an ``Int32``: widen unsigned (``extui``), then compare signed.""" + + def body(v: fx.Uint8): + _ = v > fx.Int32(3) + + text = source_ir(body, 1) + assert "arith.extui" in text and "sgt" in text and "ugt" not in text + + +@pytest.mark.parametrize( + "dtype,other,predicate", + [ + (fx.Uint16, lambda: fx.Uint16(3), "ugt"), + (fx.Int64, lambda: fx.Int64(3), "sgt"), + # Neither of these has a sign bit to read, so both are the unsigned predicate: + # ``i1``'s single bit is its value, and ``index`` has no signed spelling. + (fx.Boolean, lambda: fx.Boolean(True), "ugt"), + (fx.Index, lambda: fx.Index(3), "ugt"), + ], +) +def test_comparison_covers_the_whole_integer_domain(dtype, other, predicate): + """Every integer dtype, and the comparison predicate its signedness selects.""" + + def body(v: dtype): + _ = v > other() + + text = source_ir(body, 1) + unwanted = "sgt" if predicate == "ugt" else "ugt" + assert predicate in text and unwanted not in text + + +def test_vector_ordering_follows_the_element_signedness(): + """The same rule element-wise: a vector compares with its element dtype's predicate.""" + seen = {} + + def body(): + u = fx.Vector.filled(4, 200, fx.Uint8) + s = fx.Vector.filled(4, -100, fx.Int8) + seen["u"] = (u > fx.Uint8(5)).select(u, fx.Uint8(5)) + seen["s"] = (s > fx.Int8(5)).select(s, fx.Int8(5)) + + text = source_ir(body) + assert isinstance(seen["u"], fx.Vector) and seen["u"].dtype is fx.Uint8 + assert isinstance(seen["s"], fx.Vector) and seen["s"].dtype is fx.Int8 + assert "ugt" in text and "sgt" in text + + +# The boundary itself: torch dtype -> MLIR element type -> DSL ``Numeric``. No kernel is +# traced and no device is needed, so this runs in CPU-only CI. +_ROUND_TRIP = [("Uint8", "uint8"), ("Int8", "int8"), ("Int16", "int16"), ("Int32", "int32"), ("Int64", "int64")] + + +@pytest.mark.parametrize("dsl_name,torch_name", _ROUND_TRIP) +def test_torch_dtype_round_trips_to_the_dsl_dtype(dsl_name, torch_name): + """The dtype a kernel reconstructs must never disagree with the tensor it was given.""" + torch = pytest.importorskip("torch") + from flydsl.compiler.jit_argument import torch_dtype_to_mlir_type + + with ir.Context(): + assert Numeric.from_ir_type(torch_dtype_to_mlir_type(getattr(torch, torch_name))) is getattr(fx, dsl_name) + + +_DLPACK_ELEMENT_TYPES = [ + ("uint8", "ui8"), + ("uint16", "ui16"), + ("uint32", "ui32"), + ("uint64", "ui64"), + ("int8", "i8"), + ("int32", "i32"), + ("float32", "f32"), +] + + +@pytest.mark.parametrize("torch_name,spelling", _DLPACK_ELEMENT_TYPES) +def test_dlpack_adaptor_maps_unsigned_codes_to_unsigned_types(torch_name, spelling): + """The other entry point: ``DLTensorAdaptor`` reading a DLPack capsule directly. + + ``torch_dtype_to_mlir_type`` is a Python table; this is the C++ ``kDLUInt`` arm, which + a device launch only exercises indirectly. A host tensor is enough -- the adaptor reads + the capsule's dtype code, not its memory. It reads the width from the capsule, so every + unsigned width the framework can export is covered, not just the byte. + """ + torch = pytest.importorskip("torch") + from flydsl._mlir._mlir_libs._mlirDialectsFly import DLTensorAdaptor + + torch_dtype = getattr(torch, torch_name, None) + if torch_dtype is None: + pytest.skip(f"torch has no {torch_name}") + tensor = torch.zeros(4, dtype=torch_dtype) + with ir.Context(): + assert str(DLTensorAdaptor(tensor.__dlpack__()).dtype) == spelling + + +def test_wider_unsigned_torch_dtypes_are_still_rejected(): + """Only ``uint8`` is accepted at the boundary; a wider unsigned tensor is refused. + + Widening that set is out of scope here -- the point is that it fails loudly rather + than silently arriving as its signed counterpart, which is the bug this fixes. + """ + torch = pytest.importorskip("torch") + from flydsl.compiler.jit_argument import torch_dtype_to_mlir_type + + for name in ("uint16", "uint32", "uint64"): + torch_dtype = getattr(torch, name, None) + if torch_dtype is None: + continue + with pytest.raises(TypeError, match="unsupported torch dtype"): + with ir.Context(): + torch_dtype_to_mlir_type(torch_dtype) diff --git a/tests/unit/test_unsigned_tensor_dtype.py b/tests/unit/test_unsigned_tensor_dtype.py new file mode 100644 index 000000000..75e35c930 --- /dev/null +++ b/tests/unit/test_unsigned_tensor_dtype.py @@ -0,0 +1,417 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""An unsigned tensor keeps its signedness across the kernel boundary (ROCm/FlyDSL#701). + +A ``torch.uint8`` tensor used to arrive as ``Int8``, so every byte with the high bit set was +treated as negative: widening sign-extended and int-to-float used ``sitofp``. These tests +run both signedness families through the same kernel on the device and check the values, +not just the label -- reading unsigned memory, writing unsigned memory, and reinterpreting +signed storage as unsigned with ``recast_iter``. +""" + +import pytest + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl._mlir import ir + +try: + import torch +except ImportError: + torch = None + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +if torch is None or not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available.", allow_module_level=True) + +# Four bytes per thread, so the copy also exercises the vectorized load path. +PER_THREAD = 4 +BYTES = [0x00, 0x01, 0x7F, 0x80, 0xC0, 0xFE, 0xFF, 0x87] +NUM_THREADS = len(BYTES) + +_ATOM_FOR_WIDTH = {8: fx.UniversalCopy8b, 16: fx.UniversalCopy16b, 32: fx.UniversalCopy32b, 64: fx.UniversalCopy64b} + + +def _straddling_values(bits): + """Values on both sides of the sign bit, so signed and unsigned answers differ.""" + high = (1 << bits) - 1 + sign_bit = 1 << (bits - 1) + return [0, 1, sign_bit - 1, sign_bit, (3 << (bits - 2)) & high, high - 1, high, sign_bit + 7] + + +def _widen_and_scale(src, dst, element_type): + """dst[i] = int(float(src[i])) * 2 -- wrong for src[i] >= 0x80 if src is read as signed.""" + tid = fx.thread_idx.x + layout = fx.make_layout(PER_THREAD, 1) + frag_in = fx.make_rmem_tensor(layout, element_type) + frag_out = fx.make_rmem_tensor(layout, fx.Int32) + view_in = fx.Tensor(fx.make_view(fx.add_offset(fx.get_iter(src), fx.make_int_tuple(tid * PER_THREAD)), layout)) + view_out = fx.Tensor(fx.make_view(fx.add_offset(fx.get_iter(dst), fx.make_int_tuple(tid * PER_THREAD)), layout)) + + fx.copy(fx.make_copy_atom(fx.UniversalCopy32b(), element_type), view_in, frag_in) + values = frag_in.load() + frag_out.store((values.to(fx.Float32) * fx.Float32(2.0)).to(fx.Int32)) + fx.copy(fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32), frag_out, view_out) + + +@flyc.kernel +def _widen_kernel(src: fx.Tensor, dst: fx.Tensor): + _widen_and_scale(src, dst, src.element_type) + + +@flyc.jit +def _launch(src: fx.Tensor, dst: fx.Tensor, stream: fx.Stream = fx.Stream(None)): + _widen_kernel(src, dst).launch(grid=(1, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + +def _run(torch_dtype): + # Both dtypes get the identical bytes, so only their signedness can explain a + # difference in the result. ``torch.tensor(..., dtype=torch.int8)`` refuses 0x80..0xFF, + # hence the reinterpreting view. + values = BYTES * PER_THREAD + src = torch.tensor(values, dtype=torch.uint8, device="cuda").view(torch_dtype) + dst = torch.zeros(len(values), dtype=torch.int32, device="cuda") + _launch(src, dst, stream=torch.cuda.current_stream()) + torch.cuda.synchronize() + return src.tolist(), dst.tolist() + + +@pytest.mark.parametrize("torch_dtype", [torch.uint8, torch.int8]) +def test_widen_and_convert_follow_tensor_signedness(torch_dtype): + """Each dtype must reproduce what torch itself computes for the same bytes.""" + values, got = _run(torch_dtype) + assert got == [v * 2 for v in values] + + +def _elementwise_launcher(apply, src_bits, dst_bits, recast=None, recast_out=None, make_frag=None): + """Build a launcher for ``dst[i] = apply(src[i])``, one element per thread. + + The fragments come from :func:`make_fragment_like`, so the destination fragment -- + and therefore the store -- inherits the destination tensor's signedness, unless + *make_frag* builds the source fragment with an explicit dtype. + + *recast*, *recast_out* and *make_frag* are functions rather than dtypes on purpose: a + captured class does not participate in the JIT cache key, so two launchers differing + only in a captured dtype would share one compiled artifact. + """ + src_atom, dst_atom = _ATOM_FOR_WIDTH[src_bits], _ATOM_FOR_WIDTH[dst_bits] + + @flyc.kernel + def kernel(src: fx.Tensor, dst: fx.Tensor): + tid = fx.thread_idx.x + layout = fx.make_layout(1, 1) + src_iter = fx.get_iter(src) + dst_iter = fx.get_iter(dst) + if fx.const_expr(recast is not None): + # Reinterpret the raw storage as another element type, the way the packed-code + # kernels do; the recast element type is what decides signedness downstream. + src_iter = recast(src_iter) + if fx.const_expr(recast_out is not None): + dst_iter = recast_out(dst_iter) + view_in = fx.Tensor(fx.make_view(fx.add_offset(src_iter, fx.make_int_tuple(tid)), layout)) + view_out = fx.Tensor(fx.make_view(fx.add_offset(dst_iter, fx.make_int_tuple(tid)), layout)) + if fx.const_expr(make_frag is not None): + # An explicit dtype makes the fragment's own storage decide signedness, which is + # how a kernel reads bytes it knows are unsigned out of signed storage. + frag_in = make_frag(view_in) + else: + frag_in = fx.make_fragment_like(view_in) + frag_out = fx.make_fragment_like(view_out) + fx.copy(fx.make_copy_atom(src_atom(), frag_in.element_type), view_in, frag_in) + frag_out.store(apply(frag_in.load(), frag_in.element_type, view_out.element_type)) + fx.copy(fx.make_copy_atom(dst_atom(), view_out.element_type), frag_out, view_out) + + @flyc.jit + def launch(src: fx.Tensor, dst: fx.Tensor, stream: fx.Stream = fx.Stream(None)): + kernel(src, dst).launch(grid=(1, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + return launch + + +def _run_elementwise(launch, src, dst_torch_dtype): + dst = torch.zeros(src.numel(), dtype=dst_torch_dtype, device="cuda") + launch(src, dst, stream=torch.cuda.current_stream()) + torch.cuda.synchronize() + return dst.tolist() + + +def _shift_right(values, src_type, dst_type): + return values >> dst_type(4) + + +def _widen_to_i64(values, src_type, dst_type): + return values.to(fx.Int64) + + +def _widen_to_i32(values, src_type, dst_type): + return values.to(fx.Int32) + + +def _recast_to_unsigned_bytes(iterator): + return fx.recast_iter(fx.Uint8, iterator) + + +def _recast_to_unsigned_16(iterator): + return fx.recast_iter(fx.Uint16, iterator) + + +def _recast_to_unsigned_32(iterator): + return fx.recast_iter(fx.Uint32, iterator) + + +def _recast_to_unsigned_64(iterator): + return fx.recast_iter(fx.Uint64, iterator) + + +def _recast_to_signed_bytes(iterator): + return fx.recast_iter(fx.Int8, iterator) + + +# Only ``torch.uint8`` is accepted at the tensor boundary, so the wider unsigned dtypes are +# reached the way real kernels reach them: signed storage reinterpreted with ``recast_iter``. +_RECAST_UNSIGNED = {8: _recast_to_unsigned_bytes, 16: _recast_to_unsigned_16, 32: _recast_to_unsigned_32} +_RECAST_UNSIGNED[64] = _recast_to_unsigned_64 +_INT_WIDTHS = [(bits, getattr(torch, f"int{bits}")) for bits in (8, 16, 32, 64)] + + +def _twos_complement(value, bits): + """The signed spelling of *value*'s bits, so torch will accept it at that width.""" + return value - (1 << bits) if value >> (bits - 1) else value + + +def _straddling_tensor(bits, signed_dtype): + values = _straddling_values(bits) + return values, torch.tensor([_twos_complement(v, bits) for v in values], dtype=signed_dtype, device="cuda") + + +@pytest.mark.parametrize("bits, signed_dtype", _INT_WIDTHS) +def test_unsigned_source_and_destination_at_every_width(bits, signed_dtype): + """``v >> 4`` written back through storage of the same width, for every integer width. + + Both the load and the store go through unsigned storage, so this covers the + destination side as well as the source side. The same bits read as signed must still + shift arithmetically. + """ + values, src = _straddling_tensor(bits, signed_dtype) + recast = _RECAST_UNSIGNED[bits] + + unsigned = _elementwise_launcher(_shift_right, bits, bits, recast=recast, recast_out=recast) + assert _run_elementwise(unsigned, src, signed_dtype) == [v >> 4 for v in values] + + signed = _elementwise_launcher(_shift_right, bits, bits) + assert _run_elementwise(signed, src, signed_dtype) == (src >> 4).tolist() + + +@pytest.mark.parametrize("bits, signed_dtype", [w for w in _INT_WIDTHS if w[0] < 64]) +def test_widening_covers_the_full_unsigned_range_at_every_width(bits, signed_dtype): + """Zero- versus sign-extension into a destination wide enough to hold either answer. + + Every unsigned width has to widen correctly, not just ``uint8``: ``0xFFFFFFFF`` must + arrive as ``4294967295`` and the same bits read as signed as ``-1``. 64 bits is + excluded because there is no wider signed destination to observe it in. + """ + values, src = _straddling_tensor(bits, signed_dtype) + + unsigned = _elementwise_launcher(_widen_to_i64, bits, 64, recast=_RECAST_UNSIGNED[bits]) + assert _run_elementwise(unsigned, src, torch.int64) == values + + signed = _elementwise_launcher(_widen_to_i64, bits, 64) + assert _run_elementwise(signed, src, torch.int64) == src.tolist() + + +def test_recast_iter_reinterprets_storage_as_unsigned(): + """``recast_iter(fx.Uint8, ...)`` -- how packed codes are read -- yields unsigned bytes.""" + src = torch.tensor(BYTES, dtype=torch.uint8, device="cuda").view(torch.int8) + + as_unsigned = _elementwise_launcher(_widen_to_i32, 8, 32, recast=_recast_to_unsigned_bytes) + assert _run_elementwise(as_unsigned, src, torch.int32) == BYTES + + as_signed = _elementwise_launcher(_widen_to_i32, 8, 32, recast=_recast_to_signed_bytes) + assert _run_elementwise(as_signed, src, torch.int32) == src.tolist() + + +def _unsigned_byte_fragment(view): + return fx.make_fragment_like(view, fx.Uint8) + + +def _signed_byte_fragment(view): + return fx.make_fragment_like(view, fx.Int8) + + +def test_make_fragment_like_with_an_explicit_dtype_decides_signedness(): + """``make_fragment_like(view, fx.Uint8)`` makes the register storage unsigned. + + The source tensor is the same bytes both times, so only the fragment's declared dtype + can explain the difference. + """ + src = torch.tensor(BYTES, dtype=torch.uint8, device="cuda") + + as_unsigned = _elementwise_launcher(_widen_to_i32, 8, 32, make_frag=_unsigned_byte_fragment) + assert _run_elementwise(as_unsigned, src, torch.int32) == BYTES + + as_signed = _elementwise_launcher(_widen_to_i32, 8, 32, make_frag=_signed_byte_fragment) + assert _run_elementwise(as_signed, src, torch.int32) == src.view(torch.int8).tolist() + + +def _scalar_load_launcher(load): + """``dst[i] = int32(src[i])`` where the element is read one scalar at a time. + + *load* is a function rather than a flag so that the two index paths of + :func:`memref_load` compile to separate kernels instead of sharing a JIT cache entry. + """ + + @flyc.kernel + def kernel(src: fx.Tensor, dst: fx.Tensor): + tid = fx.thread_idx.x + layout = fx.make_layout(1, 1) + view_in = fx.Tensor(fx.make_view(fx.add_offset(fx.get_iter(src), fx.make_int_tuple(tid)), layout)) + view_out = fx.Tensor(fx.make_view(fx.add_offset(fx.get_iter(dst), fx.make_int_tuple(tid)), layout)) + frag_in = fx.make_fragment_like(view_in) + frag_out = fx.make_fragment_like(view_out) + fx.copy(fx.make_copy_atom(fx.UniversalCopy8b(), frag_in.element_type), view_in, frag_in) + frag_out[0] = load(frag_in).to(fx.Int32) + fx.copy(fx.make_copy_atom(fx.UniversalCopy32b(), view_out.element_type), frag_out, view_out) + + @flyc.jit + def launch(src: fx.Tensor, dst: fx.Tensor, stream: fx.Stream = fx.Stream(None)): + kernel(src, dst).launch(grid=(1, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + return launch + + +def _load_at_static_index(frag): + # A Python index: memref_load builds the int tuple itself and checks the profile. + return fx.memref_load(frag, 0) + + +def _load_at_dynamic_index(frag): + # An int-tuple value built from a run-time operand -- the other branch of memref_load. + # thread_idx.y is 0 for a one-dimensional block, but the compiler cannot know that. + return fx.memref_load(frag, fx.make_int_tuple(fx.thread_idx.y)) + + +@pytest.mark.parametrize("load", [_load_at_static_index, _load_at_dynamic_index]) +@pytest.mark.parametrize("torch_dtype", [torch.uint8, torch.int8]) +def test_scalar_memref_load_keeps_the_storage_signedness(load, torch_dtype): + """One element at a time, the scalar counterpart of the vector load already covered. + + ``memref_load`` yields a signless value, so widening it can only pick ``extui`` over + ``extsi`` if the dtype came from the memref it was loaded from. The same bytes are read + through unsigned and signed storage, so only that can explain a difference. + """ + src = torch.tensor(BYTES, dtype=torch.uint8, device="cuda").view(torch_dtype) + got = _run_elementwise(_scalar_load_launcher(load), src, torch.int32) + assert got == src.tolist() + + +def _ordered_against_64(values, src_type, dst_type): + bound = src_type(0x40) + return (values > bound).select(values, bound).to(fx.Int32) + + +@pytest.mark.parametrize("torch_dtype", [torch.uint8, torch.int8]) +def test_ordering_executes_with_the_dtype_signedness(torch_dtype): + """Comparison picks ``ugt`` or ``sgt`` from the dtype; checked on hardware. + + ``0x40`` sits below the sign bit, so the same bytes order differently under the two + readings and only the dtype can explain the result. + """ + src = torch.tensor(BYTES, dtype=torch.uint8, device="cuda").view(torch_dtype) + got = _run_elementwise(_elementwise_launcher(_ordered_against_64, 8, 32), src, torch.int32) + assert got == [max(v, 0x40) for v in src.tolist()] + + +def _predicate_widened(values, src_type, dst_type): + return (values > src_type(0x40)).to(fx.Int32) + + +def test_widening_a_predicate_uses_the_unsigned_domain(): + """``i1`` carries no sign bit, so widening ``true`` must give 1, not -1. + + A width-1 integer's only bit is its value. Reading it as signed would sign-extend + through ``extsi`` and turn every set predicate into -1, so ``ArithValue`` and + ``int_to_int`` both normalise width-1 signedness away. The device shows the + difference directly. + """ + src = torch.tensor(BYTES, dtype=torch.uint8, device="cuda") + got = _run_elementwise(_elementwise_launcher(_predicate_widened, 8, 32), src, torch.int32) + assert got == [int(v > 0x40) for v in src.tolist()] + + +@flyc.kernel +def _pointer_shift_kernel(src: fx.Pointer, dst: fx.Pointer, n: fx.Int32): + tid = fx.thread_idx.x + if tid < n: + dst[tid] = (src[tid] >> 4).to(fx.Int32) + + +@flyc.jit +def _launch_pointer_shift(src: fx.Pointer, dst: fx.Pointer, n: fx.Int32, stream: fx.Stream = fx.Stream(None)): + _pointer_shift_kernel(src, dst, n).launch(grid=(1, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + +@pytest.mark.parametrize("dsl_dtype, torch_dtype", [(fx.Uint8, torch.uint8), (fx.Int8, torch.int8)]) +def test_raw_pointer_argument_keeps_its_signedness(dsl_dtype, torch_dtype): + """A raw pointer argument carries a dtype too, and loading through it must honour it.""" + src = torch.tensor(BYTES, dtype=torch.uint8, device="cuda") + dst = torch.zeros(NUM_THREADS, dtype=torch.int32, device="cuda") + + _launch_pointer_shift( + flyc.from_c_void_p(dsl_dtype, src.data_ptr()), + flyc.from_c_void_p(fx.Int32, dst.data_ptr()), + NUM_THREADS, + stream=torch.cuda.current_stream(), + ) + torch.cuda.synchronize() + assert dst.tolist() == [v >> 4 for v in src.view(torch_dtype).tolist()] + + +@flyc.kernel +def _ptr_load_scalar_kernel(src: fx.Pointer, dst: fx.Pointer, n: fx.Int32): + tid = fx.thread_idx.x + if tid < n: + # The pointer's element type is signed; the ``ui8`` result type is what asks for an + # unsigned read. It names the dtype only -- the SSA value it produces stays signless. + value = fx.ptr_load(src + tid, result_type=ir.IntegerType.get_unsigned(8)) + dst[tid] = (value >> 4).to(fx.Int32) + + +@flyc.kernel +def _ptr_load_vector_kernel(src: fx.Pointer, dst: fx.Pointer, n: fx.Int32): + tid = fx.thread_idx.x + if tid < n: + values = fx.ptr_load(src + tid * PER_THREAD, result_type=ir.Type.parse("vector<4xui8>")) + for i in fx.range_constexpr(PER_THREAD): + dst[tid * PER_THREAD + i] = (values[i] >> 4).to(fx.Int32) + + +@flyc.jit +def _launch_ptr_load_scalar(src: fx.Pointer, dst: fx.Pointer, n: fx.Int32, stream: fx.Stream = fx.Stream(None)): + _ptr_load_scalar_kernel(src, dst, n).launch(grid=(1, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + +@flyc.jit +def _launch_ptr_load_vector(src: fx.Pointer, dst: fx.Pointer, n: fx.Int32, stream: fx.Stream = fx.Stream(None)): + _ptr_load_vector_kernel(src, dst, n).launch(grid=(1, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream) + + +@pytest.mark.parametrize("launch, per_thread", [(_launch_ptr_load_scalar, 1), (_launch_ptr_load_vector, PER_THREAD)]) +def test_ptr_load_with_an_unsigned_result_type(launch, per_thread): + """``ptr_load`` may be handed a raw ``uiN`` MLIR type; it must mean unsigned, and must + not leak into the SSA value, which ``arith`` would reject.""" + values = BYTES * per_thread + src = torch.tensor(values, dtype=torch.uint8, device="cuda").view(torch.int8) + dst = torch.zeros(len(values), dtype=torch.int32, device="cuda") + + launch( + flyc.from_c_void_p(fx.Int8, src.data_ptr()), + flyc.from_c_void_p(fx.Int32, dst.data_ptr()), + NUM_THREADS, + stream=torch.cuda.current_stream(), + ) + torch.cuda.synchronize() + assert dst.tolist() == [v >> 4 for v in values]