[tvm] 01/02: init

[me...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

mehrdadh pushed a commit to branch micro/dense_dsp_smaller_config
in repository https://gitbox.apache.org/repos/asf/tvm.git

commit ca34e5268ee6d09a63ea5ad1a788bed00242c746
Author: Mehrdad Hessar <mh...@octoml.ai>
AuthorDate: Thu Sep 8 09:11:29 2022 -0700

    init
---
 python/tvm/topi/arm_cpu/mprofile/dsp/dense.py | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/python/tvm/topi/arm_cpu/mprofile/dsp/dense.py b/python/tvm/topi/arm_cpu/mprofile/dsp/dense.py
index 123ad1ce3c..5630636c92 100644
--- a/python/tvm/topi/arm_cpu/mprofile/dsp/dense.py
+++ b/python/tvm/topi/arm_cpu/mprofile/dsp/dense.py
@@ -32,9 +32,22 @@ def dense_dsp_compute(cfg, data, weight, bias=None, out_dtype=None):
     M, K = get_const_tuple(data.shape)
     N, _ = get_const_tuple(weight.shape)
 
-    cfg.define_split("tile_x", M, policy="factors", num_outputs=2)
-    cfg.define_split("tile_y", N, policy="factors", num_outputs=2)
-    cfg.define_split("tile_k", K, policy="factors", num_outputs=2)
+    factor = 2
+    # import pdb; pdb.set_trace()
+    if M % factor == 0:
+        cfg.define_split("tile_x", M, policy="factors", num_outputs=2, filter=lambda x: x.size[-1] % factor == 0)
+    else:
+        cfg.define_split("tile_x", M, policy="candidate", num_outputs=2, candidate=[[1, M]])
+    
+    if N % factor == 0:
+        cfg.define_split("tile_y", N, policy="factors", num_outputs=2, filter=lambda x: x.size[-1] % factor == 0)
+    else:
+        cfg.define_split("tile_y", N, policy="candidate", num_outputs=2, candidate=[[1, N]])
+
+    if K % factor == 0:
+        cfg.define_split("tile_k", K, policy="factors", num_outputs=2, filter=lambda x: x.size[-1] % factor == 0)
+    else:
+        cfg.define_split("tile_k", K, policy="candidate", num_outputs=2, candidate=[[1, K]])
 
     k = te.reduce_axis((0, K), "k")
     C = te.compute(