More Strongly-Typed Rust Bindings #2559

kazimuth · 2019-02-04T05:32:12Z

After reading #2523 I was messing with the newly-merged Rust API and noticed some some changes that might be helpful.
The main change this PR makes is to start adding some "typed" rust APIs, which look like normal rust functions / structs, and call tvm PackedFuncs under the hood. The rust graph runtime example is now much simpler:

    let image: ndarray::Array = /* ... load image in Rust ... */;

    let input = tvm_frontend::NDArray::from_rust_ndarray(
        &image,
        TVMContext::cpu(0),
        TVMType::from("float32")
    )?;

    let mut module = GraphModule::from_paths(
        concat!(env!("CARGO_MANIFEST_DIR"), "/deploy_graph.json"),
        concat!(env!("CARGO_MANIFEST_DIR"), "/deploy_lib.so"),
        concat!(env!("CARGO_MANIFEST_DIR"), "/deploy_param.params"),
        ctx
    )?;

    let mut result: Vec<tvm_frontend::NDArray> = module.apply(&[&input]);

Of course, all the old PackedFunc functionality is still there. This just makes well-defined APIs a little easier to use.

This approach could be used for binding the non-runtime APIs as well, I think. (Which I want, because I kinda wanna make a tvm-backed rust deep learning framework at some point... that's the far future though ;))

I've also fixed up a few other random things I've noticed:

I think the old Function::get implementation would break for freshly registered globals, so I made it more dynamic
I allowed kNDArrayContainers to be returned from PackedFuncs without blowing up... probably. The example works, anyway. I don't entirely understand what's going on with the stuff in common::value though.

TODO:

more docs
more tests
run rustfmt

Open questions:

What other APIs would be good to wrap with types?

CC @ehsanmok , hope it's okay that i'm poking around in your code :)

ehsanmok · 2019-02-04T06:19:14Z

Thank you for testing out and making nice suggestions! any suggestions is more than welcome of course :)
I agree! GraphModule and GraphRuntime are more typed and ergonomic and we should definitely have them. Glad you've mentioned them. For example, you can also see the java equivalent in contrib.

Fyi @nhynes has been working on some more ergonomic and enhancements for common and few other things as part of his RFC that you linked to.

What other APIs would be good to wrap with types?

So far I've kept the compatibility with other frontends (python, java and golang) in mind, but that doesn't mean we can't add more. Off the top of my head, when #2498 is merged it could open up some good opportunities in Rust.

nhynes · 2019-02-04T07:02:28Z

Hey @kazimuth thanks a lot for the PR! This is really cool! It looks like you've made a lot of overlapping changes with my branch which unifies frontend and runtime (as per #2523). Would you mind if I asked you to postpone further work until I get that fully tested and merged? I'll prioritize it for this week. I can't stop you from contributing, but it will save us both a lot of work if we collaborated on the unified tvm crate (e.g., I already (also) fixed the globals dict). If it makes you more willing to wait, I'll apply the semantics of your commit to the updated branch! What do you think?

nhynes · 2019-02-04T07:12:23Z

rust/frontend/README.md

+    let mut module = GraphModule::from_paths(
+        concat!(env!("CARGO_MANIFEST_DIR"), "/deploy_graph.json"),
+        concat!(env!("CARGO_MANIFEST_DIR"), "/deploy_lib.so"),
+        concat!(env!("CARGO_MANIFEST_DIR"), "/deploy_param.params"),


note that the established pattern is to load the params independently of the graph lib+json since one might want to load multiple params sets (possibly over RPC, too). The Rust runtime crate already implements this pattern, so we should try to use that GraphModule but substitute in @ehsanmok's bindings when requested.

Makes sense, I was just aiming to minimize lines of example code.

I wasn't sure about providing from_paths at all because it won't work in e.g. wasm. Might make sense to put behind a #[cfg].

nhynes · 2019-02-04T07:16:36Z

rust/frontend/src/function.rs

-                .collect(),
-        )
+    static ref GLOBAL_FUNCTIONS: Mutex<HashMap<String, Function>> = {
+        Mutex::new(HashMap::new())


Can you provide a bit of motivation for not pre-allocating the map with the already known and presumed fixed set of globals?

I believe TVMFuncRegisterGlobal is meant for adding new global functions at runtime, which means that the set of names may change?

Yes, they could change, but functions can also be removed which makes caching at all ill advised. The de-facto semantics of TVM globals is static registration. Always re-wrapping pointers might be an option but I'm not sure if the performance cost would be worth the (unused) additional flexibility.

On the bright side, part of the reason for the rust runtime has been to remove these pesky globals. As we further expand the rust functionality, we won't have to think as much about C++ interop.

Once I finish up unifying the tvm crate and merging in your PR, a good next area of focus might be on creating device apis to enable use cases like GPUs (let's start with AMD because you seem to have a real use case) and VTA.

nhynes · 2019-02-04T07:24:40Z

rust/frontend/src/function.rs

+            &mut handle as *mut _
+        ));
+        if handle.is_null() {
+            bail!(ErrorKind::FunctionNotFound(name.to_string()));


just FYI can use ensure instead of if ... { bail!() }

nhynes · 2019-02-04T07:27:10Z

rust/frontend/src/function.rs

+/// 
+/// Note that only the last segment of the function path is used for the function name.
+#[macro_export]
+macro_rules! wrap_packed_globals {


I'd probably go as far as making this an actual Fn type (c.f. runtime::PackedFunc)

nhynes · 2019-02-04T07:28:18Z

rust/frontend/src/graph_runtime.rs

+        let get_output_ = module.get_function("get_output", false).expect("no get_output, invalid GraphModule");
+        let get_num_outputs_ = module.get_function("get_num_outputs", false).expect("no get_num_outputs, invalid GraphModule");
+        let run_ = module.get_function("run", false).expect("no run_, invalid GraphModule");
+        let load_params_ = module.get_function("load_params", false).expect("no load_params_, invalid GraphModule");


ideally we should have exactly one GraphModule which either uses a Rust Module implementation or shells out to the bindings. Code duplication is decidedly less than maintainable.

My understanding was that Module is a type wrapping a dynamic library that could have any number of named functions, e.g. one created with tvm.contrib.cc; whereas GraphModule is a particular API that some Modules expose. So it made sense to make GraphModule a simple typed Module wrapper.

kind of. there is really no graph api inside of the module: it's purely defined by the graph.json. For instance, a module with dense1, batch_norm, and softmax could create a graph with {dense1, `softmax} just as easily as with all three. Basically the Module is a repo of functions and the graph.json determines how those are actually called.

nhynes · 2019-02-04T07:28:51Z

rust/frontend/src/graph_runtime.rs

+        Ok(GraphModule { module, set_input_, get_input_, get_output_, get_num_outputs_, run_, load_params_ })
+    }
+
+    // TODO: should these functions return Result?


yes. for instance, loading params can fail if the user provides a bogus params file.

nhynes · 2019-02-04T07:29:41Z

rust/frontend/src/lib.rs

 pub mod bytearray;
 pub mod context;
+#[allow(deprecated)]
 pub mod errors;


we're moving to failure (@see #2523)

nhynes · 2019-02-04T07:30:44Z

rust/frontend/README.md

+    let image: ndarray::Array = /* ... load image in Rust ... */;
+
+    // convert to a tvm NDArray
+    // could be on a GPU, FPGA, ...


mind if I ask which exactly you're trying to use? ;)

my dream is to run deep learning stuff on my amd-gpu desktop, which I've generally found to be nigh-impossible without tvm ;-;

nhynes · 2019-02-04T07:46:43Z

also, rustfmt-on-save using rust.vim is 👌

kazimuth · 2019-02-04T13:49:06Z

Thanks for the advice @nhynes, I'll delay until you're done with that change -- didn't see your PR :) and no rush, I was just poking at things and figured I could chip in.

Are you planning on unifying frontend::Function and common::PackedFunc? They seem to do the same thing

kazimuth · 2019-02-15T21:30:53Z

Hey @nhynes, any update on your PR? No rush, just curious

nhynes · 2019-02-15T22:31:22Z

Hey @kazimuth super sorry for the long delay. I was planning on taking a TVM day tomorrow and Sunday. I don't want to block your contribution since it contains a number of useful improvements, so if I don't have the merge done by Monday, I'll merge this and rebase. Thanks!

ehsanmok · 2019-04-03T19:08:21Z

@kazimuth Would you want to still work on adding GraphModule ? it's a nice feature to have.

nhynes · 2019-04-03T20:52:19Z

I think that what we actually want is a GraphExecutor for the frontend. We've already established that there's no such thing as a "GraphModule" per se, but rather a module that exports functions that can be assembled into an executor when provided along with a graph.json.

@kazimuth if you're still interested in working on improving the frontend graph runtime, we should discuss the design of this component. It'd be great if we could start using conditional compilation to make a single runtime that uses either bindings or Rust.

ehsanmok · 2019-04-03T21:09:10Z

@nhynes In case of frontend, I think GraphModule would still be useful and we should be compatible with other frontends such as java which has it. At least it makes writing the application example more succinct.

nhynes · 2019-04-03T21:18:51Z

I'm not disagreeing with the utility, but the name leaves to be desired since it's decisively not a module in the TVM sense. Regardless of nomenclature, we only need one thing that provides the graph runtime/module/executor abstraction.

kazimuth · 2019-04-04T21:54:57Z

If you can give me a broad-strokes design (like some type names + general APIs) I'd be happy to implement it.

It'd be great if we could start using conditional compilation to make a single runtime that uses either bindings or Rust.

I'm not sure I understand what this means, could you elaborate?

nhynes · 2019-04-05T08:05:19Z

So from a high level I mean something like

impl GraphExecule { // TODO: name
	fn set_input(&mut self, input: Tensor) {
        
		#[cfg(feature = "bindings")]
		{
            module.get_function("set_input", ...);
		}
		#[cfg(not(feature = "bindings")]
        {
            // pure rust impl
        }
	}

	// ...
}

kazimuth · 2019-04-05T14:34:41Z

ah, i guess i'm still a little confused about the distinction between bindings and pure-rust; i thought that distinction had been eliminated. Is pure-rust just for, like, the wasm and sgx runtimes, or is it going to have a bigger application area?

Also, looking at common it seems like the bindings feature doesn't change much of the API - what is the difference between feature = "bindings" code and not(feature = "bindings") code?

Side note, I think I'm going to close this PR and start a new one for the new changes.

nhynes · 2019-04-05T16:35:01Z

distinction had been eliminated

It's a constant WIP. The two crates now share common types, but I suspect that there's a decent amount of duplicated logic. I'd also like for users to be able to do tvm::GraphModule instead of tvm::frontime::GraphSomething. The end goal, of course, would be load the dylib using Rust libloading and never call C++.

Actually, if you're feeling ambitious, you could actually go and implement a shared object module in tvm_runtime and update the GraphExecutor UI. I guess it also depends on whether you actually want to use the C++ impl.

Also, I remember you saying that you wanted an easy-to-use Rust ML library. I had a neat idea that you might like: an option to build the entire GraphExecutor at compile-time from a pre-defined TVM module and graph json. Like there'll be a macro that generates all of the extern fns, closures calling them, and allocations statically. These can only be added to the Rust runtime, though.

kazimuth · 2019-04-05T17:49:25Z

The end goal, of course, would be load the dylib using Rust libloading and never call C++.

Is it possible to call a PackedFunc from a deployed library without using C++? Don't you need to go through the C runtime API to call any PackedFunc? Or is that actually implemented somewhere in the rust codebase?

Also, I remember you saying that you wanted an easy-to-use Rust ML library. I had a neat idea that you might like: an option to build the entire GraphExecutor at compile-time from a pre-defined TVM module and graph json. Like there'll be a macro that generates all of the extern fns, closures calling them, and allocations statically. These can only be added to the Rust runtime, though.

Oh that is an interesting idea. It would be easier if GraphExecutor was already fully-rust too.

Question: Is the current deploy_graph.json / deploy_lib.so format stable, or is it going to be replaced by a new format to support Relay? I was under the impression it was based on the old NNVM graphs, not sure it's worth the effort to port if it's going to be immediately deprecated.

nhynes · 2019-04-05T23:25:14Z

Is it possible to call a PackedFunc from a deployed library without using C++?

Yes, if there were a way to load the functions from a dynamic library. The standalone runtime already does this for statically linked module libraries.

GraphExecutor was already fully-rust too.

Well, strictly speaking, the GraphExecutor only exist in the standalone (Rust) runtime :)

Is the current deploy_graph.json / deploy_lib.so format stable

Yes, Relay produces the same output as NNVM. At some point, we might need to add support for the more interesting control flow operators, though. That's still WIP on the C++ side AFAIK.

kazimuth · 2019-04-07T17:05:32Z

Closing this because it's bitrotted, will re-start this work later

ehsanmok · 2019-04-30T04:23:20Z

Hi @kazimuth , just pinging to see if you're still interested and would like to propose a deadline for it?

kazimuth · 2019-04-30T16:06:35Z

I am but my workload's picked up a bit, might be a few weeks before i get back to this :)

kazimuth added 3 commits February 3, 2019 17:52

Correctly handle new functions being registered at runtime

bca1ee2

Fix conversion of TVMRetValue of type kNDArrayContainer

dfe5b0c

Start adding typed APIs to rust frontend, greatly simplify example code

e412a1b

nhynes self-requested a review February 4, 2019 07:08

nhynes reviewed Feb 4, 2019

View reviewed changes

icemelon added the status: WIP label Mar 2, 2019

kazimuth closed this Apr 7, 2019

kazimuth deleted the rustpolish branch April 30, 2019 16:06

More Strongly-Typed Rust Bindings #2559

More Strongly-Typed Rust Bindings #2559

Uh oh!

Conversation

kazimuth commented Feb 4, 2019

Uh oh!

ehsanmok commented Feb 4, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nhynes commented Feb 4, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nhynes Feb 4, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nhynes Feb 4, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nhynes commented Feb 4, 2019

Uh oh!

kazimuth commented Feb 4, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kazimuth commented Feb 15, 2019

Uh oh!

nhynes commented Feb 15, 2019

Uh oh!

ehsanmok commented Apr 3, 2019

Uh oh!

nhynes commented Apr 3, 2019

Uh oh!

ehsanmok commented Apr 3, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nhynes commented Apr 3, 2019

Uh oh!

kazimuth commented Apr 4, 2019

Uh oh!

nhynes commented Apr 5, 2019

Uh oh!

kazimuth commented Apr 5, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nhynes commented Apr 5, 2019

Uh oh!

kazimuth commented Apr 5, 2019

Uh oh!

nhynes commented Apr 5, 2019

Uh oh!

kazimuth commented Apr 7, 2019

Uh oh!

ehsanmok commented Apr 30, 2019

ehsanmok commented Feb 4, 2019 •

edited

Loading

nhynes commented Feb 4, 2019 •

edited

Loading

nhynes Feb 4, 2019 •

edited

Loading

nhynes Feb 4, 2019 •

edited

Loading

kazimuth commented Feb 4, 2019 •

edited

Loading

ehsanmok commented Apr 3, 2019 •

edited

Loading

kazimuth commented Apr 5, 2019 •

edited

Loading