From de526e0aeafdf4d42896bbf635d2dc14fa599233 Mon Sep 17 00:00:00 2001 From: Arist12 Date: Tue, 28 Jul 2026 19:32:27 +0000 Subject: [PATCH] [DSL] Preserve logical signedness of unsigned integer dtypes A torch.uint8 tensor arrived in a kernel with element type Int8, because both Int8.ir_type and Uint8.ir_type lower to a signless i8 and the DSL wrapper was rebuilt from that signless type. Every operation whose signed and unsigned forms differ then took the signed one: widening 0xF0 gave -16 instead of 240, and integer-to-float sign-extended. No error, no warning, just wrong values for every byte with the high bit set. Carry signedness where it can be represented and drop it where it cannot. Storage descriptors spell it -- !fly.memref, !fly.ptr -- so the DSL can rebuild the logical dtype, while SSA values stay signless because arith accepts nothing else and encodes signedness in the opcode instead. toSSAValueType maps one to the other at the load boundary, with the same function mirrored on the Python side. Width-1 integers are normalised to unsigned: i1 has no bit to spare for a sign, so reading it as signed makes true compare as -1 and inverts every ordering. Signed dtypes are unaffected: they were already signless and stay that way. The set of accepted boundary dtypes is unchanged. Fixes #701 Signed-off-by: Arist12 --- include/flydsl/Dialect/Fly/Utils/TypeUtils.h | 30 ++ lib/Bindings/Python/DLTensorAdaptor.h | 4 +- lib/Dialect/Fly/CMakeLists.txt | 1 + lib/Dialect/Fly/IR/FlyOps.cpp | 6 +- lib/Dialect/Fly/Utils/PointerUtils.cpp | 5 +- lib/Dialect/Fly/Utils/TypeUtils.cpp | 17 + python/flydsl/compiler/jit_argument.py | 6 +- python/flydsl/expr/derived.py | 2 +- python/flydsl/expr/numeric.py | 44 ++ python/flydsl/expr/primitive.py | 44 +- python/flydsl/expr/utils/arith.py | 2 + tests/language/test_unsigned_semantics.py | 265 ++++++++++++ tests/unit/test_unsigned_tensor_dtype.py | 417 +++++++++++++++++++ 13 files changed, 829 insertions(+), 14 deletions(-) create mode 100644 include/flydsl/Dialect/Fly/Utils/TypeUtils.h create mode 100644 lib/Dialect/Fly/Utils/TypeUtils.cpp create mode 100644 tests/language/test_unsigned_semantics.py create mode 100644 tests/unit/test_unsigned_tensor_dtype.py 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]