From dd3b3de3932d4e3f6cbe8796fa14375a732a4d1c Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 1 Apr 2022 10:04:28 +0900 Subject: [PATCH 1/9] add test mod --- .../test_meta_schedule_multi_anchor.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 tests/python/unittest/test_meta_schedule_multi_anchor.py diff --git a/tests/python/unittest/test_meta_schedule_multi_anchor.py b/tests/python/unittest/test_meta_schedule_multi_anchor.py new file mode 100644 index 000000000000..b6975a381504 --- /dev/null +++ b/tests/python/unittest/test_meta_schedule_multi_anchor.py @@ -0,0 +1,51 @@ +# 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 numpy as np + +import tvm +from tvm import relay + + +def get_dense_dense(): + M, N, K = 128, 128, 128 + data_shape = (M, K) + weight_shape = (N, K) + + def multi_dense(): + p_data = relay.var("p_data", shape=data_shape, dtype="uint8") + p_weight1 = relay.var("p_weight1", shape=weight_shape, dtype="int8") + p_weight2 = relay.var("p_weight2", shape=weight_shape, dtype="int8") + + dense1 = relay.nn.dense(p_data, p_weight1, out_dtype="int32") + dense2 = relay.nn.dense(relay.cast(dense1, "uint8"), + p_weight2, out_dtype="int32") + + f = relay.Function([p_data, p_weight1, p_weight2], dense2) + f = f.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) + return f + + data = relay.var("data", shape=data_shape, dtype="uint8") + weight1 = relay.var("weight1", shape=weight_shape, dtype="int8") + weight2 = relay.var("weight2", shape=weight_shape, dtype="int8") + + out = relay.Call(multi_dense(), [data, weight1, weight2]) + return relay.Function([data, weight1, weight2], out) + + +mod = tvm.IRModule.from_expr(get_dense_dense()) + +print(relay.transform.InferType()(mod)) From cec5da18ce823d94a32cefcc7972c385cee1758e Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 1 Apr 2022 10:09:16 +0900 Subject: [PATCH 2/9] task extraction works --- .../test_meta_schedule_multi_anchor.py | 30 ++++++++++++++----- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/tests/python/unittest/test_meta_schedule_multi_anchor.py b/tests/python/unittest/test_meta_schedule_multi_anchor.py index b6975a381504..c6d8a9ee2392 100644 --- a/tests/python/unittest/test_meta_schedule_multi_anchor.py +++ b/tests/python/unittest/test_meta_schedule_multi_anchor.py @@ -18,13 +18,10 @@ import tvm from tvm import relay +from tvm.meta_schedule.tune import Parse, extract_task_from_relay -def get_dense_dense(): - M, N, K = 128, 128, 128 - data_shape = (M, K) - weight_shape = (N, K) - +def get_dense_dense(data_shape, weight_shape): def multi_dense(): p_data = relay.var("p_data", shape=data_shape, dtype="uint8") p_weight1 = relay.var("p_weight1", shape=weight_shape, dtype="int8") @@ -46,6 +43,25 @@ def multi_dense(): return relay.Function([data, weight1, weight2], out) -mod = tvm.IRModule.from_expr(get_dense_dense()) -print(relay.transform.InferType()(mod)) +M, N, K = 128, 128, 128 +data_shape = (M, K) +weight_shape = (N, K) + + +relay_mod = tvm.IRModule.from_expr(get_dense_dense(data_shape, weight_shape)) + +# print(relay.transform.InferType()(relay_mod)) + +target = "llvm -mcpu=cascadelake -num-cores 4" + +weight1_np = np.random.uniform(1, 10, size=weight_shape).astype("int8") +weight2_np = np.random.uniform(1, 10, size=weight_shape).astype("int8") + +params = {"weight1": weight1_np, "weight2": weight2_np} + +extracted_tasks = extract_task_from_relay(relay_mod, target, params) + +for task in extracted_tasks: + mod = Parse._mod(task.dispatched[0]) + print(mod) From 9b4ea12bb69535457521431a750ba9905b0f9504 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 1 Apr 2022 10:12:55 +0900 Subject: [PATCH 3/9] trying relay.build --- .../test_meta_schedule_multi_anchor.py | 36 +++++++++++++++---- 1 file changed, 30 insertions(+), 6 deletions(-) diff --git a/tests/python/unittest/test_meta_schedule_multi_anchor.py b/tests/python/unittest/test_meta_schedule_multi_anchor.py index c6d8a9ee2392..2743d6dd30cb 100644 --- a/tests/python/unittest/test_meta_schedule_multi_anchor.py +++ b/tests/python/unittest/test_meta_schedule_multi_anchor.py @@ -19,6 +19,8 @@ import tvm from tvm import relay from tvm.meta_schedule.tune import Parse, extract_task_from_relay +from tvm.meta_schedule.database import TuningRecord, JSONDatabase +from tvm.meta_schedule.integration import ApplyHistoryBest def get_dense_dense(data_shape, weight_shape): @@ -28,8 +30,7 @@ def multi_dense(): p_weight2 = relay.var("p_weight2", shape=weight_shape, dtype="int8") dense1 = relay.nn.dense(p_data, p_weight1, out_dtype="int32") - dense2 = relay.nn.dense(relay.cast(dense1, "uint8"), - p_weight2, out_dtype="int32") + dense2 = relay.nn.dense(relay.cast(dense1, "uint8"), p_weight2, out_dtype="int32") f = relay.Function([p_data, p_weight1, p_weight2], dense2) f = f.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) @@ -43,7 +44,6 @@ def multi_dense(): return relay.Function([data, weight1, weight2], out) - M, N, K = 128, 128, 128 data_shape = (M, K) weight_shape = (N, K) @@ -62,6 +62,30 @@ def multi_dense(): extracted_tasks = extract_task_from_relay(relay_mod, target, params) -for task in extracted_tasks: - mod = Parse._mod(task.dispatched[0]) - print(mod) +assert len(extracted_tasks) == 1 + +task = extracted_tasks[0] + +database = JSONDatabase( + path_workload="database_workload.json", + path_tuning_record="database_tuning_record.json", +) + +mod = Parse._mod(task.dispatched[0]) + +workload = database.commit_workload(mod) + +sch = tvm.tir.Schedule(mod) +# print(sch.mod.script()) + +tune_rec = TuningRecord(sch.trace, [0.0], workload, tvm.target.Target(target), []) + +database.commit_tuning_record(tune_rec) + + +with ApplyHistoryBest(database): + with tvm.transform.PassContext( + opt_level=3, + config={"relay.backend.use_meta_schedule": True}, + ): + lib = relay.build(relay_mod, target=target, params=params) From 5f8dffd9bc8e87a75ba5f5535492e8c0c4152525 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 1 Apr 2022 10:24:36 +0900 Subject: [PATCH 4/9] test runs but result not correct --- .../test_meta_schedule_multi_anchor.py | 78 ++++++++++++------- 1 file changed, 50 insertions(+), 28 deletions(-) diff --git a/tests/python/unittest/test_meta_schedule_multi_anchor.py b/tests/python/unittest/test_meta_schedule_multi_anchor.py index 2743d6dd30cb..3d547850a461 100644 --- a/tests/python/unittest/test_meta_schedule_multi_anchor.py +++ b/tests/python/unittest/test_meta_schedule_multi_anchor.py @@ -44,48 +44,70 @@ def multi_dense(): return relay.Function([data, weight1, weight2], out) -M, N, K = 128, 128, 128 -data_shape = (M, K) -weight_shape = (N, K) +def get_ref(data_np, weight1_np, weight2_np): + dense1 = np.dot(data_np, np.transpose(weight1_np)) + return np.dot(dense1.astype("uint8"), np.transpose(weight2_np)) -relay_mod = tvm.IRModule.from_expr(get_dense_dense(data_shape, weight_shape)) +def test(): + M, N, K = 128, 128, 128 + data_shape = (M, K) + weight_shape = (N, K) -# print(relay.transform.InferType()(relay_mod)) + relay_mod = tvm.IRModule.from_expr(get_dense_dense(data_shape, weight_shape)) -target = "llvm -mcpu=cascadelake -num-cores 4" + # print(relay.transform.InferType()(relay_mod)) -weight1_np = np.random.uniform(1, 10, size=weight_shape).astype("int8") -weight2_np = np.random.uniform(1, 10, size=weight_shape).astype("int8") + target = "llvm" -params = {"weight1": weight1_np, "weight2": weight2_np} + data_np = np.random.uniform(1, 10, size=data_shape).astype("uint8") + weight1_np = np.random.uniform(1, 10, size=weight_shape).astype("int8") + weight2_np = np.random.uniform(1, 10, size=weight_shape).astype("int8") -extracted_tasks = extract_task_from_relay(relay_mod, target, params) + params = {"weight1": weight1_np, "weight2": weight2_np} -assert len(extracted_tasks) == 1 + extracted_tasks = extract_task_from_relay(relay_mod, target, params) -task = extracted_tasks[0] + assert len(extracted_tasks) == 1 -database = JSONDatabase( - path_workload="database_workload.json", - path_tuning_record="database_tuning_record.json", -) + task = extracted_tasks[0] -mod = Parse._mod(task.dispatched[0]) + database = JSONDatabase( + path_workload="database_workload.json", + path_tuning_record="database_tuning_record.json", + ) -workload = database.commit_workload(mod) + mod = Parse._mod(task.dispatched[0]) -sch = tvm.tir.Schedule(mod) -# print(sch.mod.script()) + workload = database.commit_workload(mod) -tune_rec = TuningRecord(sch.trace, [0.0], workload, tvm.target.Target(target), []) + sch = tvm.tir.Schedule(mod) + print(sch.mod.script()) -database.commit_tuning_record(tune_rec) + tune_rec = TuningRecord(sch.trace, [0.0], workload, tvm.target.Target(target), []) + database.commit_tuning_record(tune_rec) -with ApplyHistoryBest(database): - with tvm.transform.PassContext( - opt_level=3, - config={"relay.backend.use_meta_schedule": True}, - ): - lib = relay.build(relay_mod, target=target, params=params) + with ApplyHistoryBest(database): + with tvm.transform.PassContext( + opt_level=3, + config={"relay.backend.use_meta_schedule": True}, + ): + lib = relay.build(relay_mod, target=target, params=params) + + dev = tvm.device(target, 0) + + runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev)) + + runtime.set_input("data", data_np) + runtime.run() + + out = runtime.get_output(0).numpy() + + ref = get_ref(data_np, weight1_np, weight2_np) + + np.testing.assert_equal(out, ref) + + +if __name__ == "__main__": + test() From c27a48167582ea45020be1828424edde66597654 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 1 Apr 2022 10:29:02 +0900 Subject: [PATCH 5/9] test worked --- .../test_meta_schedule_multi_anchor.py | 34 ++++++++++++------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/tests/python/unittest/test_meta_schedule_multi_anchor.py b/tests/python/unittest/test_meta_schedule_multi_anchor.py index 3d547850a461..0ebcd8d56673 100644 --- a/tests/python/unittest/test_meta_schedule_multi_anchor.py +++ b/tests/python/unittest/test_meta_schedule_multi_anchor.py @@ -17,6 +17,7 @@ import numpy as np import tvm +import tvm.testing from tvm import relay from tvm.meta_schedule.tune import Parse, extract_task_from_relay from tvm.meta_schedule.database import TuningRecord, JSONDatabase @@ -25,20 +26,20 @@ def get_dense_dense(data_shape, weight_shape): def multi_dense(): - p_data = relay.var("p_data", shape=data_shape, dtype="uint8") - p_weight1 = relay.var("p_weight1", shape=weight_shape, dtype="int8") - p_weight2 = relay.var("p_weight2", shape=weight_shape, dtype="int8") + p_data = relay.var("p_data", shape=data_shape, dtype="float32") + p_weight1 = relay.var("p_weight1", shape=weight_shape, dtype="float32") + p_weight2 = relay.var("p_weight2", shape=weight_shape, dtype="float32") - dense1 = relay.nn.dense(p_data, p_weight1, out_dtype="int32") - dense2 = relay.nn.dense(relay.cast(dense1, "uint8"), p_weight2, out_dtype="int32") + dense1 = relay.nn.dense(p_data, p_weight1) + dense2 = relay.nn.dense(dense1, p_weight2) f = relay.Function([p_data, p_weight1, p_weight2], dense2) f = f.with_attr("Primitive", tvm.tir.IntImm("int32", 1)) return f - data = relay.var("data", shape=data_shape, dtype="uint8") - weight1 = relay.var("weight1", shape=weight_shape, dtype="int8") - weight2 = relay.var("weight2", shape=weight_shape, dtype="int8") + data = relay.var("data", shape=data_shape, dtype="float32") + weight1 = relay.var("weight1", shape=weight_shape, dtype="float32") + weight2 = relay.var("weight2", shape=weight_shape, dtype="float32") out = relay.Call(multi_dense(), [data, weight1, weight2]) return relay.Function([data, weight1, weight2], out) @@ -46,7 +47,11 @@ def multi_dense(): def get_ref(data_np, weight1_np, weight2_np): dense1 = np.dot(data_np, np.transpose(weight1_np)) - return np.dot(dense1.astype("uint8"), np.transpose(weight2_np)) + return np.dot(dense1, np.transpose(weight2_np)) + + +def schedule_dense_dense(sch): + pass def test(): @@ -60,9 +65,9 @@ def test(): target = "llvm" - data_np = np.random.uniform(1, 10, size=data_shape).astype("uint8") - weight1_np = np.random.uniform(1, 10, size=weight_shape).astype("int8") - weight2_np = np.random.uniform(1, 10, size=weight_shape).astype("int8") + data_np = np.random.randn(*data_shape).astype("float32") + weight1_np = np.random.randn(*weight_shape).astype("float32") + weight2_np = np.random.randn(*weight_shape).astype("float32") params = {"weight1": weight1_np, "weight2": weight2_np} @@ -82,6 +87,9 @@ def test(): workload = database.commit_workload(mod) sch = tvm.tir.Schedule(mod) + + schedule_dense_dense(sch) + print(sch.mod.script()) tune_rec = TuningRecord(sch.trace, [0.0], workload, tvm.target.Target(target), []) @@ -106,7 +114,7 @@ def test(): ref = get_ref(data_np, weight1_np, weight2_np) - np.testing.assert_equal(out, ref) + tvm.testing.assert_allclose(out, ref, atol=1e-4, rtol=1e-4) if __name__ == "__main__": From b3a3a7c7496c50a7de29037952d903dc8c27732a Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 1 Apr 2022 12:15:56 +0900 Subject: [PATCH 6/9] update te_compiler_cache --- src/relay/backend/te_compiler_cache.cc | 6 ++++-- .../unittest/test_meta_schedule_multi_anchor.py | 12 +++++++++--- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/relay/backend/te_compiler_cache.cc b/src/relay/backend/te_compiler_cache.cc index 3534697beca3..dc31d12118cb 100644 --- a/src/relay/backend/te_compiler_cache.cc +++ b/src/relay/backend/te_compiler_cache.cc @@ -299,6 +299,7 @@ class ScheduleBuilder : public ExprVisitor { explicit ScheduleBuilder(Target target) : target_(target) { // Whether to use auto_scheduler schedule. use_auto_scheduler_ = backend::IsAutoSchedulerEnabled(); + use_meta_scheduler_ = backend::IsMetaScheduleEnabled(); } CachedFunc Create(const Function& relay_func, std::function renamer) { @@ -336,7 +337,7 @@ class ScheduleBuilder : public ExprVisitor { schedule = Downcast(obj); } } - if (backend::IsMetaScheduleEnabled()) { + if (use_auto_scheduler_) { IRModule relay_mod({{prim_fn_var, relay_func}}); IRModule tir_mod({{prim_fn_var, tir::CreatePrimFunc(Concat(fn_inputs, tensor_outs))}}); Optional scheduled_mod = meta_schedule::MetaScheduleContext::QueryInsideWithScope( @@ -377,7 +378,7 @@ class ScheduleBuilder : public ExprVisitor { } int op_pattern = fpattern[op]; - if (!use_auto_scheduler_ && op_pattern >= kCommReduce) { + if (!use_auto_scheduler_ && !use_meta_scheduler_ && op_pattern >= kCommReduce) { ICHECK(!anchor_op_.defined() || anchor_op_pattern_ < kCommReduce) << "Cannot apply TOPI schedule to a primitive function with two complicated ops" << " anchor=" << anchor_op_ << " current=" << op; @@ -395,6 +396,7 @@ class ScheduleBuilder : public ExprVisitor { Attrs anchor_attrs_; int anchor_op_pattern_{0}; bool use_auto_scheduler_; + bool use_meta_scheduler_; }; /*! diff --git a/tests/python/unittest/test_meta_schedule_multi_anchor.py b/tests/python/unittest/test_meta_schedule_multi_anchor.py index 0ebcd8d56673..bbc57c0d5723 100644 --- a/tests/python/unittest/test_meta_schedule_multi_anchor.py +++ b/tests/python/unittest/test_meta_schedule_multi_anchor.py @@ -51,10 +51,16 @@ def get_ref(data_np, weight1_np, weight2_np): def schedule_dense_dense(sch): - pass + dense1 = sch.get_block("T_matmul_NT") + dense2 = sch.get_block("T_matmul_NT_1") + y1, x1, k1 = sch.get_loops(dense1) + y2, x2, k2 = sch.get_loops(dense2) -def test(): + # ... + + +def test_dense_dense(): M, N, K = 128, 128, 128 data_shape = (M, K) weight_shape = (N, K) @@ -118,4 +124,4 @@ def test(): if __name__ == "__main__": - test() + test_dense_dense() From cbea61ecb49a8111b02c9d62f0ad47799d2c34e6 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 1 Apr 2022 12:40:08 +0900 Subject: [PATCH 7/9] use temp dir for database json --- .../test_meta_schedule_multi_anchor.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/tests/python/unittest/test_meta_schedule_multi_anchor.py b/tests/python/unittest/test_meta_schedule_multi_anchor.py index bbc57c0d5723..9f5280b543dc 100644 --- a/tests/python/unittest/test_meta_schedule_multi_anchor.py +++ b/tests/python/unittest/test_meta_schedule_multi_anchor.py @@ -14,6 +14,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import os +import tempfile + import numpy as np import tvm @@ -83,24 +86,25 @@ def test_dense_dense(): task = extracted_tasks[0] - database = JSONDatabase( - path_workload="database_workload.json", - path_tuning_record="database_tuning_record.json", - ) - mod = Parse._mod(task.dispatched[0]) - workload = database.commit_workload(mod) + with tempfile.TemporaryDirectory() as work_dir: + database = JSONDatabase( + path_workload=os.path.join(work_dir, "database_workload.json"), + path_tuning_record=os.path.join(work_dir, "database_tuning_record.json"), + ) + + workload = database.commit_workload(mod) - sch = tvm.tir.Schedule(mod) + sch = tvm.tir.Schedule(mod) - schedule_dense_dense(sch) + schedule_dense_dense(sch) - print(sch.mod.script()) + print(sch.mod.script()) - tune_rec = TuningRecord(sch.trace, [0.0], workload, tvm.target.Target(target), []) + tune_rec = TuningRecord(sch.trace, [0.0], workload, tvm.target.Target(target), []) - database.commit_tuning_record(tune_rec) + database.commit_tuning_record(tune_rec) with ApplyHistoryBest(database): with tvm.transform.PassContext( From a8ac4b04e670311f34fe318983ec85bbc19b45bf Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 1 Apr 2022 14:46:45 +0900 Subject: [PATCH 8/9] comment out schedule dump --- tests/python/unittest/test_meta_schedule_multi_anchor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/unittest/test_meta_schedule_multi_anchor.py b/tests/python/unittest/test_meta_schedule_multi_anchor.py index 9f5280b543dc..e59639170d0f 100644 --- a/tests/python/unittest/test_meta_schedule_multi_anchor.py +++ b/tests/python/unittest/test_meta_schedule_multi_anchor.py @@ -100,7 +100,7 @@ def test_dense_dense(): schedule_dense_dense(sch) - print(sch.mod.script()) + # print(sch.mod.script()) tune_rec = TuningRecord(sch.trace, [0.0], workload, tvm.target.Target(target), []) From b9b5e7cc79f66166dd9d0f89aaf3ce6d6a8a66ad Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Fri, 1 Apr 2022 15:33:25 +0900 Subject: [PATCH 9/9] Update src/relay/backend/te_compiler_cache.cc Co-authored-by: Junru Shao --- src/relay/backend/te_compiler_cache.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/relay/backend/te_compiler_cache.cc b/src/relay/backend/te_compiler_cache.cc index dc31d12118cb..cd3ce80459b8 100644 --- a/src/relay/backend/te_compiler_cache.cc +++ b/src/relay/backend/te_compiler_cache.cc @@ -337,7 +337,7 @@ class ScheduleBuilder : public ExprVisitor { schedule = Downcast(obj); } } - if (use_auto_scheduler_) { + if (use_meta_scheduler_) { IRModule relay_mod({{prim_fn_var, relay_func}}); IRModule tir_mod({{prim_fn_var, tir::CreatePrimFunc(Concat(fn_inputs, tensor_outs))}}); Optional scheduled_mod = meta_schedule::MetaScheduleContext::QueryInsideWithScope(