From 0aade2053d77065589d4aa7055e185561a13f80b Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Mon, 8 May 2023 04:48:53 +0900 Subject: [PATCH 1/2] Allow arbitrary stride when the corresponding shape is 1 --- src/tir/transforms/arg_binder.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tir/transforms/arg_binder.cc b/src/tir/transforms/arg_binder.cc index c785b732abc3..571acf8e092f 100644 --- a/src/tir/transforms/arg_binder.cc +++ b/src/tir/transforms/arg_binder.cc @@ -230,7 +230,7 @@ void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type, for (size_t i = buffer->shape.size(); i != 0; --i) { size_t k = i - 1; PrimExpr svalue = cast(stype, BufferLoad(buf_strides, {IntImm(DataType::Int(32), k)})); - conds.push_back(expect_stride == svalue); + conds.push_back(buffer->shape[k] == 1 || expect_stride == svalue); expect_stride = expect_stride * buffer->shape[k]; } std::ostringstream stride_err_msg; From a3f0320bf16b44043a0cad3420038f973a1bdb59 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Mon, 8 May 2023 04:57:01 +0900 Subject: [PATCH 2/2] add test --- tests/python/unittest/test_runtime_dlpack.py | 52 ++++++++++++++++++++ 1 file changed, 52 insertions(+) create mode 100644 tests/python/unittest/test_runtime_dlpack.py diff --git a/tests/python/unittest/test_runtime_dlpack.py b/tests/python/unittest/test_runtime_dlpack.py new file mode 100644 index 000000000000..3f13e2e5fed5 --- /dev/null +++ b/tests/python/unittest/test_runtime_dlpack.py @@ -0,0 +1,52 @@ +# 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 tvm +import tvm.testing +from tvm import te +import numpy as np + + +@tvm.testing.requires_package("torch") +def test_from_dlpack_shape_one(): + # A test case for the issue https://github.com/pytorch/pytorch/issues/99803 + import torch + from torch.utils.dlpack import to_dlpack + + tgt = tvm.target.Target(target="llvm", host="llvm") + + rows = 1 + a = tvm.runtime.ndarray.from_dlpack(to_dlpack(torch.randn(rows, 16))) + + A = te.placeholder((rows, 16), name="A") + B = te.placeholder((rows, 16), name="B") + C = te.compute(A.shape, lambda i, j: A[i, j] + B[i, j], name="C") + + s = te.create_schedule(C.op) + + fadd = tvm.build(s, [A, B, C], tgt) + + dev = tvm.device(tgt.kind.name, 0) + + b = tvm.nd.array(np.random.uniform(size=(rows, 16)).astype(B.dtype), dev) + c = tvm.nd.array(np.zeros((rows, 16), dtype=C.dtype), dev) + fadd(a, b, c) + + tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy()) + + +if __name__ == "__main__": + tvm.testing.main()