From a2c01ac6b660e3dbb21ac6c8e32cd1cfc4f8183f Mon Sep 17 00:00:00 2001 From: tqchen Date: Mon, 20 Feb 2023 11:16:47 -0500 Subject: [PATCH] [BugFix][UMA] Protect target registration (#13624) This PR address fixes for UMA target registration. * Fix the doc issue #13304 * Continues stalled PR #12731 Changes: * Incorporates all proposed fixes from mentioned [PR #12731](https://github.com/apache/tvm/pull/12731) * Address test case concerns and discussions from [PR #12731](https://github.com/apache/tvm/pull/12731) * **NEW:** Already exiting target cannot be created, explicit error on this. * **NEW:** Attributes having special/reserved scope cannot be created explicitly. It also address proper test cases for all the above. Co-authored-by: Balint Cristian --- gallery/tutorial/uma.py | 2 +- .../tvm/relay/backend/contrib/uma/backend.py | 7 +++--- src/relay/backend/contrib/uma/targets.cc | 24 +++++++++++------- tests/python/contrib/test_uma/test_target.py | 25 ++++++++++++++++--- 4 files changed, 42 insertions(+), 16 deletions(-) diff --git a/gallery/tutorial/uma.py b/gallery/tutorial/uma.py index ed4fc4cf805cf..ea38813a7acef 100644 --- a/gallery/tutorial/uma.py +++ b/gallery/tutorial/uma.py @@ -57,7 +57,7 @@ # ###################################################################### -# .. image:: https://raw.githubusercontent.com/apache/tvm-site/main/images/tutorial/uma_vanilla_block_diagram.png +# .. image:: https://raw.githubusercontent.com/tlc-pack/web-data/main/images/tutorial/uma_vanilla_block_diagram.png # :width: 100% # :alt: A block diagram of Vanilla # diff --git a/python/tvm/relay/backend/contrib/uma/backend.py b/python/tvm/relay/backend/contrib/uma/backend.py index 40ec06e45367a..550109f1700d3 100644 --- a/python/tvm/relay/backend/contrib/uma/backend.py +++ b/python/tvm/relay/backend/contrib/uma/backend.py @@ -278,11 +278,12 @@ def register(self) -> None: """ registration_func = tvm.get_global_func("relay.backend.contrib.uma.RegisterTarget") - for name, attr in self._target_attrs: + for name, attr in self._target_attrs.items(): if attr is None: raise ValueError("Target attribute None is not supported.") - - if registration_func(self.target_name, self._target_attrs): + # skip if target is already registered + if self.target_name not in tvm.target.Target.list_kinds(): + registration_func(self.target_name, self._target_attrs) self._relay_to_relay.register() self._relay_to_tir.register() self._tir_to_runtime.register() diff --git a/src/relay/backend/contrib/uma/targets.cc b/src/relay/backend/contrib/uma/targets.cc index ed2cc047cf2f4..e2fe644cb9bf0 100644 --- a/src/relay/backend/contrib/uma/targets.cc +++ b/src/relay/backend/contrib/uma/targets.cc @@ -31,7 +31,7 @@ namespace tvm { namespace relay { namespace contrib { namespace uma { -tvm::transform::Pass RelayToTIR(String target_name); +transform::Pass RelayToTIR(String target_name); runtime::Module TIRToRuntime(IRModule mod, Target target); } // namespace uma } // namespace contrib @@ -39,16 +39,15 @@ runtime::Module TIRToRuntime(IRModule mod, Target target); TVM_REGISTER_GLOBAL("relay.backend.contrib.uma.RegisterTarget") .set_body_typed([](String target_name, Map attr_options) -> bool { - // @todo(cgerum): We probably should get rid of target.register rather sooner than later - // And use a proper registry for uma backends - for (const String registered_target_name : ::tvm::TargetKindRegEntry::ListTargetKinds()) { + // create only new target and init only once + for (const String registered_target_name : TargetKindRegEntry::ListTargetKinds()) { if (registered_target_name == target_name) { - return false; + LOG(FATAL) << "TVM UMA Error: Target is already registered: " << target_name; } } auto target_kind = - ::tvm::TargetKindRegEntry::RegisterOrGet(target_name) + TargetKindRegEntry::RegisterOrGet(target_name) .set_name() .set_default_device_type(kDLCPU) .add_attr_option>("keys") @@ -58,20 +57,27 @@ TVM_REGISTER_GLOBAL("relay.backend.contrib.uma.RegisterTarget") .add_attr_option>("libs") .add_attr_option("host") .add_attr_option("from_device") - .set_attr(tvm::attr::kRelayToTIR, + .set_attr(attr::kRelayToTIR, relay::contrib::uma::RelayToTIR(target_name)) .set_attr("TIRToRuntime", relay::contrib::uma::TIRToRuntime); + // target kind attrs inventory + auto kind = TargetKind::Get(target_name).value(); + auto list_attrs = TargetKindRegEntry::ListTargetKindOptions(kind); + for (auto& attr_option : attr_options) { auto option_name = attr_option.first; auto default_value = attr_option.second; + if (list_attrs.find(option_name) != list_attrs.end()) { + LOG(FATAL) << "TVM UMA Error: Attribute is already registered: " << option_name; + } if (default_value->IsInstance()) { target_kind.add_attr_option(option_name, Downcast(default_value)); } else if (default_value->IsInstance()) { target_kind.add_attr_option(option_name, Downcast(default_value)); } else { - LOG(FATAL) << "Only String, Integer, or Bool are supported. Given attribute option type: " - << attr_option.second->GetTypeKey(); + LOG(FATAL) << "TypeError: Only String, Integer, or Bool are supported. " + << "Given attribute option type: " << attr_option.second->GetTypeKey(); } } return true; diff --git a/tests/python/contrib/test_uma/test_target.py b/tests/python/contrib/test_uma/test_target.py index 558c4e5182304..1662becf088dc 100644 --- a/tests/python/contrib/test_uma/test_target.py +++ b/tests/python/contrib/test_uma/test_target.py @@ -63,23 +63,42 @@ def test_uma_target(target_name, target_attrs, target_args): [ ("float_attr", 3.14), ("none_attr", None), + ("model", "my_model"), ], ) def test_invalid_attr_option(attr_name: str, target_attr: Union[str, int, bool, float, None]): + registration_func = tvm.get_global_func("relay.backend.contrib.uma.RegisterTarget") if target_attr is None: # None cannot be caught as TVMError, as it causes a SIGKILL, therefore it must be prevented to be # entered into relay.backend.contrib.uma.RegisterTarget at Python level. - with pytest.raises(ValueError): + with pytest.raises(ValueError, match=r"Target attribute None is not supported."): uma_backend = VanillaAcceleratorBackend() uma_backend._target_attrs = {attr_name: target_attr} uma_backend.register() + elif "model" in attr_name: + target_name = f"{attr_name}_{target_attr}" + target_attr = {attr_name: target_attr} + with pytest.raises(tvm.TVMError, match=r"Attribute is already registered: .*"): + registration_func(target_name, target_attr) else: - registration_func = tvm.get_global_func("relay.backend.contrib.uma.RegisterTarget") target_name = f"{attr_name}_{target_attr}" target_attr = {attr_name: target_attr} - with pytest.raises(tvm.TVMError, match=r"Only String, Integer, or Bool are supported. .*"): + with pytest.raises(TypeError, match=r"Only String, Integer, or Bool are supported. .*"): registration_func(target_name, target_attr) +@pytest.mark.parametrize( + "target_name", + [ + "llvm", + "c", + ], +) +def test_target_duplication(target_name: str): + with pytest.raises(tvm.TVMError, match=r"TVM UMA Error: Target is already registered: .*"): + registration_func = tvm.get_global_func("relay.backend.contrib.uma.RegisterTarget") + registration_func(target_name, {}) + + if __name__ == "__main__": tvm.testing.main()