ashutosh-arm commented on a change in pull request #10100: URL: https://github.com/apache/tvm/pull/10100#discussion_r796561427
########## File path: src/relay/backend/contrib/cmsisnn/scalar_to_tensor_constant.cc ########## @@ -0,0 +1,188 @@ +/* + * 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. + */ +/*! + * \file scalar_to_tensor_constant.cc + * \brief Converts scalar constant into tensor constant for binary ops of CMSIS-NN + */ + +#include <tvm/relay/attrs/nn.h> +#include <tvm/relay/attrs/transform.h> +#include <tvm/relay/expr_functor.h> +#include <tvm/relay/transform.h> +#include <tvm/runtime/ndarray.h> + +#include "../../../op/make_op.h" +#include "../../../qnn/utils.h" +#include "../../../transforms/pattern_utils.h" + +namespace tvm { +namespace relay { +namespace contrib { +namespace cmsisnn { + +/*! + * \brief This Mutator finds all partitioned functions meant for CMSIS-NN binary ops. + * Then, it substitutes the scalar constants with tensor constants. It makes the shape of this + * new constant same as that of the neighbouring constant of the other binary operand. The + * expectation is that the ExtractConstant pass would later kick this tensor constant out of the + * global partitioned function, thus making the entire global partitioned and its composite function + * constant free. This makes the TIR generation for binary ops via CMSIS-NN independent of + * constants. + */ +class ScalarToTensorConstantMutator : public MixedModeMutator { + public: + explicit ScalarToTensorConstantMutator(const IRModule& mod) : mod_(mod) {} + + private: + using MixedModeMutator::VisitExpr_; + + // Here is an example with the annotated scalar constant: + // def @tvmgen_default_cmsis_nn_main_1(%cmsis_nn_input: Tensor[], Inline=1, Compiler="cmsis-nn", + // global_symbol="tvmgen_default_cmsis_nn_main", + // Primitive=1) -> Tensor[] { + // %56 = fn (%input0: _scalar_constant_, %input1: Tensor[], + // PartitionedFromPattern="qnn.mul_", Composite="cmsis-nn.qnn_mul") -> Tensor[] { + // qnn.mul(%input0, %input1, scale0, zero_point0, + // scale1, zero_point_1, output_scale, output_zero_point) + // }; + // %56(meta[relay.Constant] /* _scalar constant_ */, %cmsis-nn_input) + // } + Expr Rewrite_(const CallNode* call, const Expr& post) final { + Expr final_call = post; + call = post.as<CallNode>(); + + // Create a new variable argument that is of the same shape as the neibouring argument + // in the binary op. This needs to be done only when one of the arguments is a scalar. + if (auto* opnode = call->op.as<OpNode>()) { + String op_name = opnode->name; + Array<Expr> new_args; + for (uint32_t i = 0; i < call->args.size(); ++i) { + Expr arg = call->args[i]; + new_args.push_back(arg); + if (!arg->checked_type_.defined()) { + continue; + } + auto* arg_type = arg->type_as<TensorTypeNode>(); + if (arg_type->shape.size() != 0 || arg.as<ConstantNode>()) { + continue; + } + String arg_name = arg.as<VarNode>()->name_hint(); + int tensor_arg_id = (i + 1) % 2; + Expr tensor_arg = call->args[tensor_arg_id]; + if (!tensor_arg->checked_type_.defined()) { + continue; + } + TensorType tensor_type = GetRef<TensorType>(tensor_arg->type_as<TensorTypeNode>()); + new_args.Set(i, Var(arg_name, tensor_type)); + } + final_call = Call(call->op, new_args, call->attrs, {}); Review comment: Sure, that would make sense -- 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]
