[REFACTOR][IR] Unify StructInfo and Type#19852
Conversation
The Relax dependent type refactor needs the helper and analysis surfaces to use the consolidated type vocabulary instead of keeping StructInfo compatibility shims around. This removes the remaining renamed helper wrappers, updates the public analysis and transform entry points to Type names, and tightens final stale-name cleanup in the nested-message and analysis internals.
tqchen
changed the title
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors TVM's Relax compiler framework by replacing the StructInfo abstraction with a unified Type system, renaming related classes (e.g., TensorStructInfo to TensorType, ShapeStructInfo to ShapeType, and DTensorStructInfo to DTensorType) and updating fields from struct_info to ty across both C++ and Python APIs. The review comments identify several critical issues, including a C++ compilation error in the Adreno backend due to incorrect type casting and pointer usage, multiple incorrect Python type annotations in the parser proxies, and a syntax error in a docstring example.
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.
| if (auto tensor_ty = ty.as<TensorType>()) { | ||
| bool is_texture = | ||
| i < is_texture_supported.size() ? is_texture_supported[i] : is_texture_supported[0]; | ||
| auto scope = | ||
| is_texture ? Scope(GetShapeFromTensorStructInfo(tensor_sinfo.value())) : "global"; | ||
| auto scope = is_texture ? Scope(GetShapeFromTensorType(tensor_ty.value())) : "global"; |
There was a problem hiding this comment.
Using ty.as<TensorType>() will result in a compilation error because TensorType is an ObjectRef subclass, not an Object subclass. It should be ty.as<TensorTypeNode>(). Additionally, tensor_ty is a pointer (const TensorTypeNode*), so calling .value() on it is invalid. Instead, use ffi::GetRef<TensorType>(tensor_ty) to pass the reference to GetShapeFromTensorType.
| if (auto tensor_ty = ty.as<TensorType>()) { | |
| bool is_texture = | |
| i < is_texture_supported.size() ? is_texture_supported[i] : is_texture_supported[0]; | |
| auto scope = | |
| is_texture ? Scope(GetShapeFromTensorStructInfo(tensor_sinfo.value())) : "global"; | |
| auto scope = is_texture ? Scope(GetShapeFromTensorType(tensor_ty.value())) : "global"; | |
| if (auto tensor_ty = ty.as<TensorTypeNode>()) { | |
| bool is_texture = | |
| i < is_texture_supported.size() ? is_texture_supported[i] : is_texture_supported[0]; | |
| auto scope = is_texture ? Scope(GetShapeFromTensorType(ffi::GetRef<TensorType>(tensor_ty))) : "global"; |
|
|
||
| def as_struct_info(self, dict_globals: dict[str, Any] | None = None) -> ShapeStructInfo: | ||
| return ObjectStructInfo() | ||
| def as_ty(self, dict_globals: dict[str, Any] | None = None) -> ShapeType: |
There was a problem hiding this comment.
The return type annotation of as_ty in ObjectProxy is incorrectly specified as ShapeType instead of ObjectType. This appears to be a copy-paste error from ShapeProxy.
| def as_ty(self, dict_globals: dict[str, Any] | None = None) -> ShapeType: | |
| def as_ty(self, dict_globals: dict[str, Any] | None = None) -> ObjectType: |
| return set() | ||
|
|
||
| def as_struct_info(self, dict_globals: dict[str, Any] | None = None) -> ShapeStructInfo: | ||
| def as_ty(self, dict_globals: dict[str, Any] | None = None) -> ShapeType: |
There was a problem hiding this comment.
The return type annotation of as_ty in PrimProxy is incorrectly specified as ShapeType instead of PrimType. This appears to be a copy-paste error from ShapeProxy.
| def as_ty(self, dict_globals: dict[str, Any] | None = None) -> ShapeType: | |
| def as_ty(self, dict_globals: dict[str, Any] | None = None) -> PrimType: |
| def as_struct_info(self, dict_globals: dict[str, Any] | None = None) -> TensorStructInfo: | ||
| return DTensorStructInfo( | ||
| self.tensor_sinfo_proxy.as_struct_info(dict_globals), | ||
| def as_ty(self, dict_globals: dict[str, Any] | None = None) -> TensorType: |
There was a problem hiding this comment.
The return type annotation of as_ty in DTensorProxy is incorrectly specified as TensorType instead of DTensorType. Since DTensorType inherits from Type (not TensorType), this is a type annotation mismatch.
| def as_ty(self, dict_globals: dict[str, Any] | None = None) -> TensorType: | |
| def as_ty(self, dict_globals: dict[str, Any] | None = None) -> DTensorType: |
| x = rx.Var("x", rx.TensorStructInfo([m, n], "float16")) | ||
| y = rx.Var("y", rx.TensorStructInfo([n], "float16") | ||
| x = rx.Var("x", rx.TensorType([m, n], "float16")) | ||
| y = rx.Var("y", rx.TensorType([n], "float16") |
There was a problem hiding this comment.
There is a missing closing parenthesis ) at the end of the line in the docstring example.
| y = rx.Var("y", rx.TensorType([n], "float16") | |
| y = rx.Var("y", rx.TensorType([n], "float16")) |
1 participant
This PR consolidates Relax's StructInfo vocabulary into the unified Type surface and moves Relax expression type storage to
Expr.ty.Summary:
tystorage and migrates Relax type storage/users away fromstruct_info_.StructInfo/*_sinfovocabulary toType/*_ty.TupleTyperepresentation while keeping Relax-specificPrimTypeandFuncTypebehavior.Validation:
ninja -C build -j$(nproc)passed.709 passed, 1 xfailed.git diff --checkpassed.Additional broad validation recorded in the task notes includes Relax op/distributed and analysis/transform/TVMScript/VM pytest sweeps.