[DSL] Preserve logical signedness of unsigned integer dtypes

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<ui8, ...>`, `!fly.ptr<ui8, ...>`), 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

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:

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

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 <mlir/IR/Attributes.h>
 #include <mlir/IR/BuiltinAttributes.h>
@@ -1913,7 +1914,8 @@ FLY_INFER_RETURN_TYPES(DecompositionOp) {
 
 FLY_INFER_RETURN_TYPES(MemRefLoadOp) {
   if (auto memrefTy = dyn_cast<MemRefType>(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<CoordTensorType>(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();
 }
 

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) {

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<IntegerType>(elemTy);
+  if (!intTy || !intTy.isUnsigned())
+    return elemTy;
+  return IntegerType::get(elemTy.getContext(), intTy.getWidth());
+}
+
+} // namespace mlir::fly

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):

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)
 
 

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<ui8, ...>``, ``!fly.ptr<ui8, ...>``).
+        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

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

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):