[REFACTOR][IR] Unify PrimExpr type mechanism to PrimType instead of DataType

[tqchen]

In the past we have been using DataType in PrimExpr.dtype field to check type information for PrimExpr while still having BaseExpr.ty for richer type information. DataType is also used both in runtime and compiler. This PR streamlines the boundary:

• PrimExpr.ty now carries PrimType that replaces original use of DataType
• Runtime use will now favor DLPack DLDataType, removing one layer of indirection.
• Constants attributes where values are usually runtime values, will use DLDataType
• DataType will be phased out after this PR

We also brings up helper functions in PrimType, but also limits them to a more concise set so the functions do not grow with the data type codes in DLPack.

This is a major refactor that changes the IR primitive. It helps to bring possible future benefits:

• Unified type mechanism through Expr.ty
• Possibility of carry future Type nodes

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.

Performance:

Using Object type instead of DLDataType would indeed bring some performance impact to the IR. We have done the following performance optimizations:

• Make sure most of the outputs reuse one of the PrimType from inputs
• Cache a thread local PrimType based on input so we don't repeatly realloc

We did benchmarks show that rewrite simplify operation stays within +-10% overhead of original one. Which merits the refactor given the benefit the unfication brings

[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.

[tqchen]

/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]

/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.

[tqchen]

/gemini review

[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.