This is an automated email from the ASF dual-hosted git repository.
masahi 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 843310b03c [Fix,AutoScheduler] Handle 0-dim buffers in featurization
(#13718)
843310b03c is described below
commit 843310b03c601bc5324932bdce7ba0195f9e94dc
Author: Tristan Konolige <tkonol...@octoml.ai>
AuthorDate: Fri Jan 6 21:15:38 2023 -0800
[Fix,AutoScheduler] Handle 0-dim buffers in featurization (#13718)
0-dim buffers get lowered to a 0-dim buffer and a 1-dim buffer decl. We
explicitly check for this case and return a stride of 1.
---
src/auto_scheduler/feature.cc | 5 +++++
.../python/unittest/test_auto_scheduler_feature.py | 26 ++++++++++++++++++++++
2 files changed, 31 insertions(+)
diff --git a/src/auto_scheduler/feature.cc b/src/auto_scheduler/feature.cc
index 4ce7ad13bc..884215c24a 100644
--- a/src/auto_scheduler/feature.cc
+++ b/src/auto_scheduler/feature.cc
@@ -443,6 +443,11 @@ class CoefficientExtractor : public StmtExprVisitor {
// Compute stride for the accesses to a buffer
int64_t ComputeStride(const std::vector<std::vector<PrimExpr>>& indices,
const std::vector<int>& shape, const VarNode*
stride_var) {
+ // Use stride of 1 for 0-dimensional buffers. 0-dim buffers has a single
+ // index access, so we have to check here.
+ if (shape.size() == 0) {
+ return 1;
+ }
int64_t min_stride = std::numeric_limits<int64_t>::max();
bool find = false;
CoefficientExtractor extractor;
diff --git a/tests/python/unittest/test_auto_scheduler_feature.py
b/tests/python/unittest/test_auto_scheduler_feature.py
index 3f435366e1..8be6e0a8f2 100644
--- a/tests/python/unittest/test_auto_scheduler_feature.py
+++ b/tests/python/unittest/test_auto_scheduler_feature.py
@@ -273,6 +273,32 @@ def test_negative_extent():
assert features["B0.unique_bytes"] == 0
+@T.prim_func
+def zero_dim(
+ p2: T.Buffer[(), "float32"],
+ T_cast: T.Buffer[(T.int64(1), T.int64(768)), "int8"],
+):
+ # function attr dict
+ T.func_attr(
+ {
+ "tir.noalias": True,
+ "Primitive": 1,
+ }
+ )
+ # buffer definition
+ T_cast_1 = T.buffer_decl([T.int64(768)], dtype="int8", data=T_cast.data)
+ p2_1 = T.buffer_decl([1], dtype="float32", data=p2.data)
+ # body
+ for i0_i1_fused in T.serial(768):
+ T_cast_1[i0_i1_fused] = p2_1[0]
+
+
+def test_zero_dim():
+ features = auto_scheduler.feature.named_features_from_primfunc(zero_dim)
+ assert features["B1.stride"] == 1
+ assert features["B0.stride"] == 1
+
+
if __name__ == "__main__":
test_cpu_matmul()
test_cpu_fusion()
