ashutosh-arm commented on a change in pull request #9331:
URL: https://github.com/apache/tvm/pull/9331#discussion_r733917252
##########
File path: src/relay/backend/contrib/cmsisnn/relay_to_tir.cc
##########
@@ -43,23 +44,153 @@ class RelayToTIRVisitor : public MixedModeVisitor {
inline IntImm ToArg(int32_t value) { return IntImm(DataType::Int(32),
value); }
void CreatePrimFuncForExtern(Array<tir::Var> func_signature,
- tvm::Array<PrimExpr> call_extern_args) {
+ tvm::Array<PrimExpr> call_extern_args, int
context_buffer_size = 0) {
Map<String, ObjectRef> dict_attrs;
dict_attrs.Set("global_symbol", func_name_);
dict_attrs.Set("tir.noalias", Bool(true));
tir::Stmt body = tir::Evaluate(
tvm::tir::Call(DataType::Int(8), tir::builtin::call_extern(),
call_extern_args));
+ if (context_buffer_size) {
+ // TODO(@ashutosh-arm) while supporting MVE, we need to move allocation
through TVMBAW
+ tir::Var buffer_var("context_buffer",
PointerType(PrimType(DataType::Int(8)), "global"));
+ body = tir::Allocate(buffer_var, DataType::Int(8),
{context_buffer_size}, tir::const_true(),
+ body);
+ }
+
primfunc_ = tir::PrimFunc(func_signature, body, VoidType(), Map<tir::Var,
tir::Buffer>(),
DictAttrs(dict_attrs));
}
+ void EmitConv2D(const Expr& expr) {
+ const CallNode* clip_call = nullptr;
+ const CallNode* requantize_call = nullptr;
+ const CallNode* bias_add_call = nullptr;
+ const CallNode* conv2d_call = nullptr;
+ auto* final_call = expr.as<CallNode>();
+ auto* final_op = final_call->op.as<OpNode>();
+ if (final_op->name == "clip") {
+ clip_call = final_call;
+ requantize_call = clip_call->args[0].as<CallNode>();
+ } else {
+ requantize_call = final_call;
+ }
+ auto* requantize_input = requantize_call->args[0].as<CallNode>();
+ auto* requantize_input_op = requantize_input->op.as<OpNode>();
+ if (requantize_input_op->name == "nn.bias_add") {
+ bias_add_call = requantize_input;
+ conv2d_call = bias_add_call->args[0].as<CallNode>();
+ } else {
+ conv2d_call = requantize_input;
+ }
+
+ // TIR variables are created in the order they appear in the Relay
partitioned function
+ // %1 = qnn.conv2d(%input, %weight_const_0, input_zero_point_scalar,
+ // %cmsisnn_multiplier_const_1, %input_scale_scalar,
%weight_scale_const_2)
+ // %2 = nn.bias_add(%1, %bias_const_3, axis=3)
+ // %3 = qnn.requantize(%2, %input_scale_const_4, %cmsisnn_shift_const_5,
+ // %output_scale_scalar, %output_zero_point_scalar)
+ // clip(%3, a_min=%min_scalar, a_max=%max_scalar)
+ auto in_var = tir::Var("input", DataType::Handle(8));
+ auto const_var0 = tir::Var("filter", DataType::Handle(8)); //
weight
+ auto const_var1 = tir::Var("multiplier", DataType::Handle(32)); //
quant multiplier
+ auto const_var2 = tir::Var("filter_scale", DataType::Handle(32)); //
weight scale
+ auto const_var3 = tir::Var("bias", DataType::Handle(32)); // bias
+ auto const_var4 = tir::Var("input_scale", DataType::Handle(32)); //
input_scale * weight_scale
+ auto const_var5 = tir::Var("shift", DataType::Handle(32)); //
quant shift
+ auto out_var = tir::Var("output", DataType::Handle(8));
+
+ // Individual arguments to the structs arguments of the CMSIS-NN API are
filled into call_extern
+ //
https://github.com/ARM-software/CMSIS_5/blob/def6f800f95661eb3451d317f7d0dde504f6020d/CMSIS/NN/Source/ConvolutionFunctions/arm_convolve_wrapper_s8.c#L50
+
+ // prepare cmsis_nn_conv_params
+ auto* conv2d_attrs = conv2d_call->attrs.as<Conv2DAttrs>();
+ int32_t input_offset =
-GetScalarFromConstant<int32_t>(conv2d_call->args[2]);
+ int32_t output_offset =
GetScalarFromConstant<int32_t>(requantize_call->args[4]);
+ int32_t stride_w = qnn::get_const_int(conv2d_attrs->strides[1]);
+ int32_t stride_h = qnn::get_const_int(conv2d_attrs->strides[0]);
+ int32_t padding_w = qnn::get_const_int(conv2d_attrs->padding[1]);
+ int32_t padding_h = qnn::get_const_int(conv2d_attrs->padding[0]);
+ int32_t dilation_w = qnn::get_const_int(conv2d_attrs->dilation[1]);
+ int32_t dilation_h = qnn::get_const_int(conv2d_attrs->dilation[0]);
+ int32_t clip_min, clip_max;
+ if (clip_call) {
+ auto* clip_attrs = clip_call->attrs.as<ClipAttrs>();
+ clip_min = clip_attrs->a_min;
+ clip_max = clip_attrs->a_max;
+ } else {
+ clip_min = -128;
+ clip_max = 127;
+ }
+
+ // cmsis_nn_dims *input_dims
Review comment:
That looks more readable. Done!
--
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
