on_device annotation

apache · Aug 3, 2020 · e2eedb2 · e2eedb2
1 parent b70a5fe
commit e2eedb2
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/python/tvm/relay/op/_tensor.py b/python/tvm/relay/op/_tensor.py
@@ -87,6 +87,9 @@
 register_broadcast_schedule("fast_exp")
 register_broadcast_schedule("fast_tanh")
 register_broadcast_schedule("fast_erf")
+# a fake on_device schedule.
+# this will not be used in actual computation as on_device will be removed during DeviceAnnotation pass
+register_injective_schedule("on_device")
 
 
 # zeros

diff --git a/src/relay/op/annotation/annotation.cc b/src/relay/op/annotation/annotation.cc
@@ -54,7 +54,13 @@ RELAY_REGISTER_OP("on_device")
     .add_type_rel("Identity", IdentityRel)
     .set_attr<TOpPattern>("TOpPattern", kOpaque)
     .set_attr<TOpIsStateful>("TOpIsStateful", false)
-    .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ElemwiseArbitraryLayout);
+    .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ElemwiseArbitraryLayout)
+    .set_attr<FTVMCompute>("FTVMCompute",
+                           [] (const Attrs& attrs,
+                               const Array<te::Tensor>& inputs,
+                               const Type& out_type) -> Array<te::Tensor> {
+                           return {topi::identity(inputs[0])};
+                           });
 
 Expr StopFusion(Expr data) {
   static const Op& op = Op::Get("annotation.stop_fusion");