Skip to content

[REFACTOR][IR] Phase out class Integer and class Bool in Attrs and PassConfig#19614

Merged
spectrometerHBH merged 7 commits into
apache:mainfrom
tqchen:tvm-phase-out-integer-bool
May 26, 2026
Merged

[REFACTOR][IR] Phase out class Integer and class Bool in Attrs and PassConfig#19614
spectrometerHBH merged 7 commits into
apache:mainfrom
tqchen:tvm-phase-out-integer-bool

Conversation

@tqchen
Copy link
Copy Markdown
Member

@tqchen tqchen commented May 26, 2026

Summary

Now that the ffi container machinery (Array, Optional, Map, Variant) accepts bare int64_t and bool, the Integer/Bool ObjectRef wrappers add no value in attribute fields, pass-config options, function-attr flags, and OpAttrMap registries — every reader paid an extra .IntValue() / ->value unbox per access for no information gain. This PR is the first stage of phasing out class Integer and class Bool: migrate the bulk of those sites at the field-declaration and call-site level. A follow-up will rewrite the remaining IR-position Integer(N) / Bool(b) constructors to IntImm(...) / const_true() / const_false() and delete the two classes entirely.

  • Relax Attrs fields and their container forms (Array<Integer> / Optional<Array<Integer>> / Optional<Integer> / Optional<Bool>) migrated to bare int64_t / bool (manipulate.h, nn.h, op.h, statistical.h, script/builder/frame.h, target/virtual_device.h, distributed/global_info.h, relax/expr.h).
  • OpAttrMap registry (set_attr<Bool>("FPurity", Bool(true)) ↔ GetAttrMap<Bool>("FPurity")) migrated to bool across ~38 files.
  • PassContext config registrations + GetConfig<Bool> / GetConfig<Integer> readers, and function-attr GetAttr<Bool> / GetAttr<Integer> readers (~42 files), all migrated; HasNonzeroAttr in ir/attrs.h dropped its .IntValue() unbox.
  • Schedule decision arrays (SampleCategorical candidates, perfect-tile factors, autobind thread_extents, multi-level-tiling levels) migrated to Array<int64_t> / Optional<int64_t> — this is a virtual-signature change on ConcreteScheduleNode::SampleCategorical and related methods, acceptable per the phase-out intent.
  • Variant<Bool, Array<String>> for LiftTransformParams.shared_transform migrated to Variant<bool, Array<String>>.

gemini-code-assist[bot]
gemini-code-assist Bot reviewed May 26, 2026
View reviewed changes
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request refactors the codebase to replace TVM's wrapper types Integer and Bool with primitive int64_t and bool types across various Relax, S-TIR, and target codegen components. The review feedback points out a critical compilation error in the TensorRT codegen where Integer wrappers are still returned despite the return type change. Additionally, it suggests simplifying value_or defaults to primitive boolean literals in scalable_expression.cc and codegen_webgpu.cc, and recommends using .value_or() to avoid redundant double lookups for target attributes in spirv_support.cc.

Comment thread src/relax/backend/contrib/tensorrt/codegen.cc
Comment thread src/arith/scalable_expression.cc Outdated
Comment thread src/target/webgpu/codegen_webgpu.cc Outdated
Comment thread src/target/vulkan/spirv_support.cc Outdated
tqchen added a commit to tqchen/tvm that referenced this pull request May 26, 2026
@tqchen
[FIX] Address Gemini review on PR apache#19614
32e8ee5 
Fixes from Gemini review feedback on the Integer/Bool migration:

- src/relax/backend/contrib/tensorrt/codegen.cc: GetTensorRTVersion
  now returns ffi::Array<int64_t>, so drop the now-incompatible
  Integer(...) wrappers around NV_TENSORRT_{MAJOR,MINOR,PATCH}
  literals. The implicit int->int64_t conversion does the rest.
  (Critical: this was a compile error in TensorRT-enabled builds.)
- src/arith/scalable_expression.cc: GetAttr<bool>("feature.has_sve")
  now returns Optional<bool>, so value_or(Bool(false)) became a
  type mismatch (Bool is an ObjectRef, not a bool). Pass plain false.
- src/target/webgpu/codegen_webgpu.cc: same value_or(Bool(false))
  -> value_or(false) for GetAttr<bool>("supports_subgroups").
- src/target/vulkan/spirv_support.cc: collapse the repetitive
  `if (target->GetAttr<T>(k)) { x = target->GetAttr<T>(k).value(); }`
  pattern into `x = target->GetAttr<T>(k).value_or(x);` for all
  17 SPIRV capability flags. No behavior change, half the lines.
tqchen added a commit to tqchen/tvm that referenced this pull request May 26, 2026
@tqchen
[FIX] Internal audit: migrate remaining Integer/Bool sites in attrs/c…
3038fd6 
…onfigs/registries

Internal audit on the partial-landing of PR apache#19614 found several
attr/registry sites that were missed. Migrated them following the
same Integer/Bool -> int64_t/bool pattern.

OpAttrMap registries (TCallEffectKind, TScriptDtypePrintLocation)
----------------------------------------------------------------
The two operator-attr-map type aliases in
`include/tvm/tirx/op_attr_types.h`:

    using TScriptDtypePrintLocation = Integer;  -> int64_t
    using TCallEffectKind            = Integer;  -> int64_t

Updated every `set_attr<TCallEffectKind>(..., Integer(K))` and
`set_attr<TScriptDtypePrintLocation>(..., Integer(K))` writer in
`src/target/{cuda,metal,webgpu}/intrin_rule_*.cc` and
`src/tirx/op/{op,builtin,runtime,target_builtin/*}.cc` and
`src/tirx/ir/stmt.cc` to use `static_cast<int64_t>(K)`.

Updated the 4 readers (`Op::GetAttrMap<TCallEffectKind>(...)[op]`
and the `TScriptDtypePrintLocation` reader in
`src/tirx/script/printer/expr.cc`) to drop `->value` /
`.IntValue()` and explicitly cast the int64_t back to the enum.

Function-attr / DictAttrs / WithAttr / attr_map.Set
----------------------------------------------------
Replaced `tvm::Bool(true)` with bare `true` at:

  - `src/te/operation/create_primfunc.cc` (2x DictAttrs init)
  - `src/relax/transform/fuse_tir.cc` (attr_map.Set)
  - `src/relax/transform/compute_prim_value.cc` (DictAttrs init)
  - `src/tirx/script/builder/frame.cc` (insert_attr for kSTir and
    kPersistentKernel)
  - `src/relax/backend/vm/codegen_vm_tir.cc` (WithAttr)
  - `src/relax/backend/vm/vm_shape_lower.cc` (WithAttr)
  - `src/tirx/transform/split_host_device.cc` (WithAttr)

Replaced `WithAttr(..., "scoped", Integer(1))` with
`WithAttr(..., "scoped", static_cast<int64_t>(1))` in
`src/relax/transform/specialize_primfunc_based_on_callsite.cc`.

Buffer/loop/block annotations
------------------------------
`annotations.Set(kVolatile, Bool(true))` -> `annotations.Set(kVolatile,
true)` at the 7 call sites in
`src/s_tir/transform/{merge_shared_memory_allocations,lower_thread_allreduce}.cc`
and `src/tirx/transform/storage_rewrite.cc`. All readers use
`annotations.count(...)` (presence-only), so the value type does not
matter to consumers.

Python downstream fixes
------------------------
The original landing missed Python callers that still did `.value`
on what is now a bare int after the migration. Fixed:

  - `python/tvm/relax/transform/legalize_ops/statistical.py`:
    `_compute_shape_prod`, `_te_mean`, `_te_variance`, `_te_median`,
    `_te_std` and `_te_median` -- the `axis` parameter is now
    `list[int]`, drop the `_axis.value` / `axis[0].value` unboxing.
  - `python/tvm/relax/transform/legalize_ops/manipulate.py`: stack's
    `call.attrs.axis.value` -> `call.attrs.axis` (StackAttrs.axis is
    `Optional[int64_t]` now).

Test mirror fixes
------------------
Tests that constructed s_tir attr via `tirx.IntImm("bool", 1)` and
compared structurally now mismatch the bare `True` produced by
`@T.prim_func(s_tir=True)`. Updated 3 tests to use plain `True`:

  - `tests/python/tirx-base/test_tir_stmt_functor_substitute.py`
  - `tests/python/tirx-transform/test_tir_transform_convert_ssa.py`
  - `tests/python/tvmscript/test_tvmscript_ir_builder_tir.py` (2x)

Verified
---------
Local CPU build clean. Test sweeps green:
  - tests/python/ir/                         90 passed
  - tests/python/tirx-base/                  288 passed, 25 skipped
  - tests/python/tirx-transform/             349 passed (1 pre-existing
    LLVM-20-only failure in test_vectorize_llvm_pure_intrin_fail,
    unrelated)
  - tests/python/tvmscript/                  769 passed, 1 skipped
  - tests/python/relax/test_op_*.py          1164 passed, 5 skipped
  - tests/python/relax/test_transform_*.py   908 passed (1 pre-existing
    failure in test_transform_meta_schedule_tuning -- see note below)

Out-of-scope / not fixed (would require API design judgment)
------------------------------------------------------------
`src/s_tir/schedule/trace.cc` JSON decision deserializer produces
`Array<IntImm>` from JSON literals, but the migrated instruction
traits (e.g. `SamplePerfectTileTraits::UnpackedApplyToSchedule`) now
expect `Optional<Array<int64_t>>`. The FFI layer does not convert
`Array<IntImm>` to `Array<int64_t>` element-wise across containers.
This breaks trace replay from JSON in meta_schedule. Needs a separate
fix to either (a) convert decisions at parse time in `Trace::FromJSON`
or (b) widen the trait signature accept-types. Pre-existing in PR apache#19614
landing, not introduced here.
tqchen and others added 7 commits May 26, 2026 20:17
@tqchen
[REFACTOR][IR] Phase out class Integer and class Bool in Attrs and Pa…
920302d 
…ssConfig

Now that the ffi container machinery (Array, Optional, Map, Variant)
accepts bare int64_t and bool, the Integer/Bool ObjectRef wrappers add
no value in attribute fields, pass-config options, function-attr
flags, and operator-attr-map registries. Every reader paid an extra
.IntValue() / ->value unbox per access for no information gain. This
commit migrates the bulk of those sites at the field-declaration and
call-site level.

Main changes:

- Relax Attrs fields: every `Integer axis;` / `Bool keepdims;` and the
  container forms `Array<Integer>` / `Optional<Array<Integer>>` /
  `Optional<Integer>` / `Optional<Bool>` migrated to the bare type
  (`int64_t` / `bool`). Headers under `include/tvm/relax/attrs/` plus
  builder/Frame and target/VirtualDevice.
- OpAttrMap registry: `set_attr<Bool>("FPurity", Bool(true))` and the
  paired `GetAttrMap<Bool>("FPurity")` readers migrated to `bool`.
- PassContext config: `TVM_REGISTER_PASS_CONFIG_OPTION("...", Bool)`
  and `("...", Integer)` and their `GetConfig<Bool>` / `GetConfig<Integer>`
  readers migrated to bare types.
- Function-attr flags: `GetAttr<Bool>(k, Bool(false)).value()` and
  `GetAttr<Integer>(k, Integer(N)).value().IntValue()` migrated to
  `GetAttr<bool>(k, false).value()` and `GetAttr<int64_t>(k, N).value()`.
- `HasNonzeroAttr` in `ir/attrs.h` dropped its `.IntValue()` unbox.
- Schedule decision arrays (SampleCategorical candidates, perfect-tile
  factors, autobind thread_extents, multi-level-tiling levels, etc.)
  migrated to `Array<int64_t>` / `Optional<int64_t>` since they are
  integer payloads, not IR exprs. This is a virtual-signature change
  on ConcreteScheduleNode::SampleCategorical and related methods.
- `Variant<Bool, Array<String>>` for `LiftTransformParams.shared_transform`
  migrated to `Variant<bool, Array<String>>`.

IR-position `Integer(N)` / `Bool(true|false)` constructors that flow
into PrimExpr / Stmt positions are left alone in this commit; they
are valid as long as `class Integer` and `class Bool` still exist
(both still derive from `IntImm` and share `IntImmNode`). A follow-up
commit will rewrite those to `IntImm(DataType::Int(32), N)` /
`const_true()` / `const_false()` and then delete the two classes.

Build verified clean for CPU target. Smoke-tested Python suites:
relax/test_op_* (130+296 passes), relax/test_transform_legalize_ops,
test_transform_fuse_ops, test_transform_lift_transform_params (56
passes), and tests/python/ir + tests/python/tirx-base (378 passes,
25 skipped).
@tqchen
[FIX] Address Gemini review on PR apache#19614
a5c8037 
Fixes from Gemini review feedback on the Integer/Bool migration:

- src/relax/backend/contrib/tensorrt/codegen.cc: GetTensorRTVersion
  now returns ffi::Array<int64_t>, so drop the now-incompatible
  Integer(...) wrappers around NV_TENSORRT_{MAJOR,MINOR,PATCH}
  literals. The implicit int->int64_t conversion does the rest.
  (Critical: this was a compile error in TensorRT-enabled builds.)
- src/arith/scalable_expression.cc: GetAttr<bool>("feature.has_sve")
  now returns Optional<bool>, so value_or(Bool(false)) became a
  type mismatch (Bool is an ObjectRef, not a bool). Pass plain false.
- src/target/webgpu/codegen_webgpu.cc: same value_or(Bool(false))
  -> value_or(false) for GetAttr<bool>("supports_subgroups").
- src/target/vulkan/spirv_support.cc: collapse the repetitive
  `if (target->GetAttr<T>(k)) { x = target->GetAttr<T>(k).value(); }`
  pattern into `x = target->GetAttr<T>(k).value_or(x);` for all
  17 SPIRV capability flags. No behavior change, half the lines.
@tqchen
[FIX] Internal audit: migrate remaining Integer/Bool sites in attrs/c…
0bb1d65 
…onfigs/registries

Internal audit on the partial-landing of PR apache#19614 found several
attr/registry sites that were missed. Migrated them following the
same Integer/Bool -> int64_t/bool pattern.

OpAttrMap registries (TCallEffectKind, TScriptDtypePrintLocation)
----------------------------------------------------------------
The two operator-attr-map type aliases in
`include/tvm/tirx/op_attr_types.h`:

    using TScriptDtypePrintLocation = Integer;  -> int64_t
    using TCallEffectKind            = Integer;  -> int64_t

Updated every `set_attr<TCallEffectKind>(..., Integer(K))` and
`set_attr<TScriptDtypePrintLocation>(..., Integer(K))` writer in
`src/target/{cuda,metal,webgpu}/intrin_rule_*.cc` and
`src/tirx/op/{op,builtin,runtime,target_builtin/*}.cc` and
`src/tirx/ir/stmt.cc` to use `static_cast<int64_t>(K)`.

Updated the 4 readers (`Op::GetAttrMap<TCallEffectKind>(...)[op]`
and the `TScriptDtypePrintLocation` reader in
`src/tirx/script/printer/expr.cc`) to drop `->value` /
`.IntValue()` and explicitly cast the int64_t back to the enum.

Function-attr / DictAttrs / WithAttr / attr_map.Set
----------------------------------------------------
Replaced `tvm::Bool(true)` with bare `true` at:

  - `src/te/operation/create_primfunc.cc` (2x DictAttrs init)
  - `src/relax/transform/fuse_tir.cc` (attr_map.Set)
  - `src/relax/transform/compute_prim_value.cc` (DictAttrs init)
  - `src/tirx/script/builder/frame.cc` (insert_attr for kSTir and
    kPersistentKernel)
  - `src/relax/backend/vm/codegen_vm_tir.cc` (WithAttr)
  - `src/relax/backend/vm/vm_shape_lower.cc` (WithAttr)
  - `src/tirx/transform/split_host_device.cc` (WithAttr)

Replaced `WithAttr(..., "scoped", Integer(1))` with
`WithAttr(..., "scoped", static_cast<int64_t>(1))` in
`src/relax/transform/specialize_primfunc_based_on_callsite.cc`.

Buffer/loop/block annotations
------------------------------
`annotations.Set(kVolatile, Bool(true))` -> `annotations.Set(kVolatile,
true)` at the 7 call sites in
`src/s_tir/transform/{merge_shared_memory_allocations,lower_thread_allreduce}.cc`
and `src/tirx/transform/storage_rewrite.cc`. All readers use
`annotations.count(...)` (presence-only), so the value type does not
matter to consumers.

Python downstream fixes
------------------------
The original landing missed Python callers that still did `.value`
on what is now a bare int after the migration. Fixed:

  - `python/tvm/relax/transform/legalize_ops/statistical.py`:
    `_compute_shape_prod`, `_te_mean`, `_te_variance`, `_te_median`,
    `_te_std` and `_te_median` -- the `axis` parameter is now
    `list[int]`, drop the `_axis.value` / `axis[0].value` unboxing.
  - `python/tvm/relax/transform/legalize_ops/manipulate.py`: stack's
    `call.attrs.axis.value` -> `call.attrs.axis` (StackAttrs.axis is
    `Optional[int64_t]` now).

Test mirror fixes
------------------
Tests that constructed s_tir attr via `tirx.IntImm("bool", 1)` and
compared structurally now mismatch the bare `True` produced by
`@T.prim_func(s_tir=True)`. Updated 3 tests to use plain `True`:

  - `tests/python/tirx-base/test_tir_stmt_functor_substitute.py`
  - `tests/python/tirx-transform/test_tir_transform_convert_ssa.py`
  - `tests/python/tvmscript/test_tvmscript_ir_builder_tir.py` (2x)

Verified
---------
Local CPU build clean. Test sweeps green:
  - tests/python/ir/                         90 passed
  - tests/python/tirx-base/                  288 passed, 25 skipped
  - tests/python/tirx-transform/             349 passed (1 pre-existing
    LLVM-20-only failure in test_vectorize_llvm_pure_intrin_fail,
    unrelated)
  - tests/python/tvmscript/                  769 passed, 1 skipped
  - tests/python/relax/test_op_*.py          1164 passed, 5 skipped
  - tests/python/relax/test_transform_*.py   908 passed (1 pre-existing
    failure in test_transform_meta_schedule_tuning -- see note below)

Out-of-scope / not fixed (would require API design judgment)
------------------------------------------------------------
`src/s_tir/schedule/trace.cc` JSON decision deserializer produces
`Array<IntImm>` from JSON literals, but the migrated instruction
traits (e.g. `SamplePerfectTileTraits::UnpackedApplyToSchedule`) now
expect `Optional<Array<int64_t>>`. The FFI layer does not convert
`Array<IntImm>` to `Array<int64_t>` element-wise across containers.
This breaks trace replay from JSON in meta_schedule. Needs a separate
fix to either (a) convert decisions at parse time in `Trace::FromJSON`
or (b) widen the trait signature accept-types. Pre-existing in PR apache#19614
landing, not introduced here.
@tqchen
[FIX][S-TIR] Trace JSON round-trip and meta_schedule mutator decision…
fa61bf3 
… types for migrated int64_t traits

After phasing out Integer from sample-instruction trait signatures, the
JSON deserializer in src/s_tir/schedule/trace.cc still produced
Array<IntImm> for attrs and decisions while their f_apply_to_schedule
now expects Array<int64_t> / Optional<int64_t> / Optional<Array<int64_t>>.
The FFI layer can unbox a top-level IntImm Any into int64_t via the
fallback, but does not recursively unbox Array<IntImm> into
Array<int64_t>, so trace replay from JSON failed with "Expected
Array<int> but got Array[index 0: ir.IntImm]". Symmetrically AsJSON
cast decisions to Optional<ObjectRef>, which fails once decisions are
stored as POD int64_t. The four meta_schedule mutators
(mutate_tile_size, mutate_thread_binding, mutate_unroll,
mutate_compute_location) were also still writing decisions as
Integer(...) / AsArray<int64_t, IntImm>(...) and reading them via
Downcast<IntImm>(...), which triggered the same in-memory trace replay
failure under search-strategy tuning.

Added NormalizeJSONIntegers helper that recursively replaces IntImm
leaves with bare int64_t in Trace::ApplyJSONToSchedule attrs/decisions
(stopping at FloatImm, IndexMap, RVs, strings, etc.); the reverse
int64_t -> IntImm conversion where a trait still wants Integer is
auto-handled by TypeTraits<IntImm>::ConvertFallbackValue. Switched
AsJSON to emit decisions as Any. Migrated the four mutator decision
writers/readers to bare int64_t / ffi::Array<int64_t>. Updated
test_sample_categorical_serialize whose .value access on the decision
no longer applies now that the decision is a Python int.
@tqchen
[FIX] Internal audit follow-up: drop .as<IntImmNode>()/.IntValue()/->…
61c26a7 
…value on migrated int64_t container elements

Three backend codegens still accessed migrated Array<int64_t> attribute
fields using the old Array<Integer> idiom:

- src/relax/backend/contrib/clml/codegen.cc:259
  PadAttrs::pad_width is Array<int64_t>; replace
  p[i].as<IntImmNode>()->value with bare p[i]. Caught by the Apache CI
  CLML container build (only that build compiles this file).
- src/relax/backend/contrib/tensorrt/codegen.cc:186-188
  TensorRTCompilerConfigNode::tensorrt_version is Array<int64_t>;
  drop the .IntValue() chain.
- src/relax/backend/contrib/nnapi/codegen.cc:61,81
  PermuteDimsAttrs::axes and StatisticalAttrs::axis are
  Optional<Array<int64_t>>; drop .IntValue() / ->value on the
  iterated element. NNAPI does compile in the regular CPU build,
  so this was also a latent compile error there.

Verified by directly compiling all three .cc files via the worktree's
clang invocation; nnapi additionally built as part of the normal
ninja build. ir/ and relax/test_op_nn_* tests pass.
@tqchen
[FIX] pad_einsum dtype preservation and annotate_buffer_access test u…
78156e7 
…pdates for int64_t migration

After migrating Array<Integer> to Array<int64_t> in pad_einsum and
annotate_buffer_access, two issues emerged:

- pad_einsum: padding constants were hardcoded as Int(64), causing dtype
  promotion when the domain variable was Int(32). Fix by using dom->dtype.
- annotate_buffer_access tests: expected TVMScript used T.int32(0) for
  explicit_read/write_region indices, but the migrated code stores bare
  int64_t values (displayed as [0] in TVMScript). Update tests to [0].
@tqchen
[LINT] ruff format annotate_buffer_access test
ab7e0c3
@tqchen tqchen force-pushed the tvm-phase-out-integer-bool branch from c4bafe4 to ab7e0c3 Compare May 26, 2026 20:37
@spectrometerHBH spectrometerHBH merged commit 0388fd0 into apache:main May 26, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@spectrometerHBH spectrometerHBH spectrometerHBH approved these changes

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@tqchen @spectrometerHBH