[REFACTOR][IR] Use PrimType for compiler dtypes by tqchen · Pull Request #19875 · apache/tvm

tqchen · 2026-06-23T15:44:09Z

This PR removes the old runtime DataType wrapper as the compiler-facing dtype abstraction and routes dtype usage by boundary.

Rationale:

PrimType is the compiler IR type for primitive expressions, so using it as the dtype/type carrier unifies dtype information with Expr.ty.
This removes the need for a dedicated expression dtype path separate from type information.
Keeping dtype as an IR type leaves room for future expression type annotations without another parallel dtype abstraction.
Runtime, ABI, and dtype-valued attrs stay raw DLDataType, where a plain DLPack dtype value is the real boundary object.

Migration guide:

Use PrimType when code reasons about compiler expression types, tensor element compiler types, or constructs a PrimExpr/compiler type.
Use existing source types such as expr.ty(), ExprOp.expr_ty(), or TE tensor element dtype where possible instead of rebuilding a type from dtype text.
Use raw DLDataType for runtime constants, ABI paths, dtype-valued attrs, and storage/runtime helper logic.
Prefer direct PrimType equality, MatchesCode(...), MatchesElementType(...), and WithCode(...) over local wrappers or string dtype checks.
Keep public names such as GetDataType and output_dtype where they are API terminology, but align their value type with the compiler/runtime boundary.

Validation:

full branch changed-file pre-commit passed over 420 changed files
git diff --check passed
final grep found no PrimType::IsPredicate, .IsBool(, .IsInt(, .IsUInt(, removed type.cc helper names, or runtime/data_type.h references in checked paths
LLVM-enabled ninja -C build -j$(nproc) completed and linked lib/libtvm_compiler.so
focused Python pytest over TIRX ops/intrinsics, Relax qdq/manipulate legalizers, and contrib sort passed
direct TOPI smoke passed for touched math/scatter paths

The branch is intentionally stacked on #19874 (2bdedc93aa) so CI has the needed base fix.

gemini-code-assist

Code Review

This pull request refactors the codebase to transition from using DataType to PrimType (backed by DLDataType), unifying the type system across IR dialects and improving type safety. The reviewer provided valuable feedback on leveraging the new PrimType APIs (such as bits(), lanes(), and IsScalableVector()) to avoid manual DLDataType unpacking and arithmetic, particularly in tensor allocation, memory view handling, static memory planning, and GPU code verification. Additionally, a bug was identified in src/relax/transform/utils.h where kDLBool was incorrectly checked against a bit width of 1 instead of the standard 8 bits.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Use PrimType as the compiler-facing dtype/type carrier so primitive expression dtype information is unified with Expr.ty instead of flowing through a dedicated dtype path. This keeps compiler IR type information in the type system and leaves room for future expression type annotations. Use raw DLDataType at runtime, ABI, storage-helper, and dtype-valued attr boundaries where a plain DLPack dtype value is the real interface. Keep the PrimType API minimal and hot-path friendly with value equality, matching helpers, documented factories, and cached common constructors. Update TIRX, TE, TOPI, Relax, codegen, Python bindings, and tests to follow the compiler PrimType versus runtime DLDataType boundary.

tqchen · 2026-06-23T16:41:19Z

/gemini review

gemini-code-assist

Code Review

This pull request refactors the TVM codebase by replacing the usage of tvm::DataType with tvm::PrimType and DLDataType across C++ and Python files to better align with DLPack. The code review feedback is highly consistent, recommending the use of the more idiomatic MatchesCode helper on PrimType instead of directly comparing raw data type codes (such as kDLFloat, kDLInt, and kDLBool). Additionally, it is suggested to simplify redundant checks by passing multiple type codes to a single variadic MatchesCode call.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

tqchen · 2026-06-23T16:58:09Z

/gemini Please run another review pass on latest head b8c7505. Since the last pass, I applied the Gemini MatchesCode cleanups, fixed the Python PrimType-to-runtime-dtype boundary in Relax gradient constants and TE NumPy test uses, reran full LLVM build plus focused Relax/TE/TIRX validation, and resolved the previous review threads that are addressed.

gemini-code-assist

Code Review

This pull request refactors TVM's type system by replacing the runtime-specific DataType with a unified compile-time PrimType (backed by DLDataType) across the compiler, runtime, and Python bindings. This extensive refactoring updates buffer declarations, expression nodes, and codegen backends to use PrimType or DLDataType directly. Feedback on the changes highlights a potential division-by-zero bug in Hexagon's GetVectorBytes for sub-byte types, a potential AttributeError in the Python PrimExpr.dtype property when self.ty is null or a pointer, and several opportunities to simplify type checks using the newly introduced IsScalar() helper and PrimType equality operators.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-23T17:06:54Z

+TVM_FFI_INLINE int GetVectorBytes(const PrimType& dtype) {
+  TVM_FFI_ICHECK(dtype.IsFixedLengthVector() || dtype.IsScalar());
+  return dtype.bits() * dtype.lanes() / 8;
+}


In GetVectorBytes, using dtype.bits() * dtype.lanes() / 8 can return 0 for sub-byte types (such as 4-bit integers or booleans). This will lead to a division-by-zero error in VectorLookupLoad where native_vector_bytes / GetVectorBytes(buffer_type) is calculated. Using (dtype.bits() * dtype.lanes() + 7) / 8 safely computes the byte size and prevents this potential compiler crash.

TVM_FFI_INLINE int GetVectorBytes(const PrimType& dtype) { TVM_FFI_ICHECK(dtype.IsFixedLengthVector() || dtype.IsScalar()); return (dtype.bits() * dtype.lanes() + 7) / 8; }

gemini-code-assist · 2026-06-23T17:06:54Z

+    @property
+    def dtype(self):
+        """Return the runtime dtype represented by this expression's PrimType."""
+        return self.ty.dtype


The dtype property on PrimExpr is implemented as return self.ty.dtype. However, self.ty can be None or a PointerType (which does not have a dtype attribute). This will raise an AttributeError when accessing .dtype on handle variables or un-typed expressions. Adding safety checks to handle None and PointerType (returning "handle" for pointers) ensures backward compatibility and robustness.

Suggested change

@property

def dtype(self):

"""Return the runtime dtype represented by this expression's PrimType."""

return self.ty.dtype

@property

def dtype(self):

"""Return the runtime dtype represented by this expression's PrimType."""

if self.ty is None:

return None

if hasattr(self.ty, "dtype"):

return self.ty.dtype

return "handle"

gemini-code-assist · 2026-06-23T17:06:54Z

+  TVM_FFI_ICHECK(!type1.IsScalableVector() && !type2.IsScalableVector());
+  TVM_FFI_ICHECK_EQ(type1.lanes(), 1);
+  TVM_FFI_ICHECK_EQ(type2.lanes(), 1);


The checks to ensure type1 and type2 are scalar can be simplified and made more robust by using the newly introduced IsScalar() helper method on PrimType.

Suggested change

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());

gemini-code-assist · 2026-06-23T17:06:54Z

+  auto cast_if_needed = [](PrimType to_type, PrimExpr expr) {
+    return to_type->dtype == expr.ty()->dtype ? expr : cast(to_type, expr);
  };


Since PrimType has operator== defined, the check to_type->dtype == expr.ty()->dtype can be simplified to to_type == expr.ty().

Suggested change

auto cast_if_needed = [](PrimType to_type, PrimExpr expr) {

return to_type->dtype == expr.ty()->dtype ? expr : cast(to_type, expr);

};

auto cast_if_needed = [](PrimType to_type, PrimExpr expr) {

return to_type == expr.ty() ? expr : cast(to_type, expr);

};

gemini-code-assist · 2026-06-23T17:06:54Z

+  TVM_FFI_DCHECK(expr->BaseExprNode::ty.defined());
+  const auto* prim_ty = expr->BaseExprNode::ty.as<PrimTypeNode>();
+  TVM_FFI_DCHECK(prim_ty != nullptr);
+  return IsIndexType(prim_ty->dtype);


The manual checks on expr->BaseExprNode::ty are redundant because PrimExprNode::ty() already implements these checks and returns PrimType. We can simplify this to a single call to expr->ty().

Suggested change

TVM_FFI_DCHECK(expr->BaseExprNode::ty.defined());

const auto* prim_ty = expr->BaseExprNode::ty.as<PrimTypeNode>();

TVM_FFI_DCHECK(prim_ty != nullptr);

return IsIndexType(prim_ty->dtype);

return IsIndexType(expr->ty()->dtype);

gemini-code-assist · 2026-06-23T17:06:54Z


  bool IsValidScalar(const PrimExpr& expr) const {
-    return expr.defined() && expr.dtype().is_scalar();
+    if (!expr.defined()) return false;


The check !ty.IsFixedLengthVector() && !ty.IsScalableVector() can be simplified using the IsScalar() helper method on PrimType.

return expr.defined() && expr.ty().IsScalar();

tqchen · 2026-06-23T17:17:22Z

/gemini Please run another review pass on latest head 4b0c660. Since the last pass, I fixed CI compile failures in CBLAS/CUBLAS/HIPBLAS/MKL/DNNL/CUDNN/random and CUDA runtime paths after DataType helper removal: contrib backends now use direct raw DLDataType equality instead of TypeMatch/TypeEqual/runtime namespace helpers, and CUDA tensor-map runtime code uses raw DLDataType fields. Local validation: changed-file pre-commit passed, LLVM-enabled ninja build passed, and syntax-only checks passed for the reported disabled-backend sources with local CUDA/CUDNN/CBLAS/DNNL headers.

gemini-code-assist

Code Review

This pull request performs a major refactoring of TVM's type system by removing the tvm::runtime::DataType wrapper class and replacing its usage with tvm::PrimType and raw DLDataType across the entire codebase, including TIR, Relax, TOPI, and various hardware backends (CUDA, Vulkan, Metal, WebGPU, Hexagon, and Trainium). PrimExpr and PrimExprNode now expose their types via PrimType (ty()) instead of DataType (dtype()), and corresponding Python bindings have been updated to reflect these changes. Since no review comments were provided, there is no specific feedback to address, and the changes appear to successfully unify and simplify the type representation.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist Bot reviewed Jun 23, 2026

View reviewed changes

tqchen force-pushed the task-use-ty-primtype-for-source-of-dtype branch 2 times, most recently from 20c4a00 to 3e75c55 Compare June 23, 2026 16:10

tqchen force-pushed the task-use-ty-primtype-for-source-of-dtype branch from 3e75c55 to 74fba5e Compare June 23, 2026 16:32

[REFACTOR][IR] Fix PrimType follow-up review issues

b204812

gemini-code-assist Bot reviewed Jun 23, 2026

View reviewed changes

[REFACTOR][IR] Address PrimType review follow-ups

b8c7505

gemini-code-assist Bot reviewed Jun 23, 2026

View reviewed changes

[REFACTOR][Runtime] Use raw DLDataType in contrib backends

4b0c660

gemini-code-assist Bot reviewed Jun 23, 2026

View reviewed changes

Conversation

tqchen commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tqchen commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tqchen commented Jun 23, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

tqchen commented Jun 23, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

tqchen commented Jun 23, 2026 •

edited

Loading

tqchen commented Jun 23, 2026 •

edited

Loading