[tvm] 04/23: Tweak build.rs to use release version of tvm-build

[jr...@apache.org]

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

jroesch pushed a commit to branch ci-docker-staging
in repository https://gitbox.apache.org/repos/asf/tvm.git

commit 03b7019ba0304b634d3cf4d0893cc808e7729f21
Author: Jared Roesch <jr...@octoml.ai>
AuthorDate: Mon Feb 22 21:34:29 2021 -0800

    Tweak build.rs to use release version of tvm-build
---
 rust/tvm-sys/Cargo.toml |  2 +-
 rust/tvm-sys/build.rs   | 19 +++++++++++--------
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/rust/tvm-sys/Cargo.toml b/rust/tvm-sys/Cargo.toml
index 6075421..0c89023 100644
--- a/rust/tvm-sys/Cargo.toml
+++ b/rust/tvm-sys/Cargo.toml
@@ -36,4 +36,4 @@ enumn = "^0.1"
 [build-dependencies]
 bindgen = { version="0.57", default-features=false }
 anyhow = "^1.0"
-tvm-build = { git = "https://github.com/octoml/tvm-build", branch = "main" }
+tvm-build = "0.1"
diff --git a/rust/tvm-sys/build.rs b/rust/tvm-sys/build.rs
index 78ff84e..c917328 100644
--- a/rust/tvm-sys/build.rs
+++ b/rust/tvm-sys/build.rs
@@ -22,8 +22,15 @@ extern crate bindgen;
 use std::path::{Path, PathBuf};
 
 use anyhow::{Context, Result};
-use tvm_build::BuildConfig;
+use tvm_build::{BuildConfig};
 
+/// The necessary information for detecting a TVM installation.
+struct TVMInstall {
+    source_path: PathBuf,
+    build_path: PathBuf,
+}
+
+/// Find the TVM install using the provided path.
 fn find_using_tvm_path<P: AsRef<Path>>(tvm_path: P) -> Result<TVMInstall> {
     Ok(TVMInstall {
         source_path: tvm_path.as_ref().into(),
@@ -47,16 +54,12 @@ fn set_env_vars() -> Result<()> {
     if_unset("LLVM_CONFIG_PATH", "/opt/homebrew/opt/llvm/bin/llvm-config")
 }
 
-struct TVMInstall {
-    source_path: PathBuf,
-    build_path: PathBuf,
-}
-
+/// Find a TVM installation using TVM build by either first installing or detecting.
 fn find_using_tvm_build() -> Result<TVMInstall> {
     set_env_vars()?;
     let mut build_config = BuildConfig::default();
-    build_config.repository = Some("https://github.com/jroesch/tvm".to_string());
-    build_config.branch = Some("rust-tvm-build".to_string());
+    build_config.repository = Some("https://github.com/apache/tvm".to_string());
+    build_config.branch = Some(option_env!("TVM_BRANCH").unwrap_or("main").into());
     let build_result = tvm_build::build(build_config)?;
     let source_path = build_result.revision.source_path();
     let build_path = build_result.revision.build_path();