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 f7b0fe62a9 [REFACTOR][S-TIR] Remove test-only force index narrowing
(#20685)
f7b0fe62a9 is described below
commit f7b0fe62a9b2685b002efdc4e69f0960b3e30ffa
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 19:11:45 2026 -0400
[REFACTOR][S-TIR] Remove test-only force index narrowing (#20685)
Remove the test-only S-TIR ForceNarrowIndexToInt32 API and keep
sort/scan integration on the i64 index path. Localize the remaining
narrowing implementation in TIRx while preserving its public entrypoint
and WebGPU use.
---
include/tvm/s_tir/transform.h | 13 --
include/tvm/tirx/transform.h | 3 +-
python/tvm/s_tir/transform/transform.py | 19 ---
python/tvm/tirx/transform/transform.py | 3 +-
src/s_tir/transform/force_narrow_index_to_i32.cc | 78 ---------
src/tirx/transform/force_narrow_index_to_i32.cc | 70 +++++++-
src/tirx/transform/force_narrow_index_to_i32.h | 113 -------------
.../relax/test_backend_dispatch_sort_scan.py | 22 +--
...st_s_tir_transform_force_narrow_index_to_i32.py | 177 ---------------------
9 files changed, 72 insertions(+), 426 deletions(-)
diff --git a/include/tvm/s_tir/transform.h b/include/tvm/s_tir/transform.h
index 1513e12cab..1e5187a0e3 100644
--- a/include/tvm/s_tir/transform.h
+++ b/include/tvm/s_tir/transform.h
@@ -310,19 +310,6 @@ TVM_DLL Pass DecorateDeviceScope();
*/
TVM_DLL Pass UseAssumeToReduceBranches();
-/*!
- * \brief Force to narrow down indexing expressions and integer tensors to
int32 dtype in
- * functions that may still contain S-TIR blocks.
- *
- * Unlike tirx::transform::ForceNarrowIndexToInt32, this pass also rewrites
block iterators,
- * block access regions, and match tensor regions, so it can run on scheduled
functions before
- * block lowering.
- *
- * \return The pass.
- * \note This pass should not be used in default cases.
- */
-TVM_DLL Pass ForceNarrowIndexToInt32();
-
} // namespace transform
} // namespace s_tir
} // namespace tvm
diff --git a/include/tvm/tirx/transform.h b/include/tvm/tirx/transform.h
index aaf1b1a31a..3b6ede5942 100644
--- a/include/tvm/tirx/transform.h
+++ b/include/tvm/tirx/transform.h
@@ -228,8 +228,7 @@ TVM_TIRX_DLL Pass NarrowDataType(int target_bits);
/*!
* \brief Force to narrow down indexing expressions and integer tensors to
int32 dtype.
*
- * The function must not contain S-TIR blocks. Use
s_tir::transform::ForceNarrowIndexToInt32
- * before block lowering.
+ * The function must not contain S-TIR blocks.
*
* \return The pass.
* \note This pass should not be used in default cases.
diff --git a/python/tvm/s_tir/transform/transform.py
b/python/tvm/s_tir/transform/transform.py
index ec6c05e5dd..3be6850e6a 100644
--- a/python/tvm/s_tir/transform/transform.py
+++ b/python/tvm/s_tir/transform/transform.py
@@ -421,22 +421,3 @@ def UseAssumeToReduceBranches():
The result pass
"""
return _ffi_api.UseAssumeToReduceBranches() # type: ignore
-
-
-def ForceNarrowIndexToInt32():
- """Force narrow down indexing expressions and integer tensors to int32
dtype.
-
- Unlike :py:func:`tvm.tirx.transform.ForceNarrowIndexToInt32`, this pass
also rewrites block
- iterators, block access regions, and match tensor regions, so it can run
on scheduled
- functions before block lowering.
-
- Returns
- -------
- fpass : tvm.transform.Pass
- The result pass
-
- Note
- ----
- This pass should not be used in default cases.
- """
- return _ffi_api.ForceNarrowIndexToInt32() # type: ignore
diff --git a/python/tvm/tirx/transform/transform.py
b/python/tvm/tirx/transform/transform.py
index 54af2ec41d..892d945bee 100644
--- a/python/tvm/tirx/transform/transform.py
+++ b/python/tvm/tirx/transform/transform.py
@@ -388,8 +388,7 @@ def NarrowDataType(target_bits: int):
def ForceNarrowIndexToInt32():
"""Force narrow down indexing expressions and integer tensors to int32
dtype.
- The function must not contain S-TIR blocks. Use
- :py:func:`tvm.s_tir.transform.ForceNarrowIndexToInt32` before block
lowering.
+ The function must not contain S-TIR blocks.
Returns
-------
diff --git a/src/s_tir/transform/force_narrow_index_to_i32.cc
b/src/s_tir/transform/force_narrow_index_to_i32.cc
deleted file mode 100644
index e234b62662..0000000000
--- a/src/s_tir/transform/force_narrow_index_to_i32.cc
+++ /dev/null
@@ -1,78 +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 force_narrow_index_to_i32.cc
- * \brief Force narrow down indexing expressions and integer tensors to int32
dtype in functions
- * that still contain S-TIR blocks.
- * \note This pass is not used in default cases.
- */
-
-#include "../../tirx/transform/force_narrow_index_to_i32.h"
-
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/s_tir/transform.h>
-#include <tvm/tirx/transform.h>
-
-#include "../ir/data_type_rewriter.h"
-
-namespace tvm {
-namespace s_tir {
-using namespace tvm::tirx;
-
-class Int32DTypeNarrower : public
Int32DTypeNarrowerBase<IndexDataTypeNormalizer> {
- public:
- using Int32DTypeNarrowerBase::Mutate;
- using Int32DTypeNarrowerBase::Mutate_;
- static Function RewriteDataType(Function func) {
- CheckTensorParams(func);
- auto narrower = ffi::make_object<Int32DTypeNarrower>(func);
- return narrower->Rewrite(func);
- }
-
- explicit Int32DTypeNarrower(Function func) :
Int32DTypeNarrowerBase(std::move(func)) {}
-
- private:
- UnchangedOr<Stmt> Mutate_(const SBlockNode* op, InplaceMode inplace_mode)
final {
- auto result = IndexDataTypeNormalizer::Mutate_(op, inplace_mode);
- auto block =
std::move(result).ValueOrUnchanged(ffi::GetRef<Stmt>(op)).as_or_throw<SBlock>();
- for (const TensorVar& tensor : block->alloc_tensors) {
- CheckAllocatedTensor(tensor);
- }
- return block;
- }
-};
-
-namespace transform {
-
-Pass ForceNarrowIndexToInt32() {
- auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- return Int32DTypeNarrower::RewriteDataType(std::move(f));
- };
- return CreateFunctionPass(pass_func, 0, "s_tir.ForceNarrowIndexToInt32");
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("s_tir.transform.ForceNarrowIndexToInt32",
ForceNarrowIndexToInt32);
-}
-
-} // namespace transform
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/tirx/transform/force_narrow_index_to_i32.cc
b/src/tirx/transform/force_narrow_index_to_i32.cc
index 4e44c07862..daf7cd152e 100644
--- a/src/tirx/transform/force_narrow_index_to_i32.cc
+++ b/src/tirx/transform/force_narrow_index_to_i32.cc
@@ -23,25 +23,87 @@
* \note This pass is not used in default cases.
*/
-#include "force_narrow_index_to_i32.h"
-
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/prim/op.h>
+#include <tvm/tirx/op/memory.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
+#include <utility>
+
+#include "../ir/data_type_rewriter.h"
+
namespace tvm {
namespace tirx {
-class Int32DTypeNarrower : public
Int32DTypeNarrowerBase<IndexDataTypeNormalizer> {
+class Int32DTypeNarrower : public IndexDataTypeNormalizer {
public:
+ using IndexDataTypeNormalizer::Mutate;
+ using IndexDataTypeNormalizer::Mutate_;
+
static Function RewriteDataType(Function func) {
CheckTensorParams(func);
auto narrower = ffi::make_object<Int32DTypeNarrower>(func);
return narrower->Rewrite(func);
}
- explicit Int32DTypeNarrower(Function func) :
Int32DTypeNarrowerBase(std::move(func)) {}
+ explicit Int32DTypeNarrower(Function func)
+ : IndexDataTypeNormalizer(PrimType::Int(32)), func_(std::move(func)) {}
+
+ private:
+ /*! \brief Reject integer tensor parameters wider than int32. */
+ static void CheckTensorParams(const Function& func) {
+ for (const Var& param : func->params) {
+ if (auto tensor = param.as<TensorVar>();
+ tensor && tensor.value()->dtype.MatchesCode(DLDataTypeCode::kDLInt)
&&
+ tensor.value()->dtype.bits() > 32) {
+ TVM_FFI_THROW(InternalError) << "The tensor parameter " <<
tensor.value() << " has dtype "
+ << tensor.value()->dtype << ". The
function is " << func;
+ }
+ }
+ }
+
+ /*! \brief Reject allocated integer tensors wider than int32. */
+ void CheckAllocatedTensor(const TensorVar& buf) const {
+ // Scalar assignments in TVMScript use local scalar storage. Keep its
explicit
+ // dtype (e.g. an int64 opaque call result) and cast at narrowed index
uses.
+ // IsScalar checks the scalar layout contract, not merely the allocation
size.
+ bool is_local_scalar = buf.scope() == "local" && buf.IsScalar();
+ if (!is_local_scalar && buf->dtype.MatchesCode(DLDataTypeCode::kDLInt) &&
+ buf->dtype.bits() > 32) {
+ TVM_FFI_THROW(InternalError)
+ << "The tensor " << buf << " allocated in the function has dtype "
<< buf->dtype
+ << ". The function is " << func_;
+ }
+ }
+
+ bool ShouldClampShiftAmounts() const final { return true; }
+
+ UnchangedOr<PrimExpr> Mutate_(const IntImmNode* op, InplaceMode
inplace_mode) final {
+ // ignore the enabled condition and always rewrite i64
+ if (op->ty.as_or_throw<PrimType>() == PrimType::Int(64)) {
+ TVM_FFI_ICHECK_LE(op->value,
prim::max_value(target_data_type_).as_or_throw<IntImm>()->value);
+ return IntImm::Int32(op->value);
+ }
+ return ffi::Unchanged();
+ }
+
+ UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
+ if (const auto* call = op->value.as<CallNode>();
+ call && call->op.same_as(tirx::alloc_tensor_op()))
+ return MutateAllocTensor(op, inplace_mode);
+ return IndexDataTypeNormalizer::Mutate_(op, inplace_mode);
+ }
+
+ UnchangedOr<Stmt> MutateAllocTensor(const BindNode* op, InplaceMode
inplace_mode) {
+ auto result = IndexDataTypeNormalizer::Mutate_(op, inplace_mode);
+ auto alloc =
std::move(result).ValueOrUnchanged(ffi::GetRef<Stmt>(op)).as_or_throw<Bind>();
+ CheckAllocatedTensor(alloc->var.as_or_throw<TensorVar>());
+ return alloc;
+ }
+
+ Function func_;
};
Function ForceNarrowIndexToInt32(Function func) {
diff --git a/src/tirx/transform/force_narrow_index_to_i32.h
b/src/tirx/transform/force_narrow_index_to_i32.h
deleted file mode 100644
index d04cb781dd..0000000000
--- a/src/tirx/transform/force_narrow_index_to_i32.h
+++ /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 force_narrow_index_to_i32.h
- * \brief Narrowing rules shared by the TIRX and S-TIR ForceNarrowIndexToInt32
passes.
- */
-#ifndef TVM_TIR_TRANSFORM_FORCE_NARROW_INDEX_TO_I32_H_
-#define TVM_TIR_TRANSFORM_FORCE_NARROW_INDEX_TO_I32_H_
-
-#include <tvm/ir/prim/op.h>
-#include <tvm/tirx/op/memory.h>
-
-#include <utility>
-
-#include "../ir/data_type_rewriter.h"
-
-namespace tvm {
-namespace tirx {
-
-/*!
- * \brief Force index expressions and integer tensors to int32.
- * \tparam Normalizer The index normalizer that determines which statements
are traversed:
- * tirx::IndexDataTypeNormalizer for lowered TIR,
s_tir::IndexDataTypeNormalizer for
- * TIR that still contains S-TIR blocks.
- */
-template <typename Normalizer>
-class Int32DTypeNarrowerBase : public Normalizer {
- public:
- using Normalizer::Mutate;
- using Normalizer::Mutate_;
-
- protected:
- explicit Int32DTypeNarrowerBase(Function func)
- : Normalizer(PrimType::Int(32)), func_(std::move(func)) {}
-
- /*! \brief Reject integer tensor parameters wider than int32. */
- static void CheckTensorParams(const Function& func) {
- for (const Var& param : func->params) {
- if (auto tensor = param.as<TensorVar>();
- tensor && tensor.value()->dtype.MatchesCode(DLDataTypeCode::kDLInt)
&&
- tensor.value()->dtype.bits() > 32) {
- TVM_FFI_THROW(InternalError) << "The tensor parameter " <<
tensor.value() << " has dtype "
- << tensor.value()->dtype << ". The
function is " << func;
- }
- }
- }
-
- /*! \brief Reject allocated integer tensors wider than int32. */
- void CheckAllocatedTensor(const TensorVar& buf) const {
- // Scalar assignments in TVMScript use local scalar storage. Keep its
explicit
- // dtype (e.g. an int64 opaque call result) and cast at narrowed index
uses.
- // IsScalar checks the scalar layout contract, not merely the allocation
size.
- bool is_local_scalar = buf.scope() == "local" && buf.IsScalar();
- if (!is_local_scalar && buf->dtype.MatchesCode(DLDataTypeCode::kDLInt) &&
- buf->dtype.bits() > 32) {
- TVM_FFI_THROW(InternalError)
- << "The tensor " << buf << " allocated in the function has dtype "
<< buf->dtype
- << ". The function is " << func_;
- }
- }
-
- bool ShouldClampShiftAmounts() const final { return true; }
-
- UnchangedOr<PrimExpr> Mutate_(const IntImmNode* op, InplaceMode
inplace_mode) final {
- // ignore the enabled condition and always rewrite i64
- if (op->ty.as_or_throw<PrimType>() == PrimType::Int(64)) {
- TVM_FFI_ICHECK_LE(
- op->value,
- prim::max_value(this->target_data_type_).template
as_or_throw<IntImm>()->value);
- return IntImm::Int32(op->value);
- }
- return ffi::Unchanged();
- }
-
- UnchangedOr<Stmt> Mutate_(const BindNode* op, InplaceMode inplace_mode)
final {
- if (const auto* call = op->value.as<CallNode>();
- call && call->op.same_as(tirx::alloc_tensor_op()))
- return MutateAllocTensor(op, inplace_mode);
- return Normalizer::Mutate_(op, inplace_mode);
- }
-
- UnchangedOr<Stmt> MutateAllocTensor(const BindNode* op, InplaceMode
inplace_mode) {
- auto result = Normalizer::Mutate_(op, inplace_mode);
- auto alloc =
- std::move(result).ValueOrUnchanged(ffi::GetRef<Stmt>(op)).template
as_or_throw<Bind>();
- CheckAllocatedTensor(alloc->var.template as_or_throw<TensorVar>());
- return alloc;
- }
-
- Function func_;
-};
-
-} // namespace tirx
-} // namespace tvm
-
-#endif // TVM_TIR_TRANSFORM_FORCE_NARROW_INDEX_TO_I32_H_
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
index 1a71a3c1c2..809dbe25e1 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -521,7 +521,7 @@ def test_dispatch_sort_webgpu_merge_grid(op, size):
pytest.param({"kind": "vulkan", "supports_int64": True},
marks=pytest.mark.gpu),
],
)
[email protected]("index_bits", [None, 32, 64])
[email protected]("index_bits", [None, 64])
def test_dispatch_cumsum_gpu(target, index_bits):
"""Test cumsum kernel dispatch and numerical correctness"""
if not tvm.testing.device_enabled(target):
@@ -544,8 +544,6 @@ def test_dispatch_cumsum_gpu(target, index_bits):
np_cumsum = np.cumsum(np_data, axis=-1)
with tvm.target.Target(target):
mod = DispatchSortScan(index_bits=index_bits)(Module)
- if index_bits == 32:
- mod = tvm.s_tir.transform.ForceNarrowIndexToInt32()(mod)
ex = tvm.compile(mod, target)
def run_and_check():
@@ -558,8 +556,8 @@ def test_dispatch_cumsum_gpu(target, index_bits):
tvm.testing.run_with_gpu_lock(run_and_check)
[email protected]("target_kind", ["metal", "webgpu", "cuda"])
[email protected]("index_bits", [None, 32, 64])
[email protected]("target_kind", ["metal", "cuda"])
[email protected]("index_bits", [None, 64])
def test_dispatch_cumsum_index_width(target_kind, index_bits):
"""Respect the caller's index budget without restricting Metal's
default."""
from tvm.relax.backend.gpu_generic import gpu_2d_continuous_cumsum
@@ -575,22 +573,10 @@ def test_dispatch_cumsum_index_width(target_kind,
index_bits):
return gv
with tvm.target.Target(target_kind, host="llvm"):
- if target_kind == "webgpu" and index_bits == 64:
- with pytest.raises(ValueError, match="WebGPU scan kernels require
index_bits=32"):
- DispatchSortScan(index_bits=index_bits)(Module)
- return
mod = DispatchSortScan(index_bits=index_bits)(Module)
- expected_bits = (
- index_bits if index_bits is not None else (32 if target_kind ==
"webgpu" else 64)
- )
- expected = gpu_2d_continuous_cumsum(
- in_dtype="float32", out_dtype="float32", index_bits=expected_bits
- )
+ expected = gpu_2d_continuous_cumsum(in_dtype="float32",
out_dtype="float32", index_bits=64)
assert_structural_equal(mod["gpu_2d_continuous_cumsum"], expected)
- if expected_bits == 32:
- # This previously failed on Metal with a 2**35 IntImm.
- tvm.s_tir.transform.ForceNarrowIndexToInt32()(mod)
@pytest.mark.parametrize("index_bits", [0, 16, 128])
diff --git
a/tests/python/s_tir/transform/test_s_tir_transform_force_narrow_index_to_i32.py
b/tests/python/s_tir/transform/test_s_tir_transform_force_narrow_index_to_i32.py
deleted file mode 100644
index 9cbfba3a08..0000000000
---
a/tests/python/s_tir/transform/test_s_tir_transform_force_narrow_index_to_i32.py
+++ /dev/null
@@ -1,177 +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
-
-import tvm
-import tvm.testing
-from tvm.s_tir import dlight as dl
-from tvm.script import s_tir as Ts
-from tvm.script import tirx as T
-from tvm.testing import env
-
-
-def _narrow(func):
- mod = tvm.IRModule.from_expr(func)
- return tvm.s_tir.transform.ForceNarrowIndexToInt32()(mod)["main"]
-
-
-def test_block():
- @Ts.function(private=True)
- def before(A: T.Tensor((128,), "float32"), B: T.Tensor((128,), "float32")):
- for i in T.serial(0, T.int64(16)):
- for j in T.serial(0, T.int64(8)):
- with Ts.sblock():
- vi = Ts.axis.spatial(T.int64(128), i * T.int64(8) + j)
- B[vi] = A[vi] + T.float32(1)
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((128,), "float32"), B: T.Tensor((128,),
"float32")):
- for i in T.serial(0, T.int32(16)):
- for j in T.serial(0, T.int32(8)):
- with Ts.sblock():
- vi = Ts.axis.spatial(T.int32(128), i * T.int32(8) + j)
- B[vi] = A[vi] + T.float32(1)
-
- tvm.ir.assert_structural_equal(_narrow(before), expected)
-
-
-def test_block_iters_used_only_in_regions():
- """Blockized blocks use their iterators only in access and match_tensor
regions."""
-
- @Ts.function(private=True)
- def before(
- A: T.Tensor((T.int64(16), T.int64(16)), "float32"),
- B: T.Tensor((T.int64(16), T.int64(16)), "float32"),
- ):
- for i_o, j_o in T.grid(T.int64(2), T.int64(2)):
- with Ts.sblock("tile_o"):
- vi_o, vj_o = Ts.axis.remap("SS", [i_o, j_o])
- Ts.reads(
- A[
- vi_o * T.int64(8) : vi_o * T.int64(8) + T.int64(8),
- vj_o * T.int64(8) : vj_o * T.int64(8) + T.int64(8),
- ]
- )
- Ts.writes(
- B[
- vi_o * T.int64(8) : vi_o * T.int64(8) + T.int64(8),
- vj_o * T.int64(8) : vj_o * T.int64(8) + T.int64(8),
- ]
- )
- A_tile = Ts.match_tensor(
- A[
- vi_o * T.int64(8) : vi_o * T.int64(8) + T.int64(8),
- vj_o * T.int64(8) : vj_o * T.int64(8) + T.int64(8),
- ],
- (T.int64(8), T.int64(8)),
- offset_factor=1,
- )
- B_tile = Ts.match_tensor(
- B[
- vi_o * T.int64(8) : vi_o * T.int64(8) + T.int64(8),
- vj_o * T.int64(8) : vj_o * T.int64(8) + T.int64(8),
- ],
- (T.int64(8), T.int64(8)),
- offset_factor=1,
- )
- for i_i, j_i in T.grid(T.int64(8), T.int64(8)):
- with Ts.sblock("tile"):
- vi_i, vj_i = Ts.axis.remap("SS", [i_i, j_i])
- B_tile[vi_i, vj_i] = A_tile[vi_i, vj_i] + T.float32(1)
-
- @Ts.function(private=True)
- def expected(A: T.Tensor((16, 16), "float32"), B: T.Tensor((16, 16),
"float32")):
- for i_o, j_o in T.grid(2, 2):
- with Ts.sblock("tile_o"):
- vi_o, vj_o = Ts.axis.remap("SS", [i_o, j_o])
- Ts.reads(A[vi_o * 8 : vi_o * 8 + 8, vj_o * 8 : vj_o * 8 + 8])
- Ts.writes(B[vi_o * 8 : vi_o * 8 + 8, vj_o * 8 : vj_o * 8 + 8])
- A_tile = Ts.match_tensor(
- A[vi_o * 8 : vi_o * 8 + 8, vj_o * 8 : vj_o * 8 + 8], (8,
8), offset_factor=1
- )
- B_tile = Ts.match_tensor(
- B[vi_o * 8 : vi_o * 8 + 8, vj_o * 8 : vj_o * 8 + 8], (8,
8), offset_factor=1
- )
- for i_i, j_i in T.grid(8, 8):
- with Ts.sblock("tile"):
- vi_i, vj_i = Ts.axis.remap("SS", [i_i, j_i])
- B_tile[vi_i, vj_i] = A_tile[vi_i, vj_i] + T.float32(1)
-
- tvm.ir.assert_structural_equal(_narrow(before), expected)
-
-
-def test_fail_on_tensor_param():
- @Ts.function(private=True)
- def func(A: T.Tensor((128,), "int64"), B: T.Tensor((128,), "int64")):
- for i in T.serial(0, 16):
- for j in T.serial(0, 8):
- with Ts.sblock():
- vi = Ts.axis.spatial(128, i * 8 + j)
- B[vi] = A[vi] + T.int64(1)
-
- with pytest.raises(RuntimeError):
- _narrow(func)
-
-
-def test_fail_on_block_alloc_tensor():
- @Ts.function(private=True)
- def func(A: T.Tensor((128,), "int32"), B: T.Tensor((128,), "int32")):
- C = Ts.sblock_alloc_tensor((128,), "int64")
- for i in T.serial(0, 16):
- for j in T.serial(0, 8):
- with Ts.sblock():
- vi = Ts.axis.spatial(128, i * 8 + j)
- C[vi] = T.cast(A[vi], "int64") + T.int64(1)
- for i in T.serial(0, 16):
- for j in T.serial(0, 8):
- with Ts.sblock():
- vi = Ts.axis.spatial(128, i * 8 + j)
- B[vi] = T.cast(C[vi] + T.int64(1), "int32")
-
- with pytest.raises(RuntimeError):
- _narrow(func)
-
-
[email protected](not env.has_llvm(), reason="need llvm for the host module")
-def test_metal_simdgroup_matmul_builds():
- """Narrowing a DLight-scheduled Metal matmul keeps its tensorized blocks
consistent."""
-
- n = T.dynamic("n")
-
- @Ts.function
- def main(
- A: T.Tensor((T.int64(1), n, T.int64(256)), "float16"),
- B: T.Tensor((T.int64(256), T.int64(256)), "float16"),
- C: T.Tensor((T.int64(1), n, T.int64(256)), "float16"),
- ):
- for i0, i1, i2, k in T.grid(T.int64(1), n, T.int64(256), T.int64(256)):
- with Ts.sblock("NT_matmul"):
- v0, v1, v2, vk = Ts.axis.remap("SSSR", [i0, i1, i2, k])
- with Ts.init():
- C[v0, v1, v2] = T.float16(0)
- C[v0, v1, v2] = C[v0, v1, v2] + A[v0, v1, vk] * B[v2, vk]
-
- target = tvm.target.Target("metal", host="llvm")
- with target:
- mod = dl.ApplyDefaultSchedule(dl.gpu.Matmul())(tvm.IRModule({"main":
main}))
- assert "metal.simdgroup" in mod.script()
- mod = tvm.s_tir.transform.ForceNarrowIndexToInt32()(mod)
- tvm.tirx.compile(mod, target=target).mod
-
-
-if __name__ == "__main__":
- tvm.testing.main()
