[REFACTOR][TIRX] Add IntImm common scalar constructors

Common bool, int32, and int64 scalar constants show up throughout TIRX and related lowering code, and named constructors make these call sites easier to read than repeated DataType spelling.

This PR establishes the scalar-constant construction policy:

- Prefer IntImm::Bool, IntImm::Int32, and IntImm::Int64 for common known scalar bool, int32, and int64 constants.

- Prefer direct IntImm or FloatImm construction when dtype is known to be scalar integer or floating point. This makes the compiled code more compact and efficient. Keep make_const for generic overload cases where dtype can be integer, floating point, or vector-valued and the caller needs its scalar/vector dispatch.

- Phase out make_zero in favor of explicit scalar constructors, or ConstHandle(0) for null handles.

Author: tqchen

include/tvm/ir/expr.h

@@ -515,6 +515,33 @@ class IntImm : public PrimExpr {
    */
   TVM_DLL IntImm(DataType dtype, int64_t value, Span span = Span());
 
+  /*!
+   * \brief Construct a scalar boolean constant.
+   * \param value The boolean value.
+   * \param span The location of this object in the source code.
+   */
+  static IntImm Bool(bool value, Span span = Span()) {
+    return IntImm(DataType::Bool(), value, span);
+  }
+
+  /*!
+   * \brief Construct a scalar int32 constant.
+   * \param value The integer value.
+   * \param span The location of this object in the source code.
+   */
+  static IntImm Int32(int64_t value, Span span = Span()) {
+    return IntImm(DataType::Int(32), value, span);
+  }
+
+  /*!
+   * \brief Construct a scalar int64 constant.
+   * \param value The integer value.
+   * \param span The location of this object in the source code.
+   */
+  static IntImm Int64(int64_t value, Span span = Span()) {
+    return IntImm(DataType::Int(64), value, span);
+  }
+
   TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntImm, PrimExpr, IntImmNode);
   TVM_DEFINE_OBJECT_REF_COW_METHOD(IntImmNode);
 };
@@ -636,7 +663,7 @@ struct TypeTraits<FloatImm> : public ObjectRefWithFallbackTraitsBase<FloatImm, d
 
 // define automatic conversion from bool, int64_t, double to PrimExpr
 TVM_FFI_INLINE PrimExpr TypeTraits<PrimExpr>::ConvertFallbackValue(StrictBool value) {
-  return IntImm(DataType::Bool(), value, Span());
+  return IntImm::Bool(value);
 }
 
 TVM_FFI_INLINE PrimExpr TypeTraits<PrimExpr>::ConvertFallbackValue(int64_t value) {

include/tvm/script/printer/doc.h

@@ -277,15 +277,15 @@ class LiteralDoc : public ExprDoc {
    * \param p The object path
    */
   static LiteralDoc Int(int64_t v, const ffi::Optional<AccessPath>& p) {
-    return LiteralDoc(IntImm(DataType::Int(64), v), p);
+    return LiteralDoc(IntImm::Int64(v), p);
   }
   /*!
    * \brief Create a LiteralDoc to represent boolean.
    * \param v The boolean value.
    * \param p The object path
    */
   static LiteralDoc Boolean(bool v, const ffi::Optional<AccessPath>& p) {
-    return LiteralDoc(IntImm(DataType::Bool(), v), p);
+    return LiteralDoc(IntImm::Bool(v), p);
   }
   /*!
    * \brief Create a LiteralDoc to represent float.

include/tvm/tirx/buffer.h

@@ -206,7 +206,7 @@ class Buffer : public ffi::ObjectRef {
    * \param input_extent The extent of ptr.
    */
   TVM_DLL PrimExpr access_ptr(int access_mask, DataType ptr_type = DataType::Handle(),
-                              int content_lanes = 1, PrimExpr offset = IntImm(DataType::Int(32), 0),
+                              int content_lanes = 1, PrimExpr offset = IntImm::Int32(0),
                               ffi::Optional<PrimExpr> input_extent = std::nullopt) const;
   /*!
    * \brief Create an Expr that does a vector load at begin index.

include/tvm/tirx/op.h

@@ -816,7 +816,14 @@ inline bool IsPointerType(const Type& type, const DataType& element_type) {
 
 /*!
  * \brief Make a const value with certain data type.
- * \param t The target type.
+ *
+ * Prefer direct IntImm or FloatImm construction when dtype is known to be
+ * scalar integer or floating point. This makes the compiled code more compact
+ * and efficient. Keep make_const for generic overload cases where dtype can be
+ * integer, floating point, or vector-valued and the caller needs its
+ * scalar/vector dispatch.
+ *
+ * \param dtype The target type.
  * \param value The input value
  * \return the result expression.
  * \tparam ValueType The constant value type
@@ -825,32 +832,14 @@ inline bool IsPointerType(const Type& type, const DataType& element_type) {
 template <typename ValueType,
           typename = typename std::enable_if<std::is_standard_layout<ValueType>::value &&
                                              std::is_trivial<ValueType>::value>::type>
-inline PrimExpr make_const(DataType t, ValueType value, Span span = Span());
-/*!
- * \brief Make a const zero expr.
- * \param t The target type.
- * \param span The location of this operation in the source.
- * \return the result expression.
- */
-inline PrimExpr make_zero(DataType t, Span span = Span());
-/*!
- * \brief Make a constant true expression.
- * \param lanes The number of lanes in the bool
- * \param span The location of this operation in the source.
- * \return The result expression.
- */
-inline PrimExpr const_true(int lanes = 1, Span span = Span()) {
-  return make_const(DataType::Bool(lanes), 1);
-}
+inline PrimExpr make_const(DataType dtype, ValueType value, Span span = Span());
 /*!
- * \brief Make a constant false expression.
- * \param lanes The number of lanes in the bool
+ * \brief Make a constant handle value.
+ * \param value The integer payload to reinterpret as a handle.
  * \param span The location of this operation in the source.
  * \return The result expression.
  */
-inline PrimExpr const_false(int lanes = 1, Span span = Span()) {
-  return make_const(DataType::Bool(lanes), 0);
-}
+inline PrimExpr ConstHandle(int64_t value, Span span = Span());
 /*!
  * \brief Get x as constant int expression.
  * \param x The expression
@@ -981,53 +970,52 @@ inline bool is_no_op(const tirx::Stmt& stmt) {
 }
 
 template <typename ValueType>
-inline PrimExpr MakeConstScalar(DataType t, ValueType value, Span span = Span()) {
-  if (t.is_int() || t.is_bool()) return IntImm(t, static_cast<int64_t>(value), span);
-  if (t.is_uint()) {
+inline PrimExpr MakeConstScalar(DataType dtype, ValueType value, Span span = Span()) {
+  if (dtype.is_int() || dtype.is_bool()) return IntImm(dtype, static_cast<int64_t>(value), span);
+  if (dtype.is_uint()) {
     // Use IntImm if it is a small integer
     uint64_t uval = static_cast<uint64_t>(value);
     if (value < static_cast<ValueType>(0)) {
       TVM_FFI_THROW(InternalError) << "cannot make uint from negative value " << value;
     } else if (uval <= static_cast<uint64_t>(std::numeric_limits<int64_t>::max())) {
-      return IntImm(t, static_cast<int64_t>(value), span);
+      return IntImm(dtype, static_cast<int64_t>(value), span);
     } else {
       uint64_t mask = (static_cast<uint64_t>(1) << 32U) - 1U;
       uint64_t low = uval & mask;
       uint64_t high = uval >> 32U;
-      return LargeUIntImm(t, static_cast<int64_t>(low), static_cast<int64_t>(high), span);
+      return LargeUIntImm(dtype, static_cast<int64_t>(low), static_cast<int64_t>(high), span);
     }
   }
-  if (t.is_float() || t.is_bfloat16() || t.is_float8() || t.is_float6() || t.is_float4())
-    return FloatImm(t, static_cast<double>(value), span);
-  TVM_FFI_THROW(InternalError) << "cannot make const for type " << t;
+  if (dtype.is_float() || dtype.is_bfloat16() || dtype.is_float8() || dtype.is_float6() ||
+      dtype.is_float4()) {
+    return FloatImm(dtype, static_cast<double>(value), span);
+  }
+  TVM_FFI_THROW(InternalError) << "cannot make const for type " << dtype;
   throw;
 }
 
 template <>
-inline PrimExpr MakeConstScalar(DataType t, bool value, Span span) {
-  return MakeConstScalar(t, static_cast<int>(value), span);
+inline PrimExpr MakeConstScalar(DataType dtype, bool value, Span span) {
+  return MakeConstScalar(dtype, static_cast<int>(value), span);
 }
 
 template <typename ValueType, typename>
-inline PrimExpr make_const(DataType t, ValueType value, Span span) {
-  if (t.is_scalar()) {
-    return MakeConstScalar(t, value, span);
+inline PrimExpr make_const(DataType dtype, ValueType value, Span span) {
+  if (dtype.is_scalar()) {
+    return MakeConstScalar(dtype, value, span);
   } else {
-    if (t.is_fixed_length_vector()) {
-      return tirx::Broadcast(MakeConstScalar(t.element_of(), value, span), t.lanes(), span);
+    if (dtype.is_fixed_length_vector()) {
+      return tirx::Broadcast(MakeConstScalar(dtype.element_of(), value, span), dtype.lanes(), span);
     } else {
-      PrimExpr lanes =
-          tirx::Mul(tirx::Call(DataType::Int(32), tirx::builtin::vscale(), {}), t.vscale_factor());
-      return tirx::Broadcast(MakeConstScalar(t.element_of(), value, span), lanes, span);
+      PrimExpr lanes = tirx::Mul(tirx::Call(DataType::Int(32), tirx::builtin::vscale(), {}),
+                                 dtype.vscale_factor());
+      return tirx::Broadcast(MakeConstScalar(dtype.element_of(), value, span), lanes, span);
     }
   }
 }
 
-inline PrimExpr make_zero(DataType t, Span span) {
-  if (t.is_handle()) {
-    return reinterpret(t, make_const(DataType::UInt(64), 0, span));
-  }
-  return make_const(t, 0, span);
+inline PrimExpr ConstHandle(int64_t value, Span span) {
+  return reinterpret(DataType::Handle(), IntImm(DataType::UInt(64), value, span));
 }
 
 }  // namespace tirx
@@ -1049,7 +1037,7 @@ inline PrimExpr make_zero(DataType t, Span span) {
     return Name(a, tirx::make_const(a.dtype(), b));                                 \
   }                                                                                 \
   inline PrimExpr Name(const PrimExpr& a, double b) {                               \
-    return Name(a, tirx::make_const(DataType::Float(64), b));                       \
+    return Name(a, FloatImm(DataType::Float(64), b));                               \
   }
 
 #define TVM_DEFINE_BINOP_CONST_VAL_OVERLOAD_SPANNED(Name)                 \
@@ -1066,7 +1054,7 @@ inline PrimExpr make_zero(DataType t, Span span) {
     return Name(a, tirx::make_const(a.dtype(), b), span);                 \
   }                                                                       \
   inline PrimExpr Name(const PrimExpr& a, double b, Span span = Span()) { \
-    return Name(a, tirx::make_const(DataType::Float(64), b), span);       \
+    return Name(a, FloatImm(DataType::Float(64), b), span);               \
   }
 
 #define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(Name)                             \

include/tvm/topi/detail/broadcast.h

@@ -130,7 +130,7 @@ inline tvm::ffi::Array<tvm::PrimExpr> InputIndexFromBroadcast(
     // Only inject 0 here if we have not yet reached the dimension of I
     // (i.e. this must be a 1)
     if (!found && (ovars.size() - i) <= expected_dims) {
-      ivars.push_back(tvm::tirx::make_zero(ovars[i].dtype()));
+      ivars.push_back(tvm::IntImm(ovars[i].dtype(), 0));
     }
   }
   TVM_FFI_ICHECK(expected_dims == ivars.size());

include/tvm/topi/detail/extern.h

@@ -108,13 +108,12 @@ inline PrimExpr pack_buffer(Buffer buf) {
   } else {
     strides = 0;
   }
-  ffi::Array<PrimExpr> pack_args{
-      buf->data,
-      shape,
-      strides,
-      make_const(DataType::Int(32), static_cast<int64_t>(buf->shape.size())),
-      make_const(buf->dtype, 0),
-      buf->elem_offset};
+  ffi::Array<PrimExpr> pack_args{buf->data,
+                                 shape,
+                                 strides,
+                                 IntImm::Int32(static_cast<int64_t>(buf->shape.size())),
+                                 make_const(buf->dtype, 0),
+                                 buf->elem_offset};
   return tvm::tirx::Call(DataType::Handle(), tvm::tirx::builtin::tvm_stack_make_array(), pack_args);
 }
 

include/tvm/topi/elemwise.h

@@ -209,7 +209,7 @@ inline Tensor sign(const Tensor& x, std::string name = "T_sign", std::string tag
   return compute(
       x->shape,
       [&](const ffi::Array<Var>& i) {
-        PrimExpr zero = make_zero(x->dtype);
+        PrimExpr zero = make_const(x->dtype, 0);
         PrimExpr one = make_const(x->dtype, 1);
         PrimExpr minus_one = make_const(x->dtype, -1);
         auto s1 = tvm::tirx::Select((x(i) < zero), minus_one, zero);
@@ -392,19 +392,19 @@ inline Tensor full_like(const Tensor& x, const PrimExpr fill_value,
  * y = exp(f) = 1 + 2 * P(x**2)/(Q(x**2) - P(x**2))
  */
 inline Tensor fast_exp_float32(const Tensor& _x, std::string name, std::string tag) {
-  auto x_hi = make_const(DataType::Float(32), 88.3762626647950f);
-  auto x_lo = make_const(DataType::Float(32), -88.3762626647949f);
-  auto log2e = make_const(DataType::Float(32), 1.44269504088896341f);
-  auto ln2 = make_const(DataType::Float(32), 0.6931471805599453f);
-  PrimExpr p[6] = {make_const(DataType::Float(32), 1.9875691500E-4f),
-                   make_const(DataType::Float(32), 1.3981999507E-3f),
-                   make_const(DataType::Float(32), 8.3334519073E-3f),
-                   make_const(DataType::Float(32), 4.1665795894E-2f),
-                   make_const(DataType::Float(32), 1.6666665459E-1f),
-                   make_const(DataType::Float(32), 5.0000001201E-1f)};
-  auto one = make_const(DataType::Float(32), 1.0f);
-  auto one_half = make_const(DataType::Float(32), 0.5f);
-  auto b = make_const(DataType::Float(32), 127.0f);
+  auto x_hi = FloatImm(DataType::Float(32), 88.3762626647950f);
+  auto x_lo = FloatImm(DataType::Float(32), -88.3762626647949f);
+  auto log2e = FloatImm(DataType::Float(32), 1.44269504088896341f);
+  auto ln2 = FloatImm(DataType::Float(32), 0.6931471805599453f);
+  PrimExpr p[6] = {FloatImm(DataType::Float(32), 1.9875691500E-4f),
+                   FloatImm(DataType::Float(32), 1.3981999507E-3f),
+                   FloatImm(DataType::Float(32), 8.3334519073E-3f),
+                   FloatImm(DataType::Float(32), 4.1665795894E-2f),
+                   FloatImm(DataType::Float(32), 1.6666665459E-1f),
+                   FloatImm(DataType::Float(32), 5.0000001201E-1f)};
+  auto one = FloatImm(DataType::Float(32), 1.0f);
+  auto one_half = FloatImm(DataType::Float(32), 0.5f);
+  auto b = FloatImm(DataType::Float(32), 127.0f);
 
   return compute(
       _x->shape,

include/tvm/topi/nn.h

@@ -232,12 +232,12 @@ inline tvm::te::Tensor pad(
       if (pad_mode == "constant") {
         return tvm::if_then_else(
             foldl([](PrimExpr a, PrimExpr b, Span span) { return tvm::logical_and(a, b, span); },
-                  const_true(1), sel),
+                  IntImm::Bool(true), sel),
             t(indices), pad_value);
       } else if (pad_mode == "edge" || pad_mode == "reflect") {
         return tvm::if_then_else(
             foldl([](PrimExpr a, PrimExpr b, Span span) { return tvm::logical_and(a, b, span); },
-                  const_true(1), sel),
+                  IntImm::Bool(true), sel),
             t(indices), t(pad_idx));
       }
     }
@@ -534,7 +534,7 @@ inline tvm::te::Tensor space_to_batch_nd(const tvm::te::Tensor& data,
         << padded_input << ")"
         << " must be divisible by its block size (" << block_size << ")";
 
-    PrimExpr bs = IntImm(DataType::Int(64), block_shape[i - 1]);
+    PrimExpr bs = IntImm::Int64(block_shape[i - 1]);
     r_shape.push_back(div(padded_shape[i], bs));
     r_shape.push_back(bs);
     block_shape_prod *= bs;
@@ -549,7 +549,7 @@ inline tvm::te::Tensor space_to_batch_nd(const tvm::te::Tensor& data,
   }
   o_shape.push_back(tvm::PrimExpr(batch) * block_shape_prod);
   for (size_t i = 1; i <= num_block_dims; i++) {
-    PrimExpr bs = IntImm(DataType::Int(64), block_shape[i - 1]);
+    PrimExpr bs = IntImm::Int64(block_shape[i - 1]);
     o_shape.push_back(div(padded_shape[i], bs));
   }
   // append remaining shape
@@ -595,7 +595,7 @@ inline tvm::te::Tensor batch_to_space_nd(const tvm::te::Tensor& data,
   int batch = static_cast<int>(GetConstInt(in_shape[0]));
 
   for (size_t i = 0; i < num_block_dims; i++) {
-    PrimExpr bs = IntImm(DataType::Int(64), block_shape[i]);
+    PrimExpr bs = IntImm::Int64(block_shape[i]);
     r_shape.push_back(bs);
     block_shape_prod *= bs;
   }
@@ -614,7 +614,7 @@ inline tvm::te::Tensor batch_to_space_nd(const tvm::te::Tensor& data,
   ffi::Array<PrimExpr> r_p_shape;
   r_p_shape.push_back(batch / block_shape_prod);
   for (size_t i = 1; i <= num_block_dims; i++) {
-    PrimExpr bs = IntImm(DataType::Int(64), block_shape[i - 1]);
+    PrimExpr bs = IntImm::Int64(block_shape[i - 1]);
     r_p_shape.push_back(in_shape[i] * bs);
   }
   for (size_t i = num_block_dims + 1; i < num_input_dims; i++) {

include/tvm/topi/nn/bnn.h

@@ -71,7 +71,7 @@ inline tvm::te::Tensor binarize_pack(const tvm::te::Tensor& data, int axis,
           start_idx.push_back(i == static_cast<size_t>(axis) ? indices[i] * 32
                                                              : static_cast<PrimExpr>(indices[i]));
         }
-        auto packed = make_const(DataType::UInt(32), 0);
+        PrimExpr packed = IntImm(DataType::UInt(32), 0);
         for (size_t j = 0; j < 32; ++j) {
           ffi::Array<PrimExpr> idx;
           for (size_t i = 0; i < n; ++i) {

include/tvm/topi/nn/group_norm.h

@@ -126,7 +126,7 @@ inline Tensor group_norm(const Tensor& data, const Tensor& gamma, const Tensor&
 
   auto temp_x = temp_x_x2[0];
   auto temp_x2 = temp_x_x2[1];
-  auto reduce_extent = make_const(DataType::Float(32), 1);
+  PrimExpr reduce_extent = FloatImm(DataType::Float(32), 1);
   for (auto axis : new_axes) {
     reduce_extent *= data_reshaped->shape[axis];
   }