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 b5113ec5f4 [REFACTOR][TIRx] Keep device-scope decoration in TIRx
(#20691)
b5113ec5f4 is described below
commit b5113ec5f4bfeb18b37fbf0383a508b0bb9cd7f4
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 20:38:46 2026 -0400
[REFACTOR][TIRx] Keep device-scope decoration in TIRx (#20691)
Keep device-scope decoration local to its public TIRx pass. Remove the
redundant S-TIR entrypoint and shared implementation header, and move
the existing device-scope test to TIRx.
---
include/tvm/s_tir/transform.h | 6 ---
python/tvm/s_tir/transform/transform.py | 11 -----
src/s_tir/transform/decorate_device_scope.cc | 50 ----------------------
src/tirx/transform/decorate_device_scope.cc | 16 +++----
src/tirx/transform/decorate_device_scope.h | 34 ---------------
.../test_transform_decorate_device_scope.py} | 2 +-
6 files changed, 6 insertions(+), 113 deletions(-)
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 1e5187a0e3..a1ce06e6e5 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -298,12 +298,6 @@ TVM_DLL Pass RemoveWeightLayoutRewriteBlock(bool
skip_tensor_rewrite = false);
*/
TVM_DLL Pass RemoveStoreUndef();
-/*!
- * \brief Decorate all the function's body as device function.
- * \return The pass.
- */
-TVM_DLL Pass DecorateDeviceScope();
-
/*!
* \brief Eliminate branches by leveraging tensor assumptions (T.assume).
* \return The pass.
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index 3be6850e6a..9130ebcee2 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -401,17 +401,6 @@ def RemoveStoreUndef():
return _ffi_api.RemoveStoreUndef() # type: ignore
-def DecorateDeviceScope():
- """Decorate all the function's body as device function.
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
- """
- return _ffi_api.DecorateDeviceScope() # type: ignore
-
-
def UseAssumeToReduceBranches():
"""Eliminate layout specific pad branch by overcomputing values for padded
region.
diff --git a/src/s_tir/transform/decorate_device_scope.cc
b/src/s_tir/transform/decorate_device_scope.cc
deleted file mode 100644
index c8333539f0..0000000000
--- a/src/s_tir/transform/decorate_device_scope.cc
+++ /dev/null
@@ -1,50 +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 decorate_device_scope.cc
- */
-#include "../../tirx/transform/decorate_device_scope.h"
-
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/prim/op.h>
-#include <tvm/s_tir/transform.h>
-
-namespace tvm {
-namespace s_tir {
-using namespace tvm::tirx;
-
-namespace transform {
-
-Pass DecorateDeviceScope() {
- auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- return tirx::DecorateDeviceScopeImpl(std::move(f));
- };
- return CreateFunctionPass(pass_func, 0, "s_tir.DecorateDeviceScope");
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.DecorateDeviceScope",
DecorateDeviceScope);
-}
-
-} // namespace transform
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/tirx/transform/decorate_device_scope.cc
b/src/tirx/transform/decorate_device_scope.cc
index ce83759602..2cbad33c0d 100644
--- a/src/tirx/transform/decorate_device_scope.cc
+++ b/src/tirx/transform/decorate_device_scope.cc
@@ -20,8 +20,6 @@
/*!
* \file decorate_device_scope.cc
*/
-#include "decorate_device_scope.h"
-
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/op.h>
@@ -31,19 +29,15 @@
namespace tvm {
namespace tirx {
-Function DecorateDeviceScopeImpl(Function f) {
- if (!f->body.has_value()) return f;
- static const Op device_scope = Op::Get("tirx.device_scope");
- auto* n = f.CopyOnWrite();
- n->body = RegionStmt(device_scope, {}, {}, DictAttrs(),
std::move(n->body).value());
- return f;
-}
-
namespace transform {
Pass DecorateDeviceScope() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- return DecorateDeviceScopeImpl(std::move(f));
+ if (!f->body.has_value()) return f;
+ static const Op device_scope = Op::Get("tirx.device_scope");
+ auto* n = f.CopyOnWrite();
+ n->body = RegionStmt(device_scope, {}, {}, DictAttrs(),
std::move(n->body).value());
+ return f;
};
return CreateFunctionPass(pass_func, 0, "tirx.DecorateDeviceScope");
}
diff --git a/src/tirx/transform/decorate_device_scope.h
b/src/tirx/transform/decorate_device_scope.h
deleted file mode 100644
index f1590ee669..0000000000
--- a/src/tirx/transform/decorate_device_scope.h
+++ /dev/null
@@ -1,34 +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 decorate_device_scope.h
- * \brief Shared implementation of native device-scope decoration.
- */
-#ifndef TVM_TIRX_TRANSFORM_DECORATE_DEVICE_SCOPE_H_
-#define TVM_TIRX_TRANSFORM_DECORATE_DEVICE_SCOPE_H_
-
-#include <tvm/tirx/function.h>
-
-namespace tvm {
-namespace tirx {
-Function DecorateDeviceScopeImpl(Function f);
-} // namespace tirx
-} // namespace tvm
-
-#endif // TVM_TIRX_TRANSFORM_DECORATE_DEVICE_SCOPE_H_
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_decorate_device_scope.py
b/tests/python/tirx/transform/test_transform_decorate_device_scope.py
similarity index 93%
rename from
tests/python/s_tir/transform/test_s_tir_transform_decorate_device_scope.py
rename to tests/python/tirx/transform/test_transform_decorate_device_scope.py
index f83340f85f..9e794cae5d 100644
--- a/tests/python/s_tir/transform/test_s_tir_transform_decorate_device_scope.py
+++ b/tests/python/tirx/transform/test_transform_decorate_device_scope.py
@@ -21,7 +21,7 @@ def test_decorate_device():
x = tvm.tirx.Var("x", "int32")
mod = tvm.IRModule.from_expr(tvm.tirx.Function([x], tvm.ir.Evaluate(x)))
- stmt = tvm.s_tir.transform.DecorateDeviceScope()(mod)["main"].body[0]
+ stmt = tvm.tirx.transform.DecorateDeviceScope()(mod)["main"].body[0]
assert stmt.op.same_as(tvm.ir.Op.get("tirx.device_scope"))
