yongwww commented on code in PR #16254: URL: https://github.com/apache/tvm/pull/16254#discussion_r1429284940
########## tests/python/relax/test_backend_dispatch_sort_scan.py: ########## @@ -0,0 +1,415 @@ +# 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.script +import tvm.testing +from tvm.script import relax as R, tir as T, ir as I + +from tvm.relax.backend import DispatchSortScan +from tvm.ir.base import assert_structural_equal + + +def test_dispatch_cumsum(): + @I.ir_module + class Before: + I.module_global_infos({"vdevice": [I.vdevice("cuda", 0), I.vdevice("llvm", 0)]}) + + @R.function + def foo(x: R.Tensor((2, 3), "float32", "llvm")): + with R.dataflow(): + gv = R.cumsum(x, axis=1, dtype="float64") + R.output(gv) + return gv + + @I.ir_module + class Expected: + I.module_global_infos({"vdevice": [I.vdevice("cuda", 0), I.vdevice("llvm", 0)]}) + + @T.prim_func(private=True) + def cumsum(var_A: T.handle, out_buf: T.Buffer((T.int64(2), T.int64(3)), "float64")): + T.func_attr({"tir.noalias": T.bool(True)}) + A = T.match_buffer(var_A, (T.int64(2), T.int64(3)), offset_factor=1) + with T.block("cumsum_generic"): + T.reads(A[T.int64(0) : T.int64(2), T.int64(0) : T.int64(3)]) + T.writes(out_buf[T.int64(0) : T.int64(2), T.int64(0) : T.int64(3)]) + for fused in T.parallel(T.int64(2)): + out_buf[ + fused * T.int64(3) // T.int64(3), fused * T.int64(3) % T.int64(3) + ] = T.Cast( + "float64", + A[fused * T.int64(3) // T.int64(3), fused * T.int64(3) % T.int64(3)], + ) + for _k in range(T.int64(2)): + out_buf[ + (fused * T.int64(3) + (_k + T.int64(1))) // T.int64(3), + (fused * T.int64(3) + (_k + T.int64(1))) % T.int64(3), + ] = out_buf[ + (fused * T.int64(3) + (_k + T.int64(1) - T.int64(1))) // T.int64(3), + (fused * T.int64(3) + (_k + T.int64(1) - T.int64(1))) % T.int64(3), + ] + T.Cast( + "float64", + A[ + (fused * T.int64(3) + (_k + T.int64(1))) // T.int64(3), + (fused * T.int64(3) + (_k + T.int64(1))) % T.int64(3), + ], + ) + + @R.function + def foo( + x: R.Tensor((2, 3), dtype="float32", vdevice="llvm") + ) -> R.Tensor((2, 3), dtype="float64", vdevice="llvm"): + cls = Expected + with R.dataflow(): + gv = R.call_tir(cls.cumsum, (x,), out_sinfo=R.Tensor((2, 3), dtype="float64")) + R.output(gv) + return gv + + mod = DispatchSortScan()(Before) + assert_structural_equal(mod, Expected, map_free_vars=True) + + [email protected]("The emitted primfunc is not roundtripable, failed in build.") +def test_dispatch_cumsum_cuda(): + @I.ir_module + class Before: Review Comment: Ran into error when compile this mod (topi.cuda.cumsum). cc: @tqchen @jinhongyii target = tvm.target.Target("cuda", host="llvm") from tvm import tir, relax mod = Before with target: mod = DispatchSortScan()(Before) mod = tir.transform.DefaultGPUSchedule()(mod) ex = relax.build(mod, target) vm = relax.VirtualMachine(ex, tvm.gpu()) Error message: ``` File "tvm/build/../tests/python/relax/test_backend_dispatch_sort_scan.py", line 332, in test_dispatch_cumsum_cuda ex = relax.build(mod, target) File "tvm/python/tvm/relax/vm_build.py", line 328, in build return _vmlink( File "tvm/python/tvm/relax/vm_build.py", line 241, in _vmlink lib = tvm.build( File "tvm/python/tvm/driver/build_module.py", line 236, in build input_mod = lower(inputs) File "tvm/python/tvm/driver/build_module.py", line 130, in lower return ffi.lower_module(inp, simple_mode) File "tvm/python/tvm/_ffi/_ctypes/packed_func.py", line 239, in __call__ raise_last_ffi_error() File "tvm/python/tvm/_ffi/base.py", line 481, in raise_last_ffi_error raise py_err File "tvm/src/driver/driver_api.cc", line 346, in operator() return LowerModule(std::move(mod), simple_mode); File "tvm/src/driver/driver_api.cc", line 342, in tvm::LowerModule(tvm::IRModule, bool) return LowerWithPassList(std::move(mod), pass_list); File "tvm/src/driver/driver_api.cc", line 281, in tvm::LowerWithPassList(tvm::IRModule, tvm::runtime::Array<tvm::transform::Pass, void>) mod = optimize(std::move(mod)); File "tvm/src/tir/ir/transform.cc", line 101, in tvm::tir::transform::PrimFuncPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const func = pass_func(std::move(func), mod, pass_ctx); File "tvm/src/tir/transforms/unify_thread_binding.cc", line 202, in operator() return UnifyThreadBinding(std::move(f)); File "tvm/src/tir/transforms/unify_thread_binding.cc", line 191, in tvm::tir::UnifyThreadBinding(tvm::tir::PrimFunc) fptr->body = ThreadBindingUnifier::Unify(std::move(f->body)); File "tvm/src/tir/transforms/unify_thread_binding.cc", line 44, in tvm::tir::ThreadBindingUnifier::Unify(tvm::tir::Stmt) static Stmt Unify(Stmt stmt) { return ThreadBindingUnifier()(std::move(stmt)); } File "tvm/src/tir/ir/stmt_functor.cc", line 211, in tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&) return MutateArray(self, arr, fmutate); File "tvm/src/tir/ir/stmt_functor.cc", line 184, in tvm::runtime::Array<tvm::tir::Stmt, std::enable_if<std::is_base_of<tvm::runtime::ObjectRef, tvm::tir::Stmt>::value, void>::type> tvm::tir::StmtMutator::Internal::MutateArray<tvm::tir::Stmt, tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}>(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, std::enable_if<std::is_base_of<tvm::runtime::ObjectRef, tvm::tir::Stmt>::value, void>::type> const&, tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}) Array<T> copy = arr.Map(fmutate); File "tvm/include/tvm/runtime/container/array.h", line 652, in tvm::runtime::Array<tvm::tir::Stmt, std::enable_if<std::is_base_of<tvm::runtime::ObjectRef, tvm::tir::Stmt>::value, void>::type> tvm::runtime::Array<tvm::tir::Stmt, void>::Map<tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}, tvm::tir::Stmt>(tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}) const return Array<U>(MapHelper(data_, fmap)); File "tvm/include/tvm/runtime/container/array.h", line 823, in tvm::runtime::ObjectPtr<tvm::runtime::Object> tvm::runtime::Array<tvm::tir::Stmt, void>::MapHelper<tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}, tvm::tir::Stmt>(tvm::runtime::ObjectPtr<tvm::runtime::Object>, tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}) U mapped = fmap(DowncastNoCheck<T>(*it)); File "tvm/src/tir/ir/stmt_functor.cc", line 210, in tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}::operator()(tvm::tir::Stmt const&) const auto fmutate = [self](const Stmt& s) { return self->VisitStmt(s); }; File "tvm/src/tir/transforms/unify_thread_binding.cc", line 64, in tvm::tir::ThreadBindingUnifier::VisitStmt_(tvm::tir::ForNode const*) Range::FromMinExtent(op->min, op->extent)); File "tvm/src/tir/transforms/unify_thread_binding.cc", line 131, in tvm::tir::Stmt tvm::tir::ThreadBindingUnifier::UnifyThreadBindingImpl<tvm::tir::ForNode>(tvm::tir::ForNode const*, tvm::tir::Var const&, tvm::tir::IterVar const&, tvm::Range const&) Stmt new_stmt = StmtMutator::VisitStmt_(op); File "tvm/src/tir/transforms/unify_thread_binding.cc", line 64, in tvm::tir::ThreadBindingUnifier::VisitStmt_(tvm::tir::ForNode const*) Range::FromMinExtent(op->min, op->extent)); File "tvm/src/tir/transforms/unify_thread_binding.cc", line 131, in tvm::tir::Stmt tvm::tir::ThreadBindingUnifier::UnifyThreadBindingImpl<tvm::tir::ForNode>(tvm::tir::ForNode const*, tvm::tir::Var const&, tvm::tir::IterVar const&, tvm::Range const&) Stmt new_stmt = StmtMutator::VisitStmt_(op); File "tvm/src/tir/transforms/unify_thread_binding.cc", line 54, in tvm::tir::ThreadBindingUnifier::VisitStmt_(tvm::tir::AttrStmtNode const*) Range::FromMinExtent(IntImm(op->value->dtype, 0), op->value)); File "tvm/src/tir/transforms/unify_thread_binding.cc", line 112, in tvm::tir::Stmt tvm::tir::ThreadBindingUnifier::UnifyThreadBindingImpl<tvm::tir::AttrStmtNode>(tvm::tir::AttrStmtNode const*, tvm::tir::Var const&, tvm::tir::IterVar const&, tvm::Range const&) CHECK(ana.CanProveEqual(dom->extent, new_iter_var->dom->extent)) ValueError: Traceback (most recent call last): 17: operator() at tvm/src/driver/driver_api.cc:346 16: tvm::LowerModule(tvm::IRModule, bool) at tvm/src/driver/driver_api.cc:342 15: tvm::LowerWithPassList(tvm::IRModule, tvm::runtime::Array<tvm::transform::Pass, void>) at tvm/src/driver/driver_api.cc:281 14: tvm::tir::transform::PrimFuncPassNode::operator()(tvm::IRModule, tvm::transform::PassContext const&) const at tvm/src/tir/ir/transform.cc:101 13: operator() at tvm/src/tir/transforms/unify_thread_binding.cc:202 12: tvm::tir::UnifyThreadBinding(tvm::tir::PrimFunc) at tvm/src/tir/transforms/unify_thread_binding.cc:191 11: tvm::tir::ThreadBindingUnifier::Unify(tvm::tir::Stmt) at tvm/src/tir/transforms/unify_thread_binding.cc:44 10: tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&) at tvm/src/tir/ir/stmt_functor.cc:211 9: tvm::runtime::Array<tvm::tir::Stmt, std::enable_if<std::is_base_of<tvm::runtime::ObjectRef, tvm::tir::Stmt>::value, void>::type> tvm::tir::StmtMutator::Internal::MutateArray<tvm::tir::Stmt, tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}>(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, std::enable_if<std::is_base_of<tvm::runtime::ObjectRef, tvm::tir::Stmt>::value, void>::type> const&, tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}) at tvm/src/tir/ir/stmt_functor.cc:184 8: tvm::runtime::Array<tvm::tir::Stmt, std::enable_if<std::is_base_of<tvm::runtime::ObjectRef, tvm::tir::Stmt>::value, void>::type> tvm::runtime::Array<tvm::tir::Stmt, void>::Map<tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}, tvm::tir::Stmt>(tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}) const at tvm/include/tvm/runtime/container/array.h:652 7: tvm::runtime::ObjectPtr<tvm::runtime::Object> tvm::runtime::Array<tvm::tir::Stmt, void>::MapHelper<tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}, tvm::tir::Stmt>(tvm::runtime::ObjectPtr<tvm::runtime::Object>, tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}) at tvm/include/tvm/runtime/container/array.h:823 6: tvm::tir::StmtMutator::Internal::Mutate(tvm::tir::StmtMutator*, tvm::runtime::Array<tvm::tir::Stmt, void> const&)::{lambda(tvm::tir::Stmt const&)#1}::operator()(tvm::tir::Stmt const&) const at tvm/src/tir/ir/stmt_functor.cc:210 5: tvm::tir::ThreadBindingUnifier::VisitStmt_(tvm::tir::ForNode const*) at tvm/src/tir/transforms/unify_thread_binding.cc:64 4: tvm::tir::Stmt tvm::tir::ThreadBindingUnifier::UnifyThreadBindingImpl<tvm::tir::ForNode>(tvm::tir::ForNode const*, tvm::tir::Var const&, tvm::tir::IterVar const&, tvm::Range const&) at tvm/src/tir/transforms/unify_thread_binding.cc:131 3: tvm::tir::ThreadBindingUnifier::VisitStmt_(tvm::tir::ForNode const*) at tvm/src/tir/transforms/unify_thread_binding.cc:64 2: tvm::tir::Stmt tvm::tir::ThreadBindingUnifier::UnifyThreadBindingImpl<tvm::tir::ForNode>(tvm::tir::ForNode const*, tvm::tir::Var const&, tvm::tir::IterVar const&, tvm::Range const&) at tvm/src/tir/transforms/unify_thread_binding.cc:131 1: tvm::tir::ThreadBindingUnifier::VisitStmt_(tvm::tir::AttrStmtNode const*) at tvm/src/tir/transforms/unify_thread_binding.cc:54 0: tvm::tir::Stmt tvm::tir::ThreadBindingUnifier::UnifyThreadBindingImpl<tvm::tir::AttrStmtNode>(tvm::tir::AttrStmtNode const*, tvm::tir::Var const&, tvm::tir::IterVar const&, tvm::Range const&) at tvm/src/tir/transforms/unify_thread_binding.cc:112 File "tvm/src/tir/transforms/unify_thread_binding.cc", line 112 ValueError: Check failed: (ana.CanProveEqual(dom->extent, new_iter_var->dom->extent)) is false: All loops that are bound to `blockIdx.x` should have the same extent. However, there are two loops with extent 1 and T.int64(2), which are not equal ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
