[GitHub] [tvm] hypercubestart commented on a change in pull request #7085: Fixes for using Python APIs from Rust.

[GitBox <gi...@apache.org>]

hypercubestart commented on a change in pull request #7085:
URL: https://github.com/apache/tvm/pull/7085#discussion_r587698328



##########
File path: rust/tvm/src/compiler/graph_rt.rs
##########
@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+use std::convert::TryInto;
+use std::io::Read;
+use std::path::Path;
+
+use once_cell::sync::Lazy;
+use thiserror::Error;
+
+use crate::ir::IRModule;
+use crate::python;
+use crate::runtime::{map::Map, Function, Module as RtModule, NDArray, String};
+
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error("{0}")]
+    IO(#[from] std::io::Error),
+    #[error("{0}")]
+    TVM(#[from] crate::errors::Error),
+}
+
+static TVM_BUILD: Lazy<Function> = Lazy::new(|| {
+    python::import("tvm").unwrap();
+    python::import("tvm.relay").unwrap();
+    Function::get("tvm.relay.build").unwrap()
+});
+
+fn _compile_module(
+    module: IRModule,
+    target: String,
+    target_host: String,
+    params: Map<String, NDArray>,
+    module_name: String,
+) -> Result<RtModule, Error> {
+    // The RAW API is Fn(IRModule, String, String, Map<String, NDArray>, String);
+    let module = TVM_BUILD.invoke(vec![
+        module.into(),
+        target.into(),
+        target_host.into(),
+        params.into(),
+        module_name.into(),
+    ])?;
+    let module: RtModule = module.try_into().unwrap();
+    Ok(module)
+}
+
+#[derive(Debug)]
+pub struct CompilerConfig {
+    target: Option<String>,
+    target_host: Option<String>,
+    params: Map<String, NDArray>,
+    module_name: Option<String>,
+}
+
+impl Default for CompilerConfig {
+    fn default() -> Self {
+        CompilerConfig {
+            target: None,
+            target_host: None,
+            params: Map::empty(),
+            module_name: None,
+        }
+    }
+}
+
+/// Compile a module from a configuration and IRModule.
+///
+/// # Arguments
+///
+/// * `config` - The configuration for the compiler.
+/// * `module` - The IRModule to compile.
+pub fn compile_module(config: CompilerConfig, module: IRModule) -> Result<RtModule, Error> {

Review comment:
       is there a test for this?

##########
File path: rust/tvm/src/python.rs
##########
@@ -29,19 +29,30 @@ use pyo3::prelude::*;
 pub fn load() -> Result<String, ()> {
     let gil = Python::acquire_gil();
     let py = gil.python();
+    // let main_mod = initialize();
+    //let main_mod = main_mod.as_ref(py);
     load_python_tvm_(py).map_err(|e| {
         // We can't display Python exceptions via std::fmt::Display,
         // so print the error here manually.
         e.print_and_set_sys_last_vars(py);
     })
 }
 
-// const TVMC_CODE: &'static str = include_str!("tvmc.py");
+pub fn import(mod_to_import: &str) -> PyResult<()> {
+    let gil = Python::acquire_gil();
+    let py = gil.python();
+    import_python(py, mod_to_import)?;
+    Ok(())
+}
+
+fn import_python<'p, 'b: 'p>(py: Python<'p>, to_import: &'b str) -> PyResult<&'p PyModule> {
+    let imported_mod = py.import(to_import)?;
+    Ok(imported_mod)
+}
 
 fn load_python_tvm_(py: Python) -> PyResult<String> {
-    let sys = py.import("tvm")?;
-    let version: String = sys.get("__version__")?.extract()?;
-    // py.run(TVMC_CODE, None, None)?;
+    let imported_mod = import_python(py, "tvm")?;

Review comment:
       ok, its strange to me that loading tvm into rust is flaky though

##########
File path: rust/tvm-rt/src/module.rs
##########
@@ -26,21 +26,24 @@ use std::{
     ptr,
 };
 
+use crate::object::Object;
+use tvm_macros::Object;
 use tvm_sys::ffi;
 
 use crate::errors::Error;
+use crate::String as TString;
 use crate::{errors, function::Function};
 
-const ENTRY_FUNC: &str = "__tvm_main__";
-
 /// Wrapper around TVM module handle which contains an entry function.
 /// The entry function can be applied to an imported module through [`entry_func`].
 ///
 /// [`entry_func`]:struct.Module.html#method.entry_func

Review comment:
       is this comment still true?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org