This PR converts a good chunk of TVM's global `PackedFunc`s to be
`TypedPackedFunc`s. Right now, this shouldn't actually change any behavior,
since the type information is immediately discarded once the functions are
registered in the `Registry`.
It does make things a lot shorter, using a few helper methods I added on
`Registry`: `set_body_simple`, `set_body_method`, and `set_body_node_method`
(names up for debate). These helpers use the properties of function pointers to
infer argument and return types:
```cpp
TVM_REGISTER_API("make.Realize")
.set_body([](TVMArgs args, TVMRetValue *ret) {
*ret = Realize::make(args[0], args[1], args[2], args[3], args[4], args[5]);
});
```
becomes ->
```
TVM_REGISTER_API("make.Realize")
.set_body_simple(Realize::make);
```
This is shorter, runtime-equivalent, and also now has type information!
I don't believe this change has much runtime cost. It might interfere with
inlining in some cases, would have to measure. Also, `PackedFunc`s aren't
intended to be all that fast anyway.
So, that's the refactoring part. Now, the RFC part:
### RFC: Inspectable PackedFunc signatures
What I really want to do with this change is prep for a later PR where we'd add
run-time type signatures to `PackedFunc`s. (That's the RFC part.) Each
`PackedFunc` would have an optional field describing what arguments it takes
and what type it returns. This field would be automatically populated during
conversion from a `TypedPackedFunc`.
Once we have these type signatures, we could use them to generate bindings in
new languages (i.e., Rust) and get coverage of a good chunk of TVM's API for
free.
This would result in much easier-to-maintain language bindings and might allow
us to eventually rip out a lot of the manually-written stuff in e.g. the Python
bindings.
The downside of this idea is that it would result in some fairly hairy codegen,
which can be difficult to maintain. It would also add small amounts of overhead
to the tvm runtime system; we could reduce that issue by adding a #define to
disable the type metadata.
A few extensions of this idea:
- Add more compile-time metadata to the Node heirarchy, allowing codegen to
access their methods / attributes.
- Add docstrings to PackedFuncs to allow auto-generation of documentation.
- Allow `std::variant` or equivalent in `TypedPackedFunc` signatures. Lots of
PackedFuncs have arguments that can be one of a few types (e.g. Int or Expr); a
simple extension to the PackedFunc system + runtime type system would allow
these to be described automatically.
You can view, comment on, or merge this pull request online at:
https://github.com/dmlc/tvm/pull/2981
-- Commit Summary --
* Add `set_body_simple` to Registry, refactor a lot of code to use it
* Add more types to Relay PackedFuncs
* Add Registry::set_body_method to easily make Node methods into
* Add set_body_method, set_body_node_method; start typing api_lang
* Add some docs, remove unused script
-- File Changes --
M include/tvm/runtime/registry.h (135)
M nnvm/src/compiler/compile_engine.cc (4)
M nnvm/src/compiler/graph_hash.cc (4)
M src/api/api_arith.cc (32)
M src/api/api_codegen.cc (5)
M src/api/api_ir.cc (156)
M src/api/api_lang.cc (267)
M src/api/api_pass.cc (99)
M src/api/api_schedule.cc (24)
M src/codegen/codegen_opencl.cc (4)
M src/codegen/codegen_opengl.cc (4)
M src/codegen/codegen_vhls.cc (4)
M src/codegen/llvm/codegen_amdgpu.cc (4)
M src/codegen/llvm/codegen_nvptx.cc (4)
M src/codegen/opt/build_cuda_on.cc (4)
M src/codegen/source_module.cc (4)
M src/codegen/spirv/build_vulkan.cc (4)
M src/codegen/stackvm/codegen_stackvm.cc (4)
M src/relay/backend/interpreter.cc (25)
M src/relay/ir/adt.cc (28)
M src/relay/ir/alpha_equal.cc (12)
M src/relay/ir/base.cc (12)
M src/relay/ir/expr.cc (56)
M src/relay/ir/expr_functor.cc (5)
M src/relay/ir/hash.cc (12)
M src/relay/ir/module.cc (53)
M src/relay/ir/type.cc (43)
M src/relay/op/debug.cc (4)
M src/relay/op/image/resize.cc (4)
M src/relay/op/nn/convolution.cc (36)
M src/relay/op/nn/nn.cc (71)
M src/relay/op/nn/pad.cc (4)
M src/relay/op/nn/pooling.cc (16)
M src/relay/op/nn/upsampling.cc (4)
M src/relay/op/tensor/reduce.cc (6)
M src/relay/op/tensor/transform.cc (104)
M src/relay/op/vision/multibox_op.cc (8)
M src/relay/op/vision/nms.cc (8)
M src/relay/op/vision/rcnn_op.cc (12)
M src/relay/op/vision/yolo.cc (4)
M src/relay/pass/canonicalize_ops.cc (4)
M src/relay/pass/combine_parallel_conv2d.cc (4)
M src/relay/pass/dead_code.cc (4)
M src/relay/pass/device_annotation.cc (12)
M src/relay/pass/fold_constant.cc (4)
M src/relay/pass/fuse_ops.cc (4)
M src/relay/pass/gradient.cc (10)
M src/relay/pass/mac_count.cc (4)
M src/relay/pass/pass_manager.cc (27)
M src/relay/pass/quantize.cc (13)
M src/relay/pass/simplify_inference.cc (4)
M src/relay/pass/to_a_normal_form.cc (4)
M src/relay/pass/to_graph_normal_form.cc (4)
M src/relay/pass/type_infer.cc (4)
M src/relay/pass/util.cc (12)
M src/relay/pass/well_formed.cc (5)
M src/runtime/cuda/cuda_module.cc (12)
M src/runtime/metal/metal_module.mm (8)
M src/runtime/opencl/aocl/aocl_module.cc (4)
M src/runtime/opencl/opencl_module.cc (12)
M src/runtime/opencl/sdaccel/sdaccel_module.cc (8)
M src/runtime/rocm/rocm_module.cc (8)
M src/runtime/rpc/rpc_event_impl.cc (6)
M src/runtime/rpc/rpc_socket_impl.cc (4)
M src/runtime/stackvm/stackvm_module.cc (4)
M src/runtime/vulkan/vulkan_module.cc (8)
M topi/src/topi.cc (266)
-- Patch Links --
https://github.com/dmlc/tvm/pull/2981.patch
https://github.com/dmlc/tvm/pull/2981.diff
--
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/pull/2981
