From 4fdf1d157e30bfd36ae60d7be03cabba218cf3de Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 30 May 2023 09:14:35 -0500 Subject: [PATCH 1/3] [Bugfix][TIR][VTA] Update host-side target, even without device func This resolves an issue introduced by the combination of https://github.com/apache/tvm/pull/14918 and https://github.com/apache/tvm/pull/14945. The bug occurred for targets that do not require device-side codegen, but do require a `device_type` other than `kDLCPU`. It wasn't caught by CI, as the issue only occurred with the combination of both PRs. 1. #14918 updated `SplitHostDevice` to only modify the `"target"` attribute when a device-side function has been extracted. 2. For VTA, there is no device-side function, as everything is done through host-side API calls. 3. From (1) and (2), the VTA examples kept the target `T.target("ext_dev", host="llvm")` after the `SplitHostDevice` pass, instead of being updated to `T.target("llvm")`. 4. #14945 restricted CombineContextCall to only apply to host-side passes. 5. From (4) and (5), the `CombineContextCall` pass was no longer applied to the VTA context calls. This PR fixes `SplitHostDevice`, updating the target from `T.target("ext_dev", host="llvm")` to `T.target("llvm")`, even if no device sections have been extracted from the function. --- src/tir/transforms/split_host_device.cc | 10 +++++----- .../test_tir_transform_split_host_device.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/src/tir/transforms/split_host_device.cc b/src/tir/transforms/split_host_device.cc index 9270b356ba22..2de831e8ad0c 100644 --- a/src/tir/transforms/split_host_device.cc +++ b/src/tir/transforms/split_host_device.cc @@ -108,12 +108,12 @@ PrimFunc SplitHostDevice(PrimFunc func, IRModule* device_mod, const GlobalVar& g HostDeviceSplitter splitter(device_mod, name_prefix); - auto body = splitter(func->body); - - if (!body.same_as(func->body)) { + if (auto body = splitter(func->body); !body.same_as(func->body)) { func.CopyOnWrite()->body = body; - auto target_host = target->GetHost().value_or(Target("llvm")); - func = WithAttr(std::move(func), tvm::attr::kTarget, target_host); + } + + if (auto target_host = target->GetHost()) { + func = WithAttr(std::move(func), tvm::attr::kTarget, target_host.value()); } return func; diff --git a/tests/python/unittest/test_tir_transform_split_host_device.py b/tests/python/unittest/test_tir_transform_split_host_device.py index cf866ae005c8..1599b9a031a0 100644 --- a/tests/python/unittest/test_tir_transform_split_host_device.py +++ b/tests/python/unittest/test_tir_transform_split_host_device.py @@ -168,5 +168,21 @@ def main_kernel(n: T.int32): return mod +class TestSplitHostDevice(BaseCompare): + """Like TestSplitHostDevice, but no device regions to extract + + Even if there are no device regions, the host-side function should + still have its "target" attribute updated. + """ + + def before(): + T.func_attr({"target": T.target("ext_dev", host="llvm")}) + T.evaluate(0) + + def expected(): + T.func_attr({"target": T.target("llvm")}) + T.evaluate(0) + + if __name__ == "__main__": tvm.testing.main() From f80b6b34a567172fb72c67c2812298952b945706 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 16 May 2023 16:19:10 -0500 Subject: [PATCH 2/3] [CodegenC] Updated unit test for sorted CodegenC output Previously, this unit test generated a `Map` whose default iteration order was not sorted by function name, built the `Map` of modules, then validated whether the resulting C code was a sorted list of 4 elements. However, this condition was stricter than necessary, as it depended on the number of items added to the `Map` until it was unsorted. This commit updates the test to instead validate that `std::is_sorted` returns true. --- tests/cpp/c_codegen_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cpp/c_codegen_test.cc b/tests/cpp/c_codegen_test.cc index e764d21505d4..bceac172e864 100644 --- a/tests/cpp/c_codegen_test.cc +++ b/tests/cpp/c_codegen_test.cc @@ -121,5 +121,5 @@ TEST(CCodegen, FunctionOrder) { auto module = build(inputs, Target()); Array func_array = module->GetFunction("get_func_names", false)(); std::vector functions{func_array.begin(), func_array.end()}; - EXPECT_THAT(functions, ElementsAre(StrEq("op_1"), _, StrEq("op_2"), _)); + EXPECT_TRUE(std::is_sorted(functions.begin(), functions.end())); } From 39d327e2fdc02a05ed42d71d1900930dbaaedb70 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 30 May 2023 10:43:30 -0500 Subject: [PATCH 3/3] Ignore __tvm_main__ in unit test --- tests/cpp/c_codegen_test.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/cpp/c_codegen_test.cc b/tests/cpp/c_codegen_test.cc index bceac172e864..a01921239a9f 100644 --- a/tests/cpp/c_codegen_test.cc +++ b/tests/cpp/c_codegen_test.cc @@ -121,5 +121,11 @@ TEST(CCodegen, FunctionOrder) { auto module = build(inputs, Target()); Array func_array = module->GetFunction("get_func_names", false)(); std::vector functions{func_array.begin(), func_array.end()}; + // The entry point is handled separately from the other functions. + functions.erase(std::remove_if(functions.begin(), functions.end(), + [](const std::string& name) { + return name == tvm::runtime::symbol::tvm_module_main; + }), + functions.end()); EXPECT_TRUE(std::is_sorted(functions.begin(), functions.end())); }