Implement C++ registry to back Python target.generic_func #892

alex-weaver · 2018-02-11T19:10:23Z

This is a step towards implementing NNVM issue https://github.com/dmlc/nnvm/issues/343

This PR sets up a C++ registry to back the target.generic_func mechanism in Python, which allows C++ to also register target-specialized functions (schedules and ops) in a Python compatible way. This mechanism provides compatibility in both directions: C++ can call functions registered in Python and vice versa. Python code can also override any function registered via C++ if necessary.

In the case of C++ calling Python, a target string is provided which is parsed and entered on the python side. In the case of Python calling C++, the global current target object is stringified and passed to the registered function as an extra first argument.

alex-weaver · 2018-02-12T16:31:36Z

OK I've fixed a couple of issues that came up when integrating the TOPI C++ schedules, but it should now be working. The python generic_func is still compatible with all existing usage in the python codebase.

The TOPI C++ library now registers all its schedules and dense ops through the GenericFunc registry. The Python TOPI code is configured to override all of these, so C++ clients will see the C++ schedules, and Python clients will continue to see the Python schedules as before.

tqchen · 2018-02-13T03:21:16Z

include/tvm/build_module.h

+  * these arguments.
+  * \param ret The return value
+  */
+  TVM_DLL void invoke_func(const tvm::Target& target, TVMArgs args, TVMRetValue* ret) const;


I think we would prefer hide target as a context, possibly in Thread Local storage, which can be queried. The invoke can simply be operator()().

tqchen · 2018-02-13T03:21:38Z

include/tvm/build_module.h

+  * false, an error will be logged if the call would override a previously registered function.
+  * \return reference to self.
+  */
+  TVM_DLL GenericFunc& set_generic_func(const PackedFunc value,


set_default_func

alex-weaver · 2018-02-13T11:37:03Z

Ok I've addressed those - the only thing I wasn't sure about was what the signature of operator() should look like - so I've mirrored the style of PackedFunc with invoke_packed(TVMArgs args, TVMRetValue* ret) and then a separate operator()(Args&& ...args)

alex-weaver · 2018-02-13T16:08:58Z

Also fixed 2 issues that cropped up testing integrating this upstream:

WrapSchedule in topi.cc did not correctly test whether it was passed Array<Tensor> or just Tensor.
Added an empty exported function to TOPI so that NNVM has a cross-platform way to make sure the library has been loaded, and therefore make sure the schedules have been registered.

alex-weaver · 2018-02-14T16:50:07Z

Removed the empty exported function from (2) above - that doesn't work with python. Instead NNVM can load tvm_topi.dll dynamically.

alex-weaver · 2018-02-20T23:06:05Z

@tqchen Sorry to hassle you but could this be reviewed again? I think the issues raised should all be addressed. Let me know if there's anything that needs fixing

tqchen · 2018-02-20T23:09:06Z

just get back from a conference trip, on my todo list

alex-weaver · 2018-02-20T23:17:01Z

Ah no problem, I just wondered if it had got lost somehow. No rush!

tqchen · 2018-02-22T18:17:14Z

include/tvm/build_module.h

+  * \param args The arguments to pass to the function.
+  * \param ret The return value
+  */
+  TVM_DLL void invoke_packed(TVMArgs args, TVMRetValue* ret) const;


Use CallPacked, same as PackedFunc

tqchen · 2018-02-22T18:20:53Z

Some thoughts on how we can support target context. I think the Target context query TLS should belong to tvm::Target .

Instead of using setter getter, the target can be supported via scoping backed by a stack in TLS, which looks like Target::EnterTargetScope(target); Target::ExitTargetScope().

We can have RAII style version of scoping. Like

tvm::TargetContext ctx(target);

Which set target at construction time, and call Exit at destructor

alex-weaver · 2018-02-26T18:31:52Z

Ok the python registration is updated to use a reference to a GenericFunc node instead of repeatedly using func_name

tqchen

thanks for keep improving the PR. I feel we are very close, here is another batch of comment

tqchen · 2018-03-01T04:47:03Z

src/codegen/build_module.cc

+  entry->context_stack.pop();
+}
+
+tvm::Target* Target::current_target(bool allow_null) {


why pointer instead of const reference?

This is because it needs to be able to return nullptr if the stack is empty and allow_null is true. This is used preserve the behavior of python current_target() when the allow_none argument is true.

tqchen · 2018-03-01T04:52:07Z

src/codegen/build_module.cc

@@ -344,4 +384,151 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable)
  p->stream << ")";
 });

+struct GenericFuncNode::Manager {


I think we only need one Get function, that is a member of GenericFunc

tqchen · 2018-03-01T04:52:30Z

src/codegen/build_module.cc

@@ -344,4 +384,151 @@ TVM_STATIC_IR_FUNCTOR(IRPrinter, vtable)
  p->stream << ")";
 });

+struct GenericFuncNode::Manager {
+  std::unordered_map<std::string, std::shared_ptr<GenericFuncNode>> fmap;


```>> -> > > `` (add space)

Yep I keep forgetting that one :)

tqchen · 2018-03-01T04:52:56Z

src/codegen/build_module.cc

+  return *this;
+}
+
+GenericFunc& GenericFunc::register_func(const std::vector<std::string>& tags,


rename to register, to be consistent with python API

This doesn't seem to work - register is a reserved keyword

OK, get it. sorry i forget

tqchen · 2018-03-01T04:53:07Z

src/codegen/build_module.cc

+  return GenericFunc(GenericFuncNode::Get(name));
+}
+
+GenericFunc& GenericFunc::set_default_func(const PackedFunc value,


set_default_func->set_default

tqchen · 2018-03-01T04:56:43Z

python/tvm/target.py

+            The register function is necessary.
+            """
+            def _do_reg(myf):
+                key_list = [key] if isinstance(key, str) else key


move this logic to GenericFunc.register

tqchen · 2018-03-01T05:04:23Z

python/tvm/target.py

+        _api_internal._GenericFuncRegisterFunc(self, func, key_list, allow_override)
+
+
+def generic_func(name=None, override=False):
    """Wrap a target generic function.


This logic seems to be overly complicated, due to the fact that now it has two functionalities:

Construct a generic function object

Retrieve a generic function from C++ by its name.

I think we want to simplify the case by only support one use-case. Keep generic_func as its original use-case (define a generic func), except that this time it is backed by C++.

Like get_global_func and register_func, we can add functions get_generic_func and register_generic_func to support fetch and store global registry

tqchen · 2018-03-01T05:04:58Z

python/tvm/target.py

+    def set_default(self, func, allow_override):
+        _api_internal._GenericFuncSetDefault(self, func, allow_override)
+
+    def register_func(self, func, key_list, allow_override):


make API consistent with old generic_func both in c++ and python

register

set_default

tqchen

One last batch of changes requested, let us move the target to Node system as well so we don't have to exchange the strings, only pointers are needed. It also makes current_target API more natural

tqchen · 2018-03-06T02:30:54Z

topi/src/topi.cc

+    Tensor data = args[0];
+    Tensor weight = args[1];
+    Tensor bias_val;
+    Tensor *bias;


Tensor can be None, use bias.defined() to check if it is none

do not use pointer here

tqchen · 2018-03-06T02:32:01Z

src/codegen/build_module.cc

+  *ret = GenericFunc(std::make_shared<GenericFuncNode>());
+  });
+
+TVM_REGISTER_API("_GenericFuncAddToRegistry")


GenericFuncRegisterGlobal

tqchen · 2018-03-06T02:32:12Z

src/codegen/build_module.cc

+  GenericFunc::RegisterGenericFunc(func, func_name);
+  });
+
+TVM_REGISTER_API("_GenericFuncGet")


_GenericFuncGetGlobal

tqchen · 2018-03-06T02:34:58Z

include/tvm/build_module.h

@@ -28,7 +29,7 @@ struct Target {
  /*! \brief The warp size that should be used by the LowerThreadAllreduce pass */
  int thread_warp_size = 1;
  /*! \brief Keys for this target */
-  std::unordered_set<std::string> keys;
+  std::vector<std::string> keys;
  /*! \brief Options for this target */


I think it might make sense to also bring Target to the Node system, so current Target can simply return a None which can be checked with target.defined()

tqchen · 2018-03-10T04:08:11Z

Need rebase from master and updates according to the last comments

…ters

…ble. Changed tvm::GenericFunc to mirror tvm::runtime::Registry as dmlc::Registry fails when used across .dll boundaries

…ent to invoke, and removed target as injected first argument to callbacks. Target is now a thread local context. Renamed invoke_func -> invoke_packed, and added operator() mirroring PackedFunc.

Library can be loaded dynamically by NNVM This reverts commit bc8f3d7.

…or cross-plat way to ensure TOPI is loaded and schedules/ops are registered

tqchen · 2018-03-16T02:46:46Z

@alex-weaver can you make the change to bring back python's generic_func? Thanks!

alex-weaver · 2018-03-16T12:18:25Z

@tqchen yep I'll sort that. I was thinking the C++ GenericFunc could be renamed NativeGenericFunc, and this could be used for schedules across both C++ and Python. That would allow the C++ schedules to be overridden by python if NNVM is changed to use the C++ ones as a base.

tqchen · 2018-03-16T16:27:45Z

I would say just keep the name GenericFunc for now in c++ side, but we could name it NativeGenericFunc in python side.

…hreading_backend

alex-weaver · 2018-03-16T23:38:33Z

@tqchen the previous python generic_func has been restored, and the python attribute for C++ GenericFunc has been changed to override_native_generic_func - this is used to provide python overrides for the schedules (+dense op) provided by C++. Now, if NNVM registers the C++ schedules, they will be overridden with the python versions for python clients of NNVM or TOPI.

tqchen

some final changes to be made

tqchen · 2018-03-17T02:31:26Z

python/tvm/target.py

+    """
+    return _api_internal._GenericFuncGetGlobal(name)
+
+def register_native_generic_func(func, name):


This is not necessary, as we can use get_native_generic_func(name).register. remove the correspond PackedFunc

tqchen · 2018-03-17T02:32:56Z

topi/src/topi.cc

+    Tensor data = args[0];
+    Tensor weight = args[1];
+    Tensor bias_val;
+    Tensor *bias;


do not use pointer here

tqchen · 2018-03-17T02:35:20Z

python/tvm/target.py

+
+    @property
+    def keys(self):
+        return [k.value for k in self.keys_array]


do a cache so it does not fetch each time,

if self._keys is None: self._keys = tuple(k.value for k in self.keys_array) return self._keys

tqchen · 2018-03-17T02:36:26Z

python/tvm/target.py

-
-    def __repr__(self):
-        return self.__str__()
+


add

def __init__(self, handle): super(Target, self).__init__(handle) self._keys = None ...

tqchen · 2018-03-17T02:37:20Z

python/tvm/target.py

+    def keys(self):
+        return [k.value for k in self.keys_array]
+
+    @property


add unit test cases to tests/python/unittest/test_target.py to test

parsing from string

assert each field is correct

tqchen · 2018-03-17T02:38:16Z

include/tvm/build_module.h

@@ -13,76 +13,141 @@
 #include "./tvm/lowered_func.h"


use relative path, instead of absolute one, which means we should do include "./lowered_func.h"

tqchen · 2018-03-17T17:01:44Z

I have made some final comments, please fix these

…e ops; cache properties in python Target; add test case for parsing Target from string

tqchen · 2018-03-17T19:23:26Z

topi/src/topi.cc

@@ -557,7 +526,7 @@ TVM_REGISTER_GENERIC_FUNC(dense)
 .set_default(WrapDenseOp([](const Target& target,
                                 const tvm::Tensor& data,
                                 const tvm::Tensor& weight,
-                                 tvm::Tensor* bias) {
+                                 tvm::Tensor bias) {


const tvm::Tensor&

tqchen · 2018-03-17T19:23:33Z

topi/src/topi.cc

@@ -524,7 +500,7 @@ TVM_REGISTER_GENERIC_FUNC(schedule_binary_dense)
 using FTVMDenseOpBuilder = std::function<tvm::Tensor(const Target& target,
                                                     const tvm::Tensor& data,
                                                     const tvm::Tensor& weight,
-                                                     tvm::Tensor* bias)>;
+                                                     tvm::Tensor bias)>;


const tvm::Tensor&

tqchen · 2018-03-17T19:24:43Z

topi/include/topi/rocm/dense.h

  }

  auto batch = data->shape[0];
  auto in_dim = data->shape[1];
  auto out_dim = weight->shape[0];

-  if (target.libs.count("rocblas") > 0) {
+  if (target->libs().count("rocblas") > 0) {


remove >0, simply use count(condition)

tqchen · 2018-03-17T19:25:07Z

topi/src/topi.cc

+
+TVM_REGISTER_GENERIC_FUNC(dense)
+.set_default(WrapDenseOp([](const Target& target,
+                                 const tvm::Tensor& data,


alignment, align to const Target& target,

tqchen · 2018-03-17T19:25:19Z

topi/src/topi.cc

+.set_default(WrapDenseOp([](const Target& target,
+                                 const tvm::Tensor& data,
+                                 const tvm::Tensor& weight,
+                                 tvm::Tensor bias) {


const tvm::Tensor& bias

tqchen · 2018-03-17T19:25:23Z

topi/src/topi.cc

+using FTVMDenseOpBuilder = std::function<tvm::Tensor(const Target& target,
+                                                     const tvm::Tensor& data,
+                                                     const tvm::Tensor& weight,
+                                                     tvm::Tensor bias)>;


const tvm::Tensor& bias

tqchen · 2018-03-19T07:17:45Z

Thanks for all the hard work in improving this PR, this is now merged!

alex-weaver · 2018-03-19T11:37:26Z

Excellent! Only a few more to go for a C++ NNVM build function ;)

tqchen requested changes Feb 13, 2018

View reviewed changes

tqchen reviewed Feb 22, 2018

View reviewed changes

tqchen mentioned this pull request Feb 22, 2018

[RFC] Target and Generic Function System on C++ #923

Closed

tqchen requested changes Mar 1, 2018

View reviewed changes

tqchen requested changes Mar 6, 2018

View reviewed changes

tqchen added the status: review in progress label Mar 7, 2018

alex-weaver added 14 commits March 12, 2018 15:25

Implement C++ registry to back target.generic_func

dba2cf6

Fix lint

c04f0b9

Fix pylint

b04180b

Fix pylint

889847c

Fix test error - changed target keys to std::vector because order mat…

37d6342

…ters

Fix lint

9eab0c2

Fixed issue with generic_func decorator where parameters were not usa…

0ddd292

…ble. Changed tvm::GenericFunc to mirror tvm::runtime::Registry as dmlc::Registry fails when used across .dll boundaries

Fix doc string

fabb071

Renamed set_generic_func -> set_default_func. Removed target as argum…

9af4fb4

…ent to invoke, and removed target as injected first argument to callbacks. Target is now a thread local context. Renamed invoke_func -> invoke_packed, and added operator() mirroring PackedFunc.

Fix issue where VC++ did not generate a .lib for tvm_topi

ac79207

Revert "Fix issue where VC++ did not generate a .lib for tvm_topi"

21e5957

Library can be loaded dynamically by NNVM This reverts commit bc8f3d7.

Fix incorrect test for array in WrapSchedule. Add exported function f…

97cd2fa

…or cross-plat way to ensure TOPI is loaded and schedules/ops are registered

Fix lint

17fff5a

Fix lint

75fe959

alex-weaver added 10 commits March 16, 2018 17:32

Fix compile warning. Restore pure python generic_func

09d6f4d

Ensure C++ schedules are registered first

36ad1d8

Set python schedules/dense op to override C++ versions

c90469b

Fix lint

36b3e9c

Fix docstrings

5fe2a1d

Change build() to remove use of Target pointer for target_host

d163cac

Fix lint

c780227

Merge branch 'master' into master

73b55ed

Fix device_type build error. Also fix missing algorithm header from t…

e6deb23

…hreading_backend

Fix unit test

52bd519

tqchen requested changes Mar 17, 2018

View reviewed changes

tqchen reviewed Mar 17, 2018

View reviewed changes

Removed register_native_generic_func; Removed bias pointers from dens…

73becd7

…e ops; cache properties in python Target; add test case for parsing Target from string

tqchen requested changes Mar 17, 2018

View reviewed changes

Style fixes

89046dd

tqchen approved these changes Mar 17, 2018

View reviewed changes

tqchen added status: accepted and removed status: review in progress labels Mar 17, 2018

Merge branch 'master' into master

72819d7

tqchen merged commit 4afc2f9 into apache:master Mar 19, 2018

merrymercy mentioned this pull request Mar 31, 2018

support string option when creating cuda/rocm/... target #1071

Merged

tqchen pushed a commit to tqchen/tvm that referenced this pull request Jul 6, 2018

Implement C++ registry to back Python target.generic_func (apache#892)

6292204

grwlf pushed a commit to grwlf/tvm that referenced this pull request Aug 8, 2018

Implement C++ registry to back Python target.generic_func (apache#892)

61f5435

Implement C++ registry to back Python target.generic_func #892

Implement C++ registry to back Python target.generic_func #892

Conversation

alex-weaver commented Feb 11, 2018

alex-weaver commented Feb 12, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

alex-weaver commented Feb 13, 2018 • edited

alex-weaver commented Feb 13, 2018 • edited

alex-weaver commented Feb 14, 2018

alex-weaver commented Feb 20, 2018

tqchen commented Feb 20, 2018

alex-weaver commented Feb 20, 2018

Choose a reason for hiding this comment

tqchen commented Feb 22, 2018 • edited

alex-weaver commented Feb 26, 2018

tqchen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tqchen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tqchen commented Mar 10, 2018

tqchen commented Mar 16, 2018

alex-weaver commented Mar 16, 2018

tqchen commented Mar 16, 2018

alex-weaver commented Mar 16, 2018

tqchen left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tqchen commented Mar 17, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tqchen commented Mar 19, 2018

alex-weaver commented Mar 19, 2018

alex-weaver commented Feb 13, 2018 •

edited

alex-weaver commented Feb 13, 2018 •

edited

tqchen commented Feb 22, 2018 •

edited