You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by lu...@apache.org on 2023/10/31 21:10:08 UTC
(tvm) branch main updated: [Fix][TIR]fix mul dtype mismatch (#16010)
This is an automated email from the ASF dual-hosted git repository.
lunderberg 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 878a61105e [Fix][TIR]fix mul dtype mismatch (#16010)
878a61105e is described below
commit 878a61105ea4c85f3547fe137a28d0a80b1f0e94
Author: Wei Tao <11...@qq.com>
AuthorDate: Wed Nov 1 05:09:58 2023 +0800
[Fix][TIR]fix mul dtype mismatch (#16010)
* [Fix][TIR]fix mul dtype mismatch
* [Fix]add unittest for dtype mismatch
---
src/tir/transforms/inject_ptx_async_copy.cc | 7 +++--
.../test_tir_transform_inject_ptx_async_copy.py | 34 ++++++++++++++++++++++
2 files changed, 38 insertions(+), 3 deletions(-)
diff --git a/src/tir/transforms/inject_ptx_async_copy.cc b/src/tir/transforms/inject_ptx_async_copy.cc
index a6685fe87b..5d23e854be 100644
--- a/src/tir/transforms/inject_ptx_async_copy.cc
+++ b/src/tir/transforms/inject_ptx_async_copy.cc
@@ -24,6 +24,7 @@
#include <tvm/tir/analysis.h>
#include <tvm/tir/builtin.h>
#include <tvm/tir/expr.h>
+#include <tvm/tir/op.h>
#include <tvm/tir/stmt_functor.h>
#include <tvm/tir/transform.h>
@@ -79,7 +80,7 @@ class PTXAsyncCopyInjector : public StmtMutator {
if (indices_lanes == 1) {
auto src_offset = load->indices[0];
auto dst_offset = store->indices[0];
- Array<PrimExpr> args = {store->buffer->data, tir::Mul(dst_offset, PrimExpr(index_factor)),
+ Array<PrimExpr> args = {store->buffer->data, mul(dst_offset, PrimExpr(index_factor)),
load->buffer->data, src_offset, PrimExpr(bytes)};
// use arguments size to indicate whether or not to use predicated cp.async
if (predicated) {
@@ -114,7 +115,7 @@ class PTXAsyncCopyInjector : public StmtMutator {
}();
if (src_offset.defined() && dst_offset.defined()) {
return Evaluate(Call(store->buffer->dtype, tvm::tir::builtin::ptx_cp_async(),
- {store->buffer->data, tir::Mul(dst_offset, PrimExpr(index_factor)),
+ {store->buffer->data, mul(dst_offset, PrimExpr(index_factor)),
load->buffer->data, src_offset, PrimExpr(bytes)}));
}
} else {
@@ -144,7 +145,7 @@ class PTXAsyncCopyInjector : public StmtMutator {
if (src_offset.defined() && dst_offset.defined()) {
return Evaluate(
Call(store->buffer->dtype, tvm::tir::builtin::ptx_cp_async(),
- {store->buffer->data, tir::Mul(dst_offset, PrimExpr(index_factor)),
+ {store->buffer->data, mul(dst_offset, PrimExpr(index_factor)),
load->buffer->data, src_offset, PrimExpr(bytes), predicate_value}));
}
}
diff --git a/tests/python/unittest/test_tir_transform_inject_ptx_async_copy.py b/tests/python/unittest/test_tir_transform_inject_ptx_async_copy.py
index 61f0892a9c..bf68200944 100644
--- a/tests/python/unittest/test_tir_transform_inject_ptx_async_copy.py
+++ b/tests/python/unittest/test_tir_transform_inject_ptx_async_copy.py
@@ -938,5 +938,39 @@ def test_vectorize_cp_async_in_if_then_else(postproc_if_missing_async_support):
assert "setp.ne.b32" in generated_code
+class TestMultiplicationNodesAreInligned(tvm.testing.CompareBeforeAfter):
+ transform = tvm.tir.transform.InjectPTXAsyncCopy()
+
+ def before(A: T.Buffer((32, 128), "float16")):
+ tx = T.launch_thread("threadIdx.x", T.int64(32))
+ A_flattened = T.Buffer((4096,), "float16", data=A.data)
+ A_shared = T.decl_buffer([4096], "float16", scope="shared")
+
+ T.attr("default", "async_scope", 1)
+ for i in range(16):
+ cse_var_1: T.int64 = T.Cast("int64", i)
+ A_shared[
+ T.Ramp(tx * T.int64(128) + cse_var_1 * T.int64(8), T.int64(1), 8)
+ ] = A_flattened[T.Ramp(tx * T.int64(128) + cse_var_1 * T.int64(8), T.int64(1), 8)]
+ T.ptx_commit_group()
+ T.ptx_wait_group(0)
+
+ def expected(A: T.Buffer((32, 128), "float16")):
+ tx = T.launch_thread("threadIdx.x", T.int64(32))
+ A_shared = T.decl_buffer((4096,), "float16", scope="shared")
+ for i in range(16):
+ cse_var_1: T.int64 = T.Cast("int64", i)
+ T.ptx_cp_async(
+ "float16",
+ A_shared.data,
+ T.Cast("int64", tx) * T.int64(128) + cse_var_1 * T.int64(8),
+ A.data,
+ T.Cast("int64", tx) * T.int64(128) + cse_var_1 * T.int64(8),
+ 16,
+ )
+ T.ptx_commit_group()
+ T.ptx_wait_group(0)
+
+
if __name__ == "__main__":
tvm.testing.main()