This is an automated email from the ASF dual-hosted git repository.
tqchen 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 7dc7f92403 refactor(op-dispatch): unify thread scopes as enum tags
(#20732)
7dc7f92403 is described below
commit 7dc7f92403d8da2258d2ef99ce099f3821af77e0
Author: Bohan Hou <[email protected]>
AuthorDate: Wed Oct 14 18:37:24 2026 -0400
refactor(op-dispatch): unify thread scopes as enum tags (#20732)
Thread scopes currently use a separate object registry, while tile Calls
store a string and recreate the scope during decoding and dispatch.
Replace these representations with canonical `ffi.Enum` tags shared by
Call attributes, Axis/Layout, axis fusers, and
`DispatchContext.thread_scope`.
CUDA and Trainium register their tags natively with explicit backend and
hierarchy-rank attributes. Rank comparisons are independent of enum
registration order. Instruction constructors resolve short names,
qualified names, or enum values before creating a Call: `scope="warp"`
now stores `ThreadScope("cuda.warp")`. Raw attribute construction
requires a scope tag or a registered qualified name.
Update TVMScript printing, JSON/pickle singleton round-trips, backend
helpers, documentation, and regression tests. Remove the former
`ExecScope` API and registration hooks; serialized data using the old
scope type must be regenerated. Native C++ tests cover registration,
rank ordering, typed attributes, and dispatch without Python.
Validation:
- Build and pre-commit pass.
- 304 targeted Python regressions pass with the repository-pinned
tvm-ffi: ThreadScope/Axis/Layout, Call validation and dispatch, CUDA
execution contexts, Trainium private allocation, IRBuilder void
emission, GPU script, printer snapshots, and JSON round-trips. This
includes the three printer expectations identified by the first CPU CI
run and updated for canonical enum tags.
- An earlier extended CUDA/TRN/script run had 571 passing cases and 3
skips. Its two JSON failures, together with three initial core JSON
failures, all pass with the pinned tvm-ffi. The original local
environment used the older 0.1.14.post1 release; a temporary
source-matched 0.1.15rc0 build resolves the dtype corruption.
- Two native C++ tests pass without loading Python and also passed in
CPU CI.
- Pinned benchmark workload import check passes for 93 kernels. Broader
downstream imports encounter existing references to the absent
`tvm.tirx.script.builder` and `IRBuilder._set_current_loc` APIs.
---
docs/reference/api/python/script/script.rst | 2 +-
docs/tirx/api/execution.rst | 37 ++++--
docs/tirx/api/tile_dispatch.rst | 24 ++--
docs/tirx/api/tirx.rst | 2 +-
docs/tirx/arch/tile_dispatch.rst | 8 +-
include/tvm/tirx/layout.h | 10 +-
include/tvm/tirx/script/ir_builder/frame.h | 2 +-
include/tvm/tirx/script/ir_builder/ir.h | 4 +-
include/tvm/tirx/stmt.h | 2 +-
include/tvm/tirx/tile/context.h | 18 +--
include/tvm/tirx/tile/exec_scope.h | 57 --------
include/tvm/tirx/tile/thread_scope.h | 82 ++++++++++++
python/tvm/backend/cuda/__init__.py | 7 +-
python/tvm/backend/cuda/script.py | 2 +-
.../{trn/exec_scope.py => cuda/thread_scope.py} | 18 +--
python/tvm/backend/trn/__init__.py | 11 +-
.../{cuda/exec_scope.py => trn/thread_scope.py} | 18 +--
python/tvm/backend/trn/tile/binary/default.py | 2 +-
python/tvm/backend/trn/tile/copy/default.py | 2 +-
python/tvm/backend/trn/tile/gemm/default.py | 2 +-
python/tvm/backend/trn/tile/reduction/utils.py | 2 +-
python/tvm/backend/trn/tile/select/default.py | 2 +-
python/tvm/backend/trn/tile/unary/default.py | 2 +-
.../tvm/backend/trn/tile/unary/with_bias_scale.py | 2 +-
.../backend/trn/transform/private_tensor_alloc.py | 6 +-
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/exec_scope.py | 45 -------
python/tvm/tirx/layout.py | 8 +-
python/tvm/tirx/script/ir_builder/ir.py | 4 +-
python/tvm/tirx/thread_scope.py | 64 +++++++++
python/tvm/tirx/tile/dispatch.py | 10 +-
python/tvm/tirx/tile/instruction.py | 35 +++--
src/backend/cuda/codegen/target_kind.cc | 12 ++
src/backend/cuda/layout/axis.cc | 87 +++++++------
src/backend/cuda/tile/exec_analysis.cc | 10 +-
src/backend/cuda/tile/exec_analysis.h | 4 +-
src/backend/cuda/tile/exec_context.h | 2 +-
.../cuda/tile/{exec_scope.cc => thread_scope.cc} | 2 +-
.../cuda/tile/{exec_scope.h => thread_scope.h} | 6 +-
src/backend/trn/codegen/target_kind.cc | 7 +
src/tirx/ir/layout/axis_registry.cc | 8 +-
src/tirx/ir/layout/tile_canonicalize.cc | 4 +-
src/tirx/ir/layout/tile_core.cc | 45 +++----
src/tirx/script/ir_builder/frame.cc | 2 +-
src/tirx/script/ir_builder/ir.cc | 2 +-
src/tirx/script/printer/block.cc | 15 +--
src/tirx/script/printer/expr.cc | 4 +-
src/tirx/tile/context.cc | 16 ++-
src/tirx/tile/exec_scope.cc | 68 ----------
src/tirx/tile/lower.cc | 12 +-
src/tirx/tile/op.cc | 12 +-
src/tirx/tile/thread_scope.cc | 80 ++++++++++++
tests/cpp/tirx/thread_scope_test.cc | 108 ++++++++++++++++
.../tirx/script/test_tirx_script_basic_usage.py | 18 ++-
tests/python/tirx/script/test_tirx_script_gpu.py | 2 +-
.../python/tirx/script/test_tirx_script_printer.py | 3 +-
tests/python/tirx/test_exec_scope.py | 43 -------
tests/python/tirx/test_layout.py | 24 ++--
tests/python/tirx/test_op.py | 2 +-
tests/python/tirx/test_thread_scope.py | 143 +++++++++++++++++++++
tests/python/tirx/test_tile_namespace.py | 24 ++--
tests/python/tirx/test_tirx_axis_enum.py | 4 +-
.../python/tirx/tile/cuda/test_add_sub_mul_div.py | 48 +++----
tests/python/tirx/tile/cuda/test_cp_async_bulk.py | 6 +-
.../tirx/tile/cuda/test_cp_async_bulk_tensor.py | 4 +-
tests/python/tirx/tile/cuda/test_mov.py | 18 +--
tests/python/tirx/tile/test_dispatcher.py | 27 +++-
tests/python/tirx/tile/trn/test_binary_trn.py | 6 +-
tests/python/tirx/tile/trn/test_compose_op_trn.py | 6 +-
tests/python/tirx/tile/trn/test_copy_trn.py | 6 +-
tests/python/tirx/tile/trn/test_gemm_trn.py | 6 +-
tests/python/tirx/tile/trn/test_reduction_trn.py | 6 +-
tests/python/tirx/tile/trn/test_select_trn.py | 6 +-
tests/python/tirx/tile/trn/test_unary_trn.py | 6 +-
74 files changed, 871 insertions(+), 535 deletions(-)
diff --git a/docs/reference/api/python/script/script.rst
b/docs/reference/api/python/script/script.rst
index 18a614e83a..dd287b4899 100644
--- a/docs/reference/api/python/script/script.rst
+++ b/docs/reference/api/python/script/script.rst
@@ -58,7 +58,7 @@ The canonical ``MemoryScope`` class is documented in
:mod:`tvm.runtime`.
:members:
:exclude-members: AlreadyEmitted, DataTypeImm, FuncType, GenericConst,
IRBuilder,
IRModuleFrame, MemoryScope, MissingType, Op, PrimType, StringImm,
StringType, TensorRegion,
- TensorRegionType, Type, Range, meta_var, Var, Call, CommReducer, Reduce,
FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod, LShift,
RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min, Max, EQ, NE, LT,
LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle, CallEffectKind,
IterVar, ComposeLayout, DtypeConstructor, ExecScope, Iter, Layout,
LetAnnotation, LocalVectorAnnotation, TileLayout, CallFFIKernelAttr,
TensorMapEncodeTiledAttr, function, jit, inline, macro
+ TensorRegionType, Type, Range, meta_var, Var, Call, CommReducer, Reduce,
FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod, LShift,
RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min, Max, EQ, NE, LT,
LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle, CallEffectKind,
IterVar, ComposeLayout, DtypeConstructor, ThreadScope, Iter, Layout,
LetAnnotation, LocalVectorAnnotation, TileLayout, CallFFIKernelAttr,
TensorMapEncodeTiledAttr, function, jit, inli [...]
.. autofunction:: tvm.script.tirx.function
diff --git a/docs/tirx/api/execution.rst b/docs/tirx/api/execution.rst
index 24c8273f68..6c22eebb30 100644
--- a/docs/tirx/api/execution.rst
+++ b/docs/tirx/api/execution.rst
@@ -15,32 +15,47 @@
specific language governing permissions and limitations
under the License.
-Execution Scope IR
-==================
+Thread Scope Enum Tags
+======================
-tvm.tirx.exec_scope
--------------------
-.. automodule:: tvm.tirx.exec_scope
+tvm.tirx.thread_scope
+---------------------
+.. automodule:: tvm.tirx.thread_scope
:members:
:no-index:
Backend scopes
--------------
-Execution scope identities remain part of generic layout IR. Backends register
-their finite execution hierarchies and provide the user-facing constructors:
-``T.cuda.ExecScope("warp")`` and ``T.trn.ExecScope("thread")``.
-``T.cuda.ExecScope("kernel")`` names the root used by grid axes such as ``bx``
+Thread scopes are canonical ``ffi.Enum`` identities shared by layout IR,
+tile Call attributes, and dispatch contexts. Query a registered tag with
+``T.ThreadScope("cuda.warp")``. The ``name`` property is the qualified enum
+name; ``local_name`` is ``"warp"``. ``backend`` and ``rank`` are enum
attributes,
+and registration order has no hierarchy semantics.
+
+Backends register their finite thread hierarchies natively using ``EnumDef``
+and explicit backend and rank attributes. Their user-facing lookup helpers are:
+``T.cuda.ThreadScope("warp")`` and ``T.trn.ThreadScope("thread")``.
+``T.cuda.ThreadScope("kernel")`` names the root used by grid axes such as
``bx``
in the ``kernel -> cta`` relationship. Every scope uniformly describes a
hierarchy node; each backend tile instruction declares the scopes it supports.
+Same-backend hierarchy comparisons use explicit ranks; cross-backend
comparisons
+are rejected.
For example, CUDA TMA instructions accept ``thread`` and ``tcgen05.ld/st``
accept
``warpgroup``. CUDA index bindings and predicate analysis live in
``tvm.backend.cuda.tile.exec_context``.
-.. automodule:: tvm.backend.cuda.exec_scope
+Instruction APIs continue accepting ``scope="warp"`` and also accept qualified
+names or ``ThreadScope`` values. Construction resolves the value before
+creating the Call; ``Call.attrs.scope`` always contains a ``ThreadScope`` enum.
+Omitting scope uses the instruction backend's ``thread`` tag. Direct attribute
+construction must provide the tag or its qualified name; shared attribute
+schemas have no backend-independent default.
+
+.. automodule:: tvm.backend.cuda.thread_scope
:members:
:no-index:
-.. automodule:: tvm.backend.trn.exec_scope
+.. automodule:: tvm.backend.trn.thread_scope
:members:
:no-index:
diff --git a/docs/tirx/api/tile_dispatch.rst b/docs/tirx/api/tile_dispatch.rst
index c48ba7466c..7feed70fe4 100644
--- a/docs/tirx/api/tile_dispatch.rst
+++ b/docs/tirx/api/tile_dispatch.rst
@@ -43,7 +43,9 @@ registers an adapter that decodes the semantic view and
invokes its lowerer.
Custom implementations can register the callback directly using
``tvm.ir.register_op_attr(op_name, "FTileOpDispatch", dispatch)`` or C++
``OpDef::set_attr<tile::FTileOpDispatch>``. The Call must carry reflected
-attributes with a string ``scope`` field in the target's execution hierarchy.
+attributes with a required ``ThreadScope`` ``scope`` field in the target's
+execution hierarchy. Direct construction may use a registered qualified name,
+such as ``"cuda.thread"``, which the typed FFI boundary resolves to its enum.
Named workspace operands have fixed optional slots. Declare a workspace policy
only for instructions that own such scratch storage; do not use it to hide
@@ -55,13 +57,15 @@ The common Python API is under ``tvm.tirx.tile`` (also
``T.tile``):
``Instruction``, ``Operand``, ``TensorCall``, and ``DispatchContext``. Backend
implementations live under ``tvm.backend.cuda.tile`` and
``tvm.backend.trn.tile``. Their user-facing instruction paths remain
-``T.cuda.tile`` and ``T.trn.tile``. Operator names and serialized type keys
-remain unchanged.
+``T.cuda.tile`` and ``T.trn.tile``. Instruction operator names and attribute
+schema type keys remain unchanged. Scope tags use the ``tirx.ThreadScope`` type
+key; serialized data containing the former scope type must be regenerated.
-Execution scopes are generic backend/name identities. Backend load hooks
-register their supported scopes and ordering; CUDA exposes
-``T.cuda.ExecScope("warp")`` and Trainium exposes
-``T.trn.ExecScope("thread")``. A string instruction scope is resolved by the
-instruction's backend. CUDA predicate analysis, index bindings, and active
-thread sets live in its private tile implementation; Trainium does not use
-CUDA's execution-context model.
+Thread scopes are backend-qualified enum identities. Native backends register
+their supported tags and explicit ranks with ``EnumDef``; CUDA exposes
+``T.cuda.ThreadScope("warp")`` and Trainium exposes
+``T.trn.ThreadScope("thread")``. A string instruction scope is resolved by the
+instruction's backend before constructing its Call. Lowerers consume the enum
+directly through ``DispatchContext.thread_scope``. CUDA predicate analysis,
+index bindings, and active thread sets live in its private tile implementation;
+Trainium does not use CUDA's execution-context model.
diff --git a/docs/tirx/api/tirx.rst b/docs/tirx/api/tirx.rst
index 033fdcbd4e..f02bfe7720 100644
--- a/docs/tirx/api/tirx.rst
+++ b/docs/tirx/api/tirx.rst
@@ -33,6 +33,6 @@ from ``tvm/ir/type.h``.
:members:
:imported-members:
:exclude-members: Call, CommReducer, ComposeLayout, DispatchContext,
- ExecScope, Expr, Layout, Op, PrimExpr, Reduce,
+ ThreadScope, Expr, Layout, Op, PrimExpr, Reduce,
TileLayout, Var, build, const,
get_default_tir_pipeline, get_tir_pipeline, register_tir_pipeline
diff --git a/docs/tirx/arch/tile_dispatch.rst b/docs/tirx/arch/tile_dispatch.rst
index 4e000d1e7f..8d2fa2522b 100644
--- a/docs/tirx/arch/tile_dispatch.rst
+++ b/docs/tirx/arch/tile_dispatch.rst
@@ -28,8 +28,9 @@ The public C++ tile interfaces are grouped under
``include/tvm/tirx/tile/``:
- ``op.h`` declares the instruction registration and validation contract.
- ``dispatch.h`` declares ``FTileOpDispatch`` and the dispatcher entry point.
-- ``exec_scope.h`` defines backend execution hierarchy nodes used by layouts
- and tile calls. The shared ``tirx::ExecScope`` type remains unchanged.
+- ``thread_scope.h`` defines backend-qualified ``tirx::ThreadScope`` enum tags
+ used by layouts, tile Calls, and dispatch contexts. Backends attach explicit
+ hierarchy ranks through enum attributes.
- ``context.h`` defines the backend-neutral ``tile::DispatchContext`` and its
allocation and initialization callbacks.
@@ -41,7 +42,8 @@ active-thread execution contexts, and predicate analysis
remain private to
``tirx.TileDispatch`` remains the first phase of ``LowerTIRx``.
The pass recognizes ``Evaluate(Call)`` by the presence of the operator's
``FTileOpDispatch`` attribute. It validates the call, reads its reflected
-``scope`` attribute, resolves the scope against the active thread set,
+``ThreadScope`` ``scope`` attribute, checks its backend against the target,
+resolves the scope against the active thread set,
constructs a ``DispatchContext``, and
invokes the instruction's registered lowerer. It replaces the Evaluate with
the returned function body and lowers any nested tensor calls.
diff --git a/include/tvm/tirx/layout.h b/include/tvm/tirx/layout.h
index ea8bb5bd7f..e499fa38e4 100644
--- a/include/tvm/tirx/layout.h
+++ b/include/tvm/tirx/layout.h
@@ -31,7 +31,7 @@
#include <tvm/ir/expr.h>
#include <tvm/ir/module.h>
#include <tvm/target/target.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
namespace tvm {
@@ -78,7 +78,7 @@ class Layout : public ffi::ObjectRef {
};
// target, subscope, scope, iter -> fused_iter
-using FAxisFuser = ffi::TypedFunction<ffi::Optional<Iter>(Target, ffi::String,
ffi::String, Iter)>;
+using FAxisFuser = ffi::TypedFunction<ffi::Optional<Iter>(Target, ThreadScope,
ThreadScope, Iter)>;
// Axis
class AxisNode : public ffi::EnumObj {
@@ -93,10 +93,10 @@ class AxisNode : public ffi::EnumObj {
ffi::Optional<ffi::String> GetBackend() const;
/*! \brief Get the scope of the (thread) axis. */
- ffi::Optional<ExecScope> GetScope() const;
+ ffi::Optional<ThreadScope> GetScope() const;
/*! \brief Get the subscope of the (thread) axis. */
- ffi::Optional<ExecScope> GetSubscope() const;
+ ffi::Optional<ThreadScope> GetSubscope() const;
/*! \brief Get the fuser of the (thread) axis. */
ffi::Optional<FAxisFuser> GetFuser() const;
@@ -198,7 +198,7 @@ class TileLayoutNode : public LayoutNode {
bool HasThreadAxis() const;
/*! \brief Get the scope pair of the layout */
- ffi::Optional<Tuple<ExecScope, ExecScope>> GetScope() const;
+ ffi::Optional<Tuple<ThreadScope, ThreadScope>> GetScope() const;
/*! \brief Get the default layout for the shape */
TVM_TIRX_DLL static TileLayout DefaultLayout(ffi::Array<PrimExpr> shape);
diff --git a/include/tvm/tirx/script/ir_builder/frame.h
b/include/tvm/tirx/script/ir_builder/frame.h
index c8c9c0d901..04e8cbc228 100644
--- a/include/tvm/tirx/script/ir_builder/frame.h
+++ b/include/tvm/tirx/script/ir_builder/frame.h
@@ -23,7 +23,7 @@
#include <tvm/script/ir_builder/frame.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include <functional>
#include <utility>
diff --git a/include/tvm/tirx/script/ir_builder/ir.h
b/include/tvm/tirx/script/ir_builder/ir.h
index e9b4a61e62..4c47bf960b 100644
--- a/include/tvm/tirx/script/ir_builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -26,7 +26,7 @@
#include <tvm/script/ir_builder/ir.h>
#include <tvm/tirx/script/ir_builder/frame.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include <tvm/tirx/type.h>
namespace tvm {
@@ -52,8 +52,8 @@ using tvm::Var;
using tvm::ffi::Tensor;
using tvm::ffi::Tuple;
using tvm::ffi::Variant;
-using tvm::tirx::ExecScope;
using tvm::tirx::TensorVar;
+using tvm::tirx::ThreadScope;
/*!
* \brief The primitive function statement.
diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h
index 95cb330ae7..eb1b26a4a7 100644
--- a/include/tvm/tirx/stmt.h
+++ b/include/tvm/tirx/stmt.h
@@ -24,7 +24,7 @@
#include <tvm/ir/stmt.h>
#include <tvm/tirx/expr.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
namespace tvm {
namespace tirx {
diff --git a/include/tvm/tirx/tile/context.h b/include/tvm/tirx/tile/context.h
index 23591e400a..b075f35d1c 100644
--- a/include/tvm/tirx/tile/context.h
+++ b/include/tvm/tirx/tile/context.h
@@ -22,7 +22,7 @@
#include <tvm/target/target.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
namespace tvm {
namespace tirx {
@@ -50,15 +50,15 @@ constexpr const char* kPostTensorDefStmt =
"post_tensor_def_stmt";
*/
class DispatchContextObj : public ffi::Object {
public:
- DispatchContextObj(Target target, ExecScope exec_scope)
- : target(std::move(target)), exec_scope(std::move(exec_scope)) {}
+ DispatchContextObj(Target target, ThreadScope thread_scope)
+ : target(std::move(target)), thread_scope(std::move(thread_scope)) {}
explicit DispatchContextObj(ffi::UnsafeInit)
- : target(ffi::UnsafeInit{}), exec_scope(ffi::UnsafeInit{}) {}
+ : target(ffi::UnsafeInit{}), thread_scope(ffi::UnsafeInit{}) {}
/*! \brief The target of the kernel. */
Target target;
- /*! \brief The exec scope of the operator */
- ExecScope exec_scope;
+ /*! \brief The thread scope of the operator */
+ ThreadScope thread_scope;
/*! \brief Hardware tag to (index variable, extent) launch parameters. */
ffi::Map<ffi::String, ffi::Tuple<PrimVar, PrimExpr>> launch_params;
/*! \brief A map from loop variables to their ranges. */
@@ -75,7 +75,7 @@ class DispatchContextObj : public ffi::Object {
* Maps axis name ("laneid"/"warpid"/"cta_id"/"wid_in_wg"/"wgid") to a
* 2-element [extent, offset] PrimExpr array. Empty map = no ExecContext
* tracking available (fallback for unresolved filters, pre-Phase-4 call
- * sites, etc.); dispatchers should fall back to exec_scope.name in that
+ * sites, etc.); dispatchers should fall back to thread_scope.local_name in
that
* case.
*/
ffi::Map<ffi::String, ffi::Array<PrimExpr>> inter;
@@ -88,7 +88,7 @@ class DispatchContextObj : public ffi::Object {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DispatchContextObj>()
.def_ro("target", &DispatchContextObj::target)
- .def_ro("exec_scope", &DispatchContextObj::exec_scope)
+ .def_ro("thread_scope", &DispatchContextObj::thread_scope)
.def_ro("launch_params", &DispatchContextObj::launch_params)
.def_ro("var_range_map", &DispatchContextObj::var_range_map)
.def_ro("alloc_only", &DispatchContextObj::alloc_only)
@@ -124,7 +124,7 @@ class DispatchContextObj : public ffi::Object {
class DispatchContext : public ffi::ObjectRef {
public:
TVM_TIRX_DLL DispatchContext(
- Target target, ExecScope exec_scope,
+ Target target, ThreadScope thread_scope,
ffi::Map<ffi::String, ffi::Tuple<PrimVar, PrimExpr>> launch_params = {},
ffi::Map<Var, Range> var_range_map = {}, bool alloc_only = false,
ffi::Map<ffi::String, ffi::ObjectRef> callbacks = {},
diff --git a/include/tvm/tirx/tile/exec_scope.h
b/include/tvm/tirx/tile/exec_scope.h
deleted file mode 100644
index 541f6391cc..0000000000
--- a/include/tvm/tirx/tile/exec_scope.h
+++ /dev/null
@@ -1,57 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-/*! \brief Backend-defined execution hierarchy nodes shared by layouts and
tile calls. */
-#ifndef TVM_TIRX_TILE_EXEC_SCOPE_H_
-#define TVM_TIRX_TILE_EXEC_SCOPE_H_
-#include <tvm/ffi/object.h>
-#include <tvm/ir/expr.h>
-
-namespace tvm::tirx {
-
-/*! \brief Register a backend's execution hierarchy. Smaller rank means wider
scope. */
-TVM_TIRX_DLL void RegisterExecScope(const ffi::String& backend, const
ffi::String& name, int rank);
-
-class ExecScopeNode : public ffi::Object {
- public:
- ffi::String backend;
- ffi::String scope_name;
- ffi::String name() const { return scope_name; }
- int rank() const;
-
- static void RegisterReflection() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<ExecScopeNode>()
- .def_ro("backend", &ExecScopeNode::backend)
- .def_ro("name", &ExecScopeNode::scope_name);
- }
- static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
- TVM_FFI_DECLARE_OBJECT_INFO("tirx.ExecScope", ExecScopeNode, ffi::Object);
-};
-
-class ExecScope : public ffi::ObjectRef {
- public:
- TVM_TIRX_DLL ExecScope(ffi::String backend, ffi::String name);
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExecScope, ffi::ObjectRef,
ExecScopeNode);
-};
-
-/*! \brief Compare scopes within the same backend's hierarchy. */
-TVM_TIRX_DLL bool ScopeHigher(const ExecScope& a, const ExecScope& b);
-
-} // namespace tvm::tirx
-#endif // TVM_TIRX_TILE_EXEC_SCOPE_H_
diff --git a/include/tvm/tirx/tile/thread_scope.h
b/include/tvm/tirx/tile/thread_scope.h
new file mode 100644
index 0000000000..98957211d8
--- /dev/null
+++ b/include/tvm/tirx/tile/thread_scope.h
@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*! \brief Enum identities for backend-defined thread hierarchies. */
+#ifndef TVM_TIRX_TILE_THREAD_SCOPE_H_
+#define TVM_TIRX_TILE_THREAD_SCOPE_H_
+#include <tvm/ffi/enum.h>
+#include <tvm/ffi/reflection/enum_def.h>
+#include <tvm/runtime/base.h>
+
+#include <ostream>
+
+namespace tvm::tirx {
+
+namespace thread_scope_attr {
+/*! \brief Backend owning a thread scope. */
+inline constexpr ffi::reflection::AttrKey<ffi::String> kBackend{"backend"};
+/*! \brief Explicit hierarchy rank. Smaller ranks denote wider scopes. */
+inline constexpr ffi::reflection::AttrKey<int64_t> kRank{"rank"};
+} // namespace thread_scope_attr
+
+class ThreadScopeObj : public ffi::EnumObj {
+ public:
+ const ffi::String& name() const { return _str_index; }
+ TVM_TIRX_DLL ffi::String backend() const;
+ TVM_TIRX_DLL ffi::String local_name() const;
+ TVM_TIRX_DLL int64_t rank() const;
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ThreadScope", ThreadScopeObj,
ffi::EnumObj);
+};
+
+/*!
+ * \brief Canonical thread-scope identity, named as backend.scope.
+ *
+ * Backends register tags with EnumDef<ThreadScopeObj> and explicit backend and
+ * rank attributes. Registration order has no hierarchy semantics.
+ */
+class ThreadScope : public ffi::Enum {
+ public:
+ explicit ThreadScope(const ffi::String& name) : ffi::Enum(Get(name)) {}
+ TVM_TIRX_DLL static ThreadScope Get(const ffi::String& name);
+ const ffi::String& name() const { return get()->name(); }
+ bool operator==(const ThreadScope& other) const { return same_as(other); }
+ bool operator!=(const ThreadScope& other) const { return !same_as(other); }
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ThreadScope, ffi::Enum,
ThreadScopeObj);
+};
+
+/*! \brief Compare explicit ranks within one backend's hierarchy. */
+TVM_TIRX_DLL bool ScopeHigher(const ThreadScope& a, const ThreadScope& b);
+
+inline std::ostream& operator<<(std::ostream& os, const ThreadScope& scope) {
+ return os << scope.name();
+}
+
+} // namespace tvm::tirx
+
+namespace tvm::ffi {
+template <>
+inline constexpr bool use_default_type_traits_v<tirx::ThreadScope> = false;
+template <>
+struct TypeTraits<tirx::ThreadScope>
+ : public ObjectRefWithFallbackTraitsBase<tirx::ThreadScope, ffi::String> {
+ static tirx::ThreadScope ConvertFallbackValue(ffi::String name) {
+ return tirx::ThreadScope::Get(name);
+ }
+};
+} // namespace tvm::ffi
+#endif // TVM_TIRX_TILE_THREAD_SCOPE_H_
diff --git a/python/tvm/backend/cuda/__init__.py
b/python/tvm/backend/cuda/__init__.py
index 98ad35d5a9..44d71992a8 100644
--- a/python/tvm/backend/cuda/__init__.py
+++ b/python/tvm/backend/cuda/__init__.py
@@ -33,7 +33,7 @@ _LAZY_SUBMODULES = {
"lang",
"layout",
"axis",
- "exec_scope",
+ "thread_scope",
"launch",
"op",
"ptx",
@@ -60,9 +60,8 @@ def _detect_target_from_device(dev):
def register_backend():
"""Register CUDA-owned Python semantics."""
- from . import axis, exec_scope
+ from . import axis
- exec_scope.register()
axis.register()
from tvm.target.detect_target import register_device_target_detector
from tvm.tirx.script.ir_builder import (
@@ -127,7 +126,6 @@ __all__ = [
"axis",
"codegen",
"cpp",
- "exec_scope",
"export_cuda_host",
"iket",
"lang",
@@ -140,6 +138,7 @@ __all__ = [
"script_namespace",
"script_namespaces",
"target_tags",
+ "thread_scope",
"tile",
"transforms",
]
diff --git a/python/tvm/backend/cuda/script.py
b/python/tvm/backend/cuda/script.py
index 4b6cdedec8..9db0ef5e2b 100644
--- a/python/tvm/backend/cuda/script.py
+++ b/python/tvm/backend/cuda/script.py
@@ -25,7 +25,6 @@ from tvm.script.ir_builder.base import auto_emit_void_call as
_auto_emit_void_ca
from . import layout as layout
from .axis import Axis as Axis
from .backend_config import BackendConfig as BackendConfig
-from .exec_scope import ExecScope as ExecScope
from .launch import AccessPolicyWindow as AccessPolicyWindow
from .launch import KernelAttributes as KernelAttributes
from .launch import LaunchCompletionEvent as LaunchCompletionEvent
@@ -34,6 +33,7 @@ from .launch import MemSyncDomainMap as MemSyncDomainMap
from .launch import ProgrammaticEvent as ProgrammaticEvent
from .ptx.engine import PTXNamespace as _PTXNamespace
from .ptx.engine import _InstrChain
+from .thread_scope import ThreadScope as ThreadScope
from .tile.instructions import make_namespace as _make_tensor_namespace
# pylint: disable=protected-access
diff --git a/python/tvm/backend/trn/exec_scope.py
b/python/tvm/backend/cuda/thread_scope.py
similarity index 67%
rename from python/tvm/backend/trn/exec_scope.py
rename to python/tvm/backend/cuda/thread_scope.py
index 7c69780d5f..2103fa1c28 100644
--- a/python/tvm/backend/trn/exec_scope.py
+++ b/python/tvm/backend/cuda/thread_scope.py
@@ -16,19 +16,11 @@
# under the License.
# pylint: disable=no-member, super-init-not-called
-"""TRN execution hierarchy."""
+"""CUDA thread-scope lookup."""
-from tvm.tirx.exec_scope import ExecScope as _ExecScope
-from tvm.tirx.exec_scope import register_exec_scope
+from tvm.tirx.thread_scope import ThreadScope as _ThreadScope
-SCOPES = ("kernel", "thread")
-
-def register():
- for rank, name in enumerate(SCOPES, start=-1):
- register_exec_scope("trn", name, rank)
-
-
-def ExecScope(name: str) -> _ExecScope:
- """Construct a TRN execution scope."""
- return _ExecScope("trn", name)
+def ThreadScope(name: str) -> _ThreadScope:
+ """Look up a CUDA thread scope by its backend-local name."""
+ return _ThreadScope("cuda." + name)
diff --git a/python/tvm/backend/trn/__init__.py
b/python/tvm/backend/trn/__init__.py
index 12e9d01bcf..e47553953e 100644
--- a/python/tvm/backend/trn/__init__.py
+++ b/python/tvm/backend/trn/__init__.py
@@ -21,7 +21,7 @@ from importlib import import_module
_LAZY_SUBMODULES = {
"layout",
"axis",
- "exec_scope",
+ "thread_scope",
"op",
"pipeline",
"script",
@@ -33,9 +33,8 @@ _LAZY_SUBMODULES = {
def register_backend():
"""Register Trainium-owned Python semantics."""
- from . import axis, exec_scope
+ from . import axis
- exec_scope.register()
axis.register()
from tvm.tirx import compilation_pipeline # pylint:
disable=import-outside-toplevel
from tvm.tirx.script.ir_builder import (
@@ -65,14 +64,14 @@ def script_namespaces(**_):
from . import layout
from .axis import Axis
- from .exec_scope import ExecScope
+ from .thread_scope import ThreadScope
from .tile import private_alloc # register workspace policies
from .tile.instructions import make_namespace
return {
"nki": script_namespace(),
"trn": SimpleNamespace(
- tile=make_namespace(), Axis=Axis, ExecScope=ExecScope,
layout=layout
+ tile=make_namespace(), Axis=Axis, ThreadScope=ThreadScope,
layout=layout
),
}
@@ -85,7 +84,6 @@ def __getattr__(name: str):
__all__ = [
"axis",
- "exec_scope",
"layout",
"op",
"pipeline",
@@ -94,6 +92,7 @@ __all__ = [
"script_namespace",
"script_namespaces",
"target_tags",
+ "thread_scope",
"tile",
"transform",
]
diff --git a/python/tvm/backend/cuda/exec_scope.py
b/python/tvm/backend/trn/thread_scope.py
similarity index 65%
rename from python/tvm/backend/cuda/exec_scope.py
rename to python/tvm/backend/trn/thread_scope.py
index e2cbd2d15b..9c35e31c0d 100644
--- a/python/tvm/backend/cuda/exec_scope.py
+++ b/python/tvm/backend/trn/thread_scope.py
@@ -16,19 +16,11 @@
# under the License.
# pylint: disable=no-member, super-init-not-called
-"""CUDA execution hierarchy."""
+"""TRN thread-scope lookup."""
-from tvm.tirx.exec_scope import ExecScope as _ExecScope
-from tvm.tirx.exec_scope import register_exec_scope
+from tvm.tirx.thread_scope import ThreadScope as _ThreadScope
-SCOPES = ("kernel", "cluster", "cta", "warpgroup", "warp", "thread")
-
-def register():
- for rank, name in enumerate(SCOPES, start=-1):
- register_exec_scope("cuda", name, rank)
-
-
-def ExecScope(name: str) -> _ExecScope:
- """Construct a CUDA execution scope."""
- return _ExecScope("cuda", name)
+def ThreadScope(name: str) -> _ThreadScope:
+ """Look up a TRN thread scope by its backend-local name."""
+ return _ThreadScope("trn." + name)
diff --git a/python/tvm/backend/trn/tile/binary/default.py
b/python/tvm/backend/trn/tile/binary/default.py
index 6d3ee356ce..ff3ad80174 100644
--- a/python/tvm/backend/trn/tile/binary/default.py
+++ b/python/tvm/backend/trn/tile/binary/default.py
@@ -31,7 +31,7 @@ from .utils import InstType, binary_map_ops,
try_find_inst_nary
def binary_trn(op: TensorCall, binary_op: MapOpType, sctx: DispatchContext) ->
Function | None:
"""Generate a binary operation schedule for Trainium."""
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
- fail("requires Trainium target and thread exec_scope")
+ fail("requires Trainium target and thread scope")
assert binary_op in binary_map_ops, f"Unsupported binary operation
{binary_op}"
diff --git a/python/tvm/backend/trn/tile/copy/default.py
b/python/tvm/backend/trn/tile/copy/default.py
index 5660a4f6b4..b3d9dffe10 100644
--- a/python/tvm/backend/trn/tile/copy/default.py
+++ b/python/tvm/backend/trn/tile/copy/default.py
@@ -47,7 +47,7 @@ def copy_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
"""Schedule copy operation between global and shared memory on CUDA."""
# Basic validation checks
if sctx.scope_kind != "thread":
- fail("requires thread exec_scope for TRN copy")
+ fail("requires thread scope for TRN copy")
dst_region, src_region = op.args
src, dst = src_region.source, dst_region.source
diff --git a/python/tvm/backend/trn/tile/gemm/default.py
b/python/tvm/backend/trn/tile/gemm/default.py
index e20add61b7..da8efff59f 100644
--- a/python/tvm/backend/trn/tile/gemm/default.py
+++ b/python/tvm/backend/trn/tile/gemm/default.py
@@ -121,7 +121,7 @@ def matmul_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
"""Schedule GEMM operation on Trainium."""
# Basic validation checks
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
- fail("requires Trainium target and thread exec_scope")
+ fail("requires Trainium target and thread scope")
# Extract arguments
(
diff --git a/python/tvm/backend/trn/tile/reduction/utils.py
b/python/tvm/backend/trn/tile/reduction/utils.py
index 5e64bceff5..e4c6732f30 100644
--- a/python/tvm/backend/trn/tile/reduction/utils.py
+++ b/python/tvm/backend/trn/tile/reduction/utils.py
@@ -87,7 +87,7 @@ def reduction_trn(
Optional[Function]: The scheduled function, or None if not applicable.
"""
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
- fail("requires Trainium target and thread exec_scope")
+ fail("requires Trainium target and thread scope")
op = TensorCall.decode(op)
dst_tensor_region, src_tensor_region = op.output, op.input
diff --git a/python/tvm/backend/trn/tile/select/default.py
b/python/tvm/backend/trn/tile/select/default.py
index f6506ac510..c591a51c14 100644
--- a/python/tvm/backend/trn/tile/select/default.py
+++ b/python/tvm/backend/trn/tile/select/default.py
@@ -41,7 +41,7 @@ def _scope_trn_sbuf():
def select_trn(op: TensorCall, sctx: DispatchContext) -> Function | None:
"""Generate schedule for select operation on Trainium."""
if sctx.scope_kind != "thread":
- fail("requires thread exec_scope for TRN select")
+ fail("requires thread scope for TRN select")
op = TensorCall.decode(op)
assert isinstance(op, Select), f"{op} is not a Select"
diff --git a/python/tvm/backend/trn/tile/unary/default.py
b/python/tvm/backend/trn/tile/unary/default.py
index 4b34297e4a..1cbaae9fb8 100644
--- a/python/tvm/backend/trn/tile/unary/default.py
+++ b/python/tvm/backend/trn/tile/unary/default.py
@@ -36,7 +36,7 @@ def unary_trn(op: TensorCall, unary_op: MapOpType, sctx:
DispatchContext) -> Fun
"""Schedule unary operation on Trainium."""
# Check execution environment
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
- fail("requires Trainium target and thread exec_scope")
+ fail("requires Trainium target and thread scope")
# Extract operation arguments
dst_tensor_region, _src = op.args
diff --git a/python/tvm/backend/trn/tile/unary/with_bias_scale.py
b/python/tvm/backend/trn/tile/unary/with_bias_scale.py
index 1e1e08a8c9..149f96ec64 100644
--- a/python/tvm/backend/trn/tile/unary/with_bias_scale.py
+++ b/python/tvm/backend/trn/tile/unary/with_bias_scale.py
@@ -35,7 +35,7 @@ def unary_with_bias_scale_trn(
"""Schedule unary operation with bias and scale on Trainium."""
# Check execution environment
if not (sctx.is_target("trn") and sctx.scope_kind == "thread"):
- fail("requires Trainium target and thread exec_scope")
+ fail("requires Trainium target and thread scope")
if len(op.args) == 2:
dst_tensor_region, src_tensor_region = op.args
diff --git a/python/tvm/backend/trn/transform/private_tensor_alloc.py
b/python/tvm/backend/trn/transform/private_tensor_alloc.py
index 6fbab19b3e..11ec7c486b 100644
--- a/python/tvm/backend/trn/transform/private_tensor_alloc.py
+++ b/python/tvm/backend/trn/transform/private_tensor_alloc.py
@@ -61,11 +61,11 @@ def _collect_private_allocations(stmt: Stmt, target:
Target):
return
call = op
op = TensorCall.decode(call)
- exec_scope = op.scope
- scope_kind = op.scope.name
+ thread_scope = op.scope
+ scope_kind = op.scope.local_name
sctx = DispatchContext(
target=target,
- exec_scope=exec_scope,
+ thread_scope=thread_scope,
launch_params=launch_params,
var_range_map=var_range_map,
alloc_only=True,
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 5c605f686d..b79baf8c63 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -78,7 +78,7 @@ from .op import gpu_dp4a
from .op import ignore_loop_partition
# TIRX-specific imports (must come before subpackage imports to avoid circular
imports)
-from .exec_scope import ExecScope
+from .thread_scope import ThreadScope
from .layout import TileLayout, Layout, ComposeLayout
from . import transform
from . import analysis
diff --git a/python/tvm/tirx/exec_scope.py b/python/tvm/tirx/exec_scope.py
deleted file mode 100644
index fecd51b539..0000000000
--- a/python/tvm/tirx/exec_scope.py
+++ /dev/null
@@ -1,45 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# pylint: disable=no-member, super-init-not-called
-
-"""Backend-defined execution scope identities."""
-
-from tvm_ffi import register_object
-
-from tvm.ir import Scriptable
-from tvm.runtime import Object
-
-from . import _ffi_api
-
-
-@register_object("tirx.ExecScope")
-class ExecScope(Object, Scriptable):
- """A node in one backend's registered execution hierarchy."""
-
- backend: str
- name: str
-
- def __init__(self, backend: str, name: str):
- self.__init_handle_by_constructor__(_ffi_api.ExecScope, backend, name)
-
-
-def register_exec_scope(backend: str, name: str, rank: int):
- """Register a backend's finite scope hierarchy, from outer to inner ranks.
-
- Backend load hooks register scopes before registering their layout axes.
- """
- _ffi_api.RegisterExecScope(backend, name, rank)
diff --git a/python/tvm/tirx/layout.py b/python/tvm/tirx/layout.py
index e43954aeae..e7056d0055 100644
--- a/python/tvm/tirx/layout.py
+++ b/python/tvm/tirx/layout.py
@@ -28,7 +28,7 @@ from tvm.runtime import Object
from tvm.tirx.expr import Expr
from . import _ffi_api
-from .exec_scope import ExecScope
+from .thread_scope import ThreadScope
@tvm_ffi.register_object("tirx.Layout")
@@ -312,11 +312,11 @@ class Axis(Enum, metaclass=_AxisMeta,
type_key="tirx.Axis", init=False):
"""Backend owning this axis, or None for a generic axis."""
return _ffi_api.AxisGetBackend(self)
- def get_scope(self) -> ExecScope | None:
+ def get_scope(self) -> ThreadScope | None:
"""Get the scope of the axis."""
return _ffi_api.AxisGetScope(self) # pylint: disable=no-member
- def get_subscope(self) -> ExecScope | None:
+ def get_subscope(self) -> ThreadScope | None:
"""Get the subscope of the axis."""
return _ffi_api.AxisGetSubscope(self) # pylint: disable=no-member
@@ -597,7 +597,7 @@ class TileLayout(Layout):
"""
return _ffi_api.TileLayoutGroup(self, shape) # pylint:
disable=no-member
- def get_scope(self) -> tuple[ExecScope, ExecScope] | None:
+ def get_scope(self) -> tuple[ThreadScope, ThreadScope] | None:
"""Get the scope pair of the layout."""
return _ffi_api.TileLayoutGetScope(self) # pylint: disable=no-member
diff --git a/python/tvm/tirx/script/ir_builder/ir.py
b/python/tvm/tirx/script/ir_builder/ir.py
index d4ffb16a95..da4d9b9615 100644
--- a/python/tvm/tirx/script/ir_builder/ir.py
+++ b/python/tvm/tirx/script/ir_builder/ir.py
@@ -42,7 +42,6 @@ from tvm.target import Target
# pylint: disable=unused-import
from tvm.tirx import Expr, IndexMap, Var, is_tensor_var, tile
-from tvm.tirx.exec_scope import ExecScope
# import tirx.expr for direct ir construction to pass structural_equal
comparison
from tvm.tirx.expr import (
@@ -51,6 +50,7 @@ from tvm.tirx.expr import (
TensorLoad,
)
from tvm.tirx.layout import Axis, ComposeLayout, Iter, Layout, R, S, TileLayout
+from tvm.tirx.thread_scope import ThreadScope
from . import _ffi_api
@@ -1064,7 +1064,6 @@ Tensor = tir.TensorType
__all__ = [
"Axis",
"ComposeLayout",
- "ExecScope",
"FloatImm",
"IntImm",
"Iter",
@@ -1080,6 +1079,7 @@ __all__ = [
"TensorLoad",
"TensorMap",
"TensorRegion",
+ "ThreadScope",
"TileLayout",
"Tuple",
"Var",
diff --git a/python/tvm/tirx/thread_scope.py b/python/tvm/tirx/thread_scope.py
new file mode 100644
index 0000000000..df99a6f30e
--- /dev/null
+++ b/python/tvm/tirx/thread_scope.py
@@ -0,0 +1,64 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=no-member, super-init-not-called
+
+"""Canonical enum tags for backend-defined thread hierarchies."""
+
+from tvm_ffi.dataclasses import Enum
+
+from tvm.ir import Scriptable
+
+from . import _ffi_api
+
+
+class ThreadScope(Enum, Scriptable, type_key="tirx.ThreadScope", init=False):
+ """A registered thread scope, queried by its qualified backend.scope name.
+
+ Backend and hierarchy rank are enum attributes. The registration index
+ does not determine the hierarchy.
+ """
+
+ @classmethod
+ def get(cls, name: str) -> "ThreadScope":
+ """Look up a registered tag without creating unknown scopes."""
+ return _ffi_api.ThreadScopeGet(name)
+
+ @property
+ def name(self) -> str:
+ """Canonical qualified name, such as cuda.warp."""
+ return self._str_index
+
+ @property
+ def backend(self) -> str:
+ """Backend owning this tag."""
+ return _ffi_api.ThreadScopeGetBackend(self)
+
+ @property
+ def local_name(self) -> str:
+ """Backend-local name, such as warp."""
+ return _ffi_api.ThreadScopeGetLocalName(self)
+
+ @property
+ def rank(self) -> int:
+ """Explicit hierarchy rank; smaller ranks denote wider scopes."""
+ return _ffi_api.ThreadScopeGetRank(self)
+
+ def __str__(self) -> str:
+ return self.name
+
+ def __repr__(self) -> str:
+ return f"ThreadScope({self.name!r})"
diff --git a/python/tvm/tirx/tile/dispatch.py b/python/tvm/tirx/tile/dispatch.py
index 12477fb3e8..8086040b12 100644
--- a/python/tvm/tirx/tile/dispatch.py
+++ b/python/tvm/tirx/tile/dispatch.py
@@ -27,7 +27,7 @@ from tvm.runtime import Object
from tvm.target import Target
from .. import _ffi_api
-from ..exec_scope import ExecScope
+from ..thread_scope import ThreadScope
@register_object("tirx.DispatchContext")
@@ -39,7 +39,7 @@ class DispatchContext(Object, Scriptable):
target : Target
The target of the dispatch context.
- exec_scope : ExecScope
+ thread_scope : ThreadScope
The execution scope of the dispatch context.
launch_params : Dict[str, Tuple[Var, Expr]]
@@ -56,7 +56,7 @@ class DispatchContext(Object, Scriptable):
"""
target: Target
- exec_scope: ExecScope
+ thread_scope: ThreadScope
launch_params: dict[str, tuple[Var, Expr]]
var_range_map: dict[Var, Range]
alloc_only: bool
@@ -74,7 +74,7 @@ class DispatchContext(Object, Scriptable):
def __init__(
self,
target: Target,
- exec_scope: ExecScope,
+ thread_scope: ThreadScope,
launch_params: dict[str, tuple[Var, Expr]],
var_range_map: dict[Var, Range],
alloc_only: bool = False,
@@ -87,7 +87,7 @@ class DispatchContext(Object, Scriptable):
self.__init_handle_by_constructor__(
_ffi_api.DispatchContext, # pylint: disable=no-member
target,
- exec_scope,
+ thread_scope,
launch_params,
var_range_map,
alloc_only,
diff --git a/python/tvm/tirx/tile/instruction.py
b/python/tvm/tirx/tile/instruction.py
index c24b358c6c..dd0ce83328 100644
--- a/python/tvm/tirx/tile/instruction.py
+++ b/python/tvm/tirx/tile/instruction.py
@@ -43,8 +43,8 @@ from tvm.ir import (
)
from tvm.ir.location import UNKNOWN_LOC
from tvm.ir.op import register_op_attr
-from tvm.tirx.exec_scope import ExecScope
from tvm.tirx.expr import IntImm
+from tvm.tirx.thread_scope import ThreadScope
_REQUIRED = object()
SPECS = {}
@@ -142,9 +142,6 @@ class Instruction:
unknown = set(kw) - set(allowed)
if unknown:
raise TypeError(f"{self.name}: unknown qualifier(s)
{sorted(unknown)}")
- scope = kw.get("scope")
- if isinstance(scope, ExecScope) and scope.backend != self.backend:
- raise ValueError(f"{self.name}: scope belongs to {scope.backend}")
attrs_values = {name: _static(value) for name, value in kw.items()}
if isinstance(attrs_values.get("axes"), int):
attrs_values["axes"] = [attrs_values["axes"]]
@@ -155,18 +152,34 @@ class Instruction:
raise TypeError("attrs cannot be combined with qualifier
keywords")
attrs = explicit_attrs
else:
+ attrs_values["scope"] = self._thread_scope(attrs_values["scope"])
attrs = make_node("tirx.tensor." + self.schema, **attrs_values)
call = Call(self.name, args, attrs=attrs, ty=result_ty, loc=loc)
call.validate()
return call
+ def _thread_scope(self, value):
+ if isinstance(value, str):
+ local_name = value.split(".", 1)[-1]
+ if local_name not in self.scopes:
+ raise ValueError(
+ f"{self.name}: unsupported scope {value!r}; expected
{self.scopes}"
+ )
+ value = ThreadScope(value if "." in value else
f"{self.backend}.{value}")
+ if not isinstance(value, ThreadScope):
+ raise TypeError(f"{self.name}: scope must be a string or
ThreadScope")
+ if value.backend != self.backend:
+ raise ValueError(f"{self.name}: scope belongs to {value.backend}")
+ if value.local_name not in self.scopes:
+ raise ValueError(
+ f"{self.name}: unsupported scope {value.local_name!r};
expected {self.scopes}"
+ )
+ return value
+
def validate(self, call):
if len(call.args) != len(self.operands):
raise TypeError(f"{self.name}: invalid operand count")
- if str(call.attrs.scope) not in self.scopes:
- raise ValueError(
- f"{self.name}: unsupported scope {call.attrs.scope!r};
expected {self.scopes}"
- )
+ self._thread_scope(call.attrs.scope)
for arg, value in zip(self.operands, call.args):
if _is_absent(value):
if arg.default is not None:
@@ -214,8 +227,8 @@ def _static(value):
return value.value
if isinstance(value, IntImm):
return bool(value) if str(value.dtype) == "bool" else int(value)
- if isinstance(value, ExecScope):
- return value.name
+ if isinstance(value, ThreadScope):
+ return value
if isinstance(value, Expr):
raise TypeError("Instruction qualifiers must be static; expressions
belong in operands")
return value
@@ -303,7 +316,7 @@ class TensorCall:
view.options = {k: _expr(v) for k, v in options.items() if v is not
None}
view.options.update({k: v for k, v in extras.items() if v is not None})
view.workspaces = {k: values[k] for k in spec.workspaces if values[k]
is not None}
- view.scope = ExecScope(spec.backend, str(call.attrs.scope))
+ view.scope = call.attrs.scope
return view
def get_private_tensors(self, tensor_dict, sctx):
diff --git a/src/backend/cuda/codegen/target_kind.cc
b/src/backend/cuda/codegen/target_kind.cc
index ac0732419f..5add58a024 100644
--- a/src/backend/cuda/codegen/target_kind.cc
+++ b/src/backend/cuda/codegen/target_kind.cc
@@ -31,6 +31,7 @@
#include <tvm/runtime/memory_scope.h>
#include <tvm/target/target.h>
#include <tvm/target/target_kind.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include <string>
@@ -117,6 +118,17 @@ ffi::Map<ffi::String, ffi::Any>
UpdateNVPTXAttrs(ffi::Map<ffi::String, ffi::Any>
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
+ for (const auto& [name, rank] : {std::pair{"kernel", -1},
+ {"cluster", 0},
+ {"cta", 1},
+ {"warpgroup", 2},
+ {"warp", 3},
+ {"thread", 4}}) {
+ auto qualified = tvm::ffi::String("cuda.") + name;
+ tvm::ffi::reflection::EnumDef<tvm::tirx::ThreadScopeObj>(qualified.c_str())
+ .set_attr(tvm::tirx::thread_scope_attr::kBackend,
tvm::ffi::String("cuda"))
+ .set_attr(tvm::tirx::thread_scope_attr::kRank, rank);
+ }
for (const char* name : {"tmem", "wmma.matrix_a", "wmma.matrix_b",
"wmma.accumulator",
"mma.matrixA", "mma.matrixB", "mma.matrixC"}) {
tvm::ffi::reflection::EnumDef<tvm::runtime::MemoryScopeObj>(name).set_attr(
diff --git a/src/backend/cuda/layout/axis.cc b/src/backend/cuda/layout/axis.cc
index 778d335cec..4e875a8ebd 100644
--- a/src/backend/cuda/layout/axis.cc
+++ b/src/backend/cuda/layout/axis.cc
@@ -30,45 +30,45 @@ void RegisterLayoutAxes() {
AxisDef("bx")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "kernel"))
- .set_attr("subscope", ExecScope("cuda", "cta"));
+ .set_attr("scope", ThreadScope("cuda.kernel"))
+ .set_attr("subscope", ThreadScope("cuda.cta"));
AxisDef("by")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "kernel"))
- .set_attr("subscope", ExecScope("cuda", "cta"));
+ .set_attr("scope", ThreadScope("cuda.kernel"))
+ .set_attr("subscope", ThreadScope("cuda.cta"));
AxisDef("bz")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "kernel"))
- .set_attr("subscope", ExecScope("cuda", "cta"));
+ .set_attr("scope", ThreadScope("cuda.kernel"))
+ .set_attr("subscope", ThreadScope("cuda.cta"));
AxisDef("cbx")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "cluster"))
- .set_attr("subscope", ExecScope("cuda", "cta"));
+ .set_attr("scope", ThreadScope("cuda.cluster"))
+ .set_attr("subscope", ThreadScope("cuda.cta"));
AxisDef("cby")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "cluster"))
- .set_attr("subscope", ExecScope("cuda", "cta"));
+ .set_attr("scope", ThreadScope("cuda.cluster"))
+ .set_attr("subscope", ThreadScope("cuda.cta"));
AxisDef("cbz")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "cluster"))
- .set_attr("subscope", ExecScope("cuda", "cta"));
+ .set_attr("scope", ThreadScope("cuda.cluster"))
+ .set_attr("subscope", ThreadScope("cuda.cta"));
AxisDef("tx")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "cta"))
- .set_attr("subscope", ExecScope("cuda", "thread"))
- .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ .set_attr("scope", ThreadScope("cuda.cta"))
+ .set_attr("subscope", ThreadScope("cuda.thread"))
+ .set_attr("fuser", FAxisFuser([](Target target, ThreadScope subscope,
ThreadScope scope,
Iter iter) -> ffi::Optional<Iter> {
if (target->kind->default_device_type == kDLCUDA) {
return std::nullopt;
@@ -79,13 +79,14 @@ void RegisterLayoutAxes() {
AxisDef("warpid")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "cta"))
- .set_attr("subscope", ExecScope("cuda", "warp"))
- .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ .set_attr("scope", ThreadScope("cuda.cta"))
+ .set_attr("subscope", ThreadScope("cuda.warp"))
+ .set_attr("fuser", FAxisFuser([](Target target, ThreadScope subscope,
ThreadScope scope,
Iter iter) -> ffi::Optional<Iter> {
if (target->kind->default_device_type == kDLCUDA) {
// cta->warp ===> cta->thread (tx)
- if (subscope == "thread" && scope == "cta") {
+ if (subscope.same_as(ThreadScope("cuda.thread")) &&
+ scope.same_as(ThreadScope("cuda.cta"))) {
return Iter(iter->extent, 32 * iter->stride,
Axis::Get("tx"));
}
return std::nullopt;
@@ -96,15 +97,17 @@ void RegisterLayoutAxes() {
AxisDef("laneid")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "warp"))
- .set_attr("subscope", ExecScope("cuda", "thread"))
- .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ .set_attr("scope", ThreadScope("cuda.warp"))
+ .set_attr("subscope", ThreadScope("cuda.thread"))
+ .set_attr("fuser", FAxisFuser([](Target target, ThreadScope subscope,
ThreadScope scope,
Iter iter) -> ffi::Optional<Iter> {
if (target->kind->default_device_type == kDLCUDA) {
- if (subscope == "thread" && scope == "warpgroup") {
+ if (subscope.same_as(ThreadScope("cuda.thread")) &&
+ scope.same_as(ThreadScope("cuda.warpgroup"))) {
// warp->thread ===> warpgroup->thread (tid_in_wg)
return Iter(iter->extent, iter->stride,
Axis::Get("tid_in_wg"));
- } else if (subscope == "thread" && scope == "cta") {
+ } else if (subscope.same_as(ThreadScope("cuda.thread")) &&
+ scope.same_as(ThreadScope("cuda.cta"))) {
// warp->thread ===> cta->thread (tx)
return Iter(iter->extent, iter->stride, Axis::Get("tx"));
}
@@ -116,15 +119,17 @@ void RegisterLayoutAxes() {
AxisDef("wgid")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "cta"))
- .set_attr("subscope", ExecScope("cuda", "warpgroup"))
- .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ .set_attr("scope", ThreadScope("cuda.cta"))
+ .set_attr("subscope", ThreadScope("cuda.warpgroup"))
+ .set_attr("fuser", FAxisFuser([](Target target, ThreadScope subscope,
ThreadScope scope,
Iter iter) -> ffi::Optional<Iter> {
if (target->kind->default_device_type == kDLCUDA) {
- if (subscope == "thread" && scope == "cta") {
+ if (subscope.same_as(ThreadScope("cuda.thread")) &&
+ scope.same_as(ThreadScope("cuda.cta"))) {
// cta->warpgroup ===> cta->thread (tx)
return Iter(iter->extent, iter->stride * 128,
Axis::Get("tx"));
- } else if (subscope == "warp" && scope == "cta") {
+ } else if (subscope.same_as(ThreadScope("cuda.warp")) &&
+ scope.same_as(ThreadScope("cuda.cta"))) {
// cta->warpgroup ===> cta->warp (warpid)
return Iter(iter->extent, iter->stride * 4,
Axis::Get("wgid"));
}
@@ -135,12 +140,13 @@ void RegisterLayoutAxes() {
AxisDef("tid_in_wg")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "warpgroup"))
- .set_attr("subscope", ExecScope("cuda", "thread"))
- .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ .set_attr("scope", ThreadScope("cuda.warpgroup"))
+ .set_attr("subscope", ThreadScope("cuda.thread"))
+ .set_attr("fuser", FAxisFuser([](Target target, ThreadScope subscope,
ThreadScope scope,
Iter iter) -> ffi::Optional<Iter> {
if (target->kind->default_device_type == kDLCUDA) {
- if (subscope == "thread" && scope == "cta") {
+ if (subscope.same_as(ThreadScope("cuda.thread")) &&
+ scope.same_as(ThreadScope("cuda.cta"))) {
// warpgroup->thread ===> cta->thread (tx)
return Iter(iter->extent, iter->stride, Axis::Get("tx"));
}
@@ -152,18 +158,21 @@ void RegisterLayoutAxes() {
AxisDef("wid_in_wg")
.set_attr("backend", ffi::String("cuda"))
.set_attr("thread", true)
- .set_attr("scope", ExecScope("cuda", "warpgroup"))
- .set_attr("subscope", ExecScope("cuda", "warp"))
- .set_attr("fuser", FAxisFuser([](Target target, ffi::String subscope,
ffi::String scope,
+ .set_attr("scope", ThreadScope("cuda.warpgroup"))
+ .set_attr("subscope", ThreadScope("cuda.warp"))
+ .set_attr("fuser", FAxisFuser([](Target target, ThreadScope subscope,
ThreadScope scope,
Iter iter) -> ffi::Optional<Iter> {
if (target->kind->default_device_type == kDLCUDA) {
- if (subscope == "thread" && scope == "warpgroup") {
+ if (subscope.same_as(ThreadScope("cuda.thread")) &&
+ scope.same_as(ThreadScope("cuda.warpgroup"))) {
// warpgroup->warp ===> warpgroup->thread (tid_in_wg)
return Iter(iter->extent, iter->stride * 32,
Axis::Get("tid_in_wg"));
- } else if (subscope == "thread" && scope == "cta") {
+ } else if (subscope.same_as(ThreadScope("cuda.thread")) &&
+ scope.same_as(ThreadScope("cuda.cta"))) {
// warpgroup->warp ===> cta->thread (tx)
return Iter(iter->extent, iter->stride * 32,
Axis::Get("tx"));
- } else if (subscope == "warp" && scope == "cta") {
+ } else if (subscope.same_as(ThreadScope("cuda.warp")) &&
+ scope.same_as(ThreadScope("cuda.cta"))) {
// warpgroup->warp ===> cta->warp (warpid)
return Iter(iter->extent, iter->stride,
Axis::Get("warpid"));
}
diff --git a/src/backend/cuda/tile/exec_analysis.cc
b/src/backend/cuda/tile/exec_analysis.cc
index 7616bb1b65..2368586a40 100644
--- a/src/backend/cuda/tile/exec_analysis.cc
+++ b/src/backend/cuda/tile/exec_analysis.cc
@@ -36,7 +36,7 @@
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/tile/context.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include <tvm/tirx/transform.h>
#include <algorithm>
@@ -1039,14 +1039,16 @@ void ExecAnalysis::Pop(int count) {
PrimExpr ExecAnalysis::RewritePredicate(const PrimExpr& condition) const {
return impl_->RewriteFilterCalls(condition);
}
-std::pair<SplitMap, SplitMap> ExecAnalysis::Split(const tirx::ExecScope&
scope) const {
- TVM_FFI_CHECK(!impl_->variable_cluster_ || scope->name() != "cluster",
ValueError)
+std::pair<SplitMap, SplitMap> ExecAnalysis::Split(const tirx::ThreadScope&
scope) const {
+ TVM_FFI_CHECK(scope->backend() == "cuda", ValueError) << "Expected a CUDA
thread scope";
+ TVM_FFI_CHECK(!impl_->variable_cluster_ || scope->local_name() != "cluster",
ValueError)
<< "Cluster tensor instructions require one static cluster shape; "
"use explicit CUDA/PTX instructions for variable or preferred cluster
shapes";
if (impl_->ctx_stack_.empty()) return {};
ExecSplit split;
std::string err;
- if (ScopeSwitch(impl_->ctx_stack_.back().A,
StringToScopeKind(scope->name()), &split, &err)) {
+ if (ScopeSwitch(impl_->ctx_stack_.back().A,
StringToScopeKind(scope->local_name()), &split,
+ &err)) {
return {EncodeSplitSide(split.inter), EncodeSplitSide(split.intra)};
}
LOG(WARNING) << "ExecContext scope_switch failed: " << err;
diff --git a/src/backend/cuda/tile/exec_analysis.h
b/src/backend/cuda/tile/exec_analysis.h
index f744871751..18f193d424 100644
--- a/src/backend/cuda/tile/exec_analysis.h
+++ b/src/backend/cuda/tile/exec_analysis.h
@@ -20,7 +20,7 @@
#ifndef TVM_BACKEND_CUDA_TILE_EXEC_ANALYSIS_H_
#define TVM_BACKEND_CUDA_TILE_EXEC_ANALYSIS_H_
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include <memory>
#include <string>
@@ -38,7 +38,7 @@ class ExecAnalysis {
int PushPredicate(const PrimExpr& condition);
void Pop(int count);
PrimExpr RewritePredicate(const PrimExpr& condition) const;
- std::pair<SplitMap, SplitMap> Split(const tirx::ExecScope& scope) const;
+ std::pair<SplitMap, SplitMap> Split(const tirx::ThreadScope& scope) const;
private:
struct Impl;
diff --git a/src/backend/cuda/tile/exec_context.h
b/src/backend/cuda/tile/exec_context.h
index 31585d1f4b..4a02d1782f 100644
--- a/src/backend/cuda/tile/exec_context.h
+++ b/src/backend/cuda/tile/exec_context.h
@@ -35,7 +35,7 @@
#include <unordered_map>
#include <vector>
-#include "exec_scope.h"
+#include "thread_scope.h"
namespace tvm::backend::cuda {
using namespace tvm::tirx;
diff --git a/src/backend/cuda/tile/exec_scope.cc
b/src/backend/cuda/tile/thread_scope.cc
similarity index 99%
rename from src/backend/cuda/tile/exec_scope.cc
rename to src/backend/cuda/tile/thread_scope.cc
index 9ed4bf0baf..685efd12bd 100644
--- a/src/backend/cuda/tile/exec_scope.cc
+++ b/src/backend/cuda/tile/thread_scope.cc
@@ -16,7 +16,7 @@
* specific language governing permissions and limitations
* under the License.
*/
-#include "exec_scope.h"
+#include "thread_scope.h"
#include <tvm/runtime/logging.h>
namespace tvm::backend::cuda {
diff --git a/src/backend/cuda/tile/exec_scope.h
b/src/backend/cuda/tile/thread_scope.h
similarity index 96%
rename from src/backend/cuda/tile/exec_scope.h
rename to src/backend/cuda/tile/thread_scope.h
index 8bc4b35142..654b287a69 100644
--- a/src/backend/cuda/tile/exec_scope.h
+++ b/src/backend/cuda/tile/thread_scope.h
@@ -16,9 +16,9 @@
* specific language governing permissions and limitations
* under the License.
*/
-#ifndef TVM_BACKEND_CUDA_TILE_EXEC_SCOPE_H_
-#define TVM_BACKEND_CUDA_TILE_EXEC_SCOPE_H_
-#include <tvm/tirx/tile/exec_scope.h>
+#ifndef TVM_BACKEND_CUDA_TILE_THREAD_SCOPE_H_
+#define TVM_BACKEND_CUDA_TILE_THREAD_SCOPE_H_
+#include <tvm/tirx/tile/thread_scope.h>
#include <string>
#include <utility>
diff --git a/src/backend/trn/codegen/target_kind.cc
b/src/backend/trn/codegen/target_kind.cc
index 573e71a2ce..47d3f6384c 100644
--- a/src/backend/trn/codegen/target_kind.cc
+++ b/src/backend/trn/codegen/target_kind.cc
@@ -28,6 +28,7 @@
#include <tvm/runtime/memory_scope.h>
#include <tvm/target/target.h>
#include <tvm/target/target_kind.h>
+#include <tvm/tirx/tile/thread_scope.h>
namespace tvm {
@@ -43,6 +44,12 @@ void RegisterTRNTransforms();
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
+ for (const auto& [name, rank] : {std::pair{"kernel", -1}, {"thread", 0}}) {
+ auto qualified = tvm::ffi::String("trn.") + name;
+ tvm::ffi::reflection::EnumDef<tvm::tirx::ThreadScopeObj>(qualified.c_str())
+ .set_attr(tvm::tirx::thread_scope_attr::kBackend,
tvm::ffi::String("trn"))
+ .set_attr(tvm::tirx::thread_scope_attr::kRank, rank);
+ }
for (const char* name : {"trn.sbuf", "trn.psum"}) {
tvm::ffi::reflection::EnumDef<tvm::runtime::MemoryScopeObj>(name).set_attr(
tvm::runtime::attr::kMemoryAllocProperty,
diff --git a/src/tirx/ir/layout/axis_registry.cc
b/src/tirx/ir/layout/axis_registry.cc
index ad9b43c639..df5bae9ce9 100644
--- a/src/tirx/ir/layout/axis_registry.cc
+++ b/src/tirx/ir/layout/axis_registry.cc
@@ -68,10 +68,12 @@ ffi::Optional<ffi::String> AxisNode::GetBackend() const {
return AxisAttr<ffi::String>(this, "backend");
}
-ffi::Optional<ExecScope> AxisNode::GetScope() const { return
AxisAttr<ExecScope>(this, "scope"); }
+ffi::Optional<ThreadScope> AxisNode::GetScope() const {
+ return AxisAttr<ThreadScope>(this, "scope");
+}
-ffi::Optional<ExecScope> AxisNode::GetSubscope() const {
- return AxisAttr<ExecScope>(this, "subscope");
+ffi::Optional<ThreadScope> AxisNode::GetSubscope() const {
+ return AxisAttr<ThreadScope>(this, "subscope");
}
ffi::Optional<FAxisFuser> AxisNode::GetFuser() const { return
AxisAttr<FAxisFuser>(this, "fuser"); }
diff --git a/src/tirx/ir/layout/tile_canonicalize.cc
b/src/tirx/ir/layout/tile_canonicalize.cc
index c22a118048..376875561f 100644
--- a/src/tirx/ir/layout/tile_canonicalize.cc
+++ b/src/tirx/ir/layout/tile_canonicalize.cc
@@ -86,8 +86,8 @@ TileLayout FuseAxesByScope(TileLayout layout) {
if (!scope_pair_opt.has_value() || !target.has_value()) {
return layout;
}
- auto subscope = scope_pair_opt.value().get<0>()->name();
- auto scope = scope_pair_opt.value().get<1>()->name();
+ auto subscope = scope_pair_opt.value().get<0>();
+ auto scope = scope_pair_opt.value().get<1>();
// Step 2: Create vectors for the new layout components
std::vector<Iter> shard;
diff --git a/src/tirx/ir/layout/tile_core.cc b/src/tirx/ir/layout/tile_core.cc
index 682fb51f1b..744c399718 100644
--- a/src/tirx/ir/layout/tile_core.cc
+++ b/src/tirx/ir/layout/tile_core.cc
@@ -327,57 +327,50 @@ bool TileLayoutNode::HasThreadAxis() const {
[](const Iter& iter) { return iter->axis->IsThreadAxis();
});
}
-ffi::Optional<ffi::Tuple<ExecScope, ExecScope>> TileLayoutNode::GetScope()
const {
+ffi::Optional<ffi::Tuple<ThreadScope, ThreadScope>> TileLayoutNode::GetScope()
const {
if (!HasThreadAxis()) return std::nullopt;
- std::unordered_map<ffi::String, ffi::String> scope_map;
- ffi::Optional<ffi::String> inner_most;
+ std::unordered_map<ThreadScope, ThreadScope, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual> scope_map;
+ ffi::Optional<ThreadScope> inner_most;
ffi::Optional<ffi::String> backend;
auto check_axis = [&](const Axis& axis) {
if (!axis->IsThreadAxis()) return;
-
auto subscope_opt = axis->GetSubscope();
auto scope_opt = axis->GetScope();
TVM_FFI_ICHECK(subscope_opt.has_value() && scope_opt.has_value())
<< "Thread axis " << axis.name() << " has no subscope or scope";
-
+ ThreadScope subscope = subscope_opt.value();
+ ThreadScope scope = scope_opt.value();
+ TVM_FFI_CHECK(scope->backend() == subscope->backend(), ValueError)
+ << "Thread axis contains scopes from different backends";
if (backend.has_value()) {
- TVM_FFI_CHECK(backend.value() == subscope_opt.value()->backend,
ValueError)
+ TVM_FFI_CHECK(backend.value() == subscope->backend(), ValueError)
<< "Layout contains thread axes from different backends";
} else {
- backend = subscope_opt.value()->backend;
+ backend = subscope->backend();
}
- ffi::String subscope = subscope_opt.value()->name();
- ffi::String scope = scope_opt.value()->name();
-
- if (!inner_most.has_value() ||
- ScopeHigher(ExecScope(backend.value(), inner_most.value()),
subscope_opt.value()))
+ if (!inner_most.has_value() || ScopeHigher(inner_most.value(), subscope)) {
inner_most = subscope;
-
- auto it = scope_map.find(subscope);
- if (it == scope_map.end())
- scope_map[subscope] = scope;
- else
- TVM_FFI_ICHECK_EQ(it->second, scope)
- << "Ill-formed tile layout: conflicting scopes for " << subscope;
+ }
+ auto [it, inserted] = scope_map.emplace(subscope, scope);
+ TVM_FFI_ICHECK(inserted || it->second.same_as(scope))
+ << "Ill-formed tile layout: conflicting scopes for " << subscope;
};
-
for (const auto& iter : shard) check_axis(iter->axis);
for (const auto& iter : replica) check_axis(iter->axis);
for (const auto& [axis, off] : offset) check_axis(axis);
- ffi::String outer_most = inner_most.value();
+ ThreadScope outer_most = inner_most.value();
size_t count = 0;
for (auto it = scope_map.find(outer_most); it != scope_map.end();
it = scope_map.find(outer_most)) {
- count++;
+ TVM_FFI_ICHECK_LT(count, scope_map.size()) << "Ill-formed tile layout:
cyclic scope chain";
+ ++count;
outer_most = it->second;
}
-
TVM_FFI_ICHECK_EQ(count, scope_map.size()) << "Ill-formed tile layout:
disconnected scope chain";
- return Tuple<ExecScope, ExecScope>{ExecScope(backend.value(),
inner_most.value()),
- ExecScope(backend.value(), outer_most)};
+ return Tuple<ThreadScope, ThreadScope>{inner_most.value(), outer_most};
}
TileLayout TileLayoutNode::DefaultLayout(ffi::Array<PrimExpr> shape) {
@@ -393,7 +386,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
"tirx.TileLayoutGetScope",
- [](const TileLayout& layout) -> ffi::Optional<ffi::Tuple<ExecScope,
ExecScope>> {
+ [](const TileLayout& layout) -> ffi::Optional<ffi::Tuple<ThreadScope,
ThreadScope>> {
return layout->GetScope();
});
}
diff --git a/src/tirx/script/ir_builder/frame.cc
b/src/tirx/script/ir_builder/frame.cc
index 9e72399f23..d1559cc26b 100644
--- a/src/tirx/script/ir_builder/frame.cc
+++ b/src/tirx/script/ir_builder/frame.cc
@@ -24,7 +24,7 @@
#include <tvm/script/ir_builder/ir.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/script/ir_builder/frame.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include <map>
diff --git a/src/tirx/script/ir_builder/ir.cc b/src/tirx/script/ir_builder/ir.cc
index e2dd1a34df..85d008ecce 100644
--- a/src/tirx/script/ir_builder/ir.cc
+++ b/src/tirx/script/ir_builder/ir.cc
@@ -29,7 +29,7 @@
#include <tvm/tirx/layout.h>
#include <tvm/tirx/script/ir_builder/ir.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include "./utils.h"
diff --git a/src/tirx/script/printer/block.cc b/src/tirx/script/printer/block.cc
index c036f7b6b1..41ab52804a 100644
--- a/src/tirx/script/printer/block.cc
+++ b/src/tirx/script/printer/block.cc
@@ -20,7 +20,7 @@
#include <tvm/script/printer/doc_translator.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include <optional>
@@ -33,20 +33,19 @@ namespace details {
namespace {
-ffi::Optional<ExprDoc> ExecScopeDocTranslate(DocTranslatorObj*, ffi::AnyView
input,
- const ffi::Object*) {
+ffi::Optional<ExprDoc> ThreadScopeDocTranslate(DocTranslatorObj*, ffi::AnyView
input,
+ const ffi::Object*) {
const auto* scope =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
tirx::ExecScopeNode>(input);
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
tirx::ThreadScopeObj>(input);
return NamespaceDoc("tirx")
- ->Attr(scope->backend)
- ->Attr("ExecScope")
+ ->Attr("ThreadScope")
->Call({LiteralDoc::Str(scope->name(), std::nullopt)});
}
TVM_FFI_STATIC_INIT_BLOCK() {
- ffi::reflection::TypeAttrDef<tirx::ExecScopeNode>().attr(
+ ffi::reflection::TypeAttrDef<tirx::ThreadScopeObj>().attr(
tvm::script::printer::type_attr::kDocTranslate,
- FDocTranslate::FromNative<&ExecScopeDocTranslate>());
+ FDocTranslate::FromNative<&ThreadScopeDocTranslate>());
}
} // namespace
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index 6db853e485..51cc83f69c 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -28,7 +28,7 @@
#include <tvm/tirx/op/abi.h>
#include <tvm/tirx/op_attr_types.h>
#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include <algorithm>
#include <limits>
@@ -223,7 +223,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
using script::printer::RegisterScriptRepr;
RegisterScriptRepr<tirx::TensorTypeNode>();
RegisterScriptRepr<tirx::ComposeLayoutNode>();
- RegisterScriptRepr<tirx::ExecScopeNode>();
+ RegisterScriptRepr<tirx::ThreadScopeObj>();
RegisterScriptRepr<tirx::IndexMapNode>();
RegisterScriptRepr<tirx::IterNode>();
RegisterScriptRepr<tirx::FunctionNode>();
diff --git a/src/tirx/tile/context.cc b/src/tirx/tile/context.cc
index 57684739f7..c6531a4385 100644
--- a/src/tirx/tile/context.cc
+++ b/src/tirx/tile/context.cc
@@ -85,7 +85,7 @@ ffi::Optional<ffi::ObjectRef>
DispatchContextObj::SharedStateGet(ffi::String key
return ffi::Optional<ffi::ObjectRef>();
}
-DispatchContext::DispatchContext(Target target, ExecScope exec_scope,
+DispatchContext::DispatchContext(Target target, ThreadScope thread_scope,
ffi::Map<ffi::String, ffi::Tuple<PrimVar,
PrimExpr>> launch_params,
ffi::Map<Var, Range> var_range_map, bool
alloc_only,
ffi::Map<ffi::String, ffi::ObjectRef>
callbacks,
@@ -93,7 +93,13 @@ DispatchContext::DispatchContext(Target target, ExecScope
exec_scope,
ffi::Map<ffi::String, ffi::Array<PrimExpr>>
inter,
ffi::Map<ffi::String, ffi::Array<PrimExpr>>
intra,
ffi::String scope_kind) {
- auto n = ffi::make_object<DispatchContextObj>(std::move(target),
std::move(exec_scope));
+ auto n = ffi::make_object<DispatchContextObj>(std::move(target),
std::move(thread_scope));
+ TVM_FFI_CHECK(n->thread_scope->backend() == n->target->kind->name,
ValueError)
+ << "Thread scope " << n->thread_scope.name() << " does not match target "
+ << n->target->kind->name;
+ ffi::String local_name = n->thread_scope->local_name();
+ TVM_FFI_CHECK(scope_kind.empty() || scope_kind == local_name, ValueError)
+ << "scope_kind must match the thread scope's local name";
for (const auto& [tag, binding] : launch_params) {
PrimType var_ty = binding.get<0>().ty();
PrimType extent_ty = binding.get<1>().ty();
@@ -107,7 +113,7 @@ DispatchContext::DispatchContext(Target target, ExecScope
exec_scope,
n->shared_state = std::move(shared_state);
n->inter = std::move(inter);
n->intra = std::move(intra);
- n->scope_kind = std::move(scope_kind);
+ n->scope_kind = std::move(local_name);
data_ = std::move(n);
}
@@ -115,14 +121,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("tirx.DispatchContext",
- [](Target target, ExecScope exec_scope,
+ [](Target target, ThreadScope thread_scope,
ffi::Map<ffi::String, ffi::Tuple<PrimVar, PrimExpr>>
launch_params,
ffi::Map<Var, Range> var_range_map, bool alloc_only,
ffi::Map<ffi::String, ffi::ObjectRef> callbacks,
ffi::Map<ffi::String, ffi::ObjectRef> shared_state,
ffi::Map<ffi::String, ffi::Array<PrimExpr>> inter,
ffi::Map<ffi::String, ffi::Array<PrimExpr>> intra, ffi::String
scope_kind) {
- return DispatchContext(target, exec_scope, launch_params,
var_range_map, alloc_only,
+ return DispatchContext(target, thread_scope, launch_params,
var_range_map, alloc_only,
callbacks, shared_state, inter, intra,
scope_kind);
})
.def_method("tirx.DispatchContextAddAllocTensor",
&DispatchContextObj::AddAllocTensor)
diff --git a/src/tirx/tile/exec_scope.cc b/src/tirx/tile/exec_scope.cc
deleted file mode 100644
index edc289cbc0..0000000000
--- a/src/tirx/tile/exec_scope.cc
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/tirx/tile/exec_scope.h>
-
-#include <map>
-#include <string>
-#include <utility>
-
-namespace tvm::tirx {
-namespace {
-using ScopeKey = std::pair<std::string, std::string>;
-std::map<ScopeKey, int>& Scopes() {
- static std::map<ScopeKey, int> scopes;
- return scopes;
-}
-int Lookup(const ffi::String& backend, const ffi::String& name) {
- auto it = Scopes().find({std::string(backend), std::string(name)});
- TVM_FFI_CHECK(it != Scopes().end(), ValueError)
- << "Unknown execution scope: " << backend << "." << name;
- return it->second;
-}
-} // namespace
-
-void RegisterExecScope(const ffi::String& backend, const ffi::String& name,
int rank) {
- Scopes().insert_or_assign({std::string(backend), std::string(name)}, rank);
-}
-int ExecScopeNode::rank() const { return Lookup(backend, scope_name); }
-ExecScope::ExecScope(ffi::String backend, ffi::String name) {
- Lookup(backend, name);
- auto n = ffi::make_object<ExecScopeNode>();
- n->backend = std::move(backend);
- n->scope_name = std::move(name);
- data_ = std::move(n);
-}
-bool ScopeHigher(const ExecScope& a, const ExecScope& b) {
- TVM_FFI_CHECK(a->backend == b->backend, ValueError)
- << "Cannot compare execution scopes from different backends";
- return a->rank() < b->rank();
-}
-TVM_FFI_STATIC_INIT_BLOCK() {
- ExecScopeNode::RegisterReflection();
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def("tirx.ExecScope",
- [](ffi::String backend, ffi::String name) {
- return ExecScope(std::move(backend), std::move(name));
- })
- .def("tirx.RegisterExecScope", RegisterExecScope);
-}
-} // namespace tvm::tirx
diff --git a/src/tirx/tile/lower.cc b/src/tirx/tile/lower.cc
index cd0f2271e9..6f8cfabc49 100644
--- a/src/tirx/tile/lower.cc
+++ b/src/tirx/tile/lower.cc
@@ -39,7 +39,7 @@
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/tile/dispatch.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include <tvm/tirx/transform.h>
#include <limits>
@@ -487,12 +487,12 @@ class TileDispatcher : public StmtExprMutator {
op->op.as_or_throw<Op>().Validate(op);
TVM_FFI_CHECK(op->attrs.has_value(), ValueError)
<< "A tile instruction must have scope attributes";
- auto scope_name =
- ffi::reflection::FieldGetter(op->attrs->GetTypeKey(),
"scope")(op->attrs.value())
- .cast<ffi::String>();
- ExecScope scope(target_->kind->name, scope_name);
+ auto scope = ffi::reflection::FieldGetter(op->attrs->GetTypeKey(),
"scope")(op->attrs.value())
+ .cast<ThreadScope>();
+ TVM_FFI_CHECK(scope->backend() == target_->kind->name, ValueError)
+ << "Thread scope " << scope.name() << " does not match target " <<
target_->kind->name;
ffi::Map<ffi::String, ffi::Array<PrimExpr>> inter_map, intra_map;
- ffi::String scope_kind = scope->name();
+ ffi::String scope_kind = scope->local_name();
if (cuda_analysis_) {
auto split = cuda_analysis_->Split(scope);
inter_map = std::move(split.first);
diff --git a/src/tirx/tile/op.cc b/src/tirx/tile/op.cc
index 6f1f4f268f..634f2d1afd 100644
--- a/src/tirx/tile/op.cc
+++ b/src/tirx/tile/op.cc
@@ -24,23 +24,25 @@
#include <tvm/ir/prim/op.h>
#include <tvm/tirx/tile/op.h>
+#include <tvm/tirx/tile/thread_scope.h>
#include <utility>
namespace tvm::tirx::tile {
struct ScopeAttrs : public AttrsNode {
- ffi::String scope;
+ ScopeAttrs() : scope(ffi::UnsafeInit{}) {}
+ ThreadScope scope;
static void RegisterReflection() {
namespace refl = ffi::reflection;
- refl::ObjectDef<ScopeAttrs>().def_ro("scope", &ScopeAttrs::scope,
-
refl::DefaultValue(ffi::String("thread")));
+ refl::ObjectDef<ScopeAttrs>().def_ro("scope", &ScopeAttrs::scope);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.tensor.ScopeAttrs", ScopeAttrs,
AttrsNode);
};
struct TrnAttrs : public AttrsNode {
- ffi::String scope;
+ TrnAttrs() : scope(ffi::UnsafeInit{}) {}
+ ThreadScope scope;
ffi::Optional<int64_t> max_inst_size;
ffi::String opcode;
ffi::String reduce_op;
@@ -52,7 +54,7 @@ struct TrnAttrs : public AttrsNode {
static void RegisterReflection() {
namespace refl = ffi::reflection;
refl::ObjectDef<TrnAttrs>()
- .def_ro("scope", &TrnAttrs::scope,
refl::DefaultValue(ffi::String("thread")))
+ .def_ro("scope", &TrnAttrs::scope)
.def_ro("max_inst_size", &TrnAttrs::max_inst_size,
refl::DefaultValue(ffi::Optional<int64_t>(std::nullopt)))
.def_ro("opcode", &TrnAttrs::opcode,
refl::DefaultValue(ffi::String("")))
diff --git a/src/tirx/tile/thread_scope.cc b/src/tirx/tile/thread_scope.cc
new file mode 100644
index 0000000000..930c219d0c
--- /dev/null
+++ b/src/tirx/tile/thread_scope.cc
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/serialization.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/tirx/tile/thread_scope.h>
+
+#include <string>
+
+namespace tvm::tirx {
+namespace refl = ffi::reflection;
+namespace {
+template <typename T>
+T ScopeAttr(const ThreadScopeObj* scope, refl::AttrKey<T> key) {
+ auto state = refl::TypeAttrColumn(
+
ffi::type_attr::kEnumState)[ThreadScopeObj::_GetOrAllocRuntimeTypeIndex()]
+ .cast<ffi::EnumState>();
+ auto column = state->attrs.Get(key.name);
+ auto value = key.ValueFrom(
+ column ?
column->Get(ffi::GetRef<ThreadScope>(scope)).value_or(ffi::Any()) : ffi::Any());
+ TVM_FFI_CHECK(value.has_value(), ValueError)
+ << "Thread scope " << scope->name() << " has no " << key.name << "
attribute";
+ return value.value();
+}
+} // namespace
+
+ffi::String ThreadScopeObj::backend() const { return ScopeAttr(this,
thread_scope_attr::kBackend); }
+ffi::String ThreadScopeObj::local_name() const {
+ const std::string prefix = std::string(backend()) + ".";
+ const std::string name(_str_index);
+ TVM_FFI_CHECK(name.size() > prefix.size() && name.compare(0, prefix.size(),
prefix) == 0,
+ ValueError)
+ << "Thread scope " << name << " must be named as backend.scope";
+ return name.substr(prefix.size());
+}
+int64_t ThreadScopeObj::rank() const { return ScopeAttr(this,
thread_scope_attr::kRank); }
+
+ThreadScope ThreadScope::Get(const ffi::String& name) {
+ auto scope =
ffi::EnumObj::_GetByStrIndex<ThreadScopeObj>(name).as_or_throw<ThreadScope>();
+ scope->local_name();
+ scope->rank();
+ return scope;
+}
+
+bool ScopeHigher(const ThreadScope& a, const ThreadScope& b) {
+ TVM_FFI_CHECK(a->backend() == b->backend(), ValueError)
+ << "Cannot compare thread scopes from different backends";
+ return a->rank() < b->rank();
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ refl::ObjectDef<ThreadScopeObj>(refl::init(false));
+ refl::TypeAttrDef<ThreadScopeObj>()
+ .def(ffi::type_attr::kDataToJson.name,
+ [](const ThreadScopeObj* scope) -> ffi::String { return
scope->name(); })
+ .def(ffi::type_attr::kDataFromJson.name, ThreadScope::Get);
+ refl::GlobalDef()
+ .def("tirx.ThreadScopeGet", ThreadScope::Get)
+ .def("tirx.ThreadScopeGetBackend", [](ThreadScope scope) { return
scope->backend(); })
+ .def("tirx.ThreadScopeGetLocalName", [](ThreadScope scope) { return
scope->local_name(); })
+ .def("tirx.ThreadScopeGetRank", [](ThreadScope scope) { return
scope->rank(); })
+ .def("tirx.ScopeHigher", ScopeHigher);
+}
+} // namespace tvm::tirx
diff --git a/tests/cpp/tirx/thread_scope_test.cc
b/tests/cpp/tirx/thread_scope_test.cc
new file mode 100644
index 0000000000..c6530ec9e2
--- /dev/null
+++ b/tests/cpp/tirx/thread_scope_test.cc
@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#include <gtest/gtest.h>
+#include <tvm/ffi/extra/serialization.h>
+#include <tvm/ffi/reflection/creator.h>
+#include <tvm/ir/op.h>
+#include <tvm/target/target_kind.h>
+#include <tvm/tirx/tile/dispatch.h>
+#include <tvm/tirx/transform.h>
+
+#include <memory>
+
+namespace tvm::tirx {
+namespace {
+namespace refl = ffi::reflection;
+
+void RegisterTestScopes() {
+ static const bool registered = [] {
+ // Register the inner scope first; enum indices must not define ordering.
+ refl::EnumDef<ThreadScopeObj>("thread_scope_test.inner")
+ .set_attr(thread_scope_attr::kBackend,
ffi::String("thread_scope_test"))
+ .set_attr(thread_scope_attr::kRank, 10);
+ refl::EnumDef<ThreadScopeObj>("thread_scope_test.outer")
+ .set_attr(thread_scope_attr::kBackend,
ffi::String("thread_scope_test"))
+ .set_attr(thread_scope_attr::kRank, -10);
+ return true;
+ }();
+ (void)registered;
+}
+
+TEST(ThreadScope, NativeRegistrationIdentityAndRank) {
+ RegisterTestScopes();
+ ThreadScope cuda("cuda.thread");
+ ThreadScope trn("trn.thread");
+ EXPECT_EQ(cuda->backend(), "cuda");
+ EXPECT_EQ(cuda->rank(), 4);
+ EXPECT_EQ(trn->backend(), "trn");
+ EXPECT_EQ(trn->rank(), 0);
+ EXPECT_FALSE(cuda.same_as(trn));
+ EXPECT_THROW(ScopeHigher(cuda, trn), ffi::Error);
+ ThreadScope inner("thread_scope_test.inner");
+ ThreadScope outer("thread_scope_test.outer");
+ EXPECT_TRUE(inner.same_as(ThreadScope::Get(inner.name())));
+ EXPECT_EQ(inner->backend(), "thread_scope_test");
+ EXPECT_EQ(inner->local_name(), "inner");
+ EXPECT_EQ(inner->rank(), 10);
+ EXPECT_TRUE(ScopeHigher(outer, inner));
+ EXPECT_FALSE(ScopeHigher(inner, outer));
+ EXPECT_FALSE(ScopeHigher(inner, inner));
+ EXPECT_GT(outer->_int_index, inner->_int_index);
+ EXPECT_TRUE(ffi::Any(inner.name()).cast<ThreadScope>().same_as(inner));
+ EXPECT_THROW(ThreadScope("thread_scope_test.missing"), ffi::Error);
+ EXPECT_THROW(refl::EnumDef<ThreadScopeObj>(inner.name().c_str()),
ffi::Error);
+
EXPECT_TRUE(ffi::FromJSONGraph(ffi::ToJSONGraph(inner)).cast<ThreadScope>().same_as(inner));
+}
+
+TEST(ThreadScope, TypedCallAndDispatchWithoutPython) {
+ RegisterTestScopes();
+ TargetKindDef("thread_scope_test").set_default_device_type(kDLCPU);
+ Target target("thread_scope_test");
+ ThreadScope scope("thread_scope_test.inner");
+ auto attrs = refl::ObjectCreator("tirx.tensor.ScopeAttrs")({{"scope",
scope}}).cast<Attrs>();
+ auto invoked = std::make_shared<int>(0);
+ OpDef("thread_scope_test.instruction")
+ .signature()
+ .set_attr<TFixedReturnType>(op_attr::kFixedReturnType, PrimType::Void())
+ .set_attr<tile::FTileOpDispatch>(
+ tile::op_attr::kTileOpDispatch,
+ [scope, invoked](const Call& call, const tile::DispatchContext&
context) {
+ EXPECT_TRUE(context->thread_scope.same_as(scope));
+ EXPECT_EQ(context->scope_kind, "inner");
+ EXPECT_TRUE(refl::FieldGetter(call->attrs->GetTypeKey(),
"scope")(call->attrs.value())
+ .cast<ThreadScope>()
+ .same_as(scope));
+ ++*invoked;
+ return Function({}, SeqStmt(Evaluate(37)));
+ });
+ Call call(PrimType::Void(), Op::Get("thread_scope_test.instruction"), {},
attrs);
+ Function func({}, SeqStmt(Evaluate(call)), VoidType(), DictAttrs({{"target",
target}}));
+ GlobalVar main("main");
+ auto result = transform::TileDispatch()(IRModule({{main, func}}));
+ auto lowered = result->Lookup(main).as_or_throw<Function>();
+ EXPECT_EQ(*invoked, 1);
+ EXPECT_EQ(
+
lowered->body.value()->seq[0].as_or_throw<Evaluate>()->value.as_or_throw<IntImm>()->value,
+ 37);
+ EXPECT_THROW(refl::ObjectCreator("tirx.tensor.ScopeAttrs")({}), ffi::Error);
+ EXPECT_THROW(refl::ObjectCreator("tirx.tensor.ScopeAttrs")({{"scope",
ffi::String("inner")}}),
+ ffi::Error);
+}
+} // namespace
+} // namespace tvm::tirx
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 822c77d1a6..52207cd308 100644
--- a/tests/python/tirx/script/test_tirx_script_basic_usage.py
+++ b/tests/python/tirx/script/test_tirx_script_basic_usage.py
@@ -359,16 +359,23 @@ def test_roundtrip_unary_inplace():
code = test.script()
# Opcode and explicit operand aliases survive printing.
- expected = 'T.cuda.tile.apply("ex2.approx.ftz.f32", (A[0:32], A[0:32]),
(), scope="warp")'
+ expected = (
+ 'T.cuda.tile.apply("ex2.approx.ftz.f32", (A[0:32], A[0:32]), '
+ '(), scope=T.ThreadScope("cuda.warp"))'
+ )
assert expected in code, (
f"expected explicit exp2, got:\n{code}"
)
- assert 'T.cuda.tile.apply("sqrt.rn.f32", (A[32:64], A[32:64]), (),
scope="warp")' in code, (
+ expected = (
+ 'T.cuda.tile.apply("sqrt.rn.f32", (A[32:64], A[32:64]), '
+ '(), scope=T.ThreadScope("cuda.warp"))'
+ )
+ assert expected in code, (
f"expected explicit sqrt, got:\n{code}"
)
expected = (
'T.cuda.tile.apply("div.rn.f32", (A[64:96], T.float32(1.0), A[64:96]),
'
- '(), scope="warp")'
+ '(), scope=T.ThreadScope("cuda.warp"))'
)
assert expected in code, (
f"expected explicit reciprocal, got:\n{code}"
@@ -395,7 +402,10 @@ def test_roundtrip_unary_different_dst_src():
# fmt: on
code = test.script()
- expected = 'T.cuda.tile.apply("ex2.approx.ftz.f32", (A[0:32], B[0:32]),
(), scope="warp")'
+ expected = (
+ 'T.cuda.tile.apply("ex2.approx.ftz.f32", (A[0:32], B[0:32]), '
+ '(), scope=T.ThreadScope("cuda.warp"))'
+ )
assert expected in code, (
f"different dst/src should keep both:\n{code}"
)
diff --git a/tests/python/tirx/script/test_tirx_script_gpu.py
b/tests/python/tirx/script/test_tirx_script_gpu.py
index cd56a746f4..514959c0c0 100644
--- a/tests/python/tirx/script/test_tirx_script_gpu.py
+++ b/tests/python/tirx/script/test_tirx_script_gpu.py
@@ -99,7 +99,7 @@ def test_roundtrip_scopeid_deferred():
assert_structural_equal(test, from_source(code))
-def test_exec_scope_filter_guard_roundtrip():
+def test_thread_scope_filter_guard_roundtrip():
@T.function(private=True)
def test(A: T.Tensor((1,), "float32", scope="global")) -> None:
T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=(128,)))
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py
b/tests/python/tirx/script/test_tirx_script_printer.py
index 5bd62ce716..7c91c15018 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -726,7 +726,8 @@ def test_print_kwargs_schedule_op_full_code():
@T.function
def test():
A = T.alloc_tensor(ty_args=[T.Tensor((16,), T.float32)])
- T.cuda.tile.apply("mov.b32", (A[0:16], T.float32(1.25)), (),
scope="warp")"""
+ T.cuda.tile.apply("mov.b32", (A[0:16], T.float32(1.25)), (),
scope=T.ThreadScope("cuda.warp"))
+""".rstrip()
code = test.script()
assert code == expected
assert from_source(code).script() == code
diff --git a/tests/python/tirx/test_exec_scope.py
b/tests/python/tirx/test_exec_scope.py
deleted file mode 100644
index 27eabab8f9..0000000000
--- a/tests/python/tirx/test_exec_scope.py
+++ /dev/null
@@ -1,43 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-import pytest
-
-from tvm.tirx.exec_scope import ExecScope
-
-
-def test_exec_scope_create():
- def is_trivial_scope(scope, name):
- return isinstance(scope, ExecScope) and scope.name == name
-
- thread = ExecScope("cuda", "thread")
- warp = ExecScope("cuda", "warp")
- wg = ExecScope("cuda", "warpgroup")
- cta = ExecScope("cuda", "cta")
- cluster = ExecScope("cuda", "cluster")
-
- assert is_trivial_scope(thread, "thread")
- assert is_trivial_scope(warp, "warp")
- assert is_trivial_scope(wg, "warpgroup")
- assert is_trivial_scope(cta, "cta")
- assert is_trivial_scope(cluster, "cluster")
-
- with pytest.raises(Exception, match="Unknown execution scope"):
- ExecScope("cuda", "aaa")
-
-
-if __name__ == "__main__":
- test_exec_scope_create()
diff --git a/tests/python/tirx/test_layout.py b/tests/python/tirx/test_layout.py
index f8e33cb0a9..3f646d4894 100644
--- a/tests/python/tirx/test_layout.py
+++ b/tests/python/tirx/test_layout.py
@@ -76,8 +76,8 @@ def test_axis():
assert CUDAAxis.TCol.is_memory()
assert CUDAAxis.TLane.is_memory()
- assert CUDAAxis.bx.get_scope().name == "kernel"
- assert CUDAAxis.bx.get_subscope().name == "cta"
+ assert CUDAAxis.bx.get_scope().name == "cuda.kernel"
+ assert CUDAAxis.bx.get_subscope().name == "cuda.cta"
def test_constructor():
@@ -201,22 +201,22 @@ def test_verify_well_formed():
layout = TileLayout(S[(8, 4, 2) : (4 @ laneid, 1 @ laneid, 1)])
res = layout.get_scope()
assert res is not None
- assert res[0].name == "thread"
- assert res[1].name == "warp"
+ assert res[0].name == "cuda.thread"
+ assert res[1].name == "cuda.warp"
assert layout.verify_well_formed()
layout = TileLayout(S[8 : 4 @ laneid] + R[4 : 1 @ laneid])
res = layout.get_scope()
assert res is not None
- assert res[0].name == "thread"
- assert res[1].name == "warp"
+ assert res[0].name == "cuda.thread"
+ assert res[1].name == "cuda.warp"
assert layout.verify_well_formed()
layout = TileLayout(S[(8, 4, 2) : (4 @ laneid, 1 @ laneid, 1)])
res = layout.get_scope()
assert res is not None
- assert res[0].name == "thread"
- assert res[1].name == "warp"
+ assert res[0].name == "cuda.thread"
+ assert res[1].name == "cuda.warp"
assert layout.verify_well_formed()
layout = TileLayout(
@@ -224,8 +224,8 @@ def test_verify_well_formed():
)
res = layout.get_scope()
assert res is not None
- assert res[0].name == "thread"
- assert res[1].name == "cta"
+ assert res[0].name == "cuda.thread"
+ assert res[1].name == "cuda.cta"
assert layout.verify_well_formed()
layout = TileLayout(
@@ -233,8 +233,8 @@ def test_verify_well_formed():
)
res = layout.get_scope()
assert res is not None
- assert res[0].name == "thread"
- assert res[1].name == "warpgroup"
+ assert res[0].name == "cuda.thread"
+ assert res[1].name == "cuda.warpgroup"
assert layout.verify_well_formed()
layout = TileLayout(S[(2, 8, 2, 4, 2) : (2 @ wgid, 4 @ laneid, 1 @
wgid, 1 @ laneid, 1)])
diff --git a/tests/python/tirx/test_op.py b/tests/python/tirx/test_op.py
index 42b1e41065..87e5fff2d1 100644
--- a/tests/python/tirx/test_op.py
+++ b/tests/python/tirx/test_op.py
@@ -36,7 +36,7 @@ def test_instruction_is_a_void_call():
assert call.ty == tvm.ir.PrimType("void")
assert call.op.get_attr("TCallEffectKind") == 3
assert callable(call.op.get_attr("FTileOpDispatch"))
- assert call.attrs.scope == "warp"
+ assert call.attrs.scope.same_as(T.cuda.ThreadScope("warp"))
assert call.args[0].value == "ld.global.v4.b32"
assert not hasattr(call, "config")
assert not hasattr(call, "workspace")
diff --git a/tests/python/tirx/test_thread_scope.py
b/tests/python/tirx/test_thread_scope.py
new file mode 100644
index 0000000000..37e8572dba
--- /dev/null
+++ b/tests/python/tirx/test_thread_scope.py
@@ -0,0 +1,143 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import os
+import pickle
+import subprocess
+import sys
+
+import pytest
+import tvm_ffi
+from tvm_ffi.dataclasses import Enum
+
+import tvm
+from tvm.ir import Call, Evaluate, make_node
+from tvm.script import tirx as T
+from tvm.tirx import Function, decl_tensor
+from tvm.tirx.thread_scope import ThreadScope
+from tvm.tirx.tile import DispatchContext
+from tvm.tirx.tile.instruction import TensorCall
+
+
[email protected](
+ "backend,names",
+ [
+ ("cuda", ("kernel", "cluster", "cta", "warpgroup", "warp", "thread")),
+ ("trn", ("kernel", "thread")),
+ ],
+)
+def test_enum_identity_and_explicit_ranks(backend, names):
+ for rank, name in enumerate(names, start=-1):
+ scope = ThreadScope(f"{backend}.{name}")
+ assert isinstance(scope, Enum)
+ assert scope.same_as(ThreadScope.get(scope.name))
+ assert scope.same_as(getattr(T, backend).ThreadScope(name))
+ assert scope.name == f"{backend}.{name}"
+ assert scope.local_name == name
+ assert scope.backend == backend
+ assert scope.rank == rank
+ assert ThreadScope.def_attr("backend")[scope] == backend
+ assert ThreadScope.def_attr("rank")[scope] == rank
+ assert hash(scope) == hash(ThreadScope(scope.name))
+ assert not ThreadScope("cuda.thread").same_as(ThreadScope("trn.thread"))
+
+
+def test_hierarchy_comparison_and_lookup_errors():
+ higher = tvm_ffi.get_global_func("tirx.ScopeHigher")
+ assert higher(ThreadScope("cuda.kernel"), ThreadScope("cuda.thread"))
+ assert not higher(ThreadScope("cuda.thread"), ThreadScope("cuda.kernel"))
+ assert not higher(ThreadScope("cuda.warp"), ThreadScope("cuda.warp"))
+ with pytest.raises(ValueError, match="different backends"):
+ higher(ThreadScope("cuda.thread"), ThreadScope("trn.thread"))
+ for name in ("warp", "cuda.unknown", "trn.warp", ""):
+ with pytest.raises(ValueError):
+ ThreadScope(name)
+
+
[email protected]("name", ["cuda.kernel", "cuda.warp", "trn.thread"])
+def test_scope_roundtrips_keep_singleton(name):
+ scope = ThreadScope(name)
+ assert tvm.ir.load_json(tvm.ir.save_json(scope)).same_as(scope)
+ assert pickle.loads(pickle.dumps(scope)).same_as(scope)
+ text = scope.script()
+ assert f'T.ThreadScope("{name}")' in text
+ assert eval(text, {"T": T}).same_as(scope)
+
+
[email protected]("scope", ["warp", "cuda.warp",
T.cuda.ThreadScope("warp")])
+def test_instruction_normalizes_scope_before_constructing_call(scope):
+ tensor = decl_tensor((16,), "float32", scope="local", layout="default")
+ call = T.cuda.tile.apply(T.ptx["sqrt.rn.f32"], tensor, tensor, scope=scope)
+ tag = ThreadScope("cuda.warp")
+ assert call.attrs.scope.same_as(tag)
+ assert TensorCall.decode(call).scope.same_as(tag)
+ default = T.cuda.tile.apply(T.ptx["sqrt.rn.f32"], tensor, tensor)
+ assert default.attrs.scope.same_as(ThreadScope("cuda.thread"))
+ context = DispatchContext(tvm.target.Target("cuda"), tag, {}, {})
+ assert context.thread_scope.same_as(tag)
+ assert context.scope_kind == "warp"
+ with pytest.raises(ValueError, match="scope_kind"):
+ DispatchContext(tvm.target.Target("cuda"), tag, {}, {},
scope_kind="thread")
+
+
+def test_raw_scope_attrs_are_required_and_typed():
+ for schema in ("ScopeAttrs", "TrnAttrs"):
+ with pytest.raises(TypeError, match="scope"):
+ make_node(f"tirx.tensor.{schema}")
+ for invalid in (None, "thread", 1):
+ with pytest.raises((TypeError, ValueError)):
+ make_node(f"tirx.tensor.{schema}", scope=invalid)
+ attrs = make_node(f"tirx.tensor.{schema}", scope="cuda.thread")
+ assert attrs.scope.same_as(ThreadScope("cuda.thread"))
+ assert
tvm.ir.load_json(tvm.ir.save_json(attrs)).scope.same_as(attrs.scope)
+
+
+def test_cross_backend_scopes_are_rejected_at_each_boundary():
+ tensor = decl_tensor((16,), "float32", scope="local", layout="default")
+ trn_scope = ThreadScope("trn.thread")
+ for scope in (trn_scope, "trn.thread"):
+ with pytest.raises(ValueError, match="belongs to trn"):
+ T.cuda.tile.apply(T.ptx["sqrt.rn.f32"], tensor, tensor,
scope=scope)
+ call = T.cuda.tile.apply(T.ptx["sqrt.rn.f32"], tensor, tensor)
+ attrs = make_node("tirx.tensor.ScopeAttrs", scope=trn_scope)
+ raw = Call(call.op, call.args, attrs=attrs, ty=call.ty)
+ with pytest.raises(ValueError, match="belongs to trn"):
+ raw.validate()
+ with pytest.raises(ValueError, match="does not match target"):
+ DispatchContext(tvm.target.Target("cuda"), trn_scope, {}, {})
+ trn_tensor = decl_tensor((16,), "float32", scope="trn.sbuf",
layout="default")
+ trn_call = T.trn.tile.tensortensor(trn_tensor, trn_tensor, trn_tensor)
+ func = Function([], Evaluate(trn_call)).with_attr("target",
tvm.target.Target("cuda"))
+ with pytest.raises(ValueError, match="does not match target"):
+ tvm.tirx.transform.TileDispatch()(tvm.IRModule({"main": func}))
+
+
+def test_runtime_only_import_and_removed_interfaces():
+ env = os.environ.copy()
+ env["TVM_USE_RUNTIME_LIB"] = "1"
+ subprocess.run(
+ [
+ sys.executable,
+ "-c",
+ "import tvm; from tvm.tirx.thread_scope import ThreadScope; "
+ "assert tvm.base._RUNTIME_ONLY",
+ ],
+ env=env,
+ check=True,
+ )
+ assert not hasattr(tvm.tirx, "ExecScope")
+ assert not hasattr(T.cuda, "ExecScope")
+ assert tvm_ffi.get_global_func("tirx.RegisterExecScope",
allow_missing=True) is None
diff --git a/tests/python/tirx/test_tile_namespace.py
b/tests/python/tirx/test_tile_namespace.py
index 672340e8d7..543175e79b 100644
--- a/tests/python/tirx/test_tile_namespace.py
+++ b/tests/python/tirx/test_tile_namespace.py
@@ -28,9 +28,9 @@ from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
from tvm.testing import env
from tvm.tirx import decl_tensor
-from tvm.tirx.exec_scope import ExecScope
from tvm.tirx.layout import Axis, S, TileLayout
from tvm.tirx.script import ir_builder as I
+from tvm.tirx.thread_scope import ThreadScope
def test_generic_layout_without_backend_autoload():
@@ -67,8 +67,8 @@ def test_backend_axis_and_scope_identity():
assert Axis.m.backend is None
assert T.cuda.Axis.laneid.backend == "cuda"
assert T.trn.Axis.P.backend == "trn"
- assert T.cuda.Axis.laneid.get_scope().name == "warp"
- assert T.cuda.Axis.laneid.get_subscope().name == "thread"
+ assert T.cuda.Axis.laneid.get_scope().name == "cuda.warp"
+ assert T.cuda.Axis.laneid.get_subscope().name == "cuda.thread"
assert not T.trn.Axis.P.is_thread()
with pytest.raises(AttributeError):
_ = Axis.laneid
@@ -76,20 +76,20 @@ def test_backend_axis_and_scope_identity():
_ = T.cuda.Axis.P
with pytest.raises(AttributeError):
_ = T.trn.Axis.laneid
- cuda = T.cuda.ExecScope("thread")
- trn = T.trn.ExecScope("thread")
+ cuda = T.cuda.ThreadScope("thread")
+ trn = T.trn.ThreadScope("thread")
assert cuda.backend == "cuda" and trn.backend == "trn"
assert not tvm_ffi.structural_equal(cuda, trn)
assert_structural_equal(cuda, load_json(save_json(cuda)))
assert_structural_equal(trn, load_json(save_json(trn)))
- for backend, factory in (("cuda", T.cuda.ExecScope), ("trn",
T.trn.ExecScope)):
+ for backend, factory in (("cuda", T.cuda.ThreadScope), ("trn",
T.trn.ThreadScope)):
root = factory("kernel")
- assert_structural_equal(root, ExecScope(backend, "kernel"))
+ assert_structural_equal(root, ThreadScope(f"{backend}.kernel"))
assert_structural_equal(root, load_json(save_json(root)))
- assert f'T.{backend}.ExecScope("kernel")' in root.script()
- assert_structural_equal(T.cuda.Axis.bx.get_scope(),
T.cuda.ExecScope("kernel"))
- with pytest.raises(ValueError, match="Unknown execution scope"):
- T.trn.ExecScope("warp")
+ assert f'T.ThreadScope("{backend}.kernel")' in root.script()
+ assert_structural_equal(T.cuda.Axis.bx.get_scope(),
T.cuda.ThreadScope("kernel"))
+ with pytest.raises(ValueError):
+ T.trn.ThreadScope("warp")
@pytest.mark.parametrize(
@@ -212,7 +212,7 @@ def test_tile_call_rejects_cross_backend_scope():
tensor = decl_tensor((8,), "float32", scope="local", layout="default")
with pytest.raises(ValueError, match="scope belongs to trn"):
T.cuda.tile.apply(
- T.ptx["add.f32"], tensor[:], tensor[:], tensor[:],
scope=T.trn.ExecScope("thread")
+ T.ptx["add.f32"], tensor[:], tensor[:], tensor[:],
scope=T.trn.ThreadScope("thread")
)
diff --git a/tests/python/tirx/test_tirx_axis_enum.py
b/tests/python/tirx/test_tirx_axis_enum.py
index e6e22d433c..7ba649119c 100644
--- a/tests/python/tirx/test_tirx_axis_enum.py
+++ b/tests/python/tirx/test_tirx_axis_enum.py
@@ -39,8 +39,8 @@ def test_axis_enum_identity_and_attributes():
assert hash(axis) == hash(Axis.get("laneid"))
assert Axis.def_attr("thread")[axis] is True
assert Axis.m.is_memory()
- assert CUDAAxis.bx.get_scope().name == "kernel"
- assert CUDAAxis.bx.get_subscope().name == "cta"
+ assert CUDAAxis.bx.get_scope().name == "cuda.kernel"
+ assert CUDAAxis.bx.get_subscope().name == "cuda.cta"
assert TileLayout(S[4 : 1 @ axis]).shard[0].axis.same_as(axis)
diff --git a/tests/python/tirx/tile/cuda/test_add_sub_mul_div.py
b/tests/python/tirx/tile/cuda/test_add_sub_mul_div.py
index 3533258446..6e557ef241 100644
--- a/tests/python/tirx/tile/cuda/test_add_sub_mul_div.py
+++ b/tests/python/tirx/tile/cuda/test_add_sub_mul_div.py
@@ -192,17 +192,17 @@ def test_binary_non_commutative_const_lhs(op_type):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
[email protected]("exec_scope", ["warp", "warpgroup"])
[email protected]("thread_scope", ["warp", "warpgroup"])
@pytest.mark.parametrize("op_type", ["add", "mul"])
-def test_binary_op_shared_subcta_scope(exec_scope, op_type):
- n_warps = 4 if exec_scope == "warpgroup" else 1
+def test_binary_op_shared_subcta_scope(thread_scope, op_type):
+ n_warps = 4 if thread_scope == "warpgroup" else 1
shape = (n_warps * 32, 8)
size = math.prod(shape)
# The old test executes warp 5 / warpgroup 1 inside a 256-thread block.
- group_shape = (n_warps, 32, 8, 1) if exec_scope == "warpgroup" else (32,
8, 1)
+ group_shape = (n_warps, 32, 8, 1) if thread_scope == "warpgroup" else (32,
8, 1)
reg_layout = (
TileLayout(S[(n_warps, 32, 8, 1) : (1 @ wid_in_wg, 1 @ laneid, 1, 1)])
- if exec_scope == "warpgroup"
+ if thread_scope == "warpgroup"
else TileLayout(S[(32, 8, 1) : (1 @ laneid, 1, 1)])
)
@@ -222,32 +222,32 @@ def test_binary_op_shared_subcta_scope(exec_scope,
op_type):
T.cuda.cta_sync()
a = T.alloc_local(group_shape, "float16", layout=reg_layout)
b = T.alloc_local(group_shape, "float16", layout=reg_layout)
- if T.constexpr(exec_scope == "warpgroup"):
+ if T.constexpr(thread_scope == "warpgroup"):
if wg == 1:
T.cuda.tile.apply(
T.ptx["ld.shared.b16"],
a.view("uint16"),
AS.view(*group_shape).view("uint16"),
- scope=exec_scope,
+ scope=thread_scope,
)
T.cuda.tile.apply(
T.ptx["ld.shared.b16"],
b.view("uint16"),
BS.view(*group_shape).view("uint16"),
- scope=exec_scope,
+ scope=thread_scope,
)
T.cuda.tile.apply(
T.ptx[f"{op_type}.rn.f16"],
a.view("uint16"),
a.view("uint16"),
b.view("uint16"),
- scope=exec_scope,
+ scope=thread_scope,
)
T.cuda.tile.apply(
T.ptx["st.shared.b16"],
AS.view(*group_shape).view("uint16"),
a.view("uint16"),
- scope=exec_scope,
+ scope=thread_scope,
)
else:
if wid == 5:
@@ -255,26 +255,26 @@ def test_binary_op_shared_subcta_scope(exec_scope,
op_type):
T.ptx["ld.shared.b16"],
a.view("uint16"),
AS.view(*group_shape).view("uint16"),
- scope=exec_scope,
+ scope=thread_scope,
)
T.cuda.tile.apply(
T.ptx["ld.shared.b16"],
b.view("uint16"),
BS.view(*group_shape).view("uint16"),
- scope=exec_scope,
+ scope=thread_scope,
)
T.cuda.tile.apply(
T.ptx[f"{op_type}.rn.f16"],
a.view("uint16"),
a.view("uint16"),
b.view("uint16"),
- scope=exec_scope,
+ scope=thread_scope,
)
T.cuda.tile.apply(
T.ptx["st.shared.b16"],
AS.view(*group_shape).view("uint16"),
a.view("uint16"),
- scope=exec_scope,
+ scope=thread_scope,
)
T.cuda.cta_sync()
for i in T.serial(T.ceildiv(size, 256)):
@@ -289,15 +289,15 @@ def test_binary_op_shared_subcta_scope(exec_scope,
op_type):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
[email protected]("exec_scope", ["cta", "warpgroup", "warp"])
[email protected]("thread_scope", ["cta", "warpgroup", "warp"])
@pytest.mark.parametrize("rhs_kind", ["region", "broadcast", "const"])
@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
-def test_binary_op_local_subcta_trivial(exec_scope, rhs_kind, op_type):
+def test_binary_op_local_subcta_trivial(thread_scope, rhs_kind, op_type):
dtype = "float16"
m, n = 4, 8
- n_threads = 256 if exec_scope == "cta" else (128 if exec_scope ==
"warpgroup" else 32)
+ n_threads = 256 if thread_scope == "cta" else (128 if thread_scope ==
"warpgroup" else 32)
# in this test, use warp3/warpgroup1 to test
- thr_str = 0 if exec_scope == "cta" else (128 if exec_scope == "warpgroup"
else 32 * 3)
+ thr_str = 0 if thread_scope == "cta" else (128 if thread_scope ==
"warpgroup" else 32 * 3)
a_shape = (n_threads, m, n)
b_n = n if rhs_kind == "region" else 1
b_shape = (n_threads, m, b_n)
@@ -318,7 +318,7 @@ def test_binary_op_local_subcta_trivial(exec_scope,
rhs_kind, op_type):
"cta": T.cuda.linear_thread_id,
"warpgroup": T.cuda.thread_in_warpgroup,
"warp": T.cuda.lane_id,
- }[exec_scope]
+ }[thread_scope]
@T.function
def kernel(
@@ -352,12 +352,12 @@ def test_binary_op_local_subcta_trivial(exec_scope,
rhs_kind, op_type):
for j in T.serial(b_n):
B_local[i, j] = B[tid_in_scope, i, j]
- if T.constexpr(exec_scope == "cta"):
+ if T.constexpr(thread_scope == "cta"):
if T.constexpr(rhs_kind == "const"):
tx_op(C_local, A_local, const)
else:
tx_op(C_local, A_local, B_values)
- elif T.constexpr(exec_scope == "warpgroup"):
+ elif T.constexpr(thread_scope == "warpgroup"):
if wg_id == 1:
if T.constexpr(rhs_kind == "const"):
tx_op(C_local, A_local, const)
@@ -428,10 +428,10 @@ def test_binary_op_local_subcta_trivial(exec_scope,
rhs_kind, op_type):
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
@pytest.mark.parametrize("storage_scope", ["shared", "local"])
[email protected]("exec_scope", ["cta", "thread"])
[email protected]("thread_scope", ["cta", "thread"])
@pytest.mark.parametrize("op_type", ["add", "sub", "mul", "fdiv"])
@pytest.mark.parametrize("dtype", ["float16"])
-def test_binary_op_vectorized(input, storage_scope, exec_scope, op_type,
dtype):
+def test_binary_op_vectorized(input, storage_scope, thread_scope, op_type,
dtype):
a_shape, b_shape, res_shape, thread_cnt, device_id = input
row_a, row_b = a_shape[1:], b_shape[1:]
n, m = math.prod(row_a), math.prod(row_b)
@@ -473,7 +473,7 @@ def test_binary_op_vectorized(input, storage_scope,
exec_scope, op_type, dtype):
for j in range(m):
bl[j] = bf[tid, j]
bv = b.view(*row_a, layout=broadcast)
- if T.constexpr(exec_scope == "cta"):
+ if T.constexpr(thread_scope == "cta"):
# Expose the implicit per-thread register arrays as an explicit
# CTA view. These declarations share their original storage.
av = a.view(thread_cnt // 32, 32, *row_a, layout=cta_a_layout)
diff --git a/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
b/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
index c8fe230613..0818d8112e 100644
--- a/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
+++ b/tests/python/tirx/tile/cuda/test_cp_async_bulk.py
@@ -35,8 +35,8 @@ from tvm.ir import TensorRegion, TensorRegionType
from tvm.script import tirx as T
from tvm.testing import env
from tvm.tirx import IntImm, Var
-from tvm.tirx.exec_scope import ExecScope
from tvm.tirx.layout import S, TileLayout
+from tvm.tirx.thread_scope import ThreadScope
from tvm.tirx.tile.dispatch import DispatchContext, run_dispatch
BULK = "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes"
@@ -53,7 +53,7 @@ def _make_dsmem_dispatch_call(shape, dtype, src_layout,
dst_layout, chunks=1):
T.ptx[BULK], dst, src, chunk_bytes=chunk_bytes, mbar=Var("mbar",
"handle"), remote_cta_id=1
)
target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
- sctx = DispatchContext(target, ExecScope("cuda", "thread"), {}, {})
+ sctx = DispatchContext(target, ThreadScope("cuda.thread"), {}, {})
return run_dispatch(op_call, sctx)
@@ -312,7 +312,7 @@ def test_dsmem_dispatch_missing_config():
buf = tvm.tirx.decl_tensor((64,), "float16", "A", scope="shared.dyn",
layout=layout)
br = TensorRegion(buf, [Range.from_min_extent(0, 64)],
ty=TensorRegionType())
target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"})
- sctx = DispatchContext(target, ExecScope("cuda", "thread"), {}, {})
+ sctx = DispatchContext(target, ThreadScope("cuda.thread"), {}, {})
with pytest.raises(RuntimeError, match="remote_cta_id"):
run_dispatch(
diff --git a/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
b/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
index b34207bad5..482ce496f0 100644
--- a/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
+++ b/tests/python/tirx/tile/cuda/test_cp_async_bulk_tensor.py
@@ -35,8 +35,8 @@ from tvm.ir import PrimType, PtrType, Range, TensorRegion,
TensorRegionType
from tvm.script import tirx as T
from tvm.sym import Analyzer
from tvm.tirx import IntImm, Var
-from tvm.tirx.exec_scope import ExecScope
from tvm.tirx.layout import ComposeLayout, S, TileLayout
+from tvm.tirx.thread_scope import ThreadScope
from tvm.tirx.tile.dispatch import DispatchContext, run_dispatch
_TMA_OPS = {
@@ -745,7 +745,7 @@ def _region(tensor, ranges=None):
def _context(arch="sm_100a", ranges=None):
return DispatchContext(
tvm.target.Target({"kind": "cuda", "arch": arch}),
- ExecScope("cuda", "thread"),
+ ThreadScope("cuda.thread"),
{},
ranges or {},
)
diff --git a/tests/python/tirx/tile/cuda/test_mov.py
b/tests/python/tirx/tile/cuda/test_mov.py
index 01cdc4886e..6320759842 100644
--- a/tests/python/tirx/tile/cuda/test_mov.py
+++ b/tests/python/tirx/tile/cuda/test_mov.py
@@ -45,14 +45,14 @@ def _run(kernel, arrays, expected, output=0, atol=1e-3,
rtol=1e-7, target="cuda"
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
[email protected]("exec_scope", ["warp", "warpgroup"])
-def test_unary_op_shared_subcta_scope(exec_scope):
- warps = 4 if exec_scope == "warpgroup" else 1
[email protected]("thread_scope", ["warp", "warpgroup"])
+def test_unary_op_shared_subcta_scope(thread_scope):
+ warps = 4 if thread_scope == "warpgroup" else 1
shape = (warps * 32, 8)
- group = (warps, 32, 8, 1) if exec_scope == "warpgroup" else (32, 8, 1)
+ group = (warps, 32, 8, 1) if thread_scope == "warpgroup" else (32, 8, 1)
layout = (
TileLayout(S[(4, 32, 8, 1) : (1 @ wid_in_wg, 1 @ laneid, 1, 1)])
- if exec_scope == "warpgroup"
+ if thread_scope == "warpgroup"
else TileLayout(S[(32, 8, 1) : (1 @ laneid, 1, 1)])
)
size = math.prod(shape)
@@ -71,7 +71,7 @@ def test_unary_op_shared_subcta_scope(exec_scope):
sf[q] = af[q]
T.cuda.cta_sync()
r = T.alloc_local(group, "uint16", layout=layout)
- if T.constexpr(exec_scope == "warp"):
+ if T.constexpr(thread_scope == "warp"):
if warp == 5:
T.cuda.tile.apply(T.ptx["mov.b16"], r, T.uint16(0),
scope="warp")
T.cuda.tile.apply(
@@ -97,9 +97,9 @@ def test_unary_op_shared_subcta_scope(exec_scope):
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
@pytest.mark.parametrize("shape", [(128, 8), (128, 4, 16), (128, 5, 5)])
@pytest.mark.parametrize("op_type", ["fill"])
[email protected]("exec_scope", ["thread", "cta"])
[email protected]("thread_scope", ["thread", "cta"])
@pytest.mark.parametrize("storage_scope", ["local", "shared"])
-def test_unary_op_vectorized(shape, op_type, exec_scope, storage_scope):
+def test_unary_op_vectorized(shape, op_type, thread_scope, storage_scope):
n = math.prod(shape[1:])
bits = int(np.float16(7.89).view("uint16"))
group_layout = TileLayout(S[(4, 32, n, 1) : (1 @ warpid, 1 @ laneid, 1,
1)])
@@ -109,7 +109,7 @@ def test_unary_op_vectorized(shape, op_type, exec_scope,
storage_scope):
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=128))
tid = T.cuda.thread_idx("x")
r = T.alloc_local(shape[1:], "float16", layout="default")
- if T.constexpr(exec_scope == "cta"):
+ if T.constexpr(thread_scope == "cta"):
values = r.view(4, 32, n, 1, layout=group_layout).view("uint16")
T.cuda.tile.apply(T.ptx["mov.b16"], values, T.uint16(bits),
scope="cta")
else:
diff --git a/tests/python/tirx/tile/test_dispatcher.py
b/tests/python/tirx/tile/test_dispatcher.py
index 660683d1db..125dd22cb3 100644
--- a/tests/python/tirx/tile/test_dispatcher.py
+++ b/tests/python/tirx/tile/test_dispatcher.py
@@ -66,12 +66,21 @@ def test_instruction_scope_contract(backend, scope):
with pytest.raises(ValueError, match=f"unsupported scope '{scope}'"):
instruction(tensor, tensor, tensor, scope=scope)
- # Raw Calls obey the same instruction contract through Op validation.
+ # Registered scope enums still obey the instruction's allowed scope set.
call = instruction(tensor, tensor, tensor)
- attrs = tvm.ir.make_node("tirx.tensor." + SPECS[call.op.name].schema,
scope=scope)
- raw = tvm.ir.Call(call.op, call.args, attrs=attrs, ty=call.ty)
- with pytest.raises(ValueError, match=f"unsupported scope '{scope}'"):
- raw.validate()
+ if scope == "kernel":
+ attrs = tvm.ir.make_node(
+ "tirx.tensor." + SPECS[call.op.name].schema,
+ scope=T.ThreadScope(f"{backend}.kernel"),
+ )
+ raw = tvm.ir.Call(call.op, call.args, attrs=attrs, ty=call.ty)
+ with pytest.raises(ValueError, match=f"unsupported scope '{scope}'"):
+ raw.validate()
+ else:
+ with pytest.raises(ValueError):
+ tvm.ir.make_node(
+ "tirx.tensor." + SPECS[call.op.name].schema,
scope=f"{backend}.{scope}"
+ )
def _custom_call(name):
@@ -81,7 +90,11 @@ def _custom_call(name):
op = Op.get(op_name)
op.set_signature(args=["token"])
op.set_attrs_type_key("tirx.tensor.ScopeAttrs")
- return Call(op, [const(7)], attrs=make_node("tirx.tensor.ScopeAttrs"))
+ return Call(
+ op,
+ [const(7)],
+ attrs=make_node("tirx.tensor.ScopeAttrs",
scope=T.ThreadScope("cuda.thread")),
+ )
def _lower(body):
@@ -101,7 +114,7 @@ def
test_op_attribute_dispatches_without_python_instruction_spec():
assert actual_call.same_as(call)
assert isinstance(context, DispatchContext)
assert context.target.kind.name == "cuda"
- assert context.exec_scope.backend == "cuda"
+ assert context.thread_scope.backend == "cuda"
assert context.scope_kind == "thread"
seen.append(actual_call)
return Function([], Evaluate(23))
diff --git a/tests/python/tirx/tile/trn/test_binary_trn.py
b/tests/python/tirx/tile/trn/test_binary_trn.py
index 286c98ed0f..60d88f9baf 100644
--- a/tests/python/tirx/tile/trn/test_binary_trn.py
+++ b/tests/python/tirx/tile/trn/test_binary_trn.py
@@ -27,7 +27,7 @@ from tvm.tirx.layout import S, TileLayout
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
-def _strip_exec_scope_stmt(stmt):
+def _strip_thread_scope_stmt(stmt):
def _strip_region(node: tvm.ir.RegionStmt):
if node.op.same_as(tvm.ir.Op.get("tirx.device_entry")):
return node.body
@@ -42,9 +42,9 @@ def _strip_exec_scope_stmt(stmt):
def assert_structural_equal(lhs, rhs, *args, **kwargs):
if isinstance(lhs, tvm.tirx.Function):
- lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+ lhs = lhs.with_body(_strip_thread_scope_stmt(lhs.body))
if isinstance(rhs, tvm.tirx.Function):
- rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+ rhs = rhs.with_body(_strip_thread_scope_stmt(rhs.body))
_assert_structural_equal(lhs, rhs, *args, **kwargs)
diff --git a/tests/python/tirx/tile/trn/test_compose_op_trn.py
b/tests/python/tirx/tile/trn/test_compose_op_trn.py
index 172f87581e..ac4d8f61ee 100644
--- a/tests/python/tirx/tile/trn/test_compose_op_trn.py
+++ b/tests/python/tirx/tile/trn/test_compose_op_trn.py
@@ -27,7 +27,7 @@ from tvm.tirx.layout import S, TileLayout
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
-def _strip_exec_scope_stmt(stmt):
+def _strip_thread_scope_stmt(stmt):
def _strip_region(node: tvm.ir.RegionStmt):
if node.op.same_as(tvm.ir.Op.get("tirx.device_entry")):
return node.body
@@ -42,9 +42,9 @@ def _strip_exec_scope_stmt(stmt):
def assert_structural_equal(lhs, rhs, *args, **kwargs):
if isinstance(lhs, tvm.tirx.Function):
- lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+ lhs = lhs.with_body(_strip_thread_scope_stmt(lhs.body))
if isinstance(rhs, tvm.tirx.Function):
- rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+ rhs = rhs.with_body(_strip_thread_scope_stmt(rhs.body))
_assert_structural_equal(lhs, rhs, *args, **kwargs)
diff --git a/tests/python/tirx/tile/trn/test_copy_trn.py
b/tests/python/tirx/tile/trn/test_copy_trn.py
index d94fcd871c..f2a21ddc03 100644
--- a/tests/python/tirx/tile/trn/test_copy_trn.py
+++ b/tests/python/tirx/tile/trn/test_copy_trn.py
@@ -31,7 +31,7 @@ from tvm.tirx.layout import S, TileLayout
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
-def _strip_exec_scope_stmt(stmt):
+def _strip_thread_scope_stmt(stmt):
def _strip_region(node: tvm.ir.RegionStmt):
if node.op.same_as(tvm.ir.Op.get("tirx.device_entry")):
return node.body
@@ -46,9 +46,9 @@ def _strip_exec_scope_stmt(stmt):
def assert_structural_equal(lhs, rhs, *args, **kwargs):
if isinstance(lhs, tvm.tirx.Function):
- lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+ lhs = lhs.with_body(_strip_thread_scope_stmt(lhs.body))
if isinstance(rhs, tvm.tirx.Function):
- rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+ rhs = rhs.with_body(_strip_thread_scope_stmt(rhs.body))
_assert_structural_equal(lhs, rhs, *args, **kwargs)
diff --git a/tests/python/tirx/tile/trn/test_gemm_trn.py
b/tests/python/tirx/tile/trn/test_gemm_trn.py
index 5b7f9067a7..fa6c575811 100644
--- a/tests/python/tirx/tile/trn/test_gemm_trn.py
+++ b/tests/python/tirx/tile/trn/test_gemm_trn.py
@@ -28,7 +28,7 @@ from tvm.tirx.layout import S, TileLayout
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
-def _strip_exec_scope_stmt(stmt):
+def _strip_thread_scope_stmt(stmt):
def _strip_region(node: tvm.ir.RegionStmt):
if node.op.same_as(tvm.ir.Op.get("tirx.device_entry")):
return node.body
@@ -43,9 +43,9 @@ def _strip_exec_scope_stmt(stmt):
def assert_structural_equal(lhs, rhs, *args, **kwargs):
if isinstance(lhs, tvm.tirx.Function):
- lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+ lhs = lhs.with_body(_strip_thread_scope_stmt(lhs.body))
if isinstance(rhs, tvm.tirx.Function):
- rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+ rhs = rhs.with_body(_strip_thread_scope_stmt(rhs.body))
_assert_structural_equal(lhs, rhs, *args, **kwargs)
diff --git a/tests/python/tirx/tile/trn/test_reduction_trn.py
b/tests/python/tirx/tile/trn/test_reduction_trn.py
index 627fab5619..582042a753 100644
--- a/tests/python/tirx/tile/trn/test_reduction_trn.py
+++ b/tests/python/tirx/tile/trn/test_reduction_trn.py
@@ -27,7 +27,7 @@ from tvm.tirx.layout import S, TileLayout
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
-def _strip_exec_scope_stmt(stmt):
+def _strip_thread_scope_stmt(stmt):
def _strip_region(node: tvm.ir.RegionStmt):
if node.op.same_as(tvm.ir.Op.get("tirx.device_entry")):
return node.body
@@ -42,9 +42,9 @@ def _strip_exec_scope_stmt(stmt):
def assert_structural_equal(lhs, rhs, *args, **kwargs):
if isinstance(lhs, tvm.tirx.Function):
- lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+ lhs = lhs.with_body(_strip_thread_scope_stmt(lhs.body))
if isinstance(rhs, tvm.tirx.Function):
- rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+ rhs = rhs.with_body(_strip_thread_scope_stmt(rhs.body))
_assert_structural_equal(lhs, rhs, *args, **kwargs)
diff --git a/tests/python/tirx/tile/trn/test_select_trn.py
b/tests/python/tirx/tile/trn/test_select_trn.py
index 46fc6ab4af..1e8b8e747a 100644
--- a/tests/python/tirx/tile/trn/test_select_trn.py
+++ b/tests/python/tirx/tile/trn/test_select_trn.py
@@ -27,7 +27,7 @@ from tvm.tirx.layout import S, TileLayout
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
-def _strip_exec_scope_stmt(stmt):
+def _strip_thread_scope_stmt(stmt):
def _strip_region(node: tvm.ir.RegionStmt):
if node.op.same_as(tvm.ir.Op.get("tirx.device_entry")):
return node.body
@@ -42,9 +42,9 @@ def _strip_exec_scope_stmt(stmt):
def assert_structural_equal(lhs, rhs, *args, **kwargs):
if isinstance(lhs, tvm.tirx.Function):
- lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+ lhs = lhs.with_body(_strip_thread_scope_stmt(lhs.body))
if isinstance(rhs, tvm.tirx.Function):
- rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+ rhs = rhs.with_body(_strip_thread_scope_stmt(rhs.body))
_assert_structural_equal(lhs, rhs, *args, **kwargs)
diff --git a/tests/python/tirx/tile/trn/test_unary_trn.py
b/tests/python/tirx/tile/trn/test_unary_trn.py
index 9f338e0db0..eb8f9d615e 100644
--- a/tests/python/tirx/tile/trn/test_unary_trn.py
+++ b/tests/python/tirx/tile/trn/test_unary_trn.py
@@ -27,7 +27,7 @@ from tvm.tirx.layout import S, TileLayout
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
-def _strip_exec_scope_stmt(stmt):
+def _strip_thread_scope_stmt(stmt):
def _strip_region(node: tvm.ir.RegionStmt):
if node.op.same_as(tvm.ir.Op.get("tirx.device_entry")):
return node.body
@@ -42,9 +42,9 @@ def _strip_exec_scope_stmt(stmt):
def assert_structural_equal(lhs, rhs, *args, **kwargs):
if isinstance(lhs, tvm.tirx.Function):
- lhs = lhs.with_body(_strip_exec_scope_stmt(lhs.body))
+ lhs = lhs.with_body(_strip_thread_scope_stmt(lhs.body))
if isinstance(rhs, tvm.tirx.Function):
- rhs = rhs.with_body(_strip_exec_scope_stmt(rhs.body))
+ rhs = rhs.with_body(_strip_thread_scope_stmt(rhs.body))
_assert_structural_equal(lhs, rhs, *args, **kwargs)
