This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new e8267ce7ee [REFACTOR][IR] Simplify type scopes and registered
annotations (#20707)
e8267ce7ee is described below
commit e8267ce7eee4c681ba12ce797e5726a09bae5d2d
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 13 17:56:58 2026 -0400
[REFACTOR][IR] Simplify type scopes and registered annotations (#20707)
Use `scope` consistently for pointer and tensor types, and print pointer
scopes as positional strings such as `T.Ptr(T.void, "shared")`. The
existing typed MemoryScope conversion restores canonical scope identity;
standalone scope values retain constructor syntax.
Restrict special annotation lowering to registered namespace syntax,
preserving scalar and indexed declarations while keeping captured Type
values ordinary. Rename the scoped printer state to
`ir.explicit_value_type` to describe the value types it preserves.
---
include/tvm/ir/type.h | 11 +--
include/tvm/tirx/expr.h | 2 +-
include/tvm/tirx/type.h | 16 ++--
python/tvm/backend/cuda/ptx/engine.py | 2 +-
python/tvm/backend/cuda/tile/patterns/common.py | 10 +-
.../tvm/backend/cuda/tile/patterns/tcgen05_mma.py | 2 +-
python/tvm/backend/cuda/tile/patterns/tma.py | 4 +-
.../tvm/backend/trn/transform/naive_allocator.py | 4 +-
python/tvm/ir/type.py | 8 +-
python/tvm/script/parser/prescan.py | 102 ---------------------
python/tvm/script/parser/transpile.py | 24 +++--
python/tvm/tirx/tensor.py | 30 ++----
src/backend/cuda/codegen/codegen_cuda.cc | 3 +-
src/backend/cuda/codegen/llvm/codegen_nvptx.cc | 3 +-
src/backend/cuda/op/target_builtin.cc | 3 +-
src/backend/metal/codegen/codegen_metal.cc | 9 +-
src/backend/opencl/codegen/codegen_opencl.cc | 10 +-
src/backend/rocm/codegen/llvm/codegen_amdgpu.cc | 3 +-
src/backend/trn/codegen/codegen_trn.cc | 3 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 3 +-
src/backend/webgpu/codegen/codegen_webgpu.cc | 3 +-
src/ir/type.cc | 21 ++---
.../transform/lower_global_view_to_local_view.cc | 2 +-
src/relax/op/op.cc | 4 +-
src/relax/script/printer/call.cc | 2 +-
src/relax/script/printer/function.cc | 2 +-
src/relax/transform/fuse_tir.cc | 6 +-
src/s_tir/analysis/calculate_allocated_memory.cc | 4 +-
src/s_tir/analysis/verify_gpu_code.cc | 2 +-
src/s_tir/backend/adreno/inject_texture_alloc.cc | 4 +-
src/s_tir/backend/adreno/texture_flatten.cc | 2 +-
src/s_tir/schedule/primitive/cache_read_write.cc | 2 +-
src/s_tir/schedule/transform.cc | 5 +-
src/s_tir/script/ir_builder/ir.cc | 2 +-
src/s_tir/script/printer/stmt.cc | 17 ++--
src/s_tir/transform/lower_vtcm_alloc.cc | 4 +-
.../transform/merge_shared_memory_allocations.cc | 12 +--
src/s_tir/transform/storage_access.cc | 4 +-
src/s_tir/transform/tensorcore_infer_fragment.cc | 12 +--
src/script/printer/doc_translator.cc | 2 +-
src/script/printer/ir/call.cc | 2 +-
src/script/printer/ir/ir.cc | 4 +-
src/script/printer/ir/misc.cc | 4 +-
src/script/printer/ir/prim_expr.cc | 2 +-
src/script/printer/ir/prim_type.cc | 23 ++++-
src/script/printer/ir/utils.cc | 4 +-
src/target/llvm/codegen_llvm.cc | 7 +-
src/target/source/codegen_c.cc | 12 +--
src/tirx/ir/specialize.cc | 8 +-
src/tirx/ir/type.cc | 24 ++---
src/tirx/op/memory.cc | 3 +-
src/tirx/script/ir_builder/ir.cc | 2 +-
src/tirx/script/printer/for_loop.cc | 4 +-
src/tirx/script/printer/function.cc | 2 +-
src/tirx/script/printer/tensor.cc | 14 +--
src/tirx/transform/flattened_tensor.h | 9 +-
src/tirx/transform/lower_tvm_builtin.cc | 2 +-
src/tirx/transform/lower_warp_memory.cc | 4 +-
src/tirx/transform/split_host_device.cc | 2 +-
src/tirx/transform/storage_rewrite.cc | 6 +-
src/tirx/transform/unsupported_dtype_legalize.cc | 6 +-
src/tirx/transform/update_pointer_storage_scope.cc | 6 +-
tests/cpp/runtime/memory_scope_test.cc | 4 +-
...nsform_specialize_function_based_on_callsite.py | 16 ++--
tests/python/runtime/test_runtime_memory_scope.py | 31 ++++---
.../test_s_tir_analysis_verify_block_scope.py | 2 +-
.../s_tir/script/test_s_tir_script_printer.py | 2 +-
tests/python/tirx-base/test_tir_nodes.py | 2 +-
tests/python/tirx-base/test_tir_tensor.py | 4 +-
.../python/tirx/codegen/test_codegen_blackwell.py | 6 +-
tests/python/tirx/codegen/test_codegen_dsmem.py | 4 +-
.../tirx/script/test_tirx_script_basic_usage.py | 4 +-
.../python/tirx/script/test_tirx_script_printer.py | 4 +-
.../python/tirx/tile/cuda/test_ld_st_registers.py | 4 +-
74 files changed, 242 insertions(+), 355 deletions(-)
diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h
index db200726fd..14b96a443b 100644
--- a/include/tvm/ir/type.h
+++ b/include/tvm/ir/type.h
@@ -110,14 +110,13 @@ class PtrTypeNode : public TypeNode {
/*!
* \brief The storage scope of the pointer
*/
- runtime::MemoryScope storage_scope;
+ runtime::MemoryScope scope;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PtrTypeNode>()
.def_ro("element_type", &PtrTypeNode::element_type)
- .def_ro("storage_scope", &PtrTypeNode::storage_scope,
- refl::DefaultValue(runtime::MemoryScope::Global()));
+ .def_ro("scope", &PtrTypeNode::scope,
refl::DefaultValue(runtime::MemoryScope::Global()));
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("ir.PtrType", PtrTypeNode, TypeNode);
};
@@ -131,12 +130,12 @@ class PtrType : public Type {
/*!
* \brief Constructor
* \param element_type The type of the element which the pointer points to.
- * \param storage_scope The storage scope into which the pointer addresses
+ * \param scope The storage scope into which the pointer addresses
*/
- TVM_TIRX_DLL explicit PtrType(Type element_type, runtime::MemoryScope
storage_scope = {});
+ TVM_TIRX_DLL explicit PtrType(Type element_type, runtime::MemoryScope scope
= {});
/*! \brief Construct an opaque pointer with void element type. */
- TVM_TIRX_DLL static PtrType VoidPtrType(runtime::MemoryScope storage_scope =
{});
+ TVM_TIRX_DLL static PtrType VoidPtrType(runtime::MemoryScope scope = {});
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PtrType, Type, PtrTypeNode);
};
diff --git a/include/tvm/tirx/expr.h b/include/tvm/tirx/expr.h
index e12a1a458e..d2c943cbd6 100644
--- a/include/tvm/tirx/expr.h
+++ b/include/tvm/tirx/expr.h
@@ -74,7 +74,7 @@ class TensorVar : public Var {
Expr data() const { return Call(type()->DataPtrType(), tensor_data_ptr_op(),
{var()}); }
/*! \brief Return the storage scope associated with this tensor. */
- runtime::MemoryScope scope() const { return (*this)->storage_scope; }
+ runtime::MemoryScope scope() const { return (*this)->scope; }
/*!
* \brief Return a new tensor with the allocated address.
diff --git a/include/tvm/tirx/type.h b/include/tvm/tirx/type.h
index 41330b8a52..9da3305155 100644
--- a/include/tvm/tirx/type.h
+++ b/include/tvm/tirx/type.h
@@ -71,7 +71,7 @@ class TensorTypeNode : public TypeNode {
/*! \brief dtype in the content of the tensor */
PrimType dtype = PrimType::Void();
/*! \brief Storage scope/address space of the tensor. */
- runtime::MemoryScope storage_scope;
+ runtime::MemoryScope scope;
/*! \brief The type of the tensor prior to flattening
*
* This contains the shape as it is accessed by
@@ -101,8 +101,7 @@ class TensorTypeNode : public TypeNode {
refl::ObjectDef<TensorTypeNode>()
.def_ro("shape", &TensorTypeNode::shape,
refl::AttachFieldFlag::SEqHashDefSimple())
.def_ro("dtype", &TensorTypeNode::dtype)
- .def_ro("scope", &TensorTypeNode::storage_scope,
- refl::DefaultValue(runtime::MemoryScope::Global()))
+ .def_ro("scope", &TensorTypeNode::scope,
refl::DefaultValue(runtime::MemoryScope::Global()))
.def_ro("strides", &TensorTypeNode::strides,
refl::AttachFieldFlag::SEqHashDefSimple(),
refl::DefaultValue(ffi::Array<PrimExpr>()))
.def_ro("elem_offset", &TensorTypeNode::elem_offset,
@@ -123,7 +122,7 @@ class TensorTypeNode : public TypeNode {
PrimType ElementType() const { return dtype; }
/*! \return type of the physical pointer projected by tensor_data_ptr. */
- PtrType DataPtrType() const { return PtrType(dtype, storage_scope); }
+ PtrType DataPtrType() const { return PtrType(dtype, scope); }
/*! \brief Whether this type supports scalar tensor syntax. */
TVM_TIRX_DLL bool IsScalar(bool alloc_or_decl = true) const;
@@ -150,11 +149,10 @@ class TensorTypeNode : public TypeNode {
*/
class TensorType : public Type {
public:
- TVM_TIRX_DLL TensorType(ffi::Array<PrimExpr> shape, PrimType dtype,
- runtime::MemoryScope storage_scope,
ffi::Array<PrimExpr> strides,
- ffi::Optional<PrimExpr> elem_offset, int
data_alignment,
- int offset_factor, ffi::Optional<Layout> layout =
std::nullopt,
- Location loc = UnknownLoc());
+ TVM_TIRX_DLL TensorType(ffi::Array<PrimExpr> shape, PrimType dtype,
runtime::MemoryScope scope,
+ ffi::Array<PrimExpr> strides,
ffi::Optional<PrimExpr> elem_offset,
+ int data_alignment, int offset_factor,
+ ffi::Optional<Layout> layout = std::nullopt,
Location loc = UnknownLoc());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TensorType, Type,
TensorTypeNode);
diff --git a/python/tvm/backend/cuda/ptx/engine.py
b/python/tvm/backend/cuda/ptx/engine.py
index 97e0cbf667..a2680c8e9b 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -587,7 +587,7 @@ def _coerce_address(entry, slot, value, mod_map):
f"or a uint32 shared address"
)
if isinstance(ty, PtrType):
- if ty.storage_scope in (_SCOPE_SHARED, _SCOPE_SHARED_DYN):
+ if ty.scope in (_SCOPE_SHARED, _SCOPE_SHARED_DYN):
raise ValueError(
f"{entry.name}: operand '{slot.name}' is a {space or
'generic'} address "
f"but got a shared-scope pointer"
diff --git a/python/tvm/backend/cuda/tile/patterns/common.py
b/python/tvm/backend/cuda/tile/patterns/common.py
index f14e690ad5..2b43c1cf4e 100644
--- a/python/tvm/backend/cuda/tile/patterns/common.py
+++ b/python/tvm/backend/cuda/tile/patterns/common.py
@@ -235,10 +235,8 @@ class Region:
self.layout = self.ty.effective_layout()
if self.layout is None:
raise ValueError(f"{name}: operand requires a physical layout")
- if self.ty.storage_scope not in tuple(MemoryScope(scope) for scope in
scopes):
- raise ValueError(
- f"{name}: expected storage scope {scopes}, got
{self.ty.storage_scope}"
- )
+ if self.ty.scope not in tuple(MemoryScope(scope) for scope in scopes):
+ raise ValueError(f"{name}: expected storage scope {scopes}, got
{self.ty.scope}")
shape = tuple(int(ctx.analyzer.simplify(r.extent)) for r in
region.region)
if any(n <= 0 for n in shape):
raise ValueError(f"{name}: pattern requires positive static region
extents")
@@ -278,7 +276,7 @@ class Region:
def physical(self, atom, repeat=None):
result = dict(self.layout.apply(*self.indices(atom, repeat),
shape=self.ty.shape))
- if self.ty.storage_scope != MemoryScope("tmem"):
+ if self.ty.scope != MemoryScope("tmem"):
result["m"] = result.get("m", 0) + self.ty.elem_offset
return {key: self.ctx.analyzer.simplify(value) for key, value in
result.items()}
@@ -356,7 +354,7 @@ class Region:
span += (extent - 1) * stride
def value(self, atom):
- if self.ty.storage_scope != MemoryScope("local"):
+ if self.ty.scope != MemoryScope("local"):
raise ValueError(f"{self.name}: scalar values must be in
registers")
if self.raw is None:
span = self.layout.span("m") + self.ty.elem_offset
diff --git a/python/tvm/backend/cuda/tile/patterns/tcgen05_mma.py
b/python/tvm/backend/cuda/tile/patterns/tcgen05_mma.py
index 36ea0769c7..1828236b43 100644
--- a/python/tvm/backend/cuda/tile/patterns/tcgen05_mma.py
+++ b/python/tvm/backend/cuda/tile/patterns/tcgen05_mma.py
@@ -169,7 +169,7 @@ def lower(opcode, entry, modifiers, args, params, sctx):
raise ValueError("tcgen05.mma D requires M/N repeat axes before the
atom")
batch, mr, nr = d_shape[:-4], d_shape[-4], d_shape[-3]
nb = len(batch)
- a_tmem = args[1].source.ty.storage_scope == MemoryScope("tmem")
+ a_tmem = args[1].source.ty.scope == MemoryScope("tmem")
a_atom = (2, M // cta_group, K) if a_tmem and ws and M == 64 else (M //
cta_group, K)
a_shape = tuple(int(r.extent) for r in args[1].region)[ng:]
kr = a_shape[-len(a_atom) - 1]
diff --git a/python/tvm/backend/cuda/tile/patterns/tma.py
b/python/tvm/backend/cuda/tile/patterns/tma.py
index 3b13e4daec..c2deabedcd 100644
--- a/python/tvm/backend/cuda/tile/patterns/tma.py
+++ b/python/tvm/backend/cuda/tile/patterns/tma.py
@@ -87,9 +87,7 @@ def _descriptor_fields(fields, rank, box, default_dtype):
def _tensor_map(ctx, region, rank, box, encoding):
- if not isinstance(region, TensorRegion) or region.source.ty.storage_scope
!= MemoryScope(
- "global"
- ):
+ if not isinstance(region, TensorRegion) or region.source.ty.scope !=
MemoryScope("global"):
raise ValueError("TMA descriptor requires a global tensor region")
tensor = region.source
ty = tensor.ty
diff --git a/python/tvm/backend/trn/transform/naive_allocator.py
b/python/tvm/backend/trn/transform/naive_allocator.py
index d57fed7ea4..e7e135059b 100644
--- a/python/tvm/backend/trn/transform/naive_allocator.py
+++ b/python/tvm/backend/trn/transform/naive_allocator.py
@@ -74,7 +74,7 @@ def _get_alloc_pool_start(stmt) -> int:
tensor_type = op.value.ty_args[0]
shape = tensor_type.shape
dtype = tensor_type.dtype.dtype
- scope = tensor_type.storage_scope
+ scope = tensor_type.scope
tensor_size = get_tensor_size(tensor, shape, dtype, scope)
if tensor_size is None:
return
@@ -95,7 +95,7 @@ def _allocate_missing_tensors(stmt, alloc_pool_start: int):
tensor_type = op.value.ty_args[0]
shape = tensor_type.shape
dtype = tensor_type.dtype.dtype
- scope = tensor_type.storage_scope
+ scope = tensor_type.scope
tensor_size = get_tensor_size(tensor, shape, dtype, scope)
allocated_addr = op.value.args[0].fields if op.value.args else []
if len(allocated_addr) == 0 and tensor_size is not None:
diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py
index 8c74745d3d..e9e44307e5 100644
--- a/python/tvm/ir/type.py
+++ b/python/tvm/ir/type.py
@@ -189,14 +189,12 @@ class PtrType(Type):
element_type : tvm.ir.Type
The type of pointer's element.
- storage_scope : str or tvm.runtime.MemoryScope
+ scope : str or tvm.runtime.MemoryScope
The storage scope into which the pointer addresses.
"""
- def __init__(self, element_type, storage_scope=MemoryScope("global"), *,
loc=UNKNOWN_LOC):
- self.__init_handle_by_constructor__(
- _ffi_api.PtrType, _normalize_type(element_type), storage_scope
- )
+ def __init__(self, element_type, scope=MemoryScope("global"), *,
loc=UNKNOWN_LOC):
+ self.__init_handle_by_constructor__(_ffi_api.PtrType,
_normalize_type(element_type), scope)
if loc is not UNKNOWN_LOC:
# Keep the script constructor's source-context behavior on the
class alias.
from tvm.script.ir_builder.base import at # pylint:
disable=import-outside-toplevel
diff --git a/python/tvm/script/parser/prescan.py
b/python/tvm/script/parser/prescan.py
index d64c36c7b0..1d7b43f8dd 100644
--- a/python/tvm/script/parser/prescan.py
+++ b/python/tvm/script/parser/prescan.py
@@ -26,8 +26,6 @@ from enum import IntEnum, auto
from types import ModuleType
from typing import NamedTuple, NoReturn
-from tvm.ir.type import PrimType, _PrimTypeArrayAnnotation
-
from . import protocol_registry
@@ -135,32 +133,6 @@ def resolve_namespace_value(node: ast.AST | None,
environment: Mapping[str, obje
return value
-def _fixed_annotation_value(node: ast.AST | None, environment: Mapping[str,
object]):
- """Read captured values and module dictionaries without running user
expressions."""
- if isinstance(node, ast.Name):
- return environment.get(node.id)
- if isinstance(node, ast.Attribute):
- owner = _fixed_annotation_value(node.value, environment)
- if type(owner) is ModuleType:
- return vars(owner).get(node.attr)
- return None
-
-
-def _primitive_annotation(node: ast.AST | None, environment: Mapping[str,
object]):
- """Recognize primitive annotations without calling constructors or
indexing shapes."""
- value = _fixed_annotation_value(node, environment)
- if type(value) in (PrimType, _PrimTypeArrayAnnotation):
- return value
- if isinstance(node, ast.Call) and _fixed_annotation_value(node.func,
environment) is PrimType:
- return PrimType
- if (
- isinstance(node, ast.Subscript)
- and _primitive_annotation(node.value, environment) is not None
- ):
- return _PrimTypeArrayAnnotation
- return None
-
-
# Design principle: Once builtin names and language namespaces are selected
# for script syntax, they cannot be shadowed. Prescan enforces this rule and
# reports clear errors at the offending bindings. This lets the transpiler
@@ -222,7 +194,6 @@ class PrescanContext:
namespaces: dict[str, str],
with_outputs: dict[ast.With, list[str]],
recursive_functions: set[ast.FunctionDef | ast.AsyncFunctionDef],
- primitive_dtypes: dict[ast.AST, str],
) -> None:
# Collected source/entry names seed the allocator; this set stays
read-only.
self.reserved_names = reserved_names
@@ -251,13 +222,6 @@ class PrescanContext:
# Collected self-reference facts select standalone
declaration-before-body
# lowering; rewriting does not add or remove functions from this set.
self.recursive_functions = recursive_functions
- self.primitive_dtypes = primitive_dtypes
-
- def scalar_annotation_dtype(self, node: ast.AST | None):
- """Return the fixed dtype of an annotation without evaluating its
source."""
- return self.primitive_dtypes.get(
- node,
protocol_registry.SCALAR_ANNOTATION_DTYPE.get(self._match_special_func(node))
- )
def _match_special_func(self, node: ast.AST | None) -> str | None:
"""Match direct source syntax against the roots fixed by this
prescan."""
@@ -287,7 +251,6 @@ class PrescanCollector(ast.NodeVisitor):
# Lexical scope and language variant restore on function/class exit.
self.scope: ast.Module | ast.ClassDef | ast.FunctionDef |
ast.AsyncFunctionDef | None = None
self.builder: object = None
- self.mutable_scopes: set[ast.AST] = set()
def collect(self, tree: ast.Module) -> PrescanContext:
"""Collect reserved names, scoped declarations and region-result
syntax.
@@ -327,8 +290,6 @@ class PrescanCollector(ast.NodeVisitor):
self.scope = tree
self.bindings[tree] = []
self.visit(tree)
- primitive_dtypes = {}
- self._resolve_primitive_annotations(tree, tree, self.environment,
primitive_dtypes)
# Transfer the completed collections directly. Consumers keep the
facts, not
# this collector, and do not mutate the collections during AST
rewriting.
return PrescanContext(
@@ -339,69 +300,8 @@ class PrescanCollector(ast.NodeVisitor):
self.namespaces,
self.exports,
self.recursive,
- primitive_dtypes,
)
- def _replace_binding(self, scope, target, kind):
- previous = self.sites[target]
- replacement = Binding(previous.name, target, kind)
- self.sites[target] = replacement
- bindings = self.bindings[scope]
- bindings[bindings.index(previous)] = replacement
-
- @staticmethod
- def _record_primitive_dtype(annotation, environment, dtypes):
- value = _fixed_annotation_value(annotation, environment)
- if type(value) is PrimType:
- dtypes[annotation] = str(value.dtype)
- elif (
- isinstance(annotation, ast.Call)
- and _fixed_annotation_value(annotation.func, environment) is
PrimType
- and len(annotation.args) == 1
- and not annotation.keywords
- and isinstance(annotation.args[0], ast.Constant)
- and isinstance(annotation.args[0].value, str)
- ):
- dtypes[annotation] = annotation.args[0].value
-
- def _resolve_primitive_annotations(self, node, scope, environment, dtypes):
- """Refine annotation facts after all lexical bindings, including later
writes, exist."""
- if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
- parameters = [*node.args.posonlyargs, *node.args.args,
*node.args.kwonlyargs]
- header_names = {parameter.name for parameter in getattr(node,
"type_params", ())}
- for parameter in getattr(node, "type_params", ()):
- self._record_primitive_dtype(getattr(parameter, "bound",
None), environment, dtypes)
- earlier_reads = set()
- for parameter in parameters:
- annotation = parameter.annotation
- self._record_primitive_dtype(annotation, environment, dtypes)
- if (
- self.sites[parameter].kind == BindingKind.PARAMETER
- and parameter.arg in earlier_reads
- and parameter.arg not in header_names
- and annotation in dtypes
- ):
- self._replace_binding(node, parameter,
BindingKind.FORWARD_PARAMETER)
- if annotation is not None:
- earlier_reads.update(
- read.id for read in
collect_annotation_free_reads(annotation)
- )
- # Header annotations capture their definition scope; body bindings only
- # shadow captures when descending into the function's body.
- if node in self.bindings:
- scope = node
- shadowed = {item.name for item in self.bindings[node]}
- environment = {key: value for key, value in environment.items() if
key not in shadowed}
- if (
- isinstance(node, ast.AnnAssign)
- and isinstance(node.target, ast.Name)
- and scope in self.mutable_scopes
- and _primitive_annotation(node.annotation, environment) is not None
- ):
- self._replace_binding(scope, node.target, BindingKind.MUTABLE)
- for child in ast.iter_child_nodes(node):
- self._resolve_primitive_annotations(child, scope, environment,
dtypes)
-
def _raise_error(self, node: ast.AST, message: str) -> NoReturn:
raise SyntaxError(
message,
@@ -598,8 +498,6 @@ class PrescanCollector(ast.NodeVisitor):
if namespace is not None:
self.builder = namespace
break
- if getattr(self.builder, "supports_mutable_declarations", True):
- self.mutable_scopes.add(node)
for parameter in getattr(node, "type_params", ()):
# -------------------- Pattern --------------------
# Python source:
diff --git a/python/tvm/script/parser/transpile.py
b/python/tvm/script/parser/transpile.py
index 596051ae9f..d0d1b69cba 100644
--- a/python/tvm/script/parser/transpile.py
+++ b/python/tvm/script/parser/transpile.py
@@ -1698,14 +1698,15 @@ class IRBuilderTranspiler(ast.NodeTransformer):
dtype = (
"int64"
if bound is None or (isinstance(bound, ast.Name) and bound.id
== "int")
- else self.module.prescan_ctx.scalar_annotation_dtype(bound)
+ else protocol_registry.SCALAR_ANNOTATION_DTYPE.get(
+ self.module.prescan_ctx._match_special_func(bound)
+ )
)
- if bound is not None and not (isinstance(bound, ast.Name) and
bound.id == "int"):
- if self.module.prescan_ctx.scalar_annotation_dtype(bound) is
None:
- self._raise_error(
- parameter,
- "A symbolic type parameter bound must be int or a
registered scalar dtype",
- )
+ if dtype is None:
+ self._raise_error(
+ parameter,
+ "A symbolic type parameter bound must be int or a
registered scalar dtype",
+ )
if getattr(parameter, "default_value", None) is not None:
self._raise_error(parameter, "A symbolic type parameter cannot
have a default")
alias = self.module.make_fresh_name("_symbol")
@@ -1781,7 +1782,9 @@ class IRBuilderTranspiler(ast.NodeTransformer):
self.function.capture_names.add(name)
alias = self.module.make_fresh_name("_parameter")
forward[name] = alias
- dtype = self.module.prescan_ctx.scalar_annotation_dtype(annotation)
+ dtype = protocol_registry.SCALAR_ANNOTATION_DTYPE.get(
+ self.module.prescan_ctx._match_special_func(annotation)
+ )
created = self._call(
self.module.ir_prefix,
"dynamic",
@@ -1848,7 +1851,10 @@ class IRBuilderTranspiler(ast.NodeTransformer):
translated = ast.Name(forward[name], ast.Load())
elif (
name in symbols
- and
self.module.prescan_ctx.scalar_annotation_dtype(annotation) is not None
+ and protocol_registry.SCALAR_ANNOTATION_DTYPE.get(
+ self.module.prescan_ctx._match_special_func(annotation)
+ )
+ is not None
):
translated = ast.Name(name, ast.Load())
else:
diff --git a/python/tvm/tirx/tensor.py b/python/tvm/tirx/tensor.py
index 6b2ef019c2..c68746aba1 100644
--- a/python/tvm/tirx/tensor.py
+++ b/python/tvm/tirx/tensor.py
@@ -75,7 +75,6 @@ class TensorType(Type):
data=None,
byte_offset=None,
loc: Location = UNKNOWN_LOC,
- storage_scope=None,
data_alignment=None,
):
"""Construct a tensor access contract.
@@ -83,9 +82,8 @@ class TensorType(Type):
An omitted element offset is zero in the configured index dtype,
independent of shape and offset_factor. Integer offsets use that dtype;
explicitly typed expressions retain theirs. Pass a symbolic offset
explicitly
- when the contract permits a variable starting position.
``storage_scope``
- and ``align`` remain compatibility spellings for ``scope`` and
- ``data_alignment``.
+ when the contract permits a variable starting position. ``align``
remains a
+ compatibility spelling for ``data_alignment``.
"""
from tvm.ir.type import _normalize_type
from tvm.script.ir_builder.base import MISSING, at
@@ -101,12 +99,6 @@ class TensorType(Type):
raise ValueError("elem_offset and byte_offset cannot be set at the
same time")
if byte_offset is not None:
elem_offset = byte_offset * 8 // tvm.DataType(dtype.dtype).bits
- if storage_scope is None:
- storage_scope = scope
- elif MemoryScope(scope) != _SCOPE_GLOBAL and MemoryScope(scope) !=
MemoryScope(
- storage_scope
- ):
- raise ValueError("scope and storage_scope must agree")
if data_alignment is None:
data_alignment = align
elif align not in (0, data_alignment):
@@ -114,7 +106,7 @@ class TensorType(Type):
if data is not None:
if not isinstance(data, tvm.ir.Expr) or not isinstance(data.ty,
PtrType):
raise TypeError("Tensor data must be an Expr with PtrType")
- storage_scope = data.ty.storage_scope
+ scope = data.ty.scope
if layout is MISSING:
layout = None
if layout == "default":
@@ -125,7 +117,7 @@ class TensorType(Type):
_ffi_api.TensorType,
shape,
dtype,
- storage_scope,
+ scope,
() if strides is None else strides,
elem_offset,
data_alignment,
@@ -166,8 +158,8 @@ class TensorType(Type):
"""
if self.layout is not None:
return self.layout
- if self.storage_scope != _SCOPE_GLOBAL:
- raise ValueError(f"An explicit layout is required for
{self.storage_scope} tensors")
+ if self.scope != _SCOPE_GLOBAL:
+ raise ValueError(f"An explicit layout is required for {self.scope}
tensors")
from tvm.sym import Analyzer
@@ -184,11 +176,6 @@ class TensorType(Type):
expected = analyzer.simplify(expected * extent)
return TileLayout(S[tuple(self.shape)])
- @property
- def storage_scope(self):
- """Compatibility alias for the canonical scope field."""
- return self.scope
-
def _scope(self, expr):
"""Return the storage scope associated with this tensor.
Returns
@@ -348,17 +335,16 @@ def decl_tensor(
if layout == "default":
layout = TileLayout(S[tuple(shape)]) if shape else None
- storage_scope = scope
if data is not None:
if not isinstance(data, tvm.ir.Expr) or not isinstance(data.ty,
PtrType):
raise TypeError("Tensor data must be an Expr with PtrType")
if not isinstance(data.ty.element_type, PrimType):
raise TypeError("Tensor data must point to a primitive type")
- storage_scope = data.ty.storage_scope
+ scope = data.ty.scope
tensor_type = _ffi_api.TensorType( # type: ignore
shape,
dtype,
- storage_scope,
+ scope,
strides,
elem_offset,
data_alignment,
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index fa7f7361c0..c3db94bce1 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1379,8 +1379,7 @@ void CodeGenCUDA::DispatchAllocTensor(const BindNode* op,
const CallNode* buffer
static const auto scope_shared = runtime::MemoryScope::Get("shared");
tvm::Tuple shape =
tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
DLDataType dtype =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
- buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.value().as_or_throw<DictAttrs>();
TVM_FFI_ICHECK(buffer.defined());
diff --git a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
index 3a1f8f8766..b395a8dc0d 100644
--- a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
+++ b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
@@ -97,8 +97,7 @@ class CodeGenNVPTX : public CodeGenLLVM {
static const auto scope_shared = runtime::MemoryScope::Get("shared");
tvm::Tuple shape =
tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
DLDataType dtype =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
- buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations =
buffer_call->attrs.value().as_or_throw<DictAttrs>();
llvm::Value* buf = nullptr;
diff --git a/src/backend/cuda/op/target_builtin.cc
b/src/backend/cuda/op/target_builtin.cc
index bd824c1f12..d0e6e19901 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -55,8 +55,7 @@ Type InferTypeDeclTmem(const CallNode* call) {
TVM_FFI_CHECK_EQ(call->ty_args.size(), 1U, ValueError)
<< "decl_tmem expects exactly one TensorType argument";
auto type = call->ty_args[0].as_or_throw<TensorType>();
- TVM_FFI_CHECK(type->storage_scope == scope_tmem, ValueError)
- << "decl_tmem requires a tmem tensor type";
+ TVM_FFI_CHECK(type->scope == scope_tmem, ValueError) << "decl_tmem requires
a tmem tensor type";
auto addr_type = call->args[0]->ty.as_or_throw<PrimType>();
TVM_FFI_CHECK(addr_type.IsScalar() &&
addr_type.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt) &&
diff --git a/src/backend/metal/codegen/codegen_metal.cc
b/src/backend/metal/codegen/codegen_metal.cc
index 321530153a..003e737e80 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -361,7 +361,7 @@ void CodeGenMetal::Dispatch_(const BindNode* op) {
return CodeGenC::Dispatch_(op);
}
- const runtime::MemoryScope& storage_scope = pointer_type->storage_scope;
+ const runtime::MemoryScope& storage_scope = pointer_type->scope;
alloc_storage_scope_[op->var.get()] = storage_scope;
RegisterHandleTypeFromPointer(op->var, &op->value);
std::string value = PrintExpr(op->value);
@@ -384,8 +384,7 @@ void CodeGenMetal::DispatchAllocTensor(const BindNode* op,
const CallNode* buffe
static const auto scope_metal_simdgroup =
runtime::MemoryScope::Get("metal.simdgroup");
tvm::Tuple shape =
tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
DLDataType dtype =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
- buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.value().as_or_throw<DictAttrs>();
TVM_FFI_ICHECK(buffer.defined());
@@ -530,10 +529,10 @@ void CodeGenMetal::Dispatch_(const CallNode* op,
std::ostream& os) { // NOLINT(
<< "Metal pointer byte offsets must have a pointer result type, but
got " << op->ty;
os << "((";
- PrintStorageScope(pointer_type->storage_scope, os);
+ PrintStorageScope(pointer_type->scope, os);
PrintType(pointer_type->element_type, os);
os << "*)(((";
- PrintStorageScope(pointer_type->storage_scope, os);
+ PrintStorageScope(pointer_type->scope, os);
os << "char*)";
PrintExpr(op->args[0], os);
os << ") + ";
diff --git a/src/backend/opencl/codegen/codegen_opencl.cc
b/src/backend/opencl/codegen/codegen_opencl.cc
index 12e401b7ee..a6e3c4f3cf 100644
--- a/src/backend/opencl/codegen/codegen_opencl.cc
+++ b/src/backend/opencl/codegen/codegen_opencl.cc
@@ -69,7 +69,7 @@ TextureArgument UnwrapTextureArgument(const Expr& texture) {
<< "Texture arguments must be a pointer Var or a
tensor_data_ptr(TensorVar) projection";
const auto* pointer_type = texture->ty.as<PtrTypeNode>();
TVM_FFI_ICHECK(pointer_type) << "Texture arguments must have PtrType";
- TVM_FFI_ICHECK(runtime::IsTextureStorage(pointer_type->storage_scope))
+ TVM_FFI_ICHECK(runtime::IsTextureStorage(pointer_type->scope))
<< "Texture intrinsics only support texture buffers";
return {var, pointer_type};
}
@@ -137,7 +137,7 @@ void CodeGenOpenCL::InitFuncState(const Function& f) {
}
for (Var arg : f->params) {
auto ptr_type = arg->ty.as<PtrTypeNode>();
- if (ptr_type && runtime::IsTextureStorage(ptr_type->storage_scope)) {
+ if (ptr_type && runtime::IsTextureStorage(ptr_type->scope)) {
// Storage scope qualifiers for textures are inferred
// and set prior to function codegen.
continue;
@@ -153,7 +153,7 @@ void CodeGenOpenCL::PrintFuncPrefix(std::ostream& os) { os
<< "__kernel "; }
void CodeGenOpenCL::PreFunctionBody(const Function& f) {
for (Var arg : f->params) {
auto ptr_type = arg->ty.as<PtrTypeNode>();
- if (ptr_type && runtime::IsTextureStorage(ptr_type->storage_scope)) {
+ if (ptr_type && runtime::IsTextureStorage(ptr_type->scope)) {
this->stream << " const sampler_t image_sampler = "
"CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP |
CLK_FILTER_NEAREST;\n";
return;
@@ -330,7 +330,7 @@ void CodeGenOpenCL::PrintType(const Type& type,
std::ostream& os) { // NOLINT(*
if (auto* ptr = type.as<PrimTypeNode>()) {
return PrintType(ffi::GetRef<PrimType>(ptr), os);
} else if (auto* ptr = type.as<PtrTypeNode>()) {
- if (runtime::IsTextureStorage(ptr->storage_scope)) {
+ if (runtime::IsTextureStorage(ptr->scope)) {
os << "image2d_array_t";
} else {
PrintType(ptr->element_type, os);
@@ -436,7 +436,7 @@ void CodeGenOpenCL::PrintStorageScope(const
runtime::MemoryScope& scope,
void CodeGenOpenCL::PrintRestrict(const Var& v, std::ostream& os) {
// Apply restrict qualifer for non-texture types only
if (auto* ptr = v->ty.as<PtrTypeNode>()) {
- if (!runtime::IsTextureStorage(ptr->storage_scope)) {
+ if (!runtime::IsTextureStorage(ptr->scope)) {
os << ' ' << restrict_keyword_;
}
}
diff --git a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
index a6563f974a..d4a8bafe43 100644
--- a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
+++ b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
@@ -111,8 +111,7 @@ class CodeGenAMDGPU : public CodeGenLLVM {
static const auto scope_shared = runtime::MemoryScope::Get("shared");
tvm::Tuple shape =
tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
DLDataType dtype =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
- buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations =
buffer_call->attrs.value().as_or_throw<DictAttrs>();
llvm::Value* buf = nullptr;
diff --git a/src/backend/trn/codegen/codegen_trn.cc
b/src/backend/trn/codegen/codegen_trn.cc
index bc437c60d2..dc41fc230a 100644
--- a/src/backend/trn/codegen/codegen_trn.cc
+++ b/src/backend/trn/codegen/codegen_trn.cc
@@ -234,8 +234,7 @@ void CodeGenTrainium::DispatchAllocTensor(const BindNode*
op, const CallNode* bu
static const auto scope_trn_psum = runtime::MemoryScope::Get("trn.psum");
tvm::Tuple shape =
tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
DLDataType dtype =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
- buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
TVM_FFI_ICHECK(buffer.defined());
std::string vid = AllocVarID(buffer.get(), buffer.name() + "_ptr");
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index e413da5fe1..7e6b8bde72 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -928,8 +928,7 @@ void CodeGenSPIRV::DispatchAllocTensor(const BindNode* op,
const CallNode* buffe
static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn");
tvm::Tuple shape =
tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
DLDataType dtype =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
- buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.value().as_or_throw<DictAttrs>();
TVM_FFI_ICHECK(!PrimType(dtype).IsVoid());
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index e0d3f5c818..d2441addaf 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -787,8 +787,7 @@ void CodeGenWebGPU::DispatchAllocTensor(const BindNode* op,
const CallNode* buff
static const auto scope_local = runtime::MemoryScope::Get("local");
tvm::Tuple shape =
tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
DLDataType dtype =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
- buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
TVM_FFI_ICHECK(buffer.defined());
std::string vid = AllocVarID(buffer.get());
diff --git a/src/ir/type.cc b/src/ir/type.cc
index 149925b558..38315311ea 100644
--- a/src/ir/type.cc
+++ b/src/ir/type.cc
@@ -125,7 +125,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
PrimTypeMaybeInplaceMut
TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> PtrTypeVisit(
ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
- // skips: storage_scope (canonical enum identity)
+ // skips: scope (canonical enum identity)
const PtrTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
PtrTypeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->element_type));
@@ -134,7 +134,7 @@ TVM_FFI_INLINE
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> PtrTypeVisit(
TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> PtrTypeMutate(
ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- // skips: storage_scope (canonical enum identity)
+ // skips: scope (canonical enum identity)
const PtrTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
PtrTypeNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>,
mapped_element_type_u,
@@ -150,7 +150,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
PtrTypeMutate(
TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
PtrTypeMaybeInplaceMutate(
ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- // skips: storage_scope (canonical enum identity)
+ // skips: scope (canonical enum identity)
PtrTypeNode* self = const_cast<PtrTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
PtrTypeNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
@@ -414,11 +414,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
// PtrType
-PtrType::PtrType(Type element_type, runtime::MemoryScope storage_scope) :
Type(ffi::UnsafeInit{}) {
+PtrType::PtrType(Type element_type, runtime::MemoryScope scope) :
Type(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(!element_type.as<MissingType>().has_value())
<< "PtrType element_type cannot be Type::Missing()";
ffi::ObjectPtr<PtrTypeNode> n = ffi::make_object<PtrTypeNode>();
- n->storage_scope = std::move(storage_scope);
+ n->scope = std::move(scope);
n->element_type = std::move(element_type);
data_ = std::move(n);
}
@@ -433,14 +433,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.attr(refl::type_attr::kStructuralMaybeInplaceMutate,
ffi::FStructuralMutate::FromNative<&PtrTypeMaybeInplaceMutate>());
- refl::GlobalDef().def("ir.PtrType",
- [](Type element_type, runtime::MemoryScope
storage_scope = {}) {
- return PtrType(element_type, storage_scope);
- });
+ refl::GlobalDef().def("ir.PtrType", [](Type element_type,
runtime::MemoryScope scope = {}) {
+ return PtrType(element_type, scope);
+ });
}
-PtrType PtrType::VoidPtrType(runtime::MemoryScope storage_scope) {
- return PtrType(PrimType::Void(), std::move(storage_scope));
+PtrType PtrType::VoidPtrType(runtime::MemoryScope scope) {
+ return PtrType(PrimType::Void(), std::move(scope));
}
FuncType::FuncType(tvm::ffi::Array<Type> arg_types, Type ret_type, Location
loc)
diff --git a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
index 35fd63f410..77af38f10c 100644
--- a/src/relax/distributed/transform/lower_global_view_to_local_view.cc
+++ b/src/relax/distributed/transform/lower_global_view_to_local_view.cc
@@ -260,7 +260,7 @@ class DistributedTensorCompactor : public
s_tir::StmtExprMutator {
shape.push_back(tensor->shape[i]);
}
}
- TensorType new_type(std::move(shape), tensor->dtype,
tensor->storage_scope, tensor->strides,
+ TensorType new_type(std::move(shape), tensor->dtype, tensor->scope,
tensor->strides,
tensor->elem_offset, tensor->data_alignment,
tensor->offset_factor,
tensor->layout);
return TensorVar(tensor.name(), std::move(new_type), tensor.loc());
diff --git a/src/relax/op/op.cc b/src/relax/op/op.cc
index 20db35f24d..0039e39c5a 100644
--- a/src/relax/op/op.cc
+++ b/src/relax/op/op.cc
@@ -392,7 +392,7 @@ Type TIRxPackedValueType(const Type& type, bool is_result) {
TVM_FFI_CHECK(default_layout, TypeError)
<< "A Relax-to-TIRx call cannot implicitly convert tensor layout or
allocated address: "
<< type;
- TVM_FFI_CHECK(IsGlobalMemoryScope(tensor->storage_scope), TypeError)
+ TVM_FFI_CHECK(IsGlobalMemoryScope(tensor->scope), TypeError)
<< "A Relax-to-TIRx call cannot implicitly convert tensor storage
scope: " << type;
ShapeExpr shape(
tensor->shape.Map([](PrimExpr dim) { return
prim::cast(PrimType::Int(64), dim); }));
@@ -476,7 +476,7 @@ void CheckTIRxCarrier(const Type& native, const Type&
actual,
<< actual;
}
if (const auto* tensor = native.as<tirx::TensorTypeNode>()) {
- auto native_scope = tensor->storage_scope;
+ auto native_scope = tensor->scope;
const auto* argument = actual.as<TensorTypeNode>();
ffi::Optional<runtime::MemoryScope> actual_scope;
if (argument && argument->vdevice.has_value()) {
diff --git a/src/relax/script/printer/call.cc b/src/relax/script/printer/call.cc
index 4f8d16ca54..3106f5bf61 100644
--- a/src/relax/script/printer/call.cc
+++ b/src/relax/script/printer/call.cc
@@ -132,7 +132,7 @@ ffi::Optional<ExprDoc>
CallDPSPackedDocTranslate(DocTranslatorObj* d, ffi::AnyVi
callee = LiteralDoc::Str(packed->global_symbol, std::nullopt);
} else if (call->args[0].as<StringImmNode>()) {
// A bare Python string now constructs ExternFunc, so retain this explicit
IR value.
- callee = d->WithExtraStateScope<bool>("ir.explicit_scalar_dtype", true,
+ callee = d->WithExtraStateScope<bool>("ir.explicit_value_type", true,
[&] { return
d->Translate(call->args[0]).value(); });
}
d->RecordOrigin(callee, call->args[0]);
diff --git a/src/relax/script/printer/function.cc
b/src/relax/script/printer/function.cc
index 7e262e4451..a4a0824614 100644
--- a/src/relax/script/printer/function.cc
+++ b/src/relax/script/printer/function.cc
@@ -97,7 +97,7 @@ ffi::Optional<ExprDoc> FunctionDocTranslate(DocTranslatorObj*
d, ffi::AnyView in
for (const auto& [key, value] : attrs) {
keys.push_back(LiteralDoc::Str(key, std::nullopt));
values.push_back(d->WithExtraStateScope<bool>(
- "ir.explicit_scalar_dtype", true, [&] { return
d->Translate(value).value(); }));
+ "ir.explicit_value_type", true, [&] { return
d->Translate(value).value(); }));
}
d->Emit(
ExprStmtDoc(NamespaceDoc("relax")->Attr("func_attr")->Call({DictDoc(keys,
values)})),
diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index cbddbd5e3b..1b17ea4b2c 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -797,9 +797,9 @@ class FusedTIRConstructor : public ExprVisitor {
return unique_name;
};
// Update tensor with new symbolic shape according to the ty
- tirx::TensorType new_type(output_shapes[i], tensor->dtype,
tensor->storage_scope,
- tensor->strides, tensor->elem_offset,
tensor->data_alignment,
- tensor->offset_factor, tensor->layout);
+ tirx::TensorType new_type(output_shapes[i], tensor->dtype,
tensor->scope, tensor->strides,
+ tensor->elem_offset, tensor->data_alignment,
tensor->offset_factor,
+ tensor->layout);
tirx::TensorVar new_tensor(unify_name_hints(), std::move(new_type),
tensor.loc());
func_info_.alloc_tensors.push_back(new_tensor);
output_tensors.push_back(new_tensor);
diff --git a/src/s_tir/analysis/calculate_allocated_memory.cc
b/src/s_tir/analysis/calculate_allocated_memory.cc
index 88d6dc39e1..0efc9924cb 100644
--- a/src/s_tir/analysis/calculate_allocated_memory.cc
+++ b/src/s_tir/analysis/calculate_allocated_memory.cc
@@ -47,7 +47,7 @@ using namespace tvm::tirx;
runtime::MemoryScope GetStorageScope(const Var& var) {
auto* ptr = var->ty.as<PtrTypeNode>();
TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PtrType";
- return ptr->storage_scope;
+ return ptr->scope;
}
/*!
@@ -78,7 +78,7 @@ class AllocTensorCalculator : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> DispatchAllocTensor(const BindNode* op, const
CallNode* call) {
tvm::Tuple shape =
tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
DLDataType dtype =
call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->scope;
auto search = _current_size.find(scope);
if (search == _current_size.end()) {
_current_size[scope] = 0;
diff --git a/src/s_tir/analysis/verify_gpu_code.cc
b/src/s_tir/analysis/verify_gpu_code.cc
index 2a640ce984..f233a0a226 100644
--- a/src/s_tir/analysis/verify_gpu_code.cc
+++ b/src/s_tir/analysis/verify_gpu_code.cc
@@ -86,7 +86,7 @@ class GPUCodeVerifier : public StmtExprVisitor {
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
tvm::Tuple shape =
tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
DLDataType dtype =
call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->scope;
runtime::MemoryScope storage_scope = scope;
int64_t const_size = 1;
for (const Expr& e : shape->fields) {
diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc
b/src/s_tir/backend/adreno/inject_texture_alloc.cc
index b1c6ece188..2b2edcc7a3 100644
--- a/src/s_tir/backend/adreno/inject_texture_alloc.cc
+++ b/src/s_tir/backend/adreno/inject_texture_alloc.cc
@@ -75,7 +75,7 @@ class TextureAllocInjector : public
s_tir::StmtExprMutatorWithAnalyzer {
UnchangedOr<Stmt> Mutate_AllocTensor(const BindNode* op, const CallNode*
call,
InplaceMode inplace_mode) {
Stmt stmt = StmtExprMutator::Mutate_(op,
inplace_mode).ValueOrUnchanged(ffi::GetRef<Stmt>(op));
- runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->scope;
if (IsTextureStorage(scope)) {
op = stmt.as<BindNode>();
if (const auto* call = op ? op->value.as<CallNode>() : nullptr;
@@ -116,7 +116,7 @@ class TextureAllocInjector : public
s_tir::StmtExprMutatorWithAnalyzer {
runtime::MemoryScope GetStorageScope(const Var& buffer_var) {
auto* ptr = buffer_var->ty.as<PtrTypeNode>();
TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PtrType";
- return ptr->storage_scope;
+ return ptr->scope;
}
};
diff --git a/src/s_tir/backend/adreno/texture_flatten.cc
b/src/s_tir/backend/adreno/texture_flatten.cc
index a722f41018..d867cae917 100644
--- a/src/s_tir/backend/adreno/texture_flatten.cc
+++ b/src/s_tir/backend/adreno/texture_flatten.cc
@@ -76,7 +76,7 @@ class TextureLoweringBase : public StmtExprMutator {
}
protected:
- runtime::MemoryScope GetStorageScope(const TensorVar& tensor) { return
tensor->storage_scope; }
+ runtime::MemoryScope GetStorageScope(const TensorVar& tensor) { return
tensor->scope; }
// Set of all external input and output tensors
std::unordered_set<TensorVar, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
extern_tensor_;
diff --git a/src/s_tir/schedule/primitive/cache_read_write.cc
b/src/s_tir/schedule/primitive/cache_read_write.cc
index df57869935..f1a2ccadb9 100644
--- a/src/s_tir/schedule/primitive/cache_read_write.cc
+++ b/src/s_tir/schedule/primitive/cache_read_write.cc
@@ -2280,7 +2280,7 @@ void CollectReindexCacheStageInfoAndCreateTensor(
// Create new tensor
ffi::ObjectPtr<TensorTypeNode> new_tensor = CopyTensorType(old_tensor);
- new_tensor->storage_scope = runtime::MemoryScope::Get(storage_scope);
+ new_tensor->scope = runtime::MemoryScope::Get(storage_scope);
new_tensor->shape = new_shape;
TensorVar rebuilt =
RebuildTensorVar(old_tensor, std::move(new_tensor), old_tensor.name() +
"_" + storage_scope);
diff --git a/src/s_tir/schedule/transform.cc b/src/s_tir/schedule/transform.cc
index f6a090dc9b..f2be42b029 100644
--- a/src/s_tir/schedule/transform.cc
+++ b/src/s_tir/schedule/transform.cc
@@ -52,9 +52,8 @@ TensorVar WithScope(const TensorVar& tensor,
runtime::MemoryScope scope) {
}
TensorVar WithDType(const TensorVar& tensor, PrimType dtype) {
- TensorType new_type(tensor->shape, dtype, tensor->storage_scope,
tensor->strides,
- tensor->elem_offset, tensor->data_alignment,
tensor->offset_factor,
- tensor->layout);
+ TensorType new_type(tensor->shape, dtype, tensor->scope, tensor->strides,
tensor->elem_offset,
+ tensor->data_alignment, tensor->offset_factor,
tensor->layout);
return TensorVar(tensor.name(), new_type, tensor.loc());
}
diff --git a/src/s_tir/script/ir_builder/ir.cc
b/src/s_tir/script/ir_builder/ir.cc
index 7c5312ae59..2034d9042a 100644
--- a/src/s_tir/script/ir_builder/ir.cc
+++ b/src/s_tir/script/ir_builder/ir.cc
@@ -53,7 +53,7 @@ TensorVar MatchTensor(ffi::ObjectRef param,
ffi::Array<PrimExpr> shape, PrimType
int align, int offset_factor, ffi::Optional<Layout>
layout,
ffi::Array<PrimExpr> allocated_addr) {
if (data.has_value()) {
- storage_scope = data.value()->ty.as_or_throw<PtrType>()->storage_scope;
+ storage_scope = data.value()->ty.as_or_throw<PtrType>()->scope;
}
// Matching/allocation builders retain their binding contract; TensorType
itself
// has a fixed zero offset and never invents a symbolic variable.
diff --git a/src/s_tir/script/printer/stmt.cc b/src/s_tir/script/printer/stmt.cc
index df7988c0ed..a91b763a42 100644
--- a/src/s_tir/script/printer/stmt.cc
+++ b/src/s_tir/script/printer/stmt.cc
@@ -45,9 +45,9 @@ CallDoc TensorHelperCall(DocTranslatorObj* d, const Type&
type, const ffi::Strin
// Supply the fixed zero explicitly when the type printer omitted it.
if (has_factor && !has_offset) {
rhs->kwargs_keys.push_back("elem_offset");
- rhs->kwargs_values.push_back(d->WithExtraStateScope<bool>(
- "ir.explicit_scalar_dtype", true,
- [&] { return
d->Translate(type.as_or_throw<tirx::TensorType>()->elem_offset).value(); }));
+
rhs->kwargs_values.push_back(d->WithExtraStateScope<bool>("ir.explicit_value_type",
true, [&] {
+ return
d->Translate(type.as_or_throw<tirx::TensorType>()->elem_offset).value();
+ }));
}
rhs->callee = NamespaceDoc("s_tir")->Attr(helper);
return rhs;
@@ -148,13 +148,12 @@ ffi::Array<StmtDoc> SBlockBody(DocTranslatorObj* d, const
s_tir::SBlockNode* blo
const auto* tensor_type = tensor.var()->ty.as<tirx::TensorTypeNode>();
auto allocated_addr =
addresses.Get(tensor.var()).value_or(ffi::Array<PrimExpr>{});
TVM_FFI_CHECK(
- allocated_addr.empty() || (tensor_type->storage_scope !=
runtime::MemoryScope::Global() &&
- tensor_type->storage_scope !=
scope_shared &&
- tensor_type->storage_scope !=
scope_shared_dyn &&
- tensor_type->storage_scope !=
scope_local),
+ allocated_addr.empty() ||
+ (tensor_type->scope != runtime::MemoryScope::Global() &&
+ tensor_type->scope != scope_shared && tensor_type->scope !=
scope_shared_dyn &&
+ tensor_type->scope != scope_local),
TypeError)
- << "Ts.sblock_alloc_tensor does not accept allocated_addr in "
- << tensor_type->storage_scope;
+ << "Ts.sblock_alloc_tensor does not accept allocated_addr in " <<
tensor_type->scope;
if (!allocated_addr.empty()) {
rhs->kwargs_keys.push_back("allocated_addr");
rhs->kwargs_values.push_back(d->Translate(allocated_addr).value());
diff --git a/src/s_tir/transform/lower_vtcm_alloc.cc
b/src/s_tir/transform/lower_vtcm_alloc.cc
index 90d8c6891f..99d4c850cb 100644
--- a/src/s_tir/transform/lower_vtcm_alloc.cc
+++ b/src/s_tir/transform/lower_vtcm_alloc.cc
@@ -53,7 +53,7 @@ class VtcmAllocator : public StmtExprMutator {
UnchangedOr<Stmt> Mutate_AllocTensor(const BindNode* op, const CallNode*
call,
InplaceMode inplace_mode) {
- runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->scope;
if (IsVtcmStorage(scope)) {
tvm::Tuple shape =
tvm::Tuple(call->ty_args[0].as_or_throw<TensorType>()->shape);
ffi::Array<Expr> args;
@@ -75,7 +75,7 @@ class VtcmAllocator : public StmtExprMutator {
runtime::MemoryScope GetStorageScope(const Var& var) {
auto* ptr = var->ty.as<PtrTypeNode>();
TVM_FFI_ICHECK(ptr) << "Buffer Var's type annotation must be of PtrType";
- return ptr->storage_scope;
+ return ptr->scope;
}
};
diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc
b/src/s_tir/transform/merge_shared_memory_allocations.cc
index 815f41c536..3864037082 100644
--- a/src/s_tir/transform/merge_shared_memory_allocations.cc
+++ b/src/s_tir/transform/merge_shared_memory_allocations.cc
@@ -85,8 +85,8 @@ bool IsDynamicSharedMemory(Var buffer_var) {
static const auto scope_shared_dyn = runtime::MemoryScope::Get("shared.dyn");
MemoryScope storage_scope =
- (buffer_var->ty.as<TensorTypeNode>() ?
buffer_var->ty.as<TensorTypeNode>()->storage_scope
- :
buffer_var->ty.as_or_throw<PtrType>()->storage_scope);
+ (buffer_var->ty.as<TensorTypeNode>() ?
buffer_var->ty.as<TensorTypeNode>()->scope
+ :
buffer_var->ty.as_or_throw<PtrType>()->scope);
return storage_scope == scope_shared_dyn;
}
@@ -94,8 +94,8 @@ bool IsStaticSharedMemory(Var buffer_var) {
static const auto scope_shared = runtime::MemoryScope::Get("shared");
MemoryScope storage_scope =
- (buffer_var->ty.as<TensorTypeNode>() ?
buffer_var->ty.as<TensorTypeNode>()->storage_scope
- :
buffer_var->ty.as_or_throw<PtrType>()->storage_scope);
+ (buffer_var->ty.as<TensorTypeNode>() ?
buffer_var->ty.as<TensorTypeNode>()->scope
+ :
buffer_var->ty.as_or_throw<PtrType>()->scope);
return storage_scope == scope_shared;
}
@@ -151,7 +151,7 @@ class AllocateCollector : public StmtExprVisitor {
static const auto scope_shared_dyn =
runtime::MemoryScope::Get("shared.dyn");
static const auto scope_shared = runtime::MemoryScope::Get("shared");
- runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->scope;
MemoryScope storage_scope = scope;
if (is_dynamic_ && storage_scope == scope_shared_dyn) {
shmem_allocs_.insert_or_assign(op->var.as_or_throw<TensorVar>().get(),
@@ -560,7 +560,7 @@ class SharedMemoryRewriter : public StmtExprMutator {
static const auto scope_shared_dyn =
runtime::MemoryScope::Get("shared.dyn");
static const auto scope_shared = runtime::MemoryScope::Get("shared");
- runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->scope;
MemoryScope storage_scope = scope;
if (storage_scope == (is_dynamic_ ? scope_shared_dyn : scope_shared)) {
if (!scope_stack_.empty()) {
diff --git a/src/s_tir/transform/storage_access.cc
b/src/s_tir/transform/storage_access.cc
index f0080eec34..4863a04b0a 100644
--- a/src/s_tir/transform/storage_access.cc
+++ b/src/s_tir/transform/storage_access.cc
@@ -409,10 +409,10 @@ ffi::Optional<VisitInterrupt>
StorageAccessVisitor::Visit_(const CallNode* op) {
MemoryScope StorageAccessVisitor::GetScope(Var buffer_var) const {
if (auto tensor_type = buffer_var->ty.as<TensorType>()) {
- return tensor_type.value()->storage_scope;
+ return tensor_type.value()->scope;
}
if (const auto* pointer_type = buffer_var->ty.as<PtrTypeNode>()) {
- return pointer_type->storage_scope;
+ return pointer_type->scope;
}
return MemoryScope(); // global by default
}
diff --git a/src/s_tir/transform/tensorcore_infer_fragment.cc
b/src/s_tir/transform/tensorcore_infer_fragment.cc
index 1ab6c7d570..343e4cdd6e 100644
--- a/src/s_tir/transform/tensorcore_infer_fragment.cc
+++ b/src/s_tir/transform/tensorcore_infer_fragment.cc
@@ -83,9 +83,9 @@ class FragmentGetter : public s_tir::StmtExprVisitor {
TVM_FFI_ICHECK(k);
TVM_FFI_ICHECK(layout);
- runtime::MemoryScope scope = (buffer_var->ty.as<TensorTypeNode>()
- ?
buffer_var->ty.as<TensorTypeNode>()->storage_scope
- :
buffer_var->ty.as_or_throw<PtrType>()->storage_scope);
+ runtime::MemoryScope scope =
+ (buffer_var->ty.as<TensorTypeNode>() ?
buffer_var->ty.as<TensorTypeNode>()->scope
+ :
buffer_var->ty.as_or_throw<PtrType>()->scope);
if (fragments.count(buffer_var)) {
// check if the fragment has met before
FragmentInfo info = fragments[buffer_var];
@@ -120,9 +120,9 @@ class FragmentGetter : public s_tir::StmtExprVisitor {
TVM_FFI_ICHECK(n);
TVM_FFI_ICHECK(k);
- runtime::MemoryScope scope = (buffer_var->ty.as<TensorTypeNode>()
- ?
buffer_var->ty.as<TensorTypeNode>()->storage_scope
- :
buffer_var->ty.as_or_throw<PtrType>()->storage_scope);
+ runtime::MemoryScope scope =
+ (buffer_var->ty.as<TensorTypeNode>() ?
buffer_var->ty.as<TensorTypeNode>()->scope
+ :
buffer_var->ty.as_or_throw<PtrType>()->scope);
if (fragments.count(buffer_var)) {
FragmentInfo info = fragments[buffer_var];
TVM_FFI_ICHECK_EQ(m->value, info.m);
diff --git a/src/script/printer/doc_translator.cc
b/src/script/printer/doc_translator.cc
index 126fa9e208..cb3c48df17 100644
--- a/src/script/printer/doc_translator.cc
+++ b/src/script/printer/doc_translator.cc
@@ -200,7 +200,7 @@ ExprDoc DocTranslatorObj::DefaultTypeTranslate(const Type&
type) {
}
ExprDoc callee = NamespaceDoc(type_key.substr(0, separator))->Attr(name);
// A constructor's fields are values, including nested Types and container
elements.
- return WithExtraStateScope<bool>("ir.explicit_scalar_dtype", true, [&]() ->
ExprDoc {
+ return WithExtraStateScope<bool>("ir.explicit_value_type", true, [&]() ->
ExprDoc {
ffi::Array<ExprDoc> args;
ffi::Array<ffi::String> keys;
ffi::Array<ExprDoc> values;
diff --git a/src/script/printer/ir/call.cc b/src/script/printer/ir/call.cc
index 7304c8b379..9953188975 100644
--- a/src/script/printer/ir/call.cc
+++ b/src/script/printer/ir/call.cc
@@ -122,7 +122,7 @@ ffi::Optional<ExprDoc>
StandardCallDocTranslate(DocTranslatorObj* d, const CallN
return RawCallDocTranslate(d, call);
}
keys.push_back(key);
-
values.push_back(d->WithExtraStateScope<bool>("ir.explicit_scalar_dtype", true,
+ values.push_back(d->WithExtraStateScope<bool>("ir.explicit_value_type",
true,
[&] { return
d->Translate(value).value(); }));
}
}
diff --git a/src/script/printer/ir/ir.cc b/src/script/printer/ir/ir.cc
index 67689f2d7c..06ddd275b8 100644
--- a/src/script/printer/ir/ir.cc
+++ b/src/script/printer/ir/ir.cc
@@ -59,7 +59,7 @@ ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj*
d, ffi::AnyView in
for (const auto& [key, value] : attrs) {
keys.push_back(LiteralDoc::Str(key, std::nullopt));
values.push_back(
-
d->WithExtraStateScope<bool>("ir.explicit_scalar_dtype", true,
+ d->WithExtraStateScope<bool>("ir.explicit_value_type",
true,
[&] { return
d->Translate(value).value(); }));
}
d->Emit(
@@ -78,7 +78,7 @@ ffi::Optional<ExprDoc> IRModuleDocTranslate(DocTranslatorObj*
d, ffi::AnyView in
ffi::Array<ExprDoc> items;
for (const GlobalInfo& entry : entries) {
items.push_back(d->WithExtraStateScope<bool>(
- "ir.explicit_scalar_dtype", true,
+ "ir.explicit_value_type", true,
[&] { return d->Translate(entry).value(); }));
}
keys.push_back(LiteralDoc::Str(key, std::nullopt));
diff --git a/src/script/printer/ir/misc.cc b/src/script/printer/ir/misc.cc
index a4b2366088..1d5d20752c 100644
--- a/src/script/printer/ir/misc.cc
+++ b/src/script/printer/ir/misc.cc
@@ -58,7 +58,7 @@ namespace {
ffi::Optional<ExprDoc> ArrayDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
const ffi::Object*) {
return d->WithExtraStateScope<bool>(
- "ir.explicit_scalar_dtype", true, [&]() -> ffi::Optional<ExprDoc> {
+ "ir.explicit_value_type", true, [&]() -> ffi::Optional<ExprDoc> {
ffi::Array<ExprDoc> elements;
for (const auto& element : input.cast<ffi::Array<ffi::Any>>()) {
elements.push_back(d->Translate(element).value());
@@ -70,7 +70,7 @@ ffi::Optional<ExprDoc> ArrayDocTranslate(DocTranslatorObj* d,
ffi::AnyView input
ffi::Optional<ExprDoc> MapDocTranslate(DocTranslatorObj* d, ffi::AnyView input,
const ffi::Object*) {
return d->WithExtraStateScope<bool>(
- "ir.explicit_scalar_dtype", true, [&]() -> ffi::Optional<ExprDoc> {
+ "ir.explicit_value_type", true, [&]() -> ffi::Optional<ExprDoc> {
auto map = input.cast<ffi::Map<ffi::Any, ffi::Any>>();
std::vector<std::pair<ffi::Any, ffi::Any>> entries(map.begin(),
map.end());
if (std::all_of(entries.begin(), entries.end(), [](const auto& entry) {
diff --git a/src/script/printer/ir/prim_expr.cc
b/src/script/printer/ir/prim_expr.cc
index 5218d3b1f5..b6e83f854d 100644
--- a/src/script/printer/ir/prim_expr.cc
+++ b/src/script/printer/ir/prim_expr.cc
@@ -302,7 +302,7 @@ ffi::Optional<ExprDoc>
StringImmDocTranslate(DocTranslatorObj* d, ffi::AnyView i
const auto* imm =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
StringImmNode>(input);
ExprDoc literal = LiteralDoc::Str(imm->value, std::nullopt);
- if (d->GetOrCreateExtraState<bool>("ir.explicit_scalar_dtype")) {
+ if (d->GetOrCreateExtraState<bool>("ir.explicit_value_type")) {
return NamespaceDoc("ir")->Attr("StringImm")->Call({literal});
}
return literal;
diff --git a/src/script/printer/ir/prim_type.cc
b/src/script/printer/ir/prim_type.cc
index 55ab9c5ca2..ad0cdcc069 100644
--- a/src/script/printer/ir/prim_type.cc
+++ b/src/script/printer/ir/prim_type.cc
@@ -75,7 +75,7 @@ ffi::Optional<ExprDoc> IntImmDocTranslate(DocTranslatorObj*
d, ffi::AnyView inpu
ExprDoc value = LiteralDoc::Int(ffi::GetRef<IntImm>(imm), std::nullopt);
DLDataType default_dtype =
ffi::StringToDLDataType(d->GetExtraConfig<ffi::String>("ir.int_dtype",
"int32"));
- if (dtype == default_dtype &&
!d->GetOrCreateExtraState<bool>("ir.explicit_scalar_dtype"))
+ if (dtype == default_dtype &&
!d->GetOrCreateExtraState<bool>("ir.explicit_value_type"))
return value;
if (!HasPrimTypeShorthand(imm->ty.as_or_throw<PrimType>()) ||
imm->value < std::numeric_limits<int64_t>::min() ||
@@ -99,7 +99,7 @@ ffi::Optional<ExprDoc> FloatImmDocTranslate(DocTranslatorObj*
d, ffi::AnyView in
DLDataType default_dtype =
ffi::StringToDLDataType(d->GetExtraConfig<ffi::String>("ir.float_dtype",
"void"));
ExprDoc value = LiteralDoc::Float(imm->value, std::nullopt);
- if (dtype == default_dtype &&
!d->GetOrCreateExtraState<bool>("ir.explicit_scalar_dtype"))
+ if (dtype == default_dtype &&
!d->GetOrCreateExtraState<bool>("ir.explicit_value_type"))
return value;
if (!HasPrimTypeShorthand(imm->ty.as_or_throw<PrimType>())) {
if (!std::isfinite(imm->value)) {
@@ -139,6 +139,25 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FDocTranslate::FromNative<&PrimTypeDocTranslate>());
}
+ffi::Optional<ExprDoc> PtrTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
input,
+ const ffi::Object*) {
+ const auto* ty =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
PtrTypeNode>(input);
+ return d->WithExtraStateScope<bool>("ir.explicit_value_type", true, [&]() ->
ExprDoc {
+ ffi::Array<ExprDoc> args{d->Translate(ty->element_type).value()};
+ if (ty->scope != runtime::MemoryScope::Global()) {
+ args.push_back(LiteralDoc::Str(ty->scope.name(), std::nullopt));
+ }
+ return NamespaceDoc("ir")->Attr("Ptr")->Call(args);
+ });
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ ffi::reflection::TypeAttrDef<PtrTypeNode>().attr(
+ tvm::script::printer::type_attr::kDocTranslate,
+ FDocTranslate::FromNative<&PtrTypeDocTranslate>());
+}
+
// The short builder name is already the TensorRegion expression constructor.
// Preserve that export and keep the explicit Type spelling for this one
collision.
ffi::Optional<ExprDoc> TensorRegionTypeDocTranslate(DocTranslatorObj*,
ffi::AnyView,
diff --git a/src/script/printer/ir/utils.cc b/src/script/printer/ir/utils.cc
index 700dd7b86d..b55c9cb7af 100644
--- a/src/script/printer/ir/utils.cc
+++ b/src/script/printer/ir/utils.cc
@@ -83,7 +83,7 @@ ExprDoc CallAttrsValue(DocTranslatorObj* d, const Attrs&
attrs) {
for (const auto& [key, value] : fields) {
if (key == "type_key") return AddMetadata(d, attrs);
keys.push_back(key);
- values.push_back(d->WithExtraStateScope<bool>("ir.explicit_scalar_dtype",
true,
+ values.push_back(d->WithExtraStateScope<bool>("ir.explicit_value_type",
true,
[&] { return
d->Translate(value).value(); }));
}
ExprDoc inline_doc =
@@ -142,7 +142,7 @@ ExprDoc RawCallDocTranslate(DocTranslatorObj* d, const
CallNode* call,
ExprDoc callee =
op.has_value() ? LiteralDoc::Str(op.value()->name, std::nullopt)
: call->op.as<StringImmNode>()
- ? d->WithExtraStateScope<bool>("ir.explicit_scalar_dtype", true,
+ ? d->WithExtraStateScope<bool>("ir.explicit_value_type", true,
[&] { return
d->Translate(call->op).value(); })
: d->Translate(call->op).value();
ffi::Array<ffi::String> keys;
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 6fca353d96..f4fa80defb 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -714,8 +714,8 @@ void CodeGenLLVM::GetAlignment(PrimType t, const VarNode*
buf_var, const PrimExp
if (it != alloc_storage_info_.end()) {
const StorageInfo& info = it->second;
*p_native_bits = NativeVectorBits(buf_var->ty.as<TensorTypeNode>()
- ?
buf_var->ty.as<TensorTypeNode>()->storage_scope
- :
buf_var->ty.as_or_throw<PtrType>()->storage_scope);
+ ?
buf_var->ty.as<TensorTypeNode>()->scope
+ :
buf_var->ty.as_or_throw<PtrType>()->scope);
max_align_bits = info.alignment * 8;
} else {
*p_native_bits = native_vector_bits_;
@@ -2335,8 +2335,7 @@ void CodeGenLLVM::Dispatch_(const SeqStmtNode* op) {
void CodeGenLLVM::DispatchDeclTensor(const BindNode* op, const CallNode*
buffer_call) {
Expr data = buffer_call->args[0];
DLDataType dtype =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
- buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
EmitDebugLocation(op);
const VarNode* buffer_var = buffer.get();
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index f05c9c33ed..c2a6def3d4 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -845,10 +845,10 @@ void CodeGenC::Dispatch_(const CallNode* op,
std::ostream& os) { // NOLINT(*)
TVM_FFI_ICHECK_EQ(op->args.size(), 2U);
const auto& pointer_type = op->ty.as_or_throw<PtrType>();
os << "((";
- if (IsScopePartOfType()) PrintStorageScope(pointer_type->storage_scope,
os);
+ if (IsScopePartOfType()) PrintStorageScope(pointer_type->scope, os);
PrintType(pointer_type->element_type, os);
os << "*)(((";
- if (IsScopePartOfType()) PrintStorageScope(pointer_type->storage_scope,
os);
+ if (IsScopePartOfType()) PrintStorageScope(pointer_type->scope, os);
os << "char*)";
this->PrintExpr(op->args[0], os);
os << ") + ";
@@ -865,7 +865,7 @@ void CodeGenC::Dispatch_(const CallNode* op, std::ostream&
os) { // NOLINT(*)
if (const auto* pointer_type = op->ty.as<PtrTypeNode>()) {
os << "((";
if (IsScopePartOfType()) {
- PrintStorageScope(pointer_type->storage_scope, os);
+ PrintStorageScope(pointer_type->scope, os);
}
if (const auto* element_type =
pointer_type->element_type.as<PrimTypeNode>()) {
this->PrintType(ffi::GetRef<PrimType>(element_type), os);
@@ -945,8 +945,7 @@ void CodeGenC::PrintVecBinaryOp(const std::string& op,
const PrimType& t, PrimEx
void CodeGenC::DispatchDeclTensor(const BindNode* op, const CallNode*
buffer_call) {
Expr data = buffer_call->args[0];
DLDataType dtype =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
- buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
const VarNode* source = data.as<VarNode>();
if (const auto* call = data.as<CallNode>();
@@ -1286,8 +1285,7 @@ void CodeGenC::Dispatch_(const BindNode* op) {
void CodeGenC::DispatchAllocTensor(const BindNode* op, const CallNode*
buffer_call) {
tvm::Tuple shape =
tvm::Tuple(buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->shape);
DLDataType dtype =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->dtype->dtype;
- runtime::MemoryScope scope =
- buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
buffer_call->ty_args[0].as_or_throw<tirx::TensorType>()->scope;
TensorVar buffer = op->var.as_or_throw<TensorVar>();
DictAttrs annotations = buffer_call->attrs.as_or_throw<DictAttrs>();
TVM_FFI_ICHECK(buffer.defined());
diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc
index 68d1ccd2fc..319d1e2f3e 100644
--- a/src/tirx/ir/specialize.cc
+++ b/src/tirx/ir/specialize.cc
@@ -105,7 +105,7 @@ class FunctionSpecializer : public StmtExprMutator {
} else {
specializer->constrained_tensor_params_.insert(param.get());
specializer->tensor_storage_scopes_.emplace(tensor.value(),
-
replacement_tensor.value()->storage_scope);
+
replacement_tensor.value()->scope);
}
}
}
@@ -290,8 +290,8 @@ class FunctionSpecializer : public StmtExprMutator {
}
}
- bool storage_scope_changed = specialized_storage_scope.has_value() &&
- specialized_storage_scope.value() !=
tensor->storage_scope;
+ bool storage_scope_changed =
+ specialized_storage_scope.has_value() &&
specialized_storage_scope.value() != tensor->scope;
if (tensor->elem_offset.same_as(elem_offset) &&
tensor->shape.same_as(shape) &&
tensor->strides.same_as(strides) && !layout_changed &&
!storage_scope_changed) {
return tensor;
@@ -304,7 +304,7 @@ class FunctionSpecializer : public StmtExprMutator {
n->layout = std::move(layout);
}
if (storage_scope_changed) {
- n->storage_scope = specialized_storage_scope.value();
+ n->scope = specialized_storage_scope.value();
}
return RebuildTensorVar(tensor, std::move(n));
}
diff --git a/src/tirx/ir/type.cc b/src/tirx/ir/type.cc
index 09441fa721..1736e9910c 100644
--- a/src/tirx/ir/type.cc
+++ b/src/tirx/ir/type.cc
@@ -74,7 +74,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
TensorMapTypeMaybeInpla
TVM_FFI_INLINE ffi::Expected<ffi::Optional<ffi::VisitInterrupt>>
TensorTypeVisit(
ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
- // skips: storage_scope, data_alignment, offset_factor
+ // skips: scope, data_alignment, offset_factor
const TensorTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TensorTypeNode>(value);
TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->dtype));
@@ -91,7 +91,7 @@ TVM_FFI_INLINE
ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> TensorTypeVisit
TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>> TensorTypeMutate(
ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- // skips: storage_scope, data_alignment, offset_factor
+ // skips: scope, data_alignment, offset_factor
const TensorTypeNode* self =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TensorTypeNode>(value);
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimType>, mapped_dtype,
@@ -127,7 +127,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
TensorTypeMutate(
TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
TensorTypeMaybeInplaceMutate(
ffi::StructuralMutatorObj* mutator, ffi::AnyView value) noexcept {
- // skips: storage_scope, data_alignment, offset_factor
+ // skips: scope, data_alignment, offset_factor
TensorTypeNode* self = const_cast<TensorTypeNode*>(
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
TensorTypeNode>(value));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimType>, mapped_dtype,
@@ -166,15 +166,15 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
TensorTypeMaybeInplaceM
} // namespace
-TensorType::TensorType(ffi::Array<PrimExpr> shape, PrimType dtype,
- runtime::MemoryScope storage_scope,
ffi::Array<PrimExpr> strides,
- ffi::Optional<PrimExpr> elem_offset, int
data_alignment, int offset_factor,
- ffi::Optional<Layout> layout, Location loc)
+TensorType::TensorType(ffi::Array<PrimExpr> shape, PrimType dtype,
runtime::MemoryScope scope,
+ ffi::Array<PrimExpr> strides, ffi::Optional<PrimExpr>
elem_offset,
+ int data_alignment, int offset_factor,
ffi::Optional<Layout> layout,
+ Location loc)
: Type(ffi::UnsafeInit{}) {
PrimExpr offset = elem_offset.value_or(IntImm(DefaultIndexPrimType(), 0));
auto n = ffi::make_object<TensorTypeNode>(std::move(offset));
n->dtype = std::move(dtype);
- n->storage_scope = std::move(storage_scope);
+ n->scope = std::move(scope);
n->shape = std::move(shape);
n->strides = std::move(strides);
n->data_alignment =
@@ -198,7 +198,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef().def(
"tirx.TensorType",
- [](ffi::Array<PrimExpr> shape, PrimType dtype, runtime::MemoryScope
storage_scope,
+ [](ffi::Array<PrimExpr> shape, PrimType dtype, runtime::MemoryScope
scope,
ffi::Array<PrimExpr> strides, ffi::AnyView elem_offset, int
data_alignment,
int offset_factor, ffi::Optional<Layout> layout, Location loc) {
ffi::Optional<PrimExpr> offset;
@@ -207,9 +207,9 @@ TVM_FFI_STATIC_INIT_BLOCK() {
} else {
offset = elem_offset.cast<ffi::Optional<PrimExpr>>();
}
- return TensorType(std::move(shape), std::move(dtype),
std::move(storage_scope),
- std::move(strides), std::move(offset),
data_alignment, offset_factor,
- std::move(layout), std::move(loc));
+ return TensorType(std::move(shape), std::move(dtype),
std::move(scope), std::move(strides),
+ std::move(offset), data_alignment, offset_factor,
std::move(layout),
+ std::move(loc));
});
}
diff --git a/src/tirx/op/memory.cc b/src/tirx/op/memory.cc
index c305b51c31..eb56466589 100644
--- a/src/tirx/op/memory.cc
+++ b/src/tirx/op/memory.cc
@@ -41,8 +41,7 @@ Type InferTypePointerOffset(const CallNode* call) {
if (element_type.as<MissingType>()) return Type::Missing();
}
auto pointer = call->args[DataIndex]->ty.as<PtrType>();
- return PtrType(element_type,
- pointer ? pointer.value()->storage_scope :
runtime::MemoryScope::Global());
+ return PtrType(element_type, pointer ? pointer.value()->scope :
runtime::MemoryScope::Global());
}
Type InferTypeAddressOf(const CallNode* call) {
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index 5d0425bdde..2f7fea2b40 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -51,7 +51,7 @@ tvm::tirx::TensorType TensorTypeDecl(ffi::Array<PrimExpr>
shape, PrimType dtype,
int offset_factor, ffi::Optional<Layout>
layout,
ffi::Optional<Expr> data) {
if (data.has_value()) {
- storage_scope = data.value()->ty.as_or_throw<PtrType>()->storage_scope;
+ storage_scope = data.value()->ty.as_or_throw<PtrType>()->scope;
}
return tvm::tirx::TensorType(shape, dtype, storage_scope,
strides.value_or(ffi::Array<PrimExpr>()),
elem_offset, align,
diff --git a/src/tirx/script/printer/for_loop.cc
b/src/tirx/script/printer/for_loop.cc
index 6ef98ddea6..bfaf22161a 100644
--- a/src/tirx/script/printer/for_loop.cc
+++ b/src/tirx/script/printer/for_loop.cc
@@ -83,7 +83,7 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
} else if (key == tvm::tirx::attr::kPragmaUnroll &&
(boolean.value_or(false) || (integer.has_value() &&
integer.value() > 0))) {
keys.push_back("unroll");
-
values.push_back(d->WithExtraStateScope<bool>("ir.explicit_scalar_dtype", true,
+ values.push_back(d->WithExtraStateScope<bool>("ir.explicit_value_type",
true,
[&] { return
d->Translate(value).value(); }));
unroll_option = true;
}
@@ -98,7 +98,7 @@ ffi::Optional<ExprDoc> ForDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
for (const auto& [key, value] : sorted) {
annotation_keys.push_back(LiteralDoc::Str(key, std::nullopt));
annotation_values.push_back(d->WithExtraStateScope<bool>(
- "ir.explicit_scalar_dtype", true, [&] { return
d->Translate(value).value(); }));
+ "ir.explicit_value_type", true, [&] { return
d->Translate(value).value(); }));
}
keys.push_back("annotations");
values.push_back(DictDoc(annotation_keys, annotation_values));
diff --git a/src/tirx/script/printer/function.cc
b/src/tirx/script/printer/function.cc
index 5a2f20ddb2..f5b04534ab 100644
--- a/src/tirx/script/printer/function.cc
+++ b/src/tirx/script/printer/function.cc
@@ -110,7 +110,7 @@ void PrintFunction(DocTranslatorObj* d, const
tirx::FunctionNode* func, ExprDoc
for (const auto& [key, value] : attrs) {
keys.push_back(LiteralDoc::Str(key, std::nullopt));
values.push_back(d->WithExtraStateScope<bool>(
- "ir.explicit_scalar_dtype", true, [&] { return
d->Translate(value).value(); }));
+ "ir.explicit_value_type", true, [&] { return
d->Translate(value).value(); }));
}
body.insert(
body.begin(),
diff --git a/src/tirx/script/printer/tensor.cc
b/src/tirx/script/printer/tensor.cc
index 81176eeb2c..5b156f66bb 100644
--- a/src/tirx/script/printer/tensor.cc
+++ b/src/tirx/script/printer/tensor.cc
@@ -50,10 +50,10 @@ ffi::Optional<ExprDoc>
TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
->Call({LiteralDoc::Str("tirx.TensorType", std::nullopt)},
{"dtype", "scope", "shape", "strides", "elem_offset",
"data_alignment",
"offset_factor", "layout"},
- {d->Translate(tensor->dtype).value(),
d->Translate(tensor->storage_scope).value(),
+ {d->Translate(tensor->dtype).value(),
d->Translate(tensor->scope).value(),
d->Translate(tensor->shape).value(),
d->Translate(tensor->strides).value(),
d->WithExtraStateScope<bool>(
- "ir.explicit_scalar_dtype", true,
+ "ir.explicit_value_type", true,
[&] { return d->Translate(tensor->elem_offset).value(); }),
LiteralDoc::Int(tensor->data_alignment, std::nullopt),
LiteralDoc::Int(tensor->offset_factor, std::nullopt),
@@ -69,10 +69,10 @@ ffi::Optional<ExprDoc>
TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
d->RecordOrigin(literal, dim);
shape.push_back(literal);
} else if (dim.as<IntImmNode>()) {
- shape.push_back(d->WithExtraStateScope<bool>("ir.explicit_scalar_dtype",
true,
+ shape.push_back(d->WithExtraStateScope<bool>("ir.explicit_value_type",
true,
[&] { return
d->Translate(dim).value(); }));
} else {
- d->WithExtraStateScope<bool>("ir.explicit_scalar_dtype", false,
+ d->WithExtraStateScope<bool>("ir.explicit_value_type", false,
[&] {
shape.push_back(d->Translate(dim).value()); });
}
}
@@ -83,12 +83,12 @@ ffi::Optional<ExprDoc>
TensorTypeDocTranslate(DocTranslatorObj* d, ffi::AnyView
keys.push_back(key);
values.push_back(value);
};
- if (tensor->storage_scope != runtime::MemoryScope::Global()) {
- optional("scope", LiteralDoc::Str(tensor->storage_scope.name(),
std::nullopt));
+ if (tensor->scope != runtime::MemoryScope::Global()) {
+ optional("scope", LiteralDoc::Str(tensor->scope.name(), std::nullopt));
}
if (!tensor->strides.empty()) optional("strides",
d->Translate(tensor->strides).value());
if (!ffi::StructuralEqual()(tensor->elem_offset,
IntImm(tirx::DefaultIndexPrimType(), 0))) {
- optional("elem_offset",
d->WithExtraStateScope<bool>("ir.explicit_scalar_dtype", true, [&] {
+ optional("elem_offset",
d->WithExtraStateScope<bool>("ir.explicit_value_type", true, [&] {
return d->Translate(tensor->elem_offset).value();
}));
}
diff --git a/src/tirx/transform/flattened_tensor.h
b/src/tirx/transform/flattened_tensor.h
index 9cff4a9f31..59893a736d 100644
--- a/src/tirx/transform/flattened_tensor.h
+++ b/src/tirx/transform/flattened_tensor.h
@@ -58,11 +58,10 @@ TVM_FFI_INLINE TensorVar FlattenedTensor(const TensorVar&
tensor) {
return tensor;
} else {
// The original layout describes the old rank; reset it for the flattened
shape.
- return TensorVar(
- tensor.name(),
- TensorType(output_shape, self->dtype, self->storage_scope, {},
self->elem_offset,
- self->data_alignment, self->offset_factor, std::nullopt),
- tensor.loc());
+ return TensorVar(tensor.name(),
+ TensorType(output_shape, self->dtype, self->scope, {},
self->elem_offset,
+ self->data_alignment, self->offset_factor,
std::nullopt),
+ tensor.loc());
}
}
diff --git a/src/tirx/transform/lower_tvm_builtin.cc
b/src/tirx/transform/lower_tvm_builtin.cc
index 84755d5093..2151a6e84c 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -279,7 +279,7 @@ class BuiltinLower : public StmtExprMutator {
tvm::Tuple shape =
tvm::Tuple(tensor_call->ty_args[0].as_or_throw<TensorType>()->shape);
DLDataType dtype =
tensor_call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
PrimType element_type(dtype);
- auto scope =
tensor_call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
+ auto scope = tensor_call->ty_args[0].as_or_throw<TensorType>()->scope;
DictAttrs annotations =
tensor_call->attrs.value().as_or_throw<DictAttrs>();
if (annotations->dict.count(transform::kDisableLowerTVMBuiltin)) {
if
(annotations->dict[transform::kDisableLowerTVMBuiltin].as_or_throw<IntImm>()->value)
{
diff --git a/src/tirx/transform/lower_warp_memory.cc
b/src/tirx/transform/lower_warp_memory.cc
index d2806a4c23..a9ded5ceca 100644
--- a/src/tirx/transform/lower_warp_memory.cc
+++ b/src/tirx/transform/lower_warp_memory.cc
@@ -372,7 +372,7 @@ class WarpAccessRewriter : public StmtExprMutator {
auto type = CopyTensorType(op->var.as_or_throw<TensorVar>());
static const auto local = runtime::MemoryScope::Get("local");
- type->storage_scope = local;
+ type->scope = local;
type->shape = {IntImm::Int32(alloc_size / width_)};
type->strides = {};
type->elem_offset =
IntImm(op->var.as_or_throw<TensorVar>()->elem_offset.ty(), 0);
@@ -674,7 +674,7 @@ class WarpMemoryRewriter : public StmtExprMutator {
static const auto local = runtime::MemoryScope::Get("local");
if (const auto* call = alloc ? alloc->value.as<CallNode>() : nullptr;
call && call->op.same_as(alloc_tensor_op()) &&
- call->ty_args[0].as_or_throw<TensorType>()->storage_scope == warp) {
+ call->ty_args[0].as_or_throw<TensorType>()->scope == warp) {
new_storage_scopes_[alloc->var] = local;
// Gather remaining siblings as the "body" for rewriting.
ffi::Array<Stmt> remaining;
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index 2e7484df79..232a0db495 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -490,7 +490,7 @@ class DeviceInfoCollector : public StmtExprVisitor {
ffi::Optional<VisitInterrupt> DispatchAllocTensor(const BindNode* op, const
CallNode* call) {
static const auto scope_shared_dyn =
runtime::MemoryScope::Get("shared.dyn");
- runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->scope;
auto storage_scope = scope;
if (storage_scope == scope_shared_dyn) {
TVM_FFI_ICHECK(!saw_dyn_shared_alloc_)
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index a2288389ef..069b8d1776 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -799,7 +799,7 @@ class StoragePlanRewriter : public StmtExprMutator {
static bool RequiresExactAllocation(const CallNode* call) {
auto type = call->ty_args[0].as_or_throw<TensorType>();
- if (type->storage_scope.alloc_property() ==
MemoryAllocProperty::kSeparateAllocations) {
+ if (type->scope.alloc_property() ==
MemoryAllocProperty::kSeparateAllocations) {
return true;
}
// Flat packing and lifetime reuse cannot reconstruct arbitrary physical
views.
@@ -1192,7 +1192,7 @@ class StoragePlanRewriter : public StmtExprMutator {
const AllocEntry& entry = alloc_info.at(var);
const BindNode* alloc = entry.alloc;
const auto* call = alloc->value.as<CallNode>();
- runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->storage_scope;
+ runtime::MemoryScope scope =
call->ty_args[0].as_or_throw<TensorType>()->scope;
auto storage_scope = scope;
StorageEntry* dst_entry = nullptr;
// inplace detection
@@ -1875,7 +1875,7 @@ class VectorTypeRewriter : public StmtExprMutator {
return RebuildTensorVar(old_tensor, std::move(type)).var();
}
return Var(old_tensor_var->name,
- PtrType(preferred,
old_tensor_var->ty.as_or_throw<PtrType>()->storage_scope),
+ PtrType(preferred,
old_tensor_var->ty.as_or_throw<PtrType>()->scope),
old_tensor_var->loc);
}();
diff --git a/src/tirx/transform/unsupported_dtype_legalize.cc
b/src/tirx/transform/unsupported_dtype_legalize.cc
index 4d223051e7..a9afa62cb3 100644
--- a/src/tirx/transform/unsupported_dtype_legalize.cc
+++ b/src/tirx/transform/unsupported_dtype_legalize.cc
@@ -648,8 +648,8 @@ class StorageLegalizer : public StmtExprMutator {
}
Call call = ret.as_or_throw<Call>();
Type new_element_type =
GetStorageUIntDType(ffi::GetRef<PrimType>(element_type));
- return Call(PtrType(new_element_type, pointer_type->storage_scope),
call->op, call->args,
- call->attrs, call->ty_args, call->loc);
+ return Call(PtrType(new_element_type, pointer_type->scope), call->op,
call->args, call->attrs,
+ call->ty_args, call->loc);
}
if (!op->ty.as<PrimTypeNode>()) {
return StmtExprMutator::Mutate_(op, inplace_mode);
@@ -703,7 +703,7 @@ class StorageLegalizer : public StmtExprMutator {
PrimType elem_prim_type = ffi::GetRef<PrimType>(elem_type);
if (MatchType(elem_prim_type)) {
Var new_var =
- Var(var->name, PtrType(GetStorageUIntDType(elem_prim_type),
ptr_type->storage_scope));
+ Var(var->name, PtrType(GetStorageUIntDType(elem_prim_type),
ptr_type->scope));
VarRemapSet(var, new_var);
return new_var;
}
diff --git a/src/tirx/transform/update_pointer_storage_scope.cc
b/src/tirx/transform/update_pointer_storage_scope.cc
index 3ab0600d96..53e5989ed9 100644
--- a/src/tirx/transform/update_pointer_storage_scope.cc
+++ b/src/tirx/transform/update_pointer_storage_scope.cc
@@ -53,7 +53,7 @@ UpdatePointerStorageScope::UpdatePointerStorageScope(
if (kv.first->ty.as<TensorTypeNode>()) {
TensorVar tensor = GetTensorVar(kv.first.get());
auto type = CopyTensorType(tensor);
- type->storage_scope = kv.second;
+ type->scope = kv.second;
TensorVar replacement = RebuildTensorVar(tensor, std::move(type));
VarRemapSet(kv.first, replacement);
} else {
@@ -82,9 +82,9 @@ UnchangedOr<Expr> UpdatePointerStorageScope::Mutate_(const
CallNode* op, Inplace
Expr value = std::move(result).ValueOrUnchanged(ffi::GetRef<Expr>(op));
auto call = value.as_or_throw<Call>();
auto type = call->ty_args[0].as_or_throw<TensorType>();
- if (type->storage_scope != it->second) {
+ if (type->scope != it->second) {
auto updated = ffi::make_object<TensorTypeNode>(*type.get());
- updated->storage_scope = it->second;
+ updated->scope = it->second;
auto copy = ffi::make_object<CallNode>(*call.get());
copy->ty_args.Set(0, TensorType(std::move(updated)));
return ReinferMutatedCallType(Expr(std::move(copy)), op, inplace_mode);
diff --git a/tests/cpp/runtime/memory_scope_test.cc
b/tests/cpp/runtime/memory_scope_test.cc
index 17daf40c97..667f62aa4e 100644
--- a/tests/cpp/runtime/memory_scope_test.cc
+++ b/tests/cpp/runtime/memory_scope_test.cc
@@ -30,7 +30,7 @@ TEST(MemoryScope, CanonicalDefaultAndTypedBoundaries) {
MemoryScope global;
EXPECT_TRUE(global.same_as(MemoryScope::Global()));
EXPECT_TRUE(global.same_as(MemoryScope::Get("global")));
- EXPECT_EQ(PtrType(PrimType::Int(32))->storage_scope, global);
+ EXPECT_EQ(PtrType(PrimType::Int(32))->scope, global);
EXPECT_THROW(ffi::Any(ffi::String("")).cast<MemoryScope>(), ffi::Error);
EXPECT_EQ(&MemoryScope::Global(), &MemoryScope::Global());
EXPECT_EQ(ffi::Any(ffi::String("shared")).cast<MemoryScope>(),
MemoryScope::Get("shared"));
@@ -47,7 +47,7 @@ TEST(MemoryScope, ExtensionRegistrationAndAttributes) {
EXPECT_EQ(scope.name(), "memory_scope_test.scratch");
EXPECT_EQ(scope.alloc_property(), MemoryAllocProperty::kNormal);
EXPECT_EQ(scope, ffi::Any(scope).cast<MemoryScope>());
- EXPECT_EQ(scope, PtrType(PrimType::Float(32), scope)->storage_scope);
+ EXPECT_EQ(scope, PtrType(PrimType::Float(32), scope)->scope);
EXPECT_THROW(MemoryScope::Register(scope.name()), ffi::Error);
auto registered = MemoryScope::Register("memory_scope_test.registered");
EXPECT_EQ(registered, MemoryScope::Get(registered.name()));
diff --git
a/tests/python/relax/test_transform_specialize_function_based_on_callsite.py
b/tests/python/relax/test_transform_specialize_function_based_on_callsite.py
index 581a5295ad..d814300bd8 100644
--- a/tests/python/relax/test_transform_specialize_function_based_on_callsite.py
+++ b/tests/python/relax/test_transform_specialize_function_based_on_callsite.py
@@ -51,8 +51,8 @@ class ValidateBufferScopes(PyExprVisitor): # pylint:
disable=abstract-method
for buf in pfunc.params:
if not tvm.tirx.is_tensor_var(buf):
continue
- assert tvm.runtime.MemoryScope("global") ==
buf.data.ty.storage_scope, (
- f"expected to be global scoped, but got
{val.data.ty.storage_scope}"
+ assert tvm.runtime.MemoryScope("global") ==
buf.data.ty.scope, (
+ f"expected to be global scoped, but got
{val.data.ty.scope}"
)
else:
for idx, arg in enumerate(call.args[1]):
@@ -61,13 +61,13 @@ class ValidateBufferScopes(PyExprVisitor): # pylint:
disable=abstract-method
f"Expected TensorType but git {type(arg_ty)}"
)
buf = pfunc.params[idx]
- assert arg_ty.vdevice.memory_scope ==
buf.data.ty.storage_scope, (
- f"scope mismatched after specialization
{arg_ty.vdevice.memory_scope} vs {buf.data.ty.storage_scope}"
+ assert arg_ty.vdevice.memory_scope == buf.data.ty.scope, (
+ f"scope mismatched after specialization
{arg_ty.vdevice.memory_scope} vs {buf.data.ty.scope}"
)
if isinstance(call.ty_args[0], relax.TensorType):
buf = pfunc.params[-1]
- assert call.ty_args[0].vdevice.memory_scope ==
buf.data.ty.storage_scope, (
- f"scope mismatched after specialization
{call.ty_args[0].vdevice.memory_scope} vs {buf.data.ty.storage_scope}"
+ assert call.ty_args[0].vdevice.memory_scope ==
buf.data.ty.scope, (
+ f"scope mismatched after specialization
{call.ty_args[0].vdevice.memory_scope} vs {buf.data.ty.scope}"
)
else:
assert isinstance(call.ty_args[0], relax.TupleType), (
@@ -75,8 +75,8 @@ class ValidateBufferScopes(PyExprVisitor): # pylint:
disable=abstract-method
)
for idx, ty in enumerate(call.ty_args[0].fields):
buf = pfunc.params[len(call.args[1]) + idx]
- assert ty.vdevice.memory_scope ==
buf.data.ty.storage_scope, (
- f"scope mismatched after specialization
{ty.vdevice.memory_scope} vs {buf.data.ty.storage_scope}"
+ assert ty.vdevice.memory_scope == buf.data.ty.scope, (
+ f"scope mismatched after specialization
{ty.vdevice.memory_scope} vs {buf.data.ty.scope}"
)
diff --git a/tests/python/runtime/test_runtime_memory_scope.py
b/tests/python/runtime/test_runtime_memory_scope.py
index 345b00bb8e..006bbe7f99 100644
--- a/tests/python/runtime/test_runtime_memory_scope.py
+++ b/tests/python/runtime/test_runtime_memory_scope.py
@@ -77,11 +77,11 @@ def test_shared_cpp_python_registry_and_typed_conversion():
ptr_ty = tvm.ir.PtrType(tvm.ir.PrimType("float32"), scope)
tensor = tvm.tirx.decl_tensor((4,), "float32", scope=scope)
tensor_ty = tvm.tirx.TensorType((4,), "float32", scope)
- assert tensor_ty.storage_scope.same_as(custom)
- assert ptr_ty.storage_scope.same_as(custom)
- assert tensor.ty.storage_scope.same_as(custom)
+ assert tensor_ty.scope.same_as(custom)
+ assert ptr_ty.scope.same_as(custom)
+ assert tensor.ty.scope.same_as(custom)
assert tensor.scope().same_as(custom)
- assert tensor.data.ty.storage_scope.same_as(custom)
+ assert tensor.data.ty.scope.same_as(custom)
unknown = "test.memory_scope.unknown." + uuid.uuid4().hex
for construct in (
@@ -97,12 +97,12 @@ def test_shared_cpp_python_registry_and_typed_conversion():
def test_global_defaults():
global_scope = MemoryScope("global")
- assert
tvm.ir.PtrType(tvm.ir.PrimType("float32")).storage_scope.same_as(global_scope)
+ assert
tvm.ir.PtrType(tvm.ir.PrimType("float32")).scope.same_as(global_scope)
assert tvm.tirx.decl_tensor((4,), "float32").scope().same_as(global_scope)
for scope in ("global", global_scope):
- assert tvm.ir.PtrType(tvm.ir.PrimType("float32"),
scope).storage_scope.same_as(global_scope)
+ assert tvm.ir.PtrType(tvm.ir.PrimType("float32"),
scope).scope.same_as(global_scope)
assert tvm.tirx.decl_tensor((4,), "float32",
scope=scope).scope().same_as(global_scope)
- assert T.Tensor((4,), "float32",
scope=scope).storage_scope.same_as(global_scope)
+ assert T.Tensor((4,), "float32",
scope=scope).scope.same_as(global_scope)
assert tvm.relax.VDevice("llvm",
memory_scope=scope).memory_scope.same_as(global_scope)
for construct in (
lambda: MemoryScope(""),
@@ -122,7 +122,7 @@ def test_json_and_script_roundtrip():
assert tvm.ir.load_json(tvm.ir.save_json(scope)).same_as(scope)
ty = T.Tensor((4,), "float32", scope=scope)
restored = tvm.ir.load_json(tvm.ir.save_json(ty))
- assert restored.storage_scope.same_as(scope)
+ assert restored.scope.same_as(scope)
tvm.ir.assert_structural_equal(ty, restored)
@T.function
@@ -144,21 +144,24 @@ def
test_canonical_constructor_aliases_and_reflected_scopes():
for scope in (MemoryScope("global"), custom):
pointer = T.Ptr(T.float32, scope)
source = pointer.script()
- assert ("storage_scope=" in source) == (scope != MemoryScope("global"))
+ expected = (
+ "T.Ptr(T.float32)"
+ if scope == MemoryScope("global")
+ else f'T.Ptr(T.float32, "{scope.name}")'
+ )
+ assert source == expected
restored = eval(source, {"T": T, "I": I})
tvm.ir.assert_structural_equal(pointer, restored)
- assert restored.storage_scope.same_as(scope)
+ assert restored.scope.same_as(scope)
tensor = T.Tensor((4,), T.float32, scope, byte_offset=8, align=128,
layout="default")
assert tensor.scope.same_as(scope)
assert tensor.elem_offset.value == 2
assert tensor.data_alignment == 128
restored_tensor = eval(tensor.script(), {"T": T, "I": I})
tvm.ir.assert_structural_equal(tensor, restored_tensor)
- for scope, alias in ((custom, custom.name), (custom.name, custom)):
- tensor = T.Tensor((4,), T.float32, scope, storage_scope=alias)
+ for scope in (custom, custom.name):
+ tensor = T.Tensor((4,), T.float32, scope=scope)
assert tensor.scope.same_as(custom)
- with pytest.raises(ValueError, match="scope and storage_scope must agree"):
- T.Tensor((4,), T.float32, custom, storage_scope="shared")
with pytest.raises(ValueError):
T.Ptr(T.float32, "")
diff --git
a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
index e5200e2daa..d9ccc02502 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_block_scope.py
@@ -25,7 +25,7 @@ from tvm.script import tirx as T
def test_tensor_region_bounds_are_visited():
- data = tvm.tirx.Var("data", tvm.ir.PtrType(tvm.ir.PrimType("int32"),
storage_scope="global"))
+ data = tvm.tirx.Var("data", tvm.ir.PtrType(tvm.ir.PrimType("int32"),
scope="global"))
tensor = tvm.tirx.decl_tensor([4], "int32", data=data)
undefined = tvm.tirx.Var("undefined", "int32")
region = TensorRegion(
diff --git a/tests/python/s_tir/script/test_s_tir_script_printer.py
b/tests/python/s_tir/script/test_s_tir_script_printer.py
index 2e44fa8a84..3591943b3e 100644
--- a/tests/python/s_tir/script/test_s_tir_script_printer.py
+++ b/tests/python/s_tir/script/test_s_tir_script_printer.py
@@ -1708,7 +1708,7 @@ def test_void_ptr_vs_handle():
T.evaluate(out_ret_value)
scoped_script = scoped_void_ptr.script()
- assert 'out_ret_value: T.Ptr(T.void,
storage_scope=T.MemoryScope("shared"))' in scoped_script
+ assert 'out_ret_value: T.Ptr(T.void, "shared")' in scoped_script
assert 'T.handle("void"' not in scoped_script
tvm.ir.assert_structural_equal(
scoped_void_ptr,
diff --git a/tests/python/tirx-base/test_tir_nodes.py
b/tests/python/tirx-base/test_tir_nodes.py
index c9f38b261f..cdb412f644 100644
--- a/tests/python/tirx-base/test_tir_nodes.py
+++ b/tests/python/tirx-base/test_tir_nodes.py
@@ -331,7 +331,7 @@ def test_scoped_storage_vars():
ptype = tvm.ir.PtrType(tvm.ir.PrimType(dtype), storage_scope)
x = tvm.tirx.Var("xyz", ptype)
assert x.ty == ptype
- assert x.ty.storage_scope == tvm.runtime.MemoryScope(storage_scope)
+ assert x.ty.scope == tvm.runtime.MemoryScope(storage_scope)
assert isinstance(ptype.element_type, tvm.ir.PrimType)
diff --git a/tests/python/tirx-base/test_tir_tensor.py
b/tests/python/tirx-base/test_tir_tensor.py
index 3719a930c7..70d5dd8e04 100644
--- a/tests/python/tirx-base/test_tir_tensor.py
+++ b/tests/python/tirx-base/test_tir_tensor.py
@@ -78,7 +78,7 @@ def test_tensor_pointer_type_derived_from_dtype_and_scope():
tensor = tvm.tirx.decl_tensor((8,), "float16", data=data)
assert tensor.ty.dtype == tvm.ir.PrimType("float16")
- assert tensor.ty.storage_scope == tvm.runtime.MemoryScope("local")
+ assert tensor.ty.scope == tvm.runtime.MemoryScope("local")
assert tensor.data.ty == tvm.ir.PtrType(tvm.ir.PrimType("float16"),
"local")
@@ -97,7 +97,7 @@ def test_tensor_type_constructor_scope_order(scope):
)
tvm.ir.assert_structural_equal(direct, scripted)
tvm.ir.assert_structural_equal(direct, declared.ty)
- assert direct.storage_scope.same_as(tvm.runtime.MemoryScope(scope))
+ assert direct.scope.same_as(tvm.runtime.MemoryScope(scope))
@pytest.mark.parametrize("scope", ["", "tensor_type.unregistered"])
diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py
b/tests/python/tirx/codegen/test_codegen_blackwell.py
index be79411fef..6e5f76a48c 100644
--- a/tests/python/tirx/codegen/test_codegen_blackwell.py
+++ b/tests/python/tirx/codegen/test_codegen_blackwell.py
@@ -71,10 +71,10 @@ def _assert_remote_mbarrier_ir(func, arrive_op_name,
n_arrives=1):
assert len(mapa_calls) == 1
assert arrive_calls
assert isinstance(bindings[0].var.ty, tvm.ir.PtrType)
- assert bindings[0].var.ty.storage_scope ==
tvm.runtime.MemoryScope("shared")
- assert bindings[0].value.ty.storage_scope ==
tvm.runtime.MemoryScope("shared")
+ assert bindings[0].var.ty.scope == tvm.runtime.MemoryScope("shared")
+ assert bindings[0].value.ty.scope == tvm.runtime.MemoryScope("shared")
assert tensors[0].value.args[0].same_as(bindings[0].var)
- assert tensors[0].value.args[0].ty.storage_scope ==
tvm.runtime.MemoryScope("shared")
+ assert tensors[0].value.args[0].ty.scope ==
tvm.runtime.MemoryScope("shared")
assert tensors[0].var.scope() == tvm.runtime.MemoryScope("shared")
# ptx operand order is the PTX operand order: mapa writes its result into
# a destination the caller declared, so args are (d, a, b) and the arrive
diff --git a/tests/python/tirx/codegen/test_codegen_dsmem.py
b/tests/python/tirx/codegen/test_codegen_dsmem.py
index 7158f2bf4d..7fa366edc4 100644
--- a/tests/python/tirx/codegen/test_codegen_dsmem.py
+++ b/tests/python/tirx/codegen/test_codegen_dsmem.py
@@ -137,8 +137,8 @@ def test_mapa_pointer_bind_codegen():
tvm_ffi.structural_walk(main.body, collect)
assert len(binds) == 1
assert isinstance(binds[0].var.ty, PtrType)
- assert binds[0].var.ty.storage_scope == tvm.runtime.MemoryScope("shared")
- assert binds[0].value.ty.storage_scope == tvm.runtime.MemoryScope("shared")
+ assert binds[0].var.ty.scope == tvm.runtime.MemoryScope("shared")
+ assert binds[0].value.ty.scope == tvm.runtime.MemoryScope("shared")
assert_structural_equal(binds[0].var.ty, binds[0].value.ty)
assert len(decl_tensors) == 1
assert decl_tensors[0].value.args[0].same_as(binds[0].var)
diff --git a/tests/python/tirx/script/test_tirx_script_basic_usage.py
b/tests/python/tirx/script/test_tirx_script_basic_usage.py
index 1544b93abd..822c77d1a6 100644
--- a/tests/python/tirx/script/test_tirx_script_basic_usage.py
+++ b/tests/python/tirx/script/test_tirx_script_basic_usage.py
@@ -75,7 +75,7 @@ def test_tir_ptr_proxy():
isinstance(ptr_0, tirx.Var)
and isinstance(ptr_0.ty, ir.PtrType)
and ptr_0.ty.element_type == ir.PrimType("int32")
- and ptr_0.ty.storage_scope == tvm.runtime.MemoryScope("global")
+ and ptr_0.ty.scope == tvm.runtime.MemoryScope("global")
)
ptr_1 = T.handle("float32", "shared")
@@ -83,7 +83,7 @@ def test_tir_ptr_proxy():
isinstance(ptr_1, tirx.Var)
and isinstance(ptr_1.ty, ir.PtrType)
and ptr_1.ty.element_type == ir.PrimType("float32")
- and ptr_1.ty.storage_scope == tvm.runtime.MemoryScope("shared")
+ and ptr_1.ty.scope == tvm.runtime.MemoryScope("shared")
)
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py
b/tests/python/tirx/script/test_tirx_script_printer.py
index 168eb3445d..a200bd7a4c 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -429,7 +429,7 @@ a
_assert_print(
a,
"""
-a = T.Var("a", T.Ptr(T.void, storage_scope=T.MemoryScope("shared")))
+a = T.Var("a", T.Ptr(T.void, "shared"))
a
""",
)
@@ -680,7 +680,7 @@ def test_pointer_type():
_assert_print(obj, "T.Ptr(T.void)")
obj = ir.PtrType(ir.PrimType("void"), "shared")
- _assert_print(obj, 'T.Ptr(T.void, storage_scope=T.MemoryScope("shared"))')
+ _assert_print(obj, 'T.Ptr(T.void, "shared")')
def test_tuple_type():
diff --git a/tests/python/tirx/tile/cuda/test_ld_st_registers.py
b/tests/python/tirx/tile/cuda/test_ld_st_registers.py
index 05527be5b4..4cee976dc4 100644
--- a/tests/python/tirx/tile/cuda/test_ld_st_registers.py
+++ b/tests/python/tirx/tile/cuda/test_ld_st_registers.py
@@ -736,8 +736,8 @@ def
test_reg_view_preserves_shared_and_register_addresses(case):
tensors.append(node.var)
tvm_ffi.structural_walk(views.body, visit)
- r_view = next(t for t in reversed(tensors) if t.ty.storage_scope ==
MemoryScope("local"))
- s_view = next(t for t in reversed(tensors) if t.ty.storage_scope ==
MemoryScope("shared"))
+ r_view = next(t for t in reversed(tensors) if t.ty.scope ==
MemoryScope("local"))
+ s_view = next(t for t in reversed(tensors) if t.ty.scope ==
MemoryScope("shared"))
analyzer = tvm.sym.Analyzer()
# Check all scalar coordinates, including the fp16 elements inside a word.
for tid in range(128):
