You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by jr...@apache.org on 2020/11/05 21:41:00 UTC
[incubator-tvm] 04/23: Hacking on Rust inside of TVM
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 131e40afeb4bd89f0b378e0e7d8ad440d017d33a
Author: Jared Roesch <jr...@octoml.ai>
AuthorDate: Tue Oct 13 15:25:56 2020 -0700
Hacking on Rust inside of TVM
---
CMakeLists.txt | 1 +
cmake/modules/RustExt.cmake | 13 +
rust/Cargo.toml | 1 +
rust/compiler-ext/Cargo.toml | 13 +
rust/compiler-ext/src/lib.rs | 7 +
rust/tvm/src/ir/source_map.rs | 0
rust/tvm/test.rly | 2 +
tests/python/relay/test_type_infer2.py | 419 +++++++++++++++++++++++++++++++++
8 files changed, 456 insertions(+)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index abf2b56..9f82754 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -79,6 +79,7 @@ tvm_option(USE_ARM_COMPUTE_LIB "Build with Arm Compute Library" OFF)
tvm_option(USE_ARM_COMPUTE_LIB_GRAPH_RUNTIME "Build with Arm Compute Library graph runtime" OFF)
tvm_option(USE_TENSORRT_CODEGEN "Build with TensorRT Codegen support" OFF)
tvm_option(USE_TENSORRT_RUNTIME "Build with TensorRT runtime" OFF)
+tvm_option(USE_RUST_EXT "Build with Rust based compiler extensions" OFF)
# include directories
include_directories(${CMAKE_INCLUDE_PATH})
diff --git a/cmake/modules/RustExt.cmake b/cmake/modules/RustExt.cmake
new file mode 100644
index 0000000..45e46bd
--- /dev/null
+++ b/cmake/modules/RustExt.cmake
@@ -0,0 +1,13 @@
+if(USE_RUST_EXT)
+ set(RUST_SRC_DIR "rust")
+ set(CARGO_OUT_DIR "rust/target"
+ set(COMPILER_EXT_PATH "${CARGO_OUT_DIR}/target/release/libcompiler_ext.dylib")
+
+ add_custom_command(
+ OUTPUT "${COMPILER_EXT_PATH}"
+ COMMAND cargo build --release
+ MAIN_DEPENDENCY "${RUST_SRC_DIR}"
+ WORKING_DIRECTORY "${RUST_SRC_DIR}/compiler-ext")
+
+ target_link_libraries(tvm "${COMPILER_EXT_PATH}" PRIVATE)
+endif(USE_RUST_EXT)
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
index 28312a5..7c092d8 100644
--- a/rust/Cargo.toml
+++ b/rust/Cargo.toml
@@ -29,4 +29,5 @@ members = [
"tvm-graph-rt/tests/test_tvm_dso",
"tvm-graph-rt/tests/test_wasm32",
"tvm-graph-rt/tests/test_nn",
+ "compiler-ext",
]
diff --git a/rust/compiler-ext/Cargo.toml b/rust/compiler-ext/Cargo.toml
new file mode 100644
index 0000000..76d10eb
--- /dev/null
+++ b/rust/compiler-ext/Cargo.toml
@@ -0,0 +1,13 @@
+[package]
+name = "compiler-ext"
+version = "0.1.0"
+authors = ["Jared Roesch <jr...@octoml.ai>"]
+edition = "2018"
+# TODO(@jroesch): would be cool to figure out how to statically link instead.
+
+[lib]
+crate-type = ["cdylib"]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
diff --git a/rust/compiler-ext/src/lib.rs b/rust/compiler-ext/src/lib.rs
new file mode 100644
index 0000000..31e1bb2
--- /dev/null
+++ b/rust/compiler-ext/src/lib.rs
@@ -0,0 +1,7 @@
+#[cfg(test)]
+mod tests {
+ #[test]
+ fn it_works() {
+ assert_eq!(2 + 2, 4);
+ }
+}
diff --git a/rust/tvm/src/ir/source_map.rs b/rust/tvm/src/ir/source_map.rs
new file mode 100644
index 0000000..e69de29
diff --git a/rust/tvm/test.rly b/rust/tvm/test.rly
new file mode 100644
index 0000000..d8b7c69
--- /dev/null
+++ b/rust/tvm/test.rly
@@ -0,0 +1,2 @@
+#[version = "0.0.5"]
+fn @main(%x: int32) -> float32 { %x }
diff --git a/tests/python/relay/test_type_infer2.py b/tests/python/relay/test_type_infer2.py
new file mode 100644
index 0000000..6758d96
--- /dev/null
+++ b/tests/python/relay/test_type_infer2.py
@@ -0,0 +1,419 @@
+# 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.
+"""Test that type checker correcly computes types
+ for expressions.
+"""
+import pytest
+import tvm
+
+from tvm import IRModule, te, relay, parser
+from tvm.relay import op, transform, analysis
+from tvm.relay import Any
+
+
+def infer_mod(mod, annotate_spans=True):
+ if annotate_spans:
+ mod = relay.transform.AnnotateSpans()(mod)
+
+ mod = transform.InferType()(mod)
+ return mod
+
+
+def infer_expr(expr, annotate_spans=True):
+ mod = IRModule.from_expr(expr)
+ mod = infer_mod(mod, annotate_spans)
+ mod = transform.InferType()(mod)
+ entry = mod["main"]
+ return entry if isinstance(expr, relay.Function) else entry.body
+
+
+def assert_has_type(expr, typ, mod=None):
+ if not mod:
+ mod = tvm.IRModule({})
+
+ mod["main"] = expr
+ mod = infer_mod(mod)
+ checked_expr = mod["main"]
+ checked_type = checked_expr.checked_type
+ if checked_type != typ:
+ raise RuntimeError("Type mismatch %s vs %s" % (checked_type, typ))
+
+
+def initialize_box_adt(mod):
+ # initializes simple ADT for tests
+ box = relay.GlobalTypeVar("box")
+ tv = relay.TypeVar("tv")
+ constructor = relay.Constructor("constructor", [tv], box)
+ data = relay.TypeData(box, [tv], [constructor])
+ mod[box] = data
+ return box, constructor
+
+
+def test_monomorphic_let():
+ "Program: let %x = 1; %x"
+ # TODO(@jroesch): this seems whack.
+ sb = relay.ScopeBuilder()
+ x = relay.var("x", dtype="float64", shape=())
+ x = sb.let("x", relay.const(1.0, "float64"))
+ sb.ret(x)
+ xchecked = infer_expr(sb.get())
+ assert xchecked.checked_type == relay.scalar_type("float64")
+
+
+def test_single_op():
+ "Program: fn (%x : float32) { let %t1 = f(%x); %t1 }"
+ x = relay.var("x", shape=[])
+ func = relay.Function([x], op.log(x))
+ ttype = relay.TensorType([], dtype="float32")
+ assert_has_type(func, relay.FuncType([ttype], ttype))
+
+
+def test_add_broadcast_op():
+ """
+ Program:
+ fn (%x: Tensor[(10, 4), float32], %y: Tensor[(5, 10, 1), float32])
+ -> Tensor[(5, 10, 4), float32] {
+ %x + %y
+ }
+ """
+ x = relay.var("x", shape=(10, 4))
+ y = relay.var("y", shape=(5, 10, 1))
+ z = x + y
+ func = relay.Function([x, y], z)
+ t1 = relay.TensorType((10, 4), "float32")
+ t2 = relay.TensorType((5, 10, 1), "float32")
+ t3 = relay.TensorType((5, 10, 4), "float32")
+ expected_ty = relay.FuncType([t1, t2], t3)
+ assert_has_type(func, expected_ty)
+
+
+def test_dual_op():
+ """Program:
+ fn (%x : Tensor[(10, 10), float32]) {
+ let %t1 = log(x);
+ let %t2 = add(%t1, %x);
+ %t1
+ }
+ """
+ tp = relay.TensorType((10, 10), "float32")
+ x = relay.var("x", tp)
+ sb = relay.ScopeBuilder()
+ t1 = sb.let("t1", relay.log(x))
+ t2 = sb.let("t2", relay.add(t1, x))
+ sb.ret(t2)
+ f = relay.Function([x], sb.get())
+ fchecked = infer_expr(f)
+ assert fchecked.checked_type == relay.FuncType([tp], tp)
+
+
+def test_decl():
+ """Program:
+ def @f(%x : Tensor[(10, 10), float32]) {
+ log(%x)
+ }
+ """
+ tp = relay.TensorType((10, 10))
+ x = relay.var("x", tp)
+ f = relay.Function([x], relay.log(x))
+ fchecked = infer_expr(f)
+ assert fchecked.checked_type == relay.FuncType([tp], tp)
+
+
+def test_recursion():
+ """
+ Program:
+ def @f(%n: int32, %data: float32) -> float32 {
+ if (%n == 0) {
+ %data
+ } else {
+ @f(%n - 1, log(%data))
+ }
+ }
+ """
+ sb = relay.ScopeBuilder()
+ f = relay.GlobalVar("f")
+ ti32 = relay.scalar_type("int32")
+ tf32 = relay.scalar_type("float32")
+ n = relay.var("n", ti32)
+ data = relay.var("data", tf32)
+
+ with sb.if_scope(relay.equal(n, relay.const(0, ti32))):
+ sb.ret(data)
+ with sb.else_scope():
+ sb.ret(f(relay.subtract(n, relay.const(1, ti32)), relay.log(data)))
+ mod = tvm.IRModule()
+ mod[f] = relay.Function([n, data], sb.get())
+ mod = infer_mod(mod)
+ assert "@f(%1, %2)" in mod.astext()
+ assert mod["f"].checked_type == relay.FuncType([ti32, tf32], tf32)
+
+
+def test_incomplete_call():
+ tt = relay.scalar_type("int32")
+ x = relay.var("x", tt)
+ f = relay.var("f")
+ func = relay.Function([x, f], relay.Call(f, [x]), tt)
+
+ ft = infer_expr(func)
+ f_type = relay.FuncType([tt], tt)
+ assert ft.checked_type == relay.FuncType([tt, f_type], tt)
+
+
+def test_higher_order_argument():
+ a = relay.TypeVar("a")
+ x = relay.Var("x", a)
+ id_func = relay.Function([x], x, a, [a])
+
+ b = relay.TypeVar("b")
+ f = relay.Var("f", relay.FuncType([b], b))
+ y = relay.Var("y", b)
+ ho_func = relay.Function([f, y], f(y), b, [b])
+
+ # id func should be an acceptable argument to the higher-order
+ # function even though id_func takes a type parameter
+ ho_call = ho_func(id_func, relay.const(0, "int32"))
+
+ hc = infer_expr(ho_call)
+ expected = relay.scalar_type("int32")
+ assert hc.checked_type == expected
+
+
+def test_higher_order_return():
+ a = relay.TypeVar("a")
+ x = relay.Var("x", a)
+ id_func = relay.Function([x], x, a, [a])
+
+ b = relay.TypeVar("b")
+ nested_id = relay.Function([], id_func, relay.FuncType([b], b), [b])
+
+ ft = infer_expr(nested_id)
+ assert ft.checked_type == relay.FuncType([], relay.FuncType([b], b), [b])
+
+
+def test_higher_order_nested():
+ a = relay.TypeVar("a")
+ x = relay.Var("x", a)
+ id_func = relay.Function([x], x, a, [a])
+
+ choice_t = relay.FuncType([], relay.scalar_type("bool"))
+ f = relay.Var("f", choice_t)
+
+ b = relay.TypeVar("b")
+ z = relay.Var("z")
+ top = relay.Function(
+ [f], relay.If(f(), id_func, relay.Function([z], z)), relay.FuncType([b], b), [b]
+ )
+
+ expected = relay.FuncType([choice_t], relay.FuncType([b], b), [b])
+ ft = infer_expr(top)
+ assert ft.checked_type == expected
+
+
+def test_tuple():
+ tp = relay.TensorType((10,))
+ x = relay.var("x", tp)
+ res = relay.Tuple([x, x])
+ assert infer_expr(res).checked_type == relay.TupleType([tp, tp])
+
+
+def test_ref():
+ x = relay.var("x", "float32")
+ y = relay.var("y", "float32")
+ r = relay.RefCreate(x)
+ st = relay.scalar_type("float32")
+ assert infer_expr(r).checked_type == relay.RefType(st)
+ g = relay.RefRead(r)
+ assert infer_expr(g).checked_type == st
+ w = relay.RefWrite(r, y)
+ assert infer_expr(w).checked_type == relay.TupleType([])
+
+
+def test_free_expr():
+ x = relay.var("x", "float32")
+ y = relay.add(x, x)
+ yy = infer_expr(y, annotate_spans=False)
+ assert tvm.ir.structural_equal(yy.args[0], x, map_free_vars=True)
+ assert yy.checked_type == relay.scalar_type("float32")
+ assert x.vid.same_as(yy.args[0].vid)
+
+
+def test_type_args():
+ x = relay.var("x", shape=(10, 10))
+ y = relay.var("y", shape=(1, 10))
+ z = relay.add(x, y)
+ ty_z = infer_expr(z)
+ ty_args = ty_z.type_args
+ assert len(ty_args) == 2
+ assert ty_args[0].dtype == "float32"
+ assert ty_args[1].dtype == "float32"
+ sh1 = ty_args[0].shape
+ sh2 = ty_args[1].shape
+ assert sh1[0].value == 10
+ assert sh1[1].value == 10
+ assert sh2[0].value == 1
+ assert sh2[1].value == 10
+
+
+def test_global_var_recursion():
+ mod = tvm.IRModule({})
+ gv = relay.GlobalVar("main")
+ x = relay.var("x", shape=[])
+ tt = relay.scalar_type("float32")
+
+ func = relay.Function([x], relay.Call(gv, [x]), tt)
+ mod[gv] = func
+ mod = infer_mod(mod)
+ func_ty = mod["main"].checked_type
+
+ assert func_ty == relay.FuncType([tt], tt)
+
+
+def test_equal():
+ i = relay.var("i", shape=[], dtype="int32")
+ eq = op.equal(i, relay.const(0, dtype="int32"))
+ func = relay.Function([i], eq)
+ ft = infer_expr(func)
+ expected = relay.FuncType([relay.scalar_type("int32")], relay.scalar_type("bool"))
+ assert ft.checked_type == expected
+
+ assert ft.checked_type == relay.FuncType(
+ [relay.scalar_type("int32")], relay.scalar_type("bool")
+ )
+
+
+def test_constructor_type():
+ mod = tvm.IRModule()
+ box, constructor = initialize_box_adt(mod)
+
+ a = relay.TypeVar("a")
+ x = relay.Var("x", a)
+ func = relay.Function([x], constructor(x), box(a), [a])
+ mod["main"] = func
+ mod = infer_mod(mod)
+ func_ty = mod["main"].checked_type
+ box = mod.get_global_type_var("box")
+ expected = relay.FuncType([a], box(a), [a])
+ assert func_ty == expected
+
+
+def test_constructor_call():
+ mod = tvm.IRModule()
+ box, constructor = initialize_box_adt(mod)
+
+ box_unit = constructor(relay.Tuple([]))
+ box_constant = constructor(relay.const(0, "float32"))
+
+ func = relay.Function([], relay.Tuple([box_unit, box_constant]))
+ mod["main"] = func
+ mod = infer_mod(mod)
+ ret_type = mod["main"].checked_type.ret_type.fields
+ # NB(@jroesch): when we annotate spans the ast fragments before
+ # annotation the previous fragments will no longer be directly equal.
+ box = mod.get_global_type_var("box")
+ expected1 = box(relay.TupleType([]))
+ expected2 = box(relay.TensorType((), "float32"))
+ assert ret_type[0] == expected1
+ assert ret_type[1] == expected2
+
+
+def test_adt_match():
+ mod = tvm.IRModule()
+ box, constructor = initialize_box_adt(mod)
+
+ v = relay.Var("v", relay.TensorType((), "float32"))
+ match = relay.Match(
+ constructor(relay.const(0, "float32")),
+ [
+ relay.Clause(
+ relay.PatternConstructor(constructor, [relay.PatternVar(v)]), relay.Tuple([])
+ ),
+ # redundant but shouldn't matter to typechecking
+ relay.Clause(relay.PatternWildcard(), relay.Tuple([])),
+ ],
+ )
+
+ func = relay.Function([], match)
+ mod["main"] = func
+ mod = infer_mod(mod)
+ actual = mod["main"].checked_type.ret_type
+ assert actual == relay.TupleType([])
+
+
+def test_adt_match_type_annotations():
+ mod = tvm.IRModule()
+ box, constructor = initialize_box_adt(mod)
+
+ # the only type annotation is inside the match pattern var
+ # but that should be enough info
+ tt = relay.TensorType((2, 2), "float32")
+ x = relay.Var("x")
+ mv = relay.Var("mv", tt)
+ match = relay.Match(
+ constructor(x),
+ [
+ relay.Clause(
+ relay.PatternConstructor(constructor, [relay.PatternVar(mv)]), relay.Tuple([])
+ )
+ ],
+ )
+
+ mod["main"] = relay.Function([x], match)
+ mod = infer_mod(mod)
+ ft = mod["main"].checked_type
+ assert ft == relay.FuncType([tt], relay.TupleType([]))
+
+
+def test_let_polymorphism():
+ id = relay.Var("id")
+ xt = relay.TypeVar("xt")
+ x = relay.Var("x", xt)
+ body = relay.Tuple([id(relay.const(1)), id(relay.Tuple([]))])
+ body = relay.Let(id, relay.Function([x], x, xt, [xt]), body)
+ body = infer_expr(body)
+ int32 = relay.TensorType((), "int32")
+ tvm.ir.assert_structural_equal(body.checked_type, relay.TupleType([int32, relay.TupleType([])]))
+
+
+def test_if():
+ choice_t = relay.FuncType([], relay.scalar_type("bool"))
+ f = relay.Var("f", choice_t)
+ true_branch = relay.Var("True", relay.TensorType([Any(), 1], dtype="float32"))
+ false_branch = relay.Var("False", relay.TensorType([Any(), Any()], dtype="float32"))
+ top = relay.Function([f, true_branch, false_branch], relay.If(f(), true_branch, false_branch))
+ ft = infer_expr(top)
+ tvm.ir.assert_structural_equal(ft.ret_type, relay.TensorType([Any(), 1], dtype="float32"))
+
+
+def test_type_arg_infer():
+ code = """
+#[version = "0.0.5"]
+def @id[A](%x: A) -> A {
+ %x
+}
+def @main(%f: float32) -> float32 {
+ @id(%f)
+}
+"""
+ mod = tvm.parser.fromtext(code)
+ mod = transform.InferType()(mod)
+ tvm.ir.assert_structural_equal(mod["main"].body.type_args, [relay.TensorType((), "float32")])
+
+
+if __name__ == "__main__":
+ import sys
+
+ pytest.main(sys.argv)