Skip to content

Commit 4c5f0c0

Browse files
mbaretleo-blonktristan-armRamana Radhakrishnan
authored
[BYOC][ETHOSN] Introduce further operator support (apache#6355)
* [BYOC][ETHOSN] Introduce further operator support This PR introduces support for the following operators: - Quantized Fully Connected - Quantized Addition - Depth-to-space - Max/Avg Pool 2D - Quantized Relu (Clip) - Reshape - Quantized Sigmoid Co-authored-by: Leo Blonk <[email protected]> Co-authored-by: Tristan O'Connor <[email protected]> Co-authored-by: Ramana Radhakrishnan <[email protected]> * Skip tf imports if not available Change-Id: I11bcf4a78014fa63e7b8e3b0cb00eecfd6cb7760 * ethos -> ethosn Change-Id: I1fb1a11d0765f6d69f04c24b9c24e08665b8af6a * Reduce random testing in test_addition Change-Id: Id06063a0a0cf5f01356df23dc5d4bbbcb47cfa99 * Reduce random testing in test fullyconnected Change-Id: I330408dfabc4bd804373f100581ce909ff724052 * Fix dumb mistake with rename Change-Id: I2c5007be485b323116a0e8bab0f9106ea5ec834b * Added comments to update the hashes in network tests when necessary Change-Id: I13828c918c959daa492b9ed942a882c86d6690d1 * Fix github name Change-Id: Idaa70ab9c2ec8db2828d51d15e7c23f28670ec82 * Use black formatting Change-Id: I538171bd547a16395bef155a1dad28e8b3e347f2 Co-authored-by: Leo Blonk <[email protected]> Co-authored-by: Tristan O'Connor <[email protected]> Co-authored-by: Ramana Radhakrishnan <[email protected]>
1 parent 69e4f44 commit 4c5f0c0

File tree

18 files changed

+1846
-24
lines changed

18 files changed

+1846
-24
lines changed

‎python/tvm/relay/op/contrib/ethosn.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,18 +57,109 @@ def qnn_conv_pattern():
5757
)
5858
return pattern
5959

60+
def qnn_fc_pattern():
61+
pattern = is_op("qnn.dense")(
62+
wildcard(), is_constant(), is_constant(), is_constant(), is_constant(), is_constant()
63+
)
64+
pattern = is_op("nn.bias_add")(pattern, is_constant())
65+
pattern = is_op("qnn.requantize")(
66+
pattern, is_constant(), is_constant(), is_constant(), is_constant()
67+
)
68+
return pattern
69+
70+
def qnn_avg_pool2d_pattern():
71+
pattern = is_op("cast")(wildcard())
72+
pattern = is_op("nn.avg_pool2d")(pattern)
73+
pattern = is_op("cast")(pattern)
74+
return pattern
75+
76+
def qnn_sigmoid_pattern():
77+
pattern = is_op("qnn.dequantize")(wildcard(), is_constant(), is_constant())
78+
pattern = is_op("sigmoid")(pattern)
79+
pattern = is_op("qnn.quantize")(pattern, is_constant(), is_constant())
80+
return pattern
81+
6082
def check_conv2d(extract):
6183
"""Check if a conv2d is supported by Ethos-N."""
6284
if not ethosn_available():
6385
return False
6486

6587
return support.conv2d(extract)
6688

89+
def check_fc(extract):
90+
"""Check if a fully connected is supported by Ethos-N."""
91+
if not ethosn_available():
92+
return False
93+
94+
return support.fc(extract)
95+
96+
def check_avg_pool2d(extract):
97+
"""Check if a avg pool2d is supported by Ethos-N."""
98+
if not ethosn_available():
99+
return False
100+
101+
return support.avg_pool2d(extract)
102+
103+
def check_sigmoid(extract):
104+
"""Check if a sigmoid is supported by Ethos-N."""
105+
if not ethosn_available():
106+
return False
107+
108+
if extract.attrs.out_dtype != "uint8":
109+
return False
110+
111+
return support.sigmoid(extract)
112+
67113
return [
68114
("ethos-n.qnn_conv2d", qnn_conv_pattern(), check_conv2d),
115+
("ethos-n.qnn_avg_pool2d", qnn_avg_pool2d_pattern(), check_avg_pool2d),
116+
("ethos-n.qnn_sigmoid", qnn_sigmoid_pattern(), check_sigmoid),
117+
("ethos-n.qnn_fc", qnn_fc_pattern(), check_fc),
69118
]
70119

71120

121+
def _is_ethosn_composite(node):
122+
if isinstance(node, tvm.relay.expr.Call) and isinstance(node.op, tvm.relay.Function):
123+
if "Composite" in node.op.attrs:
124+
comp_name = node.op.attrs["Composite"]
125+
return comp_name.split(".")[0] == "ethos-n"
126+
127+
return False
128+
129+
130+
@tvm.ir.register_op_attr("nn.max_pool2d", "target.ethos-n")
131+
def max_pool2d(attrs, args):
132+
"""Check if a max pool2d is supported by Ethos-N."""
133+
if not ethosn_available():
134+
return False
135+
136+
pool = tvm.relay.nn.max_pool2d(*args, **attrs)
137+
return support.max_pool2d(pool)
138+
139+
140+
@tvm.ir.register_op_attr("reshape", "target.ethos-n")
141+
def reshape(attrs, args):
142+
"""Check if a reshape is supported by Ethos-N."""
143+
if not ethosn_available():
144+
return False
145+
146+
if not _is_ethosn_composite(args[0]):
147+
return False
148+
149+
rs = tvm.relay.op.reshape(*args, attrs["newshape"])
150+
return support.reshape(rs)
151+
152+
153+
@tvm.ir.register_op_attr("qnn.add", "target.ethos-n")
154+
def qnn_add(attrs, args):
155+
"""Check if an addition is supported by Ethos-N."""
156+
if not ethosn_available():
157+
return False
158+
159+
add = _qnn.op.add(*args)
160+
return support.addition(add)
161+
162+
72163
@tvm.ir.register_op_attr("qnn.concatenate", "target.ethos-n")
73164
def qnn_concatenate(attrs, args):
74165
"""Check if a concatenate is supported by Ethos-N."""
@@ -116,3 +207,29 @@ def split(attrs, args):
116207
return False
117208

118209
return True
210+
211+
212+
@tvm.ir.register_op_attr("nn.depth_to_space", "target.ethos-n")
213+
def depth_to_space(attrs, args):
214+
"""Check if a depth_to_space is supported by Ethos-N."""
215+
if not ethosn_available():
216+
return False
217+
218+
depth = tvm.relay.nn.depth_to_space(*args, **attrs)
219+
if not support.depth_to_space(depth):
220+
return False
221+
222+
return True
223+
224+
225+
@tvm.ir.register_op_attr("clip", "target.ethos-n")
226+
def clip(attrs, args):
227+
"""Check if a clip is supported by Ethos-N."""
228+
if not ethosn_available():
229+
return False
230+
231+
c = tvm.relay.clip(*args, **attrs)
232+
if not support.relu(c):
233+
return False
234+
235+
return True

‎src/relay/backend/contrib/ethosn/codegen.cc‎

Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,34 @@ void InferTensorsVisitor::InferCall(const CallNode* cn) {
8383
ConvolutionParams params;
8484
err += EthosnAPI::QnnConv2d(cn->op.as<FunctionNode>()->body, &params);
8585
tensor_table_[cn->args[0]] = {params.activation_info};
86+
} else if (IsEthosnFunc(call, "ethos-n.qnn_fc")) {
87+
FullyConnectedParams params;
88+
err += EthosnAPI::QnnFullyConnected(cn->op.as<FunctionNode>()->body, &params);
89+
tensor_table_[cn->args[0]] = {params.input_info};
90+
} else if (IsEthosnOp(call, "nn.max_pool2d")) {
91+
MaxPool2DParams params;
92+
params.input_info = GetTensorInfo(tensor_table_, call);
93+
err += EthosnAPI::MaxPool2D(call, &params);
94+
tensor_table_[cn->args[0]] = {params.input_info};
95+
} else if (IsEthosnFunc(call, "ethos-n.qnn_avg_pool2d")) {
96+
AvgPool2DParams params;
97+
params.input_info = GetTensorInfo(tensor_table_, call);
98+
err += EthosnAPI::AvgPool2D(cn->op.as<FunctionNode>()->body, &params);
99+
tensor_table_[cn->args[0]] = {params.input_info};
100+
} else if (IsEthosnOp(call, "reshape")) {
101+
ReshapeParams params;
102+
params.input_info = GetTensorInfo(tensor_table_, call);
103+
err += EthosnAPI::Reshape(call, &params);
104+
tensor_table_[cn->args[0]] = {params.input_info};
105+
} else if (IsEthosnOp(call, "qnn.add")) {
106+
AdditionParams params;
107+
err += EthosnAPI::Addition(call, &params);
108+
tensor_table_[cn->args[0]] = {params.lhs_info};
109+
tensor_table_[cn->args[1]] = {params.rhs_info};
110+
} else if (IsEthosnFunc(call, "ethos-n.qnn_sigmoid")) {
111+
SigmoidParams params;
112+
err += EthosnAPI::Sigmoid(cn->op.as<FunctionNode>()->body, &params);
113+
tensor_table_[cn->args[0]] = {params.input_info};
86114
} else if (IsEthosnOp(call, "qnn.concatenate")) {
87115
ConcatenateParams params;
88116
err = EthosnAPI::Concatenate(call, &params);
@@ -92,6 +120,16 @@ void InferTensorsVisitor::InferCall(const CallNode* cn) {
92120
params.input_info = GetTensorInfo(tensor_table_, call);
93121
err = EthosnAPI::Split(call, &params);
94122
tensor_table_[cn->args[0]] = {params.input_info};
123+
} else if (IsEthosnOp(call, "nn.depth_to_space")) {
124+
DepthToSpaceParams params;
125+
params.input_info = GetTensorInfo(tensor_table_, call);
126+
err += EthosnAPI::DepthToSpace(call, &params);
127+
tensor_table_[cn->args[0]] = {params.input_info};
128+
} else if (IsEthosnOp(call, "clip")) {
129+
ReluParams params;
130+
params.input_info = GetTensorInfo(tensor_table_, call);
131+
err = EthosnAPI::Relu(call, &params);
132+
tensor_table_[cn->args[0]] = {params.input_info};
95133
} else {
96134
err = EthosnError("unknown operator");
97135
}
@@ -198,12 +236,36 @@ sl::TensorsAndId ConstructNetworkVisitor::HandleCall(const CallNode* cn) {
198236
if (IsEthosnFunc(call, "ethos-n.qnn_conv2d")) {
199237
if ((err = MakeConvolutionLayer(call, &tensor))) ReportFatalError(call, err);
200238
return MakeOps(tensor);
239+
} else if (IsEthosnFunc(call, "ethos-n.qnn_fc")) {
240+
if ((err = MakeFullyConnectedLayer(call, &tensor))) ReportFatalError(call, err);
241+
return MakeOps(tensor);
242+
} else if (IsEthosnOp(call, "nn.max_pool2d")) {
243+
if ((err = MakeMaxPool2DLayer(call, &tensor))) ReportFatalError(call, err);
244+
return MakeOps(tensor);
245+
} else if (IsEthosnFunc(call, "ethos-n.qnn_avg_pool2d")) {
246+
if ((err = MakeAvgPool2DLayer(call, &tensor))) ReportFatalError(call, err);
247+
return MakeOps(tensor);
248+
} else if (IsEthosnOp(call, "reshape")) {
249+
if ((err = MakeReshapeLayer(call, &tensor))) ReportFatalError(call, err);
250+
return MakeOps(tensor);
251+
} else if (IsEthosnOp(call, "qnn.add")) {
252+
if ((err = MakeAdditionLayer(call, &tensor))) ReportFatalError(call, err);
253+
return MakeOps(tensor);
254+
} else if (IsEthosnFunc(call, "ethos-n.qnn_sigmoid")) {
255+
if ((err = MakeSigmoidLayer(call, &tensor))) ReportFatalError(call, err);
256+
return MakeOps(tensor);
201257
} else if (IsEthosnOp(call, "qnn.concatenate")) {
202258
if ((err = MakeConcatenateLayer(call, &tensor))) ReportFatalError(call, err);
203259
return MakeOps(tensor);
204260
} else if (IsEthosnOp(call, "split")) {
205261
if ((err = MakeSplitLayer(call, &tensors))) ReportFatalError(call, err);
206262
return tensors;
263+
} else if (IsEthosnOp(call, "nn.depth_to_space")) {
264+
if ((err = MakeDepthToSpaceLayer(call, &tensor))) ReportFatalError(call, err);
265+
return MakeOps(tensor);
266+
} else if (IsEthosnOp(call, "clip")) {
267+
if ((err = MakeReluLayer(call, &tensor))) ReportFatalError(call, err);
268+
return MakeOps(tensor);
207269
} else {
208270
ReportFatalError(call, EthosnError("unknown operator"));
209271
return {};
@@ -266,6 +328,115 @@ EthosnError ConstructNetworkVisitor::MakeConvolutionLayer(const Call& call,
266328
return EthosnError();
267329
}
268330

331+
EthosnError ConstructNetworkVisitor::MakeFullyConnectedLayer(const Call& call,
332+
sl::TensorAndId<sl::Operand>* out) {
333+
FullyConnectedParams params;
334+
if (auto err = EthosnAPI::QnnFullyConnected(call->op.as<FunctionNode>()->body, &params)) {
335+
return err;
336+
}
337+
338+
auto weights = AddConstant(network_, params.weights_info, params.raw_weights).tensor;
339+
auto bias = AddConstant(network_, params.bias_info, params.raw_bias).tensor;
340+
try {
341+
auto input =
342+
AddReshape(network_, *operand_table_[call->args[0]][0], params.input_info.m_Dimensions)
343+
.tensor;
344+
*out = AddFullyConnected(network_, *input, *bias, *weights, params.fc_info);
345+
} catch (const sl::NotSupportedException& e) {
346+
return EthosnError(e.what());
347+
}
348+
return EthosnError();
349+
}
350+
351+
EthosnError ConstructNetworkVisitor::MakeMaxPool2DLayer(const Call& call,
352+
sl::TensorAndId<sl::Operand>* out) {
353+
MaxPool2DParams params;
354+
params.input_info = GetTensorInfo(tensor_table_, call);
355+
if (auto err = EthosnAPI::MaxPool2D(call, &params)) {
356+
return err;
357+
}
358+
359+
auto input = operand_table_[call->args[0]][0];
360+
361+
try {
362+
*out = AddPooling(network_, *input, params.pool_info);
363+
} catch (const sl::NotSupportedException& e) {
364+
return EthosnError(e.what());
365+
}
366+
return EthosnError();
367+
}
368+
369+
EthosnError ConstructNetworkVisitor::MakeAvgPool2DLayer(const Call& call,
370+
sl::TensorAndId<sl::Operand>* out) {
371+
AvgPool2DParams params;
372+
params.input_info = GetTensorInfo(tensor_table_, call);
373+
if (auto err = EthosnAPI::AvgPool2D(call->op.as<FunctionNode>()->body, &params)) {
374+
return err;
375+
}
376+
377+
auto input = operand_table_[call->args[0]][0];
378+
379+
try {
380+
*out = AddPooling(network_, *input, params.pool_info);
381+
} catch (const sl::NotSupportedException& e) {
382+
return EthosnError(e.what());
383+
}
384+
return EthosnError();
385+
}
386+
387+
EthosnError ConstructNetworkVisitor::MakeReshapeLayer(const Call& call,
388+
sl::TensorAndId<sl::Operand>* out) {
389+
ReshapeParams params;
390+
params.input_info = GetTensorInfo(tensor_table_, call);
391+
if (auto err = EthosnAPI::Reshape(call, &params)) {
392+
return err;
393+
}
394+
395+
auto input = operand_table_[call->args[0]][0];
396+
397+
try {
398+
*out = AddReshape(network_, *input, params.new_shape);
399+
} catch (const sl::NotSupportedException& e) {
400+
return EthosnError(e.what());
401+
}
402+
return EthosnError();
403+
}
404+
405+
EthosnError ConstructNetworkVisitor::MakeAdditionLayer(const Call& call,
406+
sl::TensorAndId<sl::Operand>* out) {
407+
AdditionParams params;
408+
if (auto err = EthosnAPI::Addition(call, &params)) {
409+
return err;
410+
}
411+
412+
auto lhs = operand_table_[call->args[0]][0];
413+
auto rhs = operand_table_[call->args[1]][0];
414+
415+
try {
416+
*out = AddAddition(network_, *lhs, *rhs, params.output_quantization_info);
417+
} catch (const sl::NotSupportedException& e) {
418+
return EthosnError(e.what());
419+
}
420+
return EthosnError();
421+
}
422+
423+
EthosnError ConstructNetworkVisitor::MakeSigmoidLayer(const Call& call,
424+
sl::TensorAndId<sl::Operand>* out) {
425+
SigmoidParams params;
426+
if (auto err = EthosnAPI::Sigmoid(call->op.as<FunctionNode>()->body, &params)) {
427+
return err;
428+
}
429+
430+
auto input = operand_table_[call->args[0]][0];
431+
432+
try {
433+
*out = AddSigmoid(network_, *input);
434+
} catch (const sl::NotSupportedException& e) {
435+
return EthosnError(e.what());
436+
}
437+
return EthosnError();
438+
}
439+
269440
EthosnError ConstructNetworkVisitor::MakeConcatenateLayer(const Call& call,
270441
sl::TensorAndId<sl::Operand>* out) {
271442
ConcatenateParams params;
@@ -304,6 +475,42 @@ EthosnError ConstructNetworkVisitor::MakeSplitLayer(const Call& call, sl::Tensor
304475
return EthosnError();
305476
}
306477

478+
EthosnError ConstructNetworkVisitor::MakeDepthToSpaceLayer(const Call& call,
479+
sl::TensorAndId<sl::Operand>* out) {
480+
DepthToSpaceParams params;
481+
params.input_info = GetTensorInfo(tensor_table_, call);
482+
if (auto err = EthosnAPI::DepthToSpace(call, &params)) {
483+
return err;
484+
}
485+
486+
auto input = operand_table_[call->args[0]][0];
487+
488+
try {
489+
*out = AddDepthToSpace(network_, *input, params.depth_info);
490+
} catch (const sl::NotSupportedException& e) {
491+
return EthosnError(e.what());
492+
}
493+
return EthosnError();
494+
}
495+
496+
EthosnError ConstructNetworkVisitor::MakeReluLayer(const Call& call,
497+
sl::TensorAndId<sl::Operand>* out) {
498+
ReluParams params;
499+
params.input_info = GetTensorInfo(tensor_table_, call);
500+
if (auto err = EthosnAPI::Relu(call, &params)) {
501+
return err;
502+
}
503+
504+
auto input = operand_table_[call->args[0]][0];
505+
506+
try {
507+
*out = AddRelu(network_, *input, params.relu_info);
508+
} catch (const sl::NotSupportedException& e) {
509+
return EthosnError(e.what());
510+
}
511+
return EthosnError();
512+
}
513+
307514
runtime::Module EthosnCompiler::CreateRuntimeModule(const ObjectRef& ref) {
308515
std::vector<runtime::ethosn::OrderedCompiledNetwork> cmms;
309516
if (ref->IsInstance<FunctionNode>()) {

0 commit comments

Comments
 (0)