Rust Refactor Stage 4: Rewrite Rust graph runtime to use new APIs #5830

jroesch · 2020-06-16T23:58:07Z

This is the fourth stage of the Rust rewrite which ports the graph-rt to the new API. The final stage after this will be turning on CI for the new bindings, updating docs, and deprecating the old bindings. This depends on #5764.

jroesch · 2020-06-18T18:46:17Z

@robo-corg @tqchen this is like a straight copy of the old runtime crate with slight modifications to use the new tvm-sys. After this patch lands I can send a final one to switch CI over to the new bindings and drop support for the old ones.

leonwanghui · 2020-06-19T01:10:57Z

rust/tvm-graph-rt/src/array.rs

+                strides: if dlt.strides.is_null() {
+                    None
+                } else {
+                    Some(slice::from_raw_parts_mut(dlt.strides as *mut usize, size).to_vec())


Suggested change

Some(slice::from_raw_parts_mut(dlt.strides as *mut usize, size).to_vec())

Some(slice::from_raw_parts_mut(dlt.strides as *mut usize, shape.len()).to_vec())

I don't think this is correct, as the field currently is passing the number of elements not the rank of the tensor.

From this figure (https://i.stack.imgur.com/oQQVI.png), I think the length of strides should be equal to the rank of the tensor (ie. dimension).

kazum · 2020-06-19T01:51:30Z

@jroesch Can we enable derive_default for bindgen? Otherwise, test_wasm32 fails with tvm-sys because of the generated padding field.

The following change looks necessary which existed in the previous implementation.

diff --git a/rust/tvm-graph-rt/src/array.rs b/rust/tvm-graph-rt/src/array.rs
index 38519bd..8209b59 100644
--- a/rust/tvm-graph-rt/src/array.rs
+++ b/rust/tvm-graph-rt/src/array.rs
@@ -288,6 +288,7 @@ impl<'a> Tensor<'a> {
                 self.strides.as_ref().unwrap().as_ptr()
             } as *mut i64,
             byte_offset: 0,
+            ..Default::default()
         }
     }
 }
diff --git a/rust/tvm-sys/build.rs b/rust/tvm-sys/build.rs
index 85e16be..01d2934 100644
--- a/rust/tvm-sys/build.rs
+++ b/rust/tvm-sys/build.rs
@@ -54,6 +54,7 @@ fn main() {
         .layout_tests(false)
         .derive_partialeq(true)
         .derive_eq(true)
+        .derive_default(true)
         .generate()
         .expect("unable to generate bindings")
         .write_to_file(PathBuf::from("src/c_runtime_api.rs"))
diff --git a/rust/tvm-sys/src/array.rs b/rust/tvm-sys/src/array.rs
index 1627e9e..5d09d86 100644
--- a/rust/tvm-sys/src/array.rs
+++ b/rust/tvm-sys/src/array.rs
@@ -48,6 +48,7 @@ macro_rules! impl_dltensor_from_ndarray {
                     shape: arr.shape().as_ptr() as *const i64 as *mut i64,
                     strides: arr.strides().as_ptr() as *const i64 as *mut i64,
                     byte_offset: 0,
+                    ..Default::default()
                 }
             }
         }

rust/tvm-graph-rt/.travis.yml

robo-corg · 2020-06-18T19:04:08Z

rust/tvm-graph-rt/src/allocator.rs

+    pub fn new(size: usize, align: Option<usize>) -> Result<Self, LayoutErr> {
+        let alignment = align.unwrap_or(DEFAULT_ALIGN_BYTES);
+        let layout = Layout::from_size_align(size, alignment)?;
+        let ptr = unsafe { alloc::alloc(layout) };
+        if ptr.is_null() {
+            alloc::handle_alloc_error(layout);
+        }
+        Ok(Self { ptr, layout })
+    }


Should this return std::mem::MaybeUninit<Allocation> or does that matter here since it is just bytes?

robo-corg · 2020-06-18T19:09:21Z

rust/tvm-graph-rt/src/allocator.rs

+
+#[derive(PartialEq, Eq)]
+pub struct Allocation {
+    layout: Layout,


I assume I will find out why we need to track alignment?

robo-corg · 2020-06-18T19:12:16Z

rust/tvm-graph-rt/src/array.rs

+/// let mut a_nd: ndarray::ArrayD<f32> = a.try_into().unwrap();
+/// ```
+#[derive(PartialEq)]
+pub struct Tensor<'a> {


Does it make sense to have an owned Tensor and a TensorRef type? I guess that can get added later.

robo-corg · 2020-06-18T19:13:33Z

rust/tvm-graph-rt/src/array.rs

+    pub(crate) data: Storage<'a>,
+    pub(crate) ctx: Context,
+    pub(crate) dtype: DataType,
+    pub(crate) shape: Vec<i64>,
+    // ^ not usize because `typedef int64_t tvm_index_t` in c_runtime_api.h
+    /// The `Tensor` strides. Can be `None` if the `Tensor` is contiguous.
+    pub(crate) strides: Option<Vec<usize>>,
+    pub(crate) byte_offset: isize,
+    /// The number of elements in the `Tensor`.
+    pub(crate) size: usize,


I would make these pub(self) or remove pub entirely since it looks like you have unsafe code using them.

robo-corg · 2020-06-19T22:37:43Z

rust/tvm-graph-rt/src/graph.rs

+        graph: &Graph,
+        lib: &'m M,
+        tensors: &[Tensor<'t>],
+    ) -> Result<Vec<Box<dyn Fn() + 'm>>, Error> {


Maybe make this execs a function of Tensor<'t> or for <'a> Tensor<'a>?

robo-corg · 2020-06-19T22:44:39Z

rust/tvm-graph-rt/src/threading.rs

+    pending: Arc<AtomicUsize>,
+}
+
+impl Job {


Any reason not to use rayon for this? I think you can tell it to spawn using tvm's thread pool: https://docs.rs/rayon/1.3.1/rayon/struct.ThreadPoolBuilder.html

robo-corg · 2020-06-19T22:53:57Z

rust/tvm-graph-rt/src/workspace.rs

+    free: Vec<usize>,
+    in_use: Vec<usize>,


Could be a good use for https://docs.rs/hibitset/0.6.3/hibitset/

robo-corg · 2020-06-19T22:56:43Z

rust/tvm-graph-rt/src/workspace.rs

+                if ws_size < size {
+                    return cur_ws_idx;
+                }


This seems like you could end up with some really extreme over allocation if you have a combination of large and small tensors.

rust/tvm-sys/src/datatype.rs

binarybana · 2020-06-20T13:42:14Z

rust/tvm-graph-rt/Cargo.toml

+
+[dependencies]
+crossbeam = "0.7.3"
+failure = "0.1"


Not anyhow or thiserror?

binarybana · 2020-06-20T13:53:15Z

rust/tvm-graph-rt/tests/test_tvm_basic/src/build_test_lib.py

+    C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
+    s = tvm.te.create_schedule(C.op)
+    s[C].parallel(s[C].op.axis[0])
+    print(tvm.lower(s, [A, B, C], simple_mode=True))


Do you want this print statement here?

binarybana · 2020-06-20T13:55:40Z

rust/tvm-graph-rt/tests/test_wasm32/src/build_test_lib.py

+    C = te.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
+    s = tvm.te.create_schedule(C.op)
+    s[C].parallel(s[C].op.axis[0])
+    print(tvm.lower(s, [A, B, C], simple_mode=True))


Another print.

jroesch · 2020-06-23T00:21:02Z

@jroesch Can we enable derive_default for bindgen? Otherwise, test_wasm32 fails with tvm-sys because of the generated padding field.

The following change looks necessary which existed in the previous implementation.

diff --git a/rust/tvm-graph-rt/src/array.rs b/rust/tvm-graph-rt/src/array.rs
index 38519bd..8209b59 100644
--- a/rust/tvm-graph-rt/src/array.rs
+++ b/rust/tvm-graph-rt/src/array.rs
@@ -288,6 +288,7 @@ impl<'a> Tensor<'a> {
                 self.strides.as_ref().unwrap().as_ptr()
             } as *mut i64,
             byte_offset: 0,
+            ..Default::default()
         }
     }
 }
diff --git a/rust/tvm-sys/build.rs b/rust/tvm-sys/build.rs
index 85e16be..01d2934 100644
--- a/rust/tvm-sys/build.rs
+++ b/rust/tvm-sys/build.rs
@@ -54,6 +54,7 @@ fn main() {
         .layout_tests(false)
         .derive_partialeq(true)
         .derive_eq(true)
+        .derive_default(true)
         .generate()
         .expect("unable to generate bindings")
         .write_to_file(PathBuf::from("src/c_runtime_api.rs"))
diff --git a/rust/tvm-sys/src/array.rs b/rust/tvm-sys/src/array.rs
index 1627e9e..5d09d86 100644
--- a/rust/tvm-sys/src/array.rs
+++ b/rust/tvm-sys/src/array.rs
@@ -48,6 +48,7 @@ macro_rules! impl_dltensor_from_ndarray {
                     shape: arr.shape().as_ptr() as *const i64 as *mut i64,
                     strides: arr.strides().as_ptr() as *const i64 as *mut i64,
                     byte_offset: 0,
+                    ..Default::default()
                 }
             }
         }

Yeah I will turn it back on, there are just so many changes to juggle I was bound to make a mistake or two.

jroesch · 2020-06-23T00:25:01Z

@binarybana and @robo-corg many of the things you guys identified existed in the previous graph runtime, it would be my preference to land this port which just makes it compile against the new bindings, and then bring other patches to improve the code in follow up PRs. thoughts?

Co-authored-by: Andrew <amcharg@gmail.com>

binarybana · 2020-06-23T00:27:07Z

Yes, that works for me, which is why I approved the PR despite my comments.

robo-corg · 2020-06-23T00:34:13Z

existed in the previous graph runtime, it would be my preference to land this port which just makes it compile against the new bindings, and then bring other patches to improve the code in follow up PRs. thoughts?

That works great!

@kazum

…ache#5830) * Port graph-runtime to new API * --amend * Fix file lint * Remove old travis file * Add @kazum's patch * Update rust/tvm-sys/src/datatype.rs Co-authored-by: Andrew <amcharg@gmail.com> Co-authored-by: Andrew <amcharg@gmail.com>

@kazum

…ache#5830) * Port graph-runtime to new API * --amend * Fix file lint * Remove old travis file * Add @kazum's patch * Update rust/tvm-sys/src/datatype.rs Co-authored-by: Andrew <amcharg@gmail.com> Co-authored-by: Andrew <amcharg@gmail.com>

jroesch added 2 commits June 18, 2020 11:38

Port graph-runtime to new API

f262f62

--amend

6c678b8

jroesch force-pushed the rust-graph-rt branch from 1827f9c to 6c678b8 Compare June 18, 2020 18:39

Fix file lint

b0df236

Remove old travis file

c2fc8c2

tqchen approved these changes Jun 18, 2020

View reviewed changes

leonwanghui reviewed Jun 19, 2020

View reviewed changes

robo-corg suggested changes Jun 19, 2020

View reviewed changes

binarybana approved these changes Jun 20, 2020

View reviewed changes

Add @kazum's patch

43dd1dd

Update rust/tvm-sys/src/datatype.rs

6ddbba5

Co-authored-by: Andrew <amcharg@gmail.com>

robo-corg approved these changes Jun 23, 2020

View reviewed changes

jroesch mentioned this pull request Jun 23, 2020

[Rust][Runtime] Address issues with unsafe code in existing Rust graph runtime. #5889

Closed

kazum approved these changes Jun 23, 2020

View reviewed changes

jroesch merged commit aa84ee2 into apache:master Jun 23, 2020

ZihengJiang mentioned this pull request Sep 25, 2020

TVM v0.7 Release Note Candidate #6486

Closed

jroesch deleted the rust-graph-rt branch February 4, 2021 04:41

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Rust Refactor Stage 4: Rewrite Rust graph runtime to use new APIs #5830

Rust Refactor Stage 4: Rewrite Rust graph runtime to use new APIs #5830

jroesch commented Jun 16, 2020

jroesch commented Jun 18, 2020

leonwanghui Jun 19, 2020

jroesch Jun 23, 2020

leonwanghui Jun 23, 2020

kazum commented Jun 19, 2020

robo-corg Jun 18, 2020

robo-corg Jun 18, 2020

robo-corg Jun 18, 2020

robo-corg Jun 18, 2020

robo-corg Jun 19, 2020

robo-corg Jun 19, 2020

robo-corg Jun 19, 2020

robo-corg Jun 19, 2020

binarybana Jun 20, 2020

binarybana Jun 20, 2020

binarybana Jun 20, 2020

jroesch commented Jun 23, 2020

jroesch commented Jun 23, 2020

binarybana commented Jun 23, 2020

robo-corg commented Jun 23, 2020

	Some(slice::from_raw_parts_mut(dlt.strides as *mut usize, size).to_vec())
	Some(slice::from_raw_parts_mut(dlt.strides as *mut usize, shape.len()).to_vec())

Rust Refactor Stage 4: Rewrite Rust graph runtime to use new APIs #5830

Rust Refactor Stage 4: Rewrite Rust graph runtime to use new APIs #5830

Conversation

jroesch commented Jun 16, 2020

jroesch commented Jun 18, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kazum commented Jun 19, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jroesch commented Jun 23, 2020

jroesch commented Jun 23, 2020

binarybana commented Jun 23, 2020

robo-corg commented Jun 23, 2020