[incubator-tvm] 16/23: More cleanup

[jroesch]

This is an automated email from the ASF dual-hosted git repository.

jroesch pushed a commit to branch cargo-build
in repository https://gitbox.apache.org/repos/asf/incubator-tvm.git

commit 187435099c4d80be2c200791e7853f8518f421e1
Author: Jared Roesch <jroe...@octoml.ai>
AuthorDate: Fri Oct 16 16:45:12 2020 -0700

    More cleanup
---
 rust/compiler-ext/src/lib.rs | 22 ++++++++--------------
 rust/tvm/src/lib.rs          |  7 ++++---
 2 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/rust/compiler-ext/src/lib.rs b/rust/compiler-ext/src/lib.rs
index 346f40f..5f83f7b 100644
--- a/rust/compiler-ext/src/lib.rs
+++ b/rust/compiler-ext/src/lib.rs
@@ -18,25 +18,19 @@
  */
 
 use env_logger;
-use tvm;
-use tvm::runtime::function::register_override;
+use tvm::export;
 
-fn test_fn() -> Result<(), tvm::Error> {
-    println!("Hello Greg from Rust!");
-    Ok(())
+fn diagnostics() -> Result<(), tvm::Error> {
+    tvm::ir::diagnostics::codespan::init()
 }
 
-fn test_fn2(message: tvm::runtime::string::String) -> Result<(), tvm::Error> {
-    println!("The message: {}", message);
-    Ok(())
-}
-
-tvm::export!(test_fn, test_fn2);
+export!(diagnostics);
 
 #[no_mangle]
-fn compiler_ext_initialize() -> i32 {
+extern fn compiler_ext_initialize() -> i32 {
     let _ = env_logger::try_init();
-    tvm_export("rust_ext").expect("failed to initialize Rust compiler_ext");
-    log::debug!("done!");
+    tvm_export("rust_ext")
+        .expect("failed to initialize the Rust compiler extensions.");
+    log::debug!("Loaded the Rust compiler extension.");
     return 0;
 }
diff --git a/rust/tvm/src/lib.rs b/rust/tvm/src/lib.rs
index d193f09..ec80ece 100644
--- a/rust/tvm/src/lib.rs
+++ b/rust/tvm/src/lib.rs
@@ -50,10 +50,11 @@ pub use runtime::version;
 
 #[macro_export]
 macro_rules! export {
-    ($($fn_names:expr),*) => {
+    ($($fn_name:expr),*) => {
         pub fn tvm_export(ns: &str) -> Result<(), tvm::Error> {
             $(
-                register_override($fn_name, concat!($ns, stringfy!($fn_name)), 
true)?;
+                let name = String::from(ns) + ::std::stringify!($fn_name);
+                tvm::runtime::function::register_override($fn_name, name, 
true)?;
             )*
             Ok(())
         }
@@ -65,7 +66,7 @@ macro_rules! export_mod {
     ($ns:expr, $($mod_name:expr),*) => {
         pub fn tvm_mod_export() -> Result<(), tvm::Error> {
             $(
-                $mod_names::tvm_export($ns)?;
+                $mod_name::tvm_export($ns)?;
             )*
             Ok(())
         }