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 15fe7c6af9 [REFACTOR][TIRx] Simplify the transform API (#20635)
15fe7c6af9 is described below
commit 15fe7c6af9b4f7681948ed9615490e1f1e0c1d95
Author: Tianqi Chen <[email protected]>
AuthorDate: Fri Oct 9 18:42:43 2026 -0400
[REFACTOR][TIRx] Simplify the transform API (#20635)
Remove the unused TIRx private-function inlining and thread-axis
remapping APIs and the unimplemented static-memory planner declaration.
Rename `TilePrimitiveDispatch` to `TileDispatch` consistently across
native and Python interfaces, pass identity, callers, source naming, and
documentation.
The dispatch implementation and lowering order are preserved.
---
docs/tirx/arch/lowering_pipeline.rst | 4 +-
docs/tirx/arch/tile_dispatch.rst | 2 +-
include/tvm/tirx/transform.h | 29 +-
.../cuda/tile_primitive/copy/vec_auto_reg.py | 4 +-
python/tvm/backend/trn/pipeline.py | 2 +-
python/tvm/backend/trn/transform/__init__.py | 2 +-
python/tvm/tirx/transform/transform.py | 15 +-
src/tirx/analysis/filter_canonical.cc | 2 +-
src/tirx/analysis/filter_canonical.h | 8 +-
src/tirx/transform/inline_private_functions.cc | 320 ---------------------
src/tirx/transform/lower_tirx.cc | 4 +-
src/tirx/transform/remap_thread_axis.cc | 113 --------
...tile_primitive_dispatch.cc => tile_dispatch.cc} | 14 +-
.../test_tir_inline_private_functions.py | 251 ----------------
.../operator/tile_primitive/cuda/copy/test_reg.py | 2 +-
15 files changed, 25 insertions(+), 747 deletions(-)
diff --git a/docs/tirx/arch/lowering_pipeline.rst
b/docs/tirx/arch/lowering_pipeline.rst
index e4a7844560..37f7fa5ad8 100644
--- a/docs/tirx/arch/lowering_pipeline.rst
+++ b/docs/tirx/arch/lowering_pipeline.rst
@@ -126,9 +126,9 @@ Inside LowerTIRx
.. code-block:: text
- LowerTIRx = Sequential([ TilePrimitiveDispatch, LowerTIRxCleanup ])
+ LowerTIRx = Sequential([ TileDispatch, LowerTIRxCleanup ])
-- **``TilePrimitiveDispatch``** replaces every tensor ``Evaluate(Call)`` with
the body emitted by its
+- **``TileDispatch``** replaces every tensor ``Evaluate(Call)`` with the body
emitted by its
instruction lowerer, including delayed mathematical composites as described
in
:doc:`tile_dispatch`. In the same pass it removes the ``device_entry``
marker, resolves standalone scope-id definitions to ``Bind`` statements,
diff --git a/docs/tirx/arch/tile_dispatch.rst b/docs/tirx/arch/tile_dispatch.rst
index a62738bbca..fd7f04f9ef 100644
--- a/docs/tirx/arch/tile_dispatch.rst
+++ b/docs/tirx/arch/tile_dispatch.rst
@@ -24,7 +24,7 @@ in ``python/tvm/backend/{cuda,trn}/tensor_instructions.py``
using the shared
``python/tvm/tirx/tensor_instruction.py`` machinery. C++ reflected attribute
types and the native Call validator are in ``src/tirx/op/tile.cc``.
-``tirx.TilePrimitiveDispatch`` remains the first phase of ``LowerTIRx``.
+``tirx.TileDispatch`` remains the first phase of ``LowerTIRx``.
The pass recognizes ``Evaluate(Call)`` by the operator's ``TIRxOpCategory``:
``tile_primitive`` or ``tile_composite``. It validates the call, resolves the
scope against the active thread set, constructs a ``DispatchContext``, and
diff --git a/include/tvm/tirx/transform.h b/include/tvm/tirx/transform.h
index a5e68d54be..a000e82f93 100644
--- a/include/tvm/tirx/transform.h
+++ b/include/tvm/tirx/transform.h
@@ -154,18 +154,6 @@ TVM_DLL Pass LowerThreadAllreduce();
*/
TVM_DLL Pass MakePackedAPI();
-/*!
- * \brief Remap the thread axis
- *
- * This can be used to get equivalent program which uses
- * threadIdx.y in place of threadIdx.x by passing
- * {"threadIdx.x": "threadIdx.y"}
- *
- *
- * \return The pass.
- */
-TVM_DLL Pass RemapThreadAxis(ffi::Map<ffi::String, ffi::String> axis_map);
-
/*!
* \brief Annotate, split, and lower host/device functions.
*
@@ -267,13 +255,6 @@ TVM_DLL Pass BF16StorageLegalize();
*/
TVM_DLL Pass FP8StorageLegalize();
-/*!
- * \brief Inline calls to private functions
- *
- * \return The pass.
- */
-TVM_DLL Pass InlinePrivateFunctions();
-
/*!
* \brief Rewrite the pointer content type of arguments,
* as well as Alloc internal to the function to use
@@ -298,14 +279,6 @@ TVM_DLL Pass FlattenBuffer();
*/
TVM_DLL Pass CommonSubexprElim();
-/*!
- * \brief This is the unified static memory planner pass that will
- * plan for memory intra- and inter- Functions together. The pass
- * requires all the function to be Functions including the main.
- * \return The pass.
- */
-TVM_DLL Pass UnifiedStaticMemoryPlanner();
-
/*!
* \brief Annotate a Function with a given target.
* \return The pass.
@@ -332,7 +305,7 @@ TVM_DLL Pass Filter(ffi::TypedFunction<bool(Function)>
fcond);
* launch_thread RegionStmts wrapping the dispatched body.
* \return The pass.
*/
-TVM_DLL Pass TilePrimitiveDispatch();
+TVM_DLL Pass TileDispatch();
/*!
* \brief Finalize TIRx lowering by applying layout rewriters and cleanup
passes.
diff --git a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
index a411e1a502..c9c6542028 100644
--- a/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
+++ b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
@@ -364,9 +364,9 @@ def _axis_decl(axis_name: str, sctx: DispatchContext):
"""Declare the runtime Var for one thread axis (called inside impl body).
Each scope_id declarator emits a ``ScopeIdDef`` stmt at the current
- builder frame. ``TilePrimitiveDispatch`` re-gathers + resolves all
+ builder frame. ``TileDispatch`` re-gathers + resolves all
ScopeIdDefs after dispatch (see ``ResolveAllScopeBinds`` in
- ``tile_primitive_dispatch.cc``), so dispatch-introduced vars are bound
+ ``tile_dispatch.cc``), so dispatch-introduced vars are bound
alongside kernel-declared ones.
Extents are deferred: the kernel header is expected to declare the full
diff --git a/python/tvm/backend/trn/pipeline.py
b/python/tvm/backend/trn/pipeline.py
index c2679e207b..dd9b252e1c 100644
--- a/python/tvm/backend/trn/pipeline.py
+++ b/python/tvm/backend/trn/pipeline.py
@@ -32,7 +32,7 @@ def trn_pipeline():
passes = [
trn_transform.TrnPrivateBufferAlloc(),
trn_transform.TrnNaiveAllocator(),
- tirx.transform.TilePrimitiveDispatch(),
+ tirx.transform.TileDispatch(),
trn_transform.LowerTrainiumLayout(),
tvm.s_tir.transform.DecorateDeviceScope(),
tirx.transform.StmtSimplify(),
diff --git a/python/tvm/backend/trn/transform/__init__.py
b/python/tvm/backend/trn/transform/__init__.py
index b74b98b0c6..703f0c6e7f 100644
--- a/python/tvm/backend/trn/transform/__init__.py
+++ b/python/tvm/backend/trn/transform/__init__.py
@@ -36,7 +36,7 @@ def LowerTrainiumLayout():
def LowerTIRx():
"""Lower TIRx tile primitive calls for the Trainium backend."""
return tvm.transform.Sequential(
- [tirx.transform.TilePrimitiveDispatch(), LowerTrainiumLayout()],
+ [tirx.transform.TileDispatch(), LowerTrainiumLayout()],
name="tirx.backend.trn.LowerTIRx",
)
diff --git a/python/tvm/tirx/transform/transform.py
b/python/tvm/tirx/transform/transform.py
index c0b5eeea23..12e69f0551 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -85,17 +85,6 @@ def StorageRewrite():
return _ffi_api.StorageRewrite() # type: ignore
-def InlinePrivateFunctions():
- """Inline calls to private functions
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- return _ffi_api.InlinePrivateFunctions() # type: ignore
-
-
def PointerValueTypeRewrite():
"""
Rewrite the pointer content type of arguments, as well as Alloc internal
to the function to use
@@ -459,7 +448,7 @@ def Filter(fcond: Callable):
return _ffi_api.Filter(fcond) # type: ignore
-def TilePrimitiveDispatch():
+def TileDispatch():
"""Lower TIRx tile primitive calls through the active backend dispatch
table.
Returns
@@ -467,7 +456,7 @@ def TilePrimitiveDispatch():
fpass : tvm.transform.Pass
The result pass
"""
- return _ffi_api.TilePrimitiveDispatch() # type: ignore
+ return _ffi_api.TileDispatch() # type: ignore
def LowerTIRx():
diff --git a/src/tirx/analysis/filter_canonical.cc
b/src/tirx/analysis/filter_canonical.cc
index 5e23d921ee..d85c9f76a2 100644
--- a/src/tirx/analysis/filter_canonical.cc
+++ b/src/tirx/analysis/filter_canonical.cc
@@ -55,7 +55,7 @@ PrimExpr StripCast(const PrimExpr& expr) {
}
// Recognized conjunction shapes: logical-And and bitwise-And nodes.
-// Mirrors FlattenConjuncts in tile_primitive_dispatch.cc so the classifier
+// Mirrors FlattenConjuncts in tile_dispatch.cc so the classifier
// accepts the same set of "fully conjunctive" predicates that the existing
// pass-internal helpers do.
void FlattenConjuncts(const PrimExpr& pred, std::vector<PrimExpr>* out) {
diff --git a/src/tirx/analysis/filter_canonical.h
b/src/tirx/analysis/filter_canonical.h
index 200fdf6207..a2df5c02a7 100644
--- a/src/tirx/analysis/filter_canonical.h
+++ b/src/tirx/analysis/filter_canonical.h
@@ -30,7 +30,7 @@
* | Call("tirx.cuda.elect_sync")
*
* Consumers:
- * 1. tile_primitive_dispatch routes a bare `if cond:` to atom-based
+ * 1. tile_dispatch routes a bare `if cond:` to atom-based
* narrowing when cond is canonical; otherwise treats it as a regular
* data-dependent branch (no narrowing).
* 2. The `tirx.gpu_thread_filter(var, pred)` escape-hatch wrapper is
intended for the
@@ -79,7 +79,7 @@ enum class FilterAtomKind {
* For `kElectSync`:
* - `elect_sync_call`: the original `Call("tirx.cuda.elect_sync")` PrimExpr,
* preserved verbatim so downstream consumers (e.g. selector construction
- * in tile_primitive_dispatch) can reuse it without re-synthesizing.
+ * in tile_dispatch) can reuse it without re-synthesizing.
* - `scopeid_var`, `lo`, `hi` are unset.
*/
struct FilterAtom {
@@ -108,7 +108,7 @@ struct CanonicalForm {
* The classifier consults this for every variable that appears on the LHS
* of a comparison atom. The callback abstracts over how scope ids are
* tracked in the caller's context:
- * - TilePrimitiveDispatcher passes a lambda that walks its
+ * - TileDispatcher passes a lambda that walks its
* `scope_id_defs_at_level_` stack.
* - Tests may pass a simpler lambda over a fixed allow-list of vars.
*
@@ -135,7 +135,7 @@ using ScopeIdPredicate = std::function<bool(const Var&)>;
* Implementation notes:
* - Conjunction is recognized via both `tir::And` nodes and
* `prim.BitwiseAnd` nodes (matching existing FlattenConjuncts behavior
- * in tile_primitive_dispatch.cc).
+ * in tile_dispatch.cc).
* - Comparison atoms with `const <op> var` are mirrored so the
* `scopeid_var` is on the LHS of the returned atom.
* - `c1 == c2` (two constants), `v1 == v2` (two vars), and any other
diff --git a/src/tirx/transform/inline_private_functions.cc
b/src/tirx/transform/inline_private_functions.cc
deleted file mode 100644
index 919d56809c..0000000000
--- a/src/tirx/transform/inline_private_functions.cc
+++ /dev/null
@@ -1,320 +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.
- */
-
-/*!
- * \file inline_private_functions.cc
- * \brief Inline private functions to their callsite
- */
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
-#include <tvm/tirx/transform.h>
-
-namespace tvm {
-namespace tirx {
-namespace transform {
-
-namespace {
-
-template <typename T>
-using PSet = std::unordered_set<T, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>;
-
-template <typename T, typename U>
-using PMap = std::unordered_map<T, U, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>;
-
-PMap<GlobalVar, PSet<GlobalVar>> CollectCallMap(const IRModule& mod) {
- struct Visitor : StmtExprVisitor {
- GlobalVar current{ffi::UnsafeInit{}};
- PMap<GlobalVar, PSet<GlobalVar>> caller_lookup;
-
- ffi::Optional<VisitInterrupt> Visit_(const CallNode* op) final {
- if (auto gvar = op->op.as<GlobalVar>()) {
- caller_lookup[gvar.value()].insert(current);
- }
- return StmtExprVisitor::Visit_(op);
- }
- };
- auto visitor = ffi::make_object<Visitor>();
-
- for (const auto& [gvar, base_func] : mod->functions) {
- if (auto function = base_func.as<FunctionNode>()) {
- visitor->current = gvar;
- visitor->Visit(function->body);
- }
- }
-
- return visitor->caller_lookup;
-}
-
-PSet<GlobalVar> CollectRecursiveFunctions(const IRModule& mod) {
- // Collect all direct callers
- auto call_map = CollectCallMap(mod);
-
- // Propagate to find all indirect callers
- while (true) {
- bool made_change = false;
- for (const auto& [callee, callers] : call_map) {
- for (const auto& caller : callers) {
- if (auto it = call_map.find(caller); it != call_map.end()) {
- PSet<GlobalVar>& indirect_callers = it->second;
-
- auto res = indirect_callers.insert(callee);
- made_change = made_change || res.second;
- }
- }
- }
- if (!made_change) {
- break;
- }
- }
-
- // Filter all GlobalVars that can be called by themselves, either
- // directly or indirectly.
- PSet<GlobalVar> recursive_funcs;
- for (const auto& [caller, callees] : call_map) {
- if (callees.count(caller)) {
- recursive_funcs.insert(caller);
- }
- }
- return recursive_funcs;
-}
-
-bool IsInlinableFunction(const GlobalVar& gvar, const Function& function,
- const PSet<GlobalVar>& recursive_functions) {
- // Only inline private functions. Externally-exposed functions
- // must be preserved so to avoid breaking callsites outside of
- // the IRModule.
- bool is_exposed =
function->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).has_value();
- if (is_exposed) return false;
-
- // We do not currently implement any analysis for termination of
- // a function. If a recursive function requires runtime checks
- // in order to terminate, we would keep inlining until the
- // recursive visits segfault.
- bool is_recursive = recursive_functions.count(gvar);
- if (is_recursive) return false;
-
- // We do not currently support inlining of functions that accept
- // buffer arguments.
- for (const Var& param : function->params) {
- if (param->ty.as<TensorTypeNode>()) return false;
- }
-
- // Generalize the old SBlockRealize exclusion to all non-native statement
roots:
- // they may introduce binder or naming rules that this pass cannot preserve.
- // Only inline roots supported by native TIRX traversal.
- struct NativeStmtTable : StmtExprVisitor {
- static VTable Make() {
- VTable table;
- InitVTable(&table);
- table.Finalize();
- return table;
- }
- };
- static const auto native_stmts = NativeStmtTable::Make();
- if (!function->body.has_value() ||
!native_stmts.CanDispatch(function->body.get())) {
- return false;
- }
-
- return true;
-}
-
-ffi::Map<GlobalVar, Function> CollectInlinableFunctions(const IRModule& mod) {
- auto recursive_functions = CollectRecursiveFunctions(mod);
-
- ffi::Map<GlobalVar, Function> output;
- for (const auto& [gvar, base_func] : mod->functions) {
- if (auto opt = base_func.as<Function>()) {
- auto function = opt.value();
- if (IsInlinableFunction(gvar, function, recursive_functions)) {
- output.Set(gvar, function);
- }
- }
- }
-
- return output;
-}
-
-class FunctionInliner : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
- explicit FunctionInliner(ffi::Map<GlobalVar, Function> inlinable_funcs)
- : inlinable_funcs_(inlinable_funcs) {
- for (const auto& [gvar, callee] : inlinable_funcs_) {
- removable_funcs_.insert(gvar);
- }
- }
-
- Function VisitFunc(Function func) {
- current_target_ = func->GetAttr<Target>(tvm::attr::kTarget);
- auto new_body_result = Mutate(func->body, InplaceMode::kDisallow);
- bool new_body_unchanged = new_body_result.UnchangedOrSameAs(func->body);
- auto new_body = std::move(new_body_result).ValueOrUnchanged(func->body);
- current_target_ = std::nullopt;
-
- if (!new_body_unchanged) {
- func.CopyOnWrite()->body = new_body;
- }
-
- return func;
- }
-
- PSet<GlobalVar> GetRemovableFunctions() const { return removable_funcs_; }
-
- private:
- UnchangedOr<Stmt> Mutate_(const EvaluateNode* eval, InplaceMode
inplace_mode) override {
- if (auto inlined = GetInlinedFunction(eval)) {
- return inlined.value();
- } else {
- return StmtExprMutator::Mutate_(eval, inplace_mode);
- }
- }
-
- ffi::Optional<Stmt> GetInlinedFunction(const EvaluateNode* eval) {
- auto call = eval->value.as<CallNode>();
- if (!call) return std::nullopt;
-
- auto gvar = call->op.as<GlobalVar>();
- if (!gvar) return std::nullopt;
-
- auto opt_callee = inlinable_funcs_.Get(gvar.value());
- if (!opt_callee) return std::nullopt;
- auto callee = opt_callee.value();
-
- bool is_same_target = [&]() -> bool {
- auto callee_target = callee->GetAttr<Target>(tvm::attr::kTarget);
- if (current_target_ && callee_target) {
- return callee_target.value()->str() == current_target_.value()->str();
- } else {
- return true;
- }
- }();
- if (!is_same_target) return std::nullopt;
-
- Stmt inlined = InlineArguments(gvar.value(), callee, call->args);
- return Mutate(inlined, InplaceMode::kDisallow).ValueOrUnchanged(inlined);
- }
-
- UnchangedOr<Expr> Mutate_(const CallNode* call, InplaceMode inplace_mode)
override {
- // Because the current implementation inlines a subroutine inserts
- // the `Stmt` body at the point of use, replacement must
- // occur in a context where a `Stmt` can be returned. Support
- // of subroutines that are called within an expression
- // (e.g. Replacing func in `Buf[0] = func(1) + func(2)`) would
- // require hoisting preprocessing done in the subroutine to the
- // parent `Stmt`.
- //
- // See `TestInlineCallOccurringInExpression` in
- // `test_tir_inline_private_functions.py` for a test of this
- // behavior, currently marked with `pytest.mark.xfail`.
- //
- // Any callee that hasn't been inlined at this point must be kept
- // in the output IRModule.
- if (auto gvar = call->op.as<GlobalVar>()) {
- removable_funcs_.erase(gvar.value());
- }
- return StmtExprMutator::Mutate_(call, inplace_mode);
- }
-
- Stmt InlineArguments(const GlobalVar& gvar, Function callee, const
ffi::Array<Expr>& args) const {
- TVM_FFI_ICHECK_EQ(callee->params.size(), args.size())
- << "Callee " << gvar << " accepts " << callee->params.size() << "
parameters ("
- << callee->params << "), but is called with " << args.size() << "
arguments (" << args
- << ")";
-
- for (const Var& param : callee->params) {
- TVM_FFI_ICHECK(!param->ty.as<TensorTypeNode>())
- << "Inlining of Functions with buffer arguments is not yet
supported, "
- << "but callee " << gvar << " has TensorType-annotated parameter "
<< param;
- }
-
- ffi::Map<Var, ffi::Variant<tirx::TensorVar, tvm::Expr>> param_map;
- for (size_t i = 0; i < callee->params.size(); i++) {
- param_map.Set(callee->params[i], args[i]);
- }
-
- callee = Specialize(callee, param_map);
-
- return callee->body.value();
- }
-
- // Map from GlobalVar to Functions which may be inlined.
- ffi::Map<GlobalVar, Function> inlinable_funcs_;
-
- /* \brief Set of callees that may be removed
- *
- * Some constructs may not be inlined (e.g. if the call site occurs
- * outside of an Evaluate node). For these cases, the output
- * IRModule must still contain the callee.
- */
- PSet<GlobalVar> removable_funcs_;
-
- ffi::Optional<Target> current_target_ = std::nullopt;
-};
-
-} // namespace
-
-Pass InlinePrivateFunctions() {
- auto pass_func = [](IRModule mod, PassContext ctx) {
- auto inlinable_functions = CollectInlinableFunctions(mod);
-
- if (inlinable_functions.empty()) {
- // Early bail-out if the module has no inlinable Functions.
- return mod;
- }
-
- auto mutator =
ffi::make_object<FunctionInliner>(std::move(inlinable_functions));
- IRModule updates;
-
- for (const auto& [gvar, base_func] : mod->functions) {
- if (auto opt = base_func.as<Function>()) {
- auto updated = mutator->VisitFunc(opt.value());
- if (!updated.same_as(base_func)) {
- updates->Add(gvar, updated);
- }
- }
- }
-
- if (updates->functions.size()) {
- auto write_ptr = mod.CopyOnWrite();
- write_ptr->Update(updates);
- for (const auto& gvar : mutator->GetRemovableFunctions()) {
- write_ptr->Remove(gvar);
- }
- mod = ConvertSSA()(mod);
- }
-
- return mod;
- };
- return tvm::transform::CreateModulePass(pass_func, 0,
"tirx.InlinePrivateFunctions");
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.transform.InlinePrivateFunctions",
InlinePrivateFunctions);
-}
-
-} // namespace transform
-
-} // namespace tirx
-} // namespace tvm
diff --git a/src/tirx/transform/lower_tirx.cc b/src/tirx/transform/lower_tirx.cc
index c351a934e3..b093d59861 100644
--- a/src/tirx/transform/lower_tirx.cc
+++ b/src/tirx/transform/lower_tirx.cc
@@ -34,7 +34,7 @@ namespace tirx {
namespace transform {
Pass LowerTIRx() {
- std::vector<tvm::transform::Pass> passes = {TilePrimitiveDispatch()};
+ std::vector<tvm::transform::Pass> passes = {TileDispatch()};
if (std::getenv("TVM_PRINT_AFTER_TIRX_DISPATCH_OPS")) {
passes.push_back(tvm::transform::PrintIR());
}
@@ -45,7 +45,7 @@ Pass LowerTIRx() {
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
- .def("tirx.transform.TilePrimitiveDispatch", TilePrimitiveDispatch)
+ .def("tirx.transform.TileDispatch", TileDispatch)
.def("tirx.transform.LowerTIRxCleanup", LowerTIRxCleanup)
.def("tirx.transform.LowerTIRx", LowerTIRx);
}
diff --git a/src/tirx/transform/remap_thread_axis.cc
b/src/tirx/transform/remap_thread_axis.cc
deleted file mode 100644
index 2d7e21fd20..0000000000
--- a/src/tirx/transform/remap_thread_axis.cc
+++ /dev/null
@@ -1,113 +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.
- */
-
-/*!
- * \file remap_thread_axis.cc
- */
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/expr.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/tirx/op/region.h>
-#include <tvm/tirx/stmt_functor.h>
-#include <tvm/tirx/transform.h>
-
-#include <unordered_map>
-
-namespace tvm {
-namespace tirx {
-
-// Mutator to change the read pattern
-class ThreadAxisRewriter : public StmtExprMutator {
- public:
- using StmtExprMutator::Mutate;
- using StmtExprMutator::Mutate_;
- explicit ThreadAxisRewriter(const std::unordered_map<std::string,
ffi::String>& tmap)
- : tmap_(tmap) {}
-
- Stmt Rewrite(Stmt stmt) { return Mutate(stmt,
InplaceMode::kAllow).ValueOrUnchanged(stmt); }
-
- private:
- UnchangedOr<Stmt> Mutate_(const RegionStmtNode* op, InplaceMode
inplace_mode) final {
- if (op->op.same_as(tirx::launch_thread_op()) &&
-
std::string(op->args[0].as_or_throw<StringImm>()->value).rfind("vthread", 0) !=
0) {
- auto it = tmap_.find(op->args[0].as_or_throw<StringImm>()->value);
- if (it != tmap_.end()) {
- PrimExpr old_extent = op->args[1].as_or_throw<PrimExpr>();
- PrimExpr extent = Mutate(old_extent,
inplace_mode).ValueOrUnchanged(old_extent);
- PrimVar old_var = op->body_params[0].as_or_throw<PrimVar>();
- PrimVar new_var(it->second, extent.ty());
- ffi::Any previous_remap = VarRemapGet(old_var);
- VarRemapSet(old_var, new_var);
- Stmt body = Mutate(op->body, inplace_mode).ValueOrUnchanged(op->body);
- VarRemapSet(old_var, previous_remap);
- return RegionStmt(tirx::launch_thread_op(), {StringImm(it->second),
extent}, {new_var},
- DictAttrs(), body, {}, op->span);
- }
- }
- return StmtExprMutator::Mutate_(op, inplace_mode);
- }
-
- // The thread map
- const std::unordered_map<std::string, ffi::String>& tmap_;
-};
-
-Function RemapThreadAxis(Function func, ffi::Map<ffi::String, ffi::String>
thread_map) {
- std::unordered_map<std::string, ffi::String> tmap;
- for (const auto& kv : thread_map) {
- tmap[kv.first] = kv.second;
- }
-
- if (auto opt =
func->GetAttr<ffi::Array<ffi::String>>(tirx::attr::kKernelLaunchParams)) {
- TVM_FFI_ICHECK(opt != nullptr) << "Require attribute " <<
tirx::attr::kKernelLaunchParams;
- auto launch_params = opt.value();
- // replace the thread axis attribute
- for (size_t i = 0; i < launch_params.size(); ++i) {
- auto it = tmap.find(launch_params[i]);
- if (it != tmap.end()) {
- launch_params.Set(i, it->second);
- }
- }
-
- func = WithAttr(std::move(func), tirx::attr::kKernelLaunchParams,
launch_params);
- }
-
- if (!func->body.has_value()) return func;
- auto* n = func.CopyOnWrite();
- n->body =
ffi::make_object<ThreadAxisRewriter>(tmap)->Rewrite(std::move(n->body).value());
- return func;
-}
-
-namespace transform {
-
-Pass RemapThreadAxis(ffi::Map<ffi::String, ffi::String> thread_map) {
- auto pass_func = [thread_map](Function f, IRModule m, PassContext ctx) {
- return RemapThreadAxis(std::move(f), thread_map);
- };
- return CreateFunctionPass(pass_func, 0, "tirx.RemapThreadAxis");
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tirx.transform.RemapThreadAxis", RemapThreadAxis);
-}
-
-} // namespace transform
-} // namespace tirx
-} // namespace tvm
diff --git a/src/tirx/transform/tile_primitive_dispatch.cc
b/src/tirx/transform/tile_dispatch.cc
similarity index 99%
rename from src/tirx/transform/tile_primitive_dispatch.cc
rename to src/tirx/transform/tile_dispatch.cc
index 75db630df4..ac470c7da9 100644
--- a/src/tirx/transform/tile_primitive_dispatch.cc
+++ b/src/tirx/transform/tile_dispatch.cc
@@ -18,7 +18,7 @@
*/
/*!
- * \file tile_primitive_dispatch.cc
+ * \file tile_dispatch.cc
* \brief Lower tensor instruction Evaluate(Call) nodes (also resolves
ScopeIdDef
* declarations and emits launch params).
*/
@@ -214,14 +214,14 @@ class NoOpCallVerifier : public
Verifier<NoOpCallVerifier> {
}
};
-class TilePrimitiveDispatcher : public StmtExprMutator {
+class TileDispatcher : public StmtExprMutator {
public:
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
- explicit TilePrimitiveDispatcher(const Target& target) : target_(target) {}
+ explicit TileDispatcher(const Target& target) : target_(target) {}
static Stmt LowerOpCalls(const Stmt& stmt, const Target& target) {
- return ffi::make_object<TilePrimitiveDispatcher>(target)
+ return ffi::make_object<TileDispatcher>(target)
->Mutate(stmt, InplaceMode::kAllow)
.ValueOrUnchanged(stmt);
}
@@ -1546,18 +1546,18 @@ Target ResolveTarget(const Function& f) {
namespace transform {
-Pass TilePrimitiveDispatch() {
+Pass TileDispatch() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
if (!f->body.has_value()) return f;
Target target = ResolveTarget(f);
auto* n = f.CopyOnWrite();
- n->body = TilePrimitiveDispatcher::LowerOpCalls(n->body.value(), target);
+ n->body = TileDispatcher::LowerOpCalls(n->body.value(), target);
if (!NoOpCallVerifier::Verify(n->body.value(), false)) {
LOG(FATAL) << "Failed to lower the TIRx program: " << f;
}
return f;
};
- return CreateFunctionPass(pass_func, 0, "tirx.TilePrimitiveDispatch");
+ return CreateFunctionPass(pass_func, 0, "tirx.TileDispatch");
}
} // namespace transform
diff --git a/tests/python/tirx-transform/test_tir_inline_private_functions.py
b/tests/python/tirx-transform/test_tir_inline_private_functions.py
deleted file mode 100644
index d87fcc1c7f..0000000000
--- a/tests/python/tirx-transform/test_tir_inline_private_functions.py
+++ /dev/null
@@ -1,251 +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.
-# ruff: noqa: F821
-
-import pytest
-
-import tvm.testing
-from tvm.script import ir as I
-from tvm.script import tirx as T
-
-
-class BaseTestCase:
- def test_well_formed(self):
- After = tvm.tirx.transform.InlinePrivateFunctions()(self.Before)
- tvm.tirx.analysis.verify_well_formed(After)
-
- def test_produces_expected(self):
- After = tvm.tirx.transform.InlinePrivateFunctions()(self.Before)
- tvm.ir.assert_structural_equal(self.Expected, After)
-
-
-class TestSimple(BaseTestCase):
- """Simple case directly acting on Function"""
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor([80, 16], "float32"), B: T.Tensor([64, 16],
"float32")):
- for i in range(64):
- Before.subroutine(T.address_of(A[i, 0]), T.address_of(B[i, 0]))
-
- @T.function(private=True)
- def subroutine(A_data: T.handle("float32"), B_data:
T.handle("float32")):
- A = T.decl_tensor([16, 16], "float32", data=A_data)
- B = T.decl_tensor([16], "float32", data=B_data)
- for i in range(16):
- B[i] = 0.0
- for j in range(16):
- B[i] = B[i] + A[i, j]
-
- @I.ir_module
- class Expected:
- @T.function
- def main(A: T.Tensor([80, 16], "float32"), B: T.Tensor([64, 16],
"float32")):
- for i in range(64):
- Aview = T.decl_tensor([16, 16], "float32",
data=T.address_of(A[i, 0]))
- Bview = T.decl_tensor([16], "float32", data=T.address_of(B[i,
0]))
- for j in range(16):
- Bview[j] = 0.0
- for k in range(16):
- Bview[j] = Bview[j] + Aview[j, k]
-
-
-class TestRetainCrossFunctionSubroutines(BaseTestCase):
- """Do not inline functions that cross device boundaries
-
- When lowering TIR, calls for which the callsite and callee have
- different targets are used at some stages, before being further
- lowered to explicit device kernel launches. Since inlining the
- function would remove this cross-device information,
- InlinePrivateSubroutines should not inline these cases.
- """
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor([80, 16], "float32"), B: T.Tensor([64, 16],
"float32")):
- T.func_attr({"target": T.target("llvm")})
- for i in range(64):
- Before.subroutine(T.address_of(A[i, 0]), T.address_of(B[i, 0]))
-
- @T.function(private=True)
- def subroutine(A_data: T.handle("float32"), B_data:
T.handle("float32")):
- T.func_attr({"target": T.target({"kind": "cuda", "arch":
"sm_80"})})
- A = T.decl_tensor([16, 16], "float32", data=A_data)
- B = T.decl_tensor([16], "float32", data=B_data)
- for i in range(16):
- B[i] = 0.0
- for j in range(16):
- B[i] = B[i] + A[i, j]
-
- Expected = Before
-
-
-class TestRetainRecursiveSubroutines(BaseTestCase):
- """Do not inline recursive functions
-
- To avoid potentially infinite loops at compile-time, disable
- inlining of recursive functions. If inlining of these functions
- would be useful, this restriction may be relaxed with improved
- analysis of the subroutine.
- """
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor(16, "float32")):
- Before.subroutine(T.address_of(A[0]), 16)
-
- @T.function(private=True)
- def subroutine(A_data: T.handle("float32"), A_size: T.int32):
- A = T.decl_tensor(A_size, "float32", data=A_data)
- A[1] = A[0] + A[1]
-
- if A_size > 1:
- Before.subroutine(T.address_of(A[1]), A_size - 1)
-
- Expected = Before
-
-
-class TestDeduplicateBlockName(BaseTestCase):
- """Block names must be de-duplicated after inlining"""
-
- @pytest.mark.xfail(reason="Inlining of schedulable TIR not yet supported")
- def test_produces_expected(self):
- super().test_produces_expected(self)
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor([2, 16], "float32"), B: T.Tensor([2, 16],
"float32")):
- Before.subroutine(T.address_of(A[0, 0]), T.address_of(B[0, 0]))
- Before.subroutine(T.address_of(A[1, 0]), T.address_of(B[1, 0]))
-
- @T.function(private=True)
- def subroutine(A_data: T.handle("float32"), B_data:
T.handle("float32")):
- A = T.decl_tensor(16, "float32", data=A_data)
- B = T.decl_tensor(16, "float32", data=B_data)
- for i in range(16):
- B[i] = A[i] * 2.0
-
- @I.ir_module
- class Expected:
- @T.function
- def main(A: T.Tensor([80, 16], "float32"), B: T.Tensor([64, 16],
"float32")):
- A_data_1 = T.bind(T.address_of(A[0, 0]), T.handle("float32"))
- A_1 = T.decl_tensor(16, "float32", data=A_data_1)
- B_data_1: T.let[T.handle("float32")] = T.address_of(B[0, 0])
- B_1 = T.decl_tensor(16, "float32", data=B_data_1)
- for i in range(16):
- B_1[i] = A_1[i] * 2.0
-
- A_data_2 = T.bind(T.address_of(A[1, 0]), T.handle("float32"))
- A_2 = T.decl_tensor(16, "float32", data=A_data_2)
- B_data_2: T.let[T.handle("float32")] = T.address_of(B[1, 0])
- B_2 = T.decl_tensor(16, "float32", data=B_data_2)
- for i in range(16):
- B_2[i] = A_2[i] * 2.0
-
-
-class TestInlineCallOccurringInExpression(BaseTestCase):
- """Inline a Call node that is used in a function
-
- The current implementation only replaces `ir.Call` instances that
- occur in a `tvm.ir.Evaluate` context. This is the primary use case,
- used in destination-passing style.
-
- This unit test is marked as xfail. If/when the implementation
- supports inlining of function calls occurring as part of an
- expression, the annotation can be removed.
- """
-
- @pytest.mark.xfail(
- reason="Inlining of Functions outside of tvm.ir.Evaluate is not yet
supported"
- )
- def test_produces_expected(self):
- super().test_produces_expected(self)
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor(16, "float32")):
- for i in range(16):
- A[i] = Before.subroutine(i)
-
- @T.function(private=True)
- def subroutine(i: T.int32) -> T.float32:
- cos = T.cos(T.cast(i, "float32"))
- sin = T.sin(T.cast(i, "float32"))
- retval = cos * cos + sin * sin
- return retval
-
- @I.ir_module
- class Expected:
- @T.function
- def main(A: T.Tensor(16, "float32")):
- for i in range(16):
- cos = T.cos(T.cast(i, "float32"))
- sin = T.sin(T.cast(i, "float32"))
- retval = cos * cos + sin * sin
- A[i] = retval
-
-
-class TestInlineFunctionWithBufferArguments(BaseTestCase):
- """Inline a function that accepts buffer arguments
-
- The current implementation does not support this usage. This unit
- test is provided to display a possible user interaction, and is
- marked with `@pytest.mark.xfail`. If/when the implementation
- supports inlining of function calls with buffer arguments, the
- annotation can be removed.
- """
-
- @pytest.mark.xfail(reason="Inlining of Functions with buffer arguments")
- def test_produces_expected(self):
- super().test_produces_expected(self)
-
- @I.ir_module
- class Before:
- @T.function
- def main(A: T.Tensor(16, "float32")):
- Before.subroutine(
- T.stack_make_dltensor(
- A.data,
- T.stack_make_shape(*A.ty.shape),
- 0,
- len(A.ty.shape),
- 0.0,
- A.ty.elem_offset,
- )
- )
-
- @T.function(private=True)
- def subroutine(A: T.Tensor(16, "float32")):
- for i in range(16):
- A[i] = A[i] * 2.0
-
- @I.ir_module
- class Expected:
- @T.function
- def main(A: T.Tensor(16, "float32")):
- for i in range(16):
- A[i] = A[i] * 2.0
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
index b2253269f9..8ebe0fed1a 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
@@ -462,7 +462,7 @@ def test_reg_copy_linear_shared_hoists_thread_base():
@tvm.transform.pass_instrument
class CaptureDispatch:
def run_after_pass(self, mod, info):
- if info.name == "tirx.TilePrimitiveDispatch":
+ if info.name == "tirx.TileDispatch":
dispatched.append(mod["main"])
target = tvm.target.Target("cuda")
