[GitHub] [tvm] masahi commented on a change in pull request #8497: [Runtime] add set_output_zero_copy

[GitBox]

masahi commented on a change in pull request #8497:
URL: https://github.com/apache/tvm/pull/8497#discussion_r697324406



##########
File path: tests/cpp/build_module_test.cc
##########
@@ -199,3 +199,129 @@ TEST(BuildModule, Heterogeneous) {
     ICHECK_LT(std::fabs(p_out[i] - (i + (i + 1.0) - (i - 1.0))), 1e-5);
   }
 }
+
+TEST(BuildModule, ZeroCopy) {
+  /*
+   *
+   *          A    B
+   *           \  /
+   *      elemwise_add(out0)
+   *              \
+   *       C      copy
+   *        \      /
+   *      elemwise_sub(out1)
+   */
+
+  using namespace tvm;
+  using namespace tvm::te;
+
+  auto target_llvm = Target("llvm");
+
+  // The shape of input tensors.
+  const int n = 4;
+  Array<PrimExpr> shape{n};
+
+  auto A = placeholder(shape, DataType::Float(32), "A");
+  auto B = placeholder(shape, DataType::Float(32), "B");
+  auto C = placeholder(shape, DataType::Float(32), "C");
+
+  auto elemwise_add = compute(
+      A->shape, [&A, &B](PrimExpr i) { return A[i] + B[i]; }, "elemwise_add");
+
+  auto copy = placeholder(shape, DataType::Float(32), "__copy");
+  auto elemwise_sub = compute(
+      C->shape, [&copy, &C](PrimExpr i) { return copy[i] - C[i]; }, 
"elemwise_sub");
+
+  With<Target> llvm_scope(target_llvm);
+  auto s1 = create_schedule({elemwise_add->op});
+  auto s2 = create_schedule({elemwise_sub->op});
+
+  auto args1 = Array<Tensor>({A, B, elemwise_add});
+  auto args2 = Array<Tensor>({copy, C, elemwise_sub});
+
+  std::unordered_map<Tensor, Buffer> binds;
+  auto lowered_s1 = LowerSchedule(s1, args1, "elemwise_add", binds);
+  auto lowered_s2 = LowerSchedule(s2, args2, "elemwise_sub", binds);
+  Map<tvm::Target, IRModule> inputs = {{target_llvm, lowered_s1}, 
{target_llvm, lowered_s2}};
+  auto module = build(inputs, Target());
+
+  // Execute the graph and check the correctness.
+  // Setup graph json.
+  std::string json =

Review comment:
       I also don't understand why there is no Python API. It is totally 
reasonable to pass pre-allocated `ndarray` storage for `get_output`. 




-- 
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.

To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org

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