Skip to content
This repository has been archived by the owner on Oct 25, 2023. It is now read-only.

[TUZ-150] Add a simplified access point for Unity Flow #32

Merged
merged 162 commits into from
Mar 8, 2023
Merged
Changes from 1 commit
Commits
Show all changes
162 commits
Select commit Hold shift + click to select a range
f7165a1
[microTVM] Fix tvmc tutorial (#14076)
mehrdadh Feb 25, 2023
10fb8c5
[MetaSchedule] Introduce Async Pipeline in MultiLevelTiling (#14009)
cblmemo Feb 25, 2023
9fab56c
[TVMScript] Use op attribute to control whether to print dtype in TVM…
liangW-intellif Feb 25, 2023
1ad1994
[Fix][TVMScript] Fix index of metadata in printed script (#14130)
Ubospica Feb 25, 2023
f21a17b
[Pytorch] frontend full_impl fix (#14122)
Feb 26, 2023
d9b0a80
[DOCKER] Configurable NDK version support (#14000)
srkreddy1238 Feb 27, 2023
54a62c1
[Fix][TIR] SampleCategorical apply-to-schedule (#14133)
MasterJH5574 Feb 27, 2023
74603ee
[Arith] ConstIntBound was incorrectly assuming bounds were over int64…
Feb 27, 2023
0e046da
[CMSIS-NN] Reduction in code size of AOT test runner binary (#13815)
NicolaLancellotti Feb 27, 2023
77df6e8
[CMSIS-NN] Add a runtime error message (#13643)
NicolaLancellotti Feb 27, 2023
bf589f3
[CRT]Cleanup unused macros in crt_config.h.template (#14125)
mehrdadh Feb 27, 2023
663f7ae
[Fix][Relay] Fix axis transformation in squeeze shape function (#14135)
Lucien0 Feb 27, 2023
4d152fe
[Unittest] merge test_cp_async_in_if_then_else into test_tir_transfor…
cblmemo Feb 27, 2023
2feb243
[Frontend][TFLite] Fix conv2d import bug (#14124)
mehrdadh Feb 27, 2023
6097df5
[ONNX][TORCH] Replace scatter op by scatter_elements (#14019)
vvchernov Feb 28, 2023
2b2cb96
[TVMScript][Printer] Remove relax prefix for now (#14140)
tqchen Feb 28, 2023
7d67bb1
[microNPU] Sum legalization support (#13997)
Aleksei-grovety Feb 28, 2023
7c06de5
[Fix][MetaSchedule] Fix redundant stages in async pipeline for mlt (#…
cblmemo Feb 28, 2023
428400c
[COMMUNITY] Cheng Wen -> Reviewer (#14153)
Hzfengsy Mar 1, 2023
1043136
[Runtime] Fix high RAM usage when saving / loading paramters of big m…
masahi Mar 1, 2023
e9cf04e
[Relay][Frontend] Span Filling PyTorch (#14050)
chunit-quic Mar 1, 2023
6c04ac5
[TRT][BYOC] allow strided_slice ops on selected dimensions (#14142) (…
AreopagX Mar 1, 2023
69acdfb
[ONNX][TOPI] Add `DFT` operator (#13999)
KJlaccHoeUM9l Mar 1, 2023
908dc8f
[CRT][microTVM] Enable USMP by default for AoTExecutor + CRT runtime …
mehrdadh Mar 2, 2023
25f4d06
[Android] Fix using system libraries in Android apps (#14145)
echuraev Mar 2, 2023
05cbe32
[microTVM]Enable TVMC micro with AoT Executor (#14077)
mehrdadh Mar 2, 2023
bd8e7d3
[bugfix] Fix the write buffer scope of `mma_store_impl` (#14174)
yzh119 Mar 2, 2023
cb37b82
[Relay] Enhance EliminateCommonSubexpr to support Tuple argument (#14…
lixiaoquan Mar 2, 2023
91dc8ef
[TIR] Fix typo in doc (#14178)
vinx13 Mar 2, 2023
a42e98b
[microTVM] Use QNN schedules to give SOTA performance (#13752)
guberti Mar 2, 2023
bc92a3f
Add v0.11.0 docs link to site (#14181)
areusch Mar 3, 2023
df429c5
[TIR] Allow TransformLayout with non-inversible index map (#14095)
vinx13 Mar 3, 2023
c0f148a
[TIR][Analysis] Implement IdentifyMemCpy analysis function (#13947)
Lunderberg Mar 4, 2023
736ceca
[HotFix][MetaSchedule] Turn off database shash check (#14188)
MasterJH5574 Mar 4, 2023
22c47ee
[TOPI] Batch Norm Training Mode (#14190)
SiriusNEO Mar 4, 2023
baedf7f
[TOPI] Group normalization (#14193)
MasterJH5574 Mar 5, 2023
befdc4e
[Fix][TIR] LowerCrossThreadReduction with write-back predicate (#14199)
MasterJH5574 Mar 5, 2023
e7b02f2
[Unity] Relax VM (#13878)
YuchenJin Feb 1, 2023
9508a18
[Unity] Relax expressions and types (#13901)
YuchenJin Feb 2, 2023
4d46290
[Unity][IR] First-class StructInfo (#13907)
YuchenJin Feb 3, 2023
b59ad48
[Unity][CI] Unity specific jenkins setup (do not upstream to main) (#…
tqchen Feb 3, 2023
ff8bfa2
[Unity] Basic StructInfo Analysis and Expr construction (#13916)
YuchenJin Feb 5, 2023
bb0c129
[Unity] Relax BlockBuilder and ExprMutator (#13926)
YuchenJin Feb 7, 2023
1807e6f
[Unity] Relax TVMScript Parser. (#13932)
Hzfengsy Feb 8, 2023
846a2c5
[Unity] Relax TVMScript Printer (#13944)
junrushao Feb 10, 2023
850d6a4
[Unity] Relax VM codegen (#13954)
YuchenJin Feb 11, 2023
a966cf1
[Unity] Relax VM shape lowering pass (#13956)
YuchenJin Feb 11, 2023
f735d93
[Unity] e2e Relax minimum build flow (#13961)
YuchenJin Feb 11, 2023
ad4185c
[Unity][TVMScript] Use explicit `R.shape` in TVMScript (#13979)
Hzfengsy Feb 14, 2023
e8227b9
[Unity] Relax op: index (#13987)
MasterJH5574 Feb 14, 2023
2e08c8c
[Unity] Relax op: datatype (#13986)
MasterJH5574 Feb 14, 2023
68a04a8
[Unity] Relax op: set (#13990)
MasterJH5574 Feb 14, 2023
5723ebb
[Unity] Relax op: image (#13994)
MasterJH5574 Feb 14, 2023
72bca0f
[Unity] Relax op: arithmetic, comparison (#13983)
MasterJH5574 Feb 14, 2023
f491b96
[Unity] Relax op: statistical (#13991)
MasterJH5574 Feb 14, 2023
71437f7
[Unity] Relax op: neural networks (#13993)
MasterJH5574 Feb 14, 2023
9113fc9
[Unity] Relax op: creation (#13984)
MasterJH5574 Feb 14, 2023
c788135
[Unity] Relax op: linear algebra (#13988)
MasterJH5574 Feb 14, 2023
b6818bb
[Unity] Relax op: search (#13992)
MasterJH5574 Feb 14, 2023
17cf446
[Unity] Relax op: manipulation (#13989)
MasterJH5574 Feb 14, 2023
9317ec8
[Unity] NestedMsg Support utility (#13995)
tqchen Feb 14, 2023
07e0dfb
[Unity][Pass] Operator Fusion Passes (#14001)
Hzfengsy Feb 15, 2023
33c4aab
[Unity][Pass] LambdaLift pass (#14012)
yongwww Feb 16, 2023
fd5c73d
[Unity][VM] Supporting "compiled" exec mode. (#14015)
tqchen Feb 17, 2023
733fc00
[Unity][Pass] BindParams pass, FoldConstant pass (#14016)
sunggg Feb 17, 2023
88852c1
[Unity][Pass][TuningAPI] Introduce TuningAPI and MetaSchedule pass (#…
sunggg Feb 17, 2023
24470c9
[Unity] Relay -> Relax translator (#14026)
YuchenJin Feb 17, 2023
fc0540c
[Unity][Pass] Normalize Pass (#14031)
LeshengJin Feb 18, 2023
91adf7b
[Unity][BlockBuilder] CallTE convert PrimValue args (#14028)
MasterJH5574 Feb 18, 2023
53f800d
[Unity][Pass] Wellformed Analysis (#14032)
LeshengJin Feb 18, 2023
596d472
[Unity][TVMScript] Move tir/relax import in script out of __init__.py…
MasterJH5574 Feb 18, 2023
b150b1a
[Unity][Pass] Operator legalization (#14029)
MasterJH5574 Feb 18, 2023
449e094
[Unity][Op] Add ShapeExpr Tests for Reshape Op (#14035)
Ubospica Feb 18, 2023
9b5f214
[Unity] Initial PyTorch Frontend (#14037)
MasterJH5574 Feb 18, 2023
fd35d1e
[Unity][Pass] Block-level static memory planning (#14038)
MasterJH5574 Feb 18, 2023
d8fdd5c
[Unity] Disallow inline prim_func in relax IR (#14040)
yongwww Feb 18, 2023
8039f6a
[Unity] Update tests to adapt to latest TVMScript syntax (#14039)
MasterJH5574 Feb 18, 2023
782c632
[Unity] Relax dataflow pattern language (matching) (#14041)
ganler Feb 18, 2023
1b85765
[Unity] Statement rewriter for DataflowBlock (#14043)
ganler Feb 19, 2023
c1439b3
[Unity][Pass] FuseOps FuseTIR fixes (#14044)
MasterJH5574 Feb 19, 2023
180bead
[Unity][TVMScript] Overload `__neg__` for relax expr (#14045)
SiriusNEO Feb 19, 2023
fe528f6
[Unity][VM] Add per-op profiling support (#14053)
masahi Feb 20, 2023
aa55c05
[Unity][BYOC] Add pattern-based partitioning pass (#14054)
masahi Feb 20, 2023
fada709
[Unity] Relax op: collapse sum (#14059)
SiriusNEO Feb 21, 2023
06de35e
[Unity][Fix][Pass] Fix FuseOps for lack graph edges (#14058)
MasterJH5574 Feb 21, 2023
a9032d9
[Unity][Pass] Remove Unused Function (#14061)
sunggg Feb 21, 2023
ed2696a
[Unity][BYOC] Add pass to merge composite functions to offload large …
masahi Feb 21, 2023
0d58835
[Unity][Frontend] Annotate number of non-static input of FX function …
vinx13 Feb 21, 2023
246c4c1
[Unity][Transform] Add LiftTransformParams pass (#14069)
vinx13 Feb 21, 2023
e8a0c4d
[Unity][BYOC][Pass] RunCodegen and TensorRT (#14078)
sunggg Feb 22, 2023
4ad8d64
[Unity][Pass] Canonicalize Bindings (#14079)
YuchenJin Feb 22, 2023
3a64963
[Unity] Add testcases for `expr_args_converter` (#14080)
Hzfengsy Feb 22, 2023
b8460eb
[Unity][BYOC] Add CUTLASS backend (#14081)
masahi Feb 22, 2023
defc15b
[Unity][BYOC] Add DNNL backend (#14082)
masahi Feb 22, 2023
7645aa7
[Unity][Op] `log_softmax` and `cross_entropy_with_logits` (#14083)
SiriusNEO Feb 22, 2023
6e2d7bb
[Unity][Analysis] TIR pattern kind analysis for multi-buffer write bl…
MasterJH5574 Feb 22, 2023
acd0e0b
[Unity][Fix][Pass] FoldConstant with DCE in dataflow block (#14087)
MasterJH5574 Feb 22, 2023
1950940
[Unity] Refactor Relax Build JIT UX (#14088)
tqchen Feb 22, 2023
cf36b7b
[Unity][Relax] Set Shape Function to Be Host Function (#14090)
zxybazh Feb 22, 2023
b1f2d53
[Unity] Fix typo in the comment (#14096)
vinx13 Feb 22, 2023
a8338e6
[Unity] Lower `shape_of` to a builtin (#14093)
YuchenJin Feb 22, 2023
74f3007
[Unity] Relax Recursive function (#14092)
yongwww Feb 23, 2023
98d0a01
[Unity][Layout] Add layout transformation analysis for PrimFunc (#14066)
psrivas2 Feb 23, 2023
b755a6f
[Unity] Remove attributes of relax.print, assert and unique (#14101)
yongwww Feb 23, 2023
eaaa1fb
[Unity][BYOC]Add relax backend pattern registry (#14106)
yelite Feb 24, 2023
89bb68b
[Unity] Update tests again to adapt to latest TVMScript syntax (#14115)
Ubospica Feb 24, 2023
02b3a1f
[Unity][Fix] Fix bug in MergeCompositeFunctions (#14117)
Ubospica Feb 24, 2023
c7d2c38
[Unity][BlockBuilder] Add `name_hint` argument for `emit` and `emit_o…
SiriusNEO Feb 25, 2023
57c86eb
[Unity][WEB] Relax vm on web runtime (#14131)
tqchen Feb 25, 2023
61c2761
[Unity] Add Global info (#14132)
jinhongyii Feb 26, 2023
368d9f6
[Unity][BYOC] Add transposed matmul support to Relax CUTLASS BYOC (#1…
yelite Feb 27, 2023
4713b52
[Unity][TVMScript] emit_te sugar (#14123)
yongwww Feb 27, 2023
28c6825
[Unity][BYOC] Assign group to unused bindings and ignroe PrimFunc (#1…
vinx13 Feb 27, 2023
63ce37d
[Unity] Add callback to FuseOpsByPattern to check match result is acc…
vinx13 Feb 27, 2023
7a5d313
[Unity][Legalize] Fix Scalar Constant Legalization (#14127)
zxybazh Feb 28, 2023
c29ac7e
[Unity][Pass] Enhance constant folding to fold relax ops by evaluatin…
psrivas2 Feb 28, 2023
4c085d2
[Unity][Debugging] AST printer (#14152)
slyubomirsky Mar 1, 2023
8a1e623
[Unity][Pass] Support Symbolic Shape Deduction during BindParam (#14154)
Hzfengsy Mar 1, 2023
67659ac
[Unity][Analysis] Checking function return struct info in well-formed…
Hzfengsy Mar 1, 2023
e57f591
[Unity][BYOC] Use Relax legalize + CPU build for reference in tests (…
masahi Mar 1, 2023
3fa880a
[Unity] Add bind_constants option to FuseOpsByPattern (#14151)
vinx13 Mar 1, 2023
6ee79e1
[Unity][Analysis] Analysis for detecting recursion in Relax (#14149)
slyubomirsky Mar 1, 2023
8423811
[Unity][BYOC] Add batch matmul support to Relax CUTLASS BYOC (#14166)
yelite Mar 2, 2023
781bfe0
[Unity][Op] Full support of Relax op `power` (#14171)
SiriusNEO Mar 2, 2023
475f3c2
[Unity][Analysis] Restore Python bindings for var analyses (#14180)
slyubomirsky Mar 4, 2023
38315af
[Unity][OP] Add an operator for fused multi head attention (#14150)
cyx-6 Mar 4, 2023
ed5367d
[Unity][WEBGPU] Codegen improvements and WebRuntime (#14187)
tqchen Mar 4, 2023
1f04221
[Unity][Transform] LiftTransformParams handling multiple functions (#…
MasterJH5574 Mar 4, 2023
22b65bc
[Unity][Op] Group normalization (#14194)
MasterJH5574 Mar 4, 2023
88ab730
[Unity][Op] Argmax and argmin (#14195)
MasterJH5574 Mar 5, 2023
0e98e6e
[Unity][Op] Legalize `round`, `floor`, `ceil`, `sign` (#14198)
MasterJH5574 Mar 5, 2023
284b278
[Unity][Frontend] FX translator supporting more ops (#14196)
MasterJH5574 Mar 5, 2023
f7f24b7
[Unity][Frontend] FX translator returning weights with `keep_params_a…
MasterJH5574 Mar 5, 2023
d103ee2
[Unity][Fix] FX translating dtype (#14201)
MasterJH5574 Mar 5, 2023
8f7c343
[Unity][TIR][Pass] ForceNarrowIndexToInt32 (#14203)
MasterJH5574 Mar 6, 2023
6de551b
[Unity][Frontend] FX translator support torch.baddbmm (#14202)
MasterJH5574 Mar 6, 2023
bb34d97
[CI] Point cpu ci to dep with onnx (#40)
Mar 6, 2023
0937202
[Unity] Introduce Default GPU Schedule Pass (#14182)
zxybazh Mar 6, 2023
1878d7b
Merge with upstream Unity
Mar 6, 2023
b12320b
Remove now unnecessary ScheduleForTarget pass.
Mar 6, 2023
627fb0a
Go back to standard ci_cpu image
Mar 6, 2023
3d72050
Refactor importer locations for consistency and cleaner import.
Mar 7, 2023
1dd7a56
Add initial octoml utility functions.
Feb 21, 2023
d843a77
Add gpu target extraction
Feb 21, 2023
ea79707
Start compilation helper file.
Feb 22, 2023
dfe37fd
Add entrypoint compile function.
Feb 27, 2023
ab05b8c
Add OctoModel helper class and testing.
Feb 27, 2023
44f4742
Add full feature support and testing.
Feb 27, 2023
61e03a4
Chatgpt refactoring of my regular expressions.
Feb 28, 2023
e49e924
Cleanup lint issues.
Feb 28, 2023
e7ab17e
API cleanups after rebase.
Feb 28, 2023
dfc0313
Simplify cuda thread binding until Xiyous full pass lands.
Mar 2, 2023
ee05ae0
Importer improvements such that full flow works end to end
Mar 2, 2023
b4d172f
Improvements after merge with main
Mar 6, 2023
57cf2ff
Merge branch 'relax' into TUZ-150
Mar 7, 2023
f5b86c1
Improve tests
Mar 7, 2023
447bc59
Lint cleanup
Mar 8, 2023
e68e731
Add octo tests to CI
Mar 8, 2023
022042c
Incorporate feedback
Mar 8, 2023
b9813e8
Fix lint
Mar 8, 2023
b5777e1
Fix small type bug
Mar 8, 2023
dfc2931
Fix test target
Mar 8, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
[Frontend][TFLite] Fix conv2d import bug (#14124)
* Fix TFLite frontend bug and add test

* lint
mehrdadh authored Feb 27, 2023

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature. The key has expired.
commit 2feb243bb73acfd15b346a0f661b6d348ef02384
7 changes: 6 additions & 1 deletion python/tvm/relay/frontend/tflite.py
Original file line number Diff line number Diff line change
@@ -2146,7 +2146,7 @@ def convert_conv(self, op, conv_type):
_, kernel_h, kernel_w, in_channels = to_int_list(self.get_tensor_shape(weight_tensor))
assert in_channels == input_c * depth_multiplier
else:
output_channels, kernel_h, kernel_w, _ = to_int_list(
output_channels, kernel_h, kernel_w, in_channels = to_int_list(
self.get_tensor_shape(weight_tensor)
)

@@ -2170,6 +2170,11 @@ def convert_conv(self, op, conv_type):
else:
params["channels"] = int(output_channels)
params["kernel_layout"] = "HWIO"
if input_c != in_channels:
assert (
input_c % in_channels == 0
), "Input channels is not divisible of kernel in_channels."
params["groups"] = int(input_c / in_channels)

# weight tensor type should be INT8/UINT8 (quantization) or FLOAT32
weight_tensor_type = weight_tensor.tensor.Type()
175 changes: 34 additions & 141 deletions tests/python/frontend/tflite/test_forward.py
Original file line number Diff line number Diff line change
@@ -26,6 +26,7 @@

import os
import tempfile
import typing
from packaging import version as package_version
import pytest
import numpy as np
@@ -292,11 +293,11 @@ def run_tflite_graph(tflite_model_buf, input_data):


def compare_tflite_with_tvm(
in_data,
in_name,
input_tensors,
output_tensors,
init_global_variables=False,
in_data: typing.List[np.ndarray],
in_name: typing.List[str],
input_tensors: typing.List,
output_tensors: typing.List,
init_global_variables: bool = False,
out_names=None,
quantized=False,
input_range=None,
@@ -5301,140 +5302,32 @@ def _golden():
_test_reshape_span()


#######################################################################
# Main
# ----
class TestConv2d:
"""Import Conv2d operator from TFLite, build with Relay and test."""

input_shape, kernel_shape, padding = tvm.testing.parameters(
((1, 128, 256, 6), (5, 5, 6, 10), "SAME"),
((1, 128, 256, 6), (5, 5, 6, 10), "VALID"),
# conv2d_group cases
((1, 30, 40, 6), (5, 5, 1, 6), "SAME"),
((1, 30, 40, 6), (5, 5, 1, 6), "VALID"),
)

def test_conv2d(self, input_shape: tuple, kernel_shape: tuple, padding: str):
dtype = tf.float32
kernel_in = np.ones(kernel_shape)
with tf.Graph().as_default():
x = array_ops.placeholder(shape=input_shape, dtype=dtype.name, name="input")
kernel = tf.constant(kernel_in, dtype=dtype, name="filter_weight")
out = tf.nn.conv2d(x, kernel, strides=[1, 1, 1, 1], padding=padding, name="conv2d")
input_data = np.random.randn(*input_shape).astype(dtype.name)
compare_tflite_with_tvm(
[input_data],
["input"],
[x],
[out],
)


if __name__ == "__main__":
# BatchToSpaceND
test_forward_batch_to_space_nd()

# SpaceToBatchND
test_forward_space_to_batch_nd()

# Split
test_forward_split()

# Transpose
test_forward_transpose()

# Cast
test_forward_cast()

# BatchMatMul
test_forward_batch_matmul()

# Tile
test_forward_tile()

# Query
test_forward_shape()

# Transforms
test_forward_concatenation()
test_forward_pad()
test_forward_pack()
test_forward_unpack()
test_forward_reshape()
test_all_resize()
test_forward_range()
test_forward_squeeze()
test_forward_slice()
test_forward_topk()
test_forward_gather()
test_forward_gather_nd()
test_forward_stridedslice()
test_forward_depthtospace()
test_forward_spacetodepth()
test_forward_reverse_sequence()
test_forward_sparse_to_dense()
test_forward_select()
test_forward_quantize_dequantize()
test_forward_arg_min_max()
test_forward_expand_dims()
test_forward_reverse_v2()
test_forward_matrix_set_diag()
test_forward_matrix_diag()

# NN
test_forward_convolution()
test_forward_transpose_conv()
test_forward_logistic()
test_forward_pooling()
test_forward_l2_pool2d()
test_forward_softmax()
test_forward_tanh()
test_forward_relu()
test_forward_relu6()
test_forward_leaky_relu()
test_forward_relu_n1_to_1()
test_forward_log_softmax()
test_forward_fully_connected()
test_forward_l2_normalization()
test_forward_local_response_normalization()
test_forward_prelu()
test_forward_unidirectional_sequence_lstm()

# Elemwise
test_all_elemwise()
test_forward_add_n()

# Unary elemwise
test_all_unary_elemwise()
# Zeros Like
test_forward_zeros_like()

# Fill
test_forward_fill()

# Reduce
test_all_reduce()

# Logical
test_all_logical()

# Detection_PostProcess
test_detection_postprocess()

# NonMaxSuppressionV5
test_forward_nms_v5()

# Overwrite Converter
test_custom_op_converter()

# test structural_equal and span information
test_structure_and_span()

# End to End
test_forward_mobilenet_v1()
test_forward_mobilenet_v2()
test_forward_mobilenet_v3()
test_forward_inception_v3_net()
test_forward_inception_v4_net()
test_forward_inception_v4_net_batched()
test_forward_coco_ssd_mobilenet_v1()
test_forward_mediapipe_hand_landmark()

# End to End Sparse models
test_forward_sparse_mobilenet_v1()
test_forward_sparse_mobilenet_v2()

# End to End quantized
test_forward_qnn_inception_v1_net()
test_forward_qnn_mobilenet_v1_net()
test_forward_qnn_mobilenet_v2_net()
# This also fails with a segmentation fault in my run
# with Tflite 1.15.2
test_forward_qnn_mobilenet_v3_net()
test_forward_qnn_coco_ssd_mobilenet_v1()

# TFLite 2.1.0 quantized tests
test_forward_quantized_convolution()
test_forward_quantized_depthwise_convolution()
test_forward_tflite2_qnn_resnet50()
test_forward_tflite2_qnn_inception_v1()
test_forward_tflite2_qnn_mobilenet_v2()

test_forward_tflite_float16()

test_forward_tflite_int16()
test_forward_ds_cnn_int16()
tvm.testing.main()