This is an automated email from the ASF dual-hosted git repository.
lukhut 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 866bec03ad [microNPU] Disable copying weights to SRAM for
FullyConnected ops in CopyConstants scheduler (#13588)
866bec03ad is described below
commit 866bec03ade5dd7e4cb529cb56bd8e58b7f38326
Author: Alexey Yazev <[email protected]>
AuthorDate: Wed Dec 14 13:15:47 2022 +0400
[microNPU] Disable copying weights to SRAM for FullyConnected ops in
CopyConstants scheduler (#13588)
In Ethos-U, CopyConstants scheduler currently copies weights for all
operators. But in Vela, there are a number of scenarios where the weights are
not buffered in SRAM, and FullyConnected case is one of them.
---
python/tvm/relay/backend/contrib/ethosu/tir/scheduler.py | 10 +++++++++-
tests/python/contrib/test_ethosu/test_scheduler.py | 16 ++++++++++++++++
2 files changed, 25 insertions(+), 1 deletion(-)
diff --git a/python/tvm/relay/backend/contrib/ethosu/tir/scheduler.py
b/python/tvm/relay/backend/contrib/ethosu/tir/scheduler.py
index bcabe2b7c2..cee8f563ff 100644
--- a/python/tvm/relay/backend/contrib/ethosu/tir/scheduler.py
+++ b/python/tvm/relay/backend/contrib/ethosu/tir/scheduler.py
@@ -132,6 +132,12 @@ def copy_constants():
def _planner(cached_func, const_dict, sch):
planned = set() # type: ignore
+ def _is_matmul(tensor):
+ if tensor.name not in ["ethosu_conv2d"]:
+ return False
+ a, b = tensor.op.input_tensors[0:2]
+ return a.shape[1:3] == [1, 1] and b.shape[1:3] == [1, 1]
+
def _visit(tensor, reader, lut):
if tensor not in planned:
planned.add(tensor)
@@ -140,7 +146,9 @@ def copy_constants():
# ambiguity when encountering a scalar.
is_same = [var.same_as(tensor) for var in
cached_func.inputs]
index = is_same.index(True)
- if index in const_dict:
+ # Along with constants, also skip for FullyConnected to
correspond
+ # with Vela behavior
+ if index in const_dict and not _is_matmul(reader):
sch.cache_read(tensor, "global", [reader])
elif isinstance(tensor.op, tvm.te.ComputeOp):
diff --git a/tests/python/contrib/test_ethosu/test_scheduler.py
b/tests/python/contrib/test_ethosu/test_scheduler.py
index fd1e1afa60..695aed0d19 100644
--- a/tests/python/contrib/test_ethosu/test_scheduler.py
+++ b/tests/python/contrib/test_ethosu/test_scheduler.py
@@ -217,5 +217,21 @@ def test_schedule_diamond_graph():
tvm.ir.assert_structural_equal(test_mod["main"], reference_mod["main"],
True)
+def test_copy_constants_fully_connected_weights():
+ """Check that MatMul-like conv2d ops do not copy weights to SRAM."""
+ ifm = relay.var("IFM", shape=(1, 1, 1, 32), dtype="int8")
+ conv = make_ethosu_conv2d(ifm, 32, 8, (1, 1), (0, 0), (1, 1), (1, 1))
+ func = relay.Function(relay.analysis.free_vars(conv), conv)
+ func = run_opt_pass(func, relay.transform.InferType())
+
+ func, const_dict = extract_constants(func)
+ cached_func = lower_to_te(func)
+
+ sch = te.create_schedule([cached_func.outputs[0].op])
+ planner = copy_constants()
+ planner(cached_func, const_dict, sch)
+ assert True not in [".global" in s.op.name for s in sch.stages]
+
+
if __name__ == "__main__":
pytest.main([__file__])
