mbs-octoml commented on a change in pull request #9038:
URL: https://github.com/apache/tvm/pull/9038#discussion_r717093154
##########
File path: src/relay/op/annotation/annotation.cc
##########
@@ -56,12 +83,98 @@ RELAY_REGISTER_OP("on_device")
.set_attr<TOpPattern>("TOpPattern", kOpaque)
.set_attr<TOpIsStateful>("TOpIsStateful", false)
.set_attr<FInferCorrectLayout>("FInferCorrectLayout",
ElemwiseArbitraryLayout)
+ .set_attr<TNonComputational>("TNonComputational", true)
.set_attr<FTVMCompute>("FTVMCompute",
[](const Attrs& attrs, const Array<te::Tensor>&
inputs,
const Type& out_type) -> Array<te::Tensor> {
return {topi::identity(inputs[0])};
});
+OnDeviceProps GetOnDeviceProps(const CallNode* call_node) {
+ if (call_node->op == OnDeviceOp()) {
+ ICHECK_EQ(call_node->args.size(), 1) << "on_device expects one argument";
Review comment:
I've gotten all the way through to unit testing the final PR without
needing it, so how about I leave it. I tend to add them as I need them.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
