[FIX][IR] Repair PrimType dtype CI fallout

Author: tqchen

include/tvm/ir/base_expr.h

@@ -29,6 +29,7 @@
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/ir/source_map.h>
 
+#include <cstddef>
 #include <cstdint>
 
 namespace tvm {
@@ -207,6 +208,14 @@ class PrimType final : public Type {
     return static_cast<int16_t>(get()->dtype.lanes) > 1;
   }
 
+  /*!
+   * \brief Return the number of bytes needed to store one value of this type.
+   *
+   * This uses the same packed sub-byte dtype sizing rule as runtime tensors.
+   * Scalable vector types have no compile-time storage size and are rejected.
+   */
+  TVM_DLL size_t StorageBytes() const;
+
   /*! \brief Return the same type with a different dtype code, preserving bits and lanes. */
   TVM_FFI_INLINE PrimType WithCode(DLDataTypeCode code) const {
     DLDataType dtype = get()->dtype;

include/tvm/topi/detail/broadcast.h

@@ -43,9 +43,7 @@ struct BroadcastHelper {
 };
 
 static inline PrimType CommonType(const PrimType& type1, const PrimType& type2) {
-  TVM_FFI_ICHECK(!type1.IsScalableVector() && !type2.IsScalableVector());
-  TVM_FFI_ICHECK_EQ(type1.lanes(), 1);
-  TVM_FFI_ICHECK_EQ(type2.lanes(), 1);
+  TVM_FFI_ICHECK(type1.IsScalar() && type2.IsScalar());
   TVM_FFI_ICHECK(type1.code() == type2.code());
   return type1.bits() < type2.bits() ? type1.WithBits(type2.bits()) : type1;
 }
@@ -59,7 +57,7 @@ inline BroadcastHelper BroadcastShape(const tvm::ffi::Array<tvm::PrimExpr>& shap
   int i;
 
   auto cast_if_needed = [](PrimType to_type, PrimExpr expr) {
-    return to_type->dtype == expr.ty()->dtype ? expr : cast(to_type, expr);
+    return to_type == expr.ty() ? expr : cast(to_type, expr);
   };
 
   for (i = 1; i <= std::min(s1_size, s2_size); ++i) {

python/tvm/ir/expr.py

@@ -45,8 +45,16 @@ class PrimExpr(BaseExpr):
 
     @property
     def dtype(self):
-        """Return the runtime dtype represented by this expression's PrimType."""
-        return self.ty.dtype
+        """Compatibility alias for the runtime dtype of scalar PrimExpr.
+
+        New code should inspect ``expr.ty`` directly.  For scalar primitive
+        expressions, use ``expr.ty.dtype``.
+        """
+        if self.ty is None:
+            return None
+        if hasattr(self.ty, "dtype"):
+            return self.ty.dtype
+        return "handle"
 
 
 @tvm_ffi.register_object("ir.RelaxExpr")

python/tvm/script/parser/core/evaluator.py

@@ -396,7 +396,11 @@ def _eval_if_exp(self, fields: dict[str, Any]) -> Any:
         orelse = self._eval_expr(fields["orelse"])
         if isinstance(test, bool):
             return body if test else orelse
-        elif isinstance(test, tvm.tirx.PrimExpr) and test.dtype.type_code == tvm.DataTypeCode.BOOL:
+        elif (
+            isinstance(test, tvm.tirx.PrimExpr)
+            and isinstance(test.ty, tvm.ir.PrimType)
+            and test.ty.matches_code(tvm.DataTypeCode.BOOL)
+        ):
             return tvm.tirx.op.if_then_else(test, body, orelse)
         else:
             raise TypeError(f"Expected Python bool or TIR bool, but got {type(test)}")

python/tvm/tirx/buffer.py

@@ -352,7 +352,7 @@ def _infer_shape(shape):
             shape = args
             assert all(
                 isinstance(arg, int)
-                or (isinstance(arg, PrimExpr) and arg.dtype in ["int32", "int64"])
+                or (isinstance(arg, PrimExpr) and arg.ty.dtype in ["int32", "int64"])
                 for arg in shape
             ), "shape must be a list of integers or PrimExprs with dtype int32 or int64"
             # Safely get optional keyword arguments
@@ -462,7 +462,7 @@ def permute(self, *dims) -> "Buffer":
 
     def __getitem__(self, indices):
         from ..arith import Analyzer  # pylint: disable=import-outside-toplevel
-        from .expr import BufferLoad, Ramp, const  # pylint: disable=import-outside-toplevel
+        from .expr import BufferLoad, Ramp  # pylint: disable=import-outside-toplevel
         from .stmt import BufferRegion  # pylint: disable=import-outside-toplevel
 
         if not isinstance(indices, tuple | list):
@@ -483,7 +483,8 @@ def __getitem__(self, indices):
                 else:
                     region.append(
                         Range.from_min_extent(
-                            index, const(1, index.dtype) if isinstance(index, PrimExpr) else 1
+                            index,
+                            tvm.tirx.expr.IntImm(index.ty, 1) if isinstance(index, PrimExpr) else 1,
                         )
                     )
             if has_implicit_slice:
@@ -499,7 +500,7 @@ def __getitem__(self, indices):
                     step = 1 if index.step is None else index.step
                     # We should ensure the dtype of start is the same with that of step.
                     if isinstance(start, tvm.tirx.expr.PrimExpr) and isinstance(step, int):
-                        step = tvm.tirx.expr.IntImm(start.dtype, step)
+                        step = tvm.tirx.expr.IntImm(start.ty, step)
                     lanes = analyzer.simplify((stop - start + step - 1) // step)
                     if lanes == 1:
                         expr_indices.append(start)
@@ -540,8 +541,8 @@ def decl_buffer(
         layout = TileLayout(S[tuple(shape)]) if shape else None
 
     if offset_factor != 0 and elem_offset is None:
-        shape_dtype = shape[0].dtype if shape and hasattr(shape[0], "dtype") else "int32"
-        elem_offset = Var(f"{name}_elem_offset", shape_dtype)
+        shape_ty = shape[0].ty if shape and isinstance(shape[0], PrimExpr) else "int32"
+        elem_offset = Var(f"{name}_elem_offset", shape_ty)
     if data is None:
         # Bool is represented as uint1 in the IR, but stored as int8
         storage_type = dtype if isinstance(dtype, PrimType) else PrimType(dtype)

python/tvm/tirx/expr.py

@@ -132,7 +132,7 @@ def __rmod__(self, other: PrimExpr) -> PrimExpr:
         return _ffi_api._OpFloorMod(other, self, None)  # type: ignore
 
     def __neg__(self) -> PrimExpr:
-        neg_one = const(-1, self.dtype)  # type: ignore
+        neg_one = const(-1, self.expr_ty().dtype)
         return self.__mul__(neg_one)
 
     def __lshift__(self, other: PrimExpr) -> PrimExpr:
@@ -215,7 +215,7 @@ def equal(self, other: PrimExpr, span: Span | None = None) -> bool:
         """
         return _ffi_api._OpEQ(self, other, span)  # type: ignore
 
-    def astype(self, dtype: str, span: Span | None = None) -> PrimExpr:
+    def astype(self, dtype: str | ir.PrimType, span: Span | None = None) -> PrimExpr:
         """Cast the expression to other type.
 
         Parameters
@@ -477,12 +477,10 @@ def __init__(
                 raise TypeError("dom need to be Range")
 
         name = var if var is not None else "iter"
-        dtype = "int32" if dom is None else dom.extent.dtype
+        dtype = "int32" if dom is None else dom.extent.ty
         var = Var(name, dtype=dtype, span=span) if not isinstance(var, Var) else var
         if dom is not None:
-            assert var.dtype == dom.extent.dtype, (
-                "IterVar's Var dtype must match its domain's extent's dtype"
-            )
+            assert var.ty == dom.extent.ty, "IterVar's Var type must match its domain's extent type"
         self.__init_handle_by_constructor__(
             _ffi_api.IterVar,
             dom,
@@ -618,7 +616,9 @@ class FloatImm(ConstExpr):
 
     value: float
 
-    def __init__(self, dtype: str, value: float, span: Span | None = None) -> None:
+    def __init__(self, dtype: str | ir.PrimType, value: float, span: Span | None = None) -> None:
+        if isinstance(dtype, ir.PrimType):
+            dtype = dtype.dtype
         self.__init_handle_by_constructor__(
             tvm.ir._ffi_api.FloatImm,
             dtype,
@@ -648,7 +648,9 @@ class IntImm(ConstExpr):
 
     value: int
 
-    def __init__(self, dtype: str, value: int, span: Span | None = None) -> None:
+    def __init__(self, dtype: str | ir.PrimType, value: int, span: Span | None = None) -> None:
+        if isinstance(dtype, ir.PrimType):
+            dtype = dtype.dtype
         self.__init_handle_by_constructor__(
             tvm.ir._ffi_api.IntImm,
             dtype,
@@ -725,7 +727,9 @@ class Cast(PrimExprWithOp):
 
     value: PrimExpr
 
-    def __init__(self, dtype, value, span: Span | None = None) -> None:
+    def __init__(self, dtype: str | ir.PrimType, value, span: Span | None = None) -> None:
+        if isinstance(dtype, ir.PrimType):
+            dtype = dtype.dtype
         self.__init_handle_by_constructor__(_ffi_api.Cast, dtype, value, span)  # type: ignore
 
 
@@ -1336,7 +1340,7 @@ class Call(PrimExprWithOp):
 
     def __init__(
         self,
-        dtype: str,
+        dtype: str | ir.PrimType,
         op: Op | str,
         args: list[PrimExpr],
         attrs: ir.Attrs | dict | None = None,

python/tvm/tirx/layout.py

@@ -332,10 +332,10 @@ def _get_default_strides(data: list[int | PrimExpr], stride: int = 1) -> tuple:
         # produce for int64-shaped buffers (otherwise the last stride stays a
         # Python ``int`` -> int32 IntImm and breaks structural-equal).
         for t in data:
-            if isinstance(t, PrimExpr) and t.dtype != "int32":
+            if isinstance(t, PrimExpr) and t.ty.dtype != "int32":
                 from .expr import IntImm  # pylint: disable=import-outside-toplevel
 
-                stride = IntImm(t.dtype, stride)
+                stride = IntImm(t.ty, stride)
                 break
         res = list()
         for t in reversed(data):