Skip to content

Commit

Permalink
[Cherry-pick] Fixing the location of DimAnalysis in onnx-to-zhigh pas…
Browse files Browse the repository at this point in the history 
…s and some rules in zhigh-to-onnx pass

The onnx-to-zhigh pass has two phases: 1) converting multiple onnx ops into a single zhigh op, and 2) converting a single onnx op to a single zhigh op, where the second phase uses DimAnalysis (Patterns in the 1st phase at this moment does not use DimAnalysis)

The problem is DimAnalysis is currently called before the 1st phase, which is not good because the 1st  phase may change the IR so the information from DimAnalysis is obsoleted to the 2nd phase. Correct position for DimAnalysis would be just before the 2nd phase.

Other than that, this PR changes slightly the rules in zhigh-to-onnx pass so that for binary ops, only one input (instead of two) that is from stick would be enough to trigger the rule to convert a zhigh op back to an onnx op.

Resolves onnx#2789

---------

Signed-off-by: Tung D. Le <tung@jp.ibm.com>
(cherry picked from commit 80a63f2)
Signed-off-by: Charles Volzka <cjvolzka@us.ibm.com>
  • Loading branch information
@tungld @cjvolzka
tungld authored and cjvolzka committed Apr 11, 2024
1 parent 0833dac commit c585f03
Show file tree
Hide file tree
Showing 9 changed files with 103 additions and 48 deletions.
1 change: 1 addition & 0 deletions src/Accelerators/NNPA/Conversion/ONNXToZHigh/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ add_onnx_mlir_library(OMONNXToZHigh
OMNNPACompilerOptions
OMONNXOps
OMONNXToKrnl
OMShapeInferencePass
OMZHighOps

ACCEL_INCLUDE_DIRS PRIVATE
Expand Down
2 changes: 1 addition & 1 deletion src/Accelerators/NNPA/Conversion/ONNXToZHigh/DevicePlacement.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,7 @@ void DevicePlacementPass::runOnOperation() {
// Call ONNXToZHigh pass for lowering multiple ONNX ops at once to ZHigh.
// E.g. `onnx.ReLu (onnx.Conv)` to zhigh.Conv.
RewritePatternSet Patterns2(context);
getONNXToZHighOneOpPatterns(Patterns2);
getONNXToZHighMultipleOpPatterns(Patterns2);
(void)applyAnalysisConversion(module, target, std::move(Patterns2),
ConversionConfig{.legalizableOps = &legalizedOps2});

Expand Down
13 changes: 8 additions & 5 deletions src/Accelerators/NNPA/Conversion/ONNXToZHigh/ONNXToZHigh.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
#include "src/Dialect/ONNX/ONNXDimAnalysis.hpp"
#include "src/Dialect/ONNX/ONNXOps.hpp"
#include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
#include "src/Dialect/ONNX/Transforms/ShapeInference.hpp"

using namespace mlir;

Expand Down Expand Up @@ -328,16 +329,13 @@ void getONNXToZHighMultipleOpPatterns(RewritePatternSet &patterns) {
patterns.insert<replaceONNXMatMulAddPattern2>(context);
patterns.insert<replaceONNXReluConvPattern>(context);
patterns.insert<replaceONNXLogSoftmaxPattern>(context);
// Shape inference for newly-added operations.
getShapeInferencePatterns(patterns);
}

void ONNXToZHighLoweringPass::runOnOperation() {
ModuleOp module = getOperation();

// Run the unknown dimension analysis to help check equality of unknown
// dimensions at compile time.
onnx_mlir::DimAnalysis dimAnalysis(module);
dimAnalysis.analyze();

// The first thing to define is the conversion target. This will define the
// final target for this lowering.
ConversionTarget target(getContext());
Expand All @@ -363,6 +361,11 @@ void ONNXToZHighLoweringPass::runOnOperation() {
// It's ok to fail.
(void)applyPatternsAndFoldGreedily(module, std::move(combinedPatterns));

// Run the unknown dimension analysis to help check equality of unknown
// dimensions at compile time.
onnx_mlir::DimAnalysis dimAnalysis(module);
dimAnalysis.analyze();

// Single ONNX to ZHigh operation lowering.
RewritePatternSet patterns(&getContext());
onnx_mlir::getONNXToZHighOneOpPatterns(patterns);
Expand Down
2 changes: 2 additions & 0 deletions src/Accelerators/NNPA/Conversion/ONNXToZHigh/ZHighToONNX.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,8 @@ void ZHighToONNXLoweringPass::runOnOperation() {

RewritePatternSet patterns(&getContext());
populateWithGenerated(patterns);
zhigh::ZHighStickOp::getCanonicalizationPatterns(patterns, &getContext());
zhigh::ZHighUnstickOp::getCanonicalizationPatterns(patterns, &getContext());

(void)applyPatternsAndFoldGreedily(function, std::move(patterns));
}
Expand Down
95 changes: 70 additions & 25 deletions src/Accelerators/NNPA/Conversion/ONNXToZHigh/ZHighToONNX.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,62 +37,107 @@ def CreateONNXMaxOp : NativeCodeCall<"$_builder.create<ONNXMaxOp>($_loc, $0.getT
// ONNXAddOp %X = ZHighUnstickOp (ZHighAddOp (ZHighStickOp %X),
// (ZHighStickOp %Y))
//===----------------------------------------------------------------------===//
def replaceZHighAddPattern : Pat<
(ZHighUnstickOp (ZHighAddOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
(ONNXAddOp $x, $y),
[(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
def replaceZHighAddPattern1 : Pat<
(ZHighUnstickOp (ZHighAddOp (ZHighStickOp:$s_x $x, $_), $y)),
(ONNXAddOp $x, (ZHighUnstickOp $y)),
[(NotBlockArgument:$x), (HasOneUse:$s_x)]
>;

def replaceZHighAddPattern2 : Pat<
(ZHighUnstickOp (ZHighAddOp $x, (ZHighStickOp:$s_y $y, $_))),
(ONNXAddOp (ZHighUnstickOp $x), $y),
[(NotBlockArgument:$y), (HasOneUse:$s_y)]
>;

//===----------------------------------------------------------------------===//
// ONNXMulOp %X = ZHighUnstickOp (ZHighMulOp (ZHighStickOp %X),
// (ZHighStickOp %Y))
//===----------------------------------------------------------------------===//
def replaceZHighMulPattern : Pat<
(ZHighUnstickOp (ZHighMulOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
(ONNXMulOp $x, $y),
[(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
def replaceZHighMulPattern1 : Pat<
(ZHighUnstickOp (ZHighMulOp (ZHighStickOp:$s_x $x, $_), $y)),
(ONNXMulOp $x, (ZHighUnstickOp $y)),
[(NotBlockArgument:$x), (HasOneUse:$s_x)], [ ],
(addBenefit 1)
>;

def replaceZHighMulPattern2 : Pat<
(ZHighUnstickOp (ZHighMulOp $x, (ZHighStickOp:$s_y $y, $_))),
(ONNXMulOp (ZHighUnstickOp $x), $y),
[(NotBlockArgument:$y), (HasOneUse:$s_y)], [],
(addBenefit 0)
>;

//===----------------------------------------------------------------------===//
// ONNXSubOp %X = ZHighUnstickOp (ZHighSubOp (ZHighStickOp %X),
// (ZHighStickOp %Y))
//===----------------------------------------------------------------------===//
def replaceZHighSubPattern : Pat<
(ZHighUnstickOp (ZHighSubOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
(ONNXSubOp $x, $y),
[(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
def replaceZHighSubPattern1 : Pat<
(ZHighUnstickOp (ZHighSubOp (ZHighStickOp:$s_x $x, $_), $y)),
(ONNXSubOp $x, (ZHighUnstickOp $y)),
[(NotBlockArgument:$x), (HasOneUse:$s_x)], [ ],
(addBenefit 1)
>;

def replaceZHighSubPattern2 : Pat<
(ZHighUnstickOp (ZHighSubOp $x, (ZHighStickOp:$s_y $y, $_))),
(ONNXSubOp (ZHighUnstickOp $x), $y),
[(NotBlockArgument:$y), (HasOneUse:$s_y)], [ ],
(addBenefit 0)
>;

//===----------------------------------------------------------------------===//
// ONNXDivOp %X = ZHighUnstickOp (ZHighDivOp (ZHighStickOp
// %X),(ZHighStickOp %Y))
// Note: turn off this pattern since NNPA is faster at this moment.
//===----------------------------------------------------------------------===//
// def replaceZHighDivPattern : Pat<
// (ZHighUnstickOp (ZHighDivOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
// (ONNXDivOp $x, $y),
// [(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
// >;
//def replaceZHighDivPattern1 : Pat<
// (ZHighUnstickOp (ZHighDivOp (ZHighStickOp:$s_x $x, $_), $y)),
// (ONNXDivOp $x, (ZHighUnstickOp $y)),
// [(NotBlockArgument:$x), (HasOneUse:$s_x)], [ ],
// (addBenefit 1)
//>;
//
//def replaceZHighDivPattern2 : Pat<
// (ZHighUnstickOp (ZHighDivOp $x, (ZHighStickOp:$s_y $y, $_))),
// (ONNXDivOp (ZHighUnstickOp $x), $y),
// [(NotBlockArgument:$y), (HasOneUse:$s_y)], [ ],
// (addBenefit 0)
//>;

//===----------------------------------------------------------------------===//
// ONNXMinOp %X = ZHighUnstickOp (ZHighMinOp (ZHighStickOp %X),
// (ZHighStickOp %Y))
//===----------------------------------------------------------------------===//
def replaceZHighMinPattern : Pat<
(ZHighUnstickOp:$u (ZHighMinOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
(CreateONNXMinOp $u, $x, $y),
[(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
def replaceZHighMinPattern1 : Pat<
(ZHighUnstickOp:$u (ZHighMinOp (ZHighStickOp:$s_x $x, $_), $y)),
(CreateONNXMinOp $u, $x, (ZHighUnstickOp $y)),
[(NotBlockArgument:$x), (HasOneUse:$s_x)], [ ],
(addBenefit 1)
>;

def replaceZHighMinPattern2 : Pat<
(ZHighUnstickOp:$u (ZHighMinOp $x, (ZHighStickOp:$s_y $y, $_))),
(CreateONNXMinOp $u, (ZHighUnstickOp $x), $y),
[(NotBlockArgument:$y), (HasOneUse:$s_y)], [ ],
(addBenefit 0)
>;

//===----------------------------------------------------------------------===//
// ONNXMaxOp %X = ZHighUnstickOp (ZHighMaxOp (ZHighStickOp %X),
// (ZHighStickOp %Y))
//===----------------------------------------------------------------------===//
def replaceZHighMaxPattern : Pat<
(ZHighUnstickOp:$u (ZHighMaxOp (ZHighStickOp:$s_x $x, $_), (ZHighStickOp:$s_y $y, $_))),
(CreateONNXMaxOp $u, $x, $y),
[(NotBlockArgument:$x), (HasOneUse:$s_x), (HasOneUse:$s_y)]
def replaceZHighMaxPattern1 : Pat<
(ZHighUnstickOp:$u (ZHighMaxOp (ZHighStickOp:$s_x $x, $_), $y)),
(CreateONNXMaxOp $u, $x, (ZHighUnstickOp $y)),
[(NotBlockArgument:$x), (HasOneUse:$s_x)], [ ],
(addBenefit 1)
>;

def replaceZHighMaxPattern2 : Pat<
(ZHighUnstickOp:$u (ZHighMaxOp $x, (ZHighStickOp:$s_y $y, $_))),
(CreateONNXMaxOp $u, (ZHighUnstickOp $x), $y),
[(NotBlockArgument:$y), (HasOneUse:$s_y)], [ ],
(addBenefit 0)
>;

//===----------------------------------------------------------------------===//
Expand Down
6 changes: 5 additions & 1 deletion test/mlir/accelerators/nnpa/conversion/instrument/add-onnx-zhigh-level.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@
func.func @test_instrument_add_onnx_zhigh(%arg0 : tensor<10x10xf32>, %arg1 : tensor<10x1xf32>) -> tensor<*xf32> {
%0 = "onnx.Add"(%arg0, %arg1) {onnx_node_name = "onnx.Add1"} : (tensor<10x10xf32>, tensor<10x1xf32>) -> tensor<*xf32>
%1 = "onnx.Add"(%arg0, %0) {onnx_node_name = "onnx.Add2"} : (tensor<10x10xf32>, tensor<*xf32>) -> tensor<*xf32>
"onnx.Return"(%1) : (tensor<*xf32>) -> ()
%2 = "onnx.Relu"(%1) {onnx_node_name = "onnx.Relu"} : (tensor<*xf32>) -> tensor<*xf32>
"onnx.Return"(%2) : (tensor<*xf32>) -> ()
}

// CHECK-LABEL: func.func @test_instrument_add_onnx_zhigh
Expand All @@ -21,6 +22,9 @@ func.func @test_instrument_add_onnx_zhigh(%arg0 : tensor<10x10xf32>, %arg1 : ten
// CHECK: "krnl.runtime_instrument"() {opName = "zhigh.Add", tag = 5 : i64} : () -> ()
// CHECK: "zhigh.Add"
// CHECK: "krnl.runtime_instrument"() {opName = "zhigh.Add", tag = 6 : i64} : () -> ()
// CHECK: "krnl.runtime_instrument"() {opName = "zhigh.Relu", tag = 5 : i64} : () -> ()
// CHECK: "zhigh.Relu"
// CHECK: "krnl.runtime_instrument"() {opName = "zhigh.Relu", tag = 6 : i64} : () -> ()
// CHECK: "krnl.runtime_instrument"() {opName = "zhigh.Unstick", tag = 5 : i64} : () -> ()
// CHECK: "zhigh.Unstick"
// CHECK: "krnl.runtime_instrument"() {opName = "zhigh.Unstick", tag = 6 : i64} : () -> ()
Expand Down
4 changes: 2 additions & 2 deletions test/mlir/accelerators/nnpa/conversion/onnx-to-zhigh/conv.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,8 +107,8 @@ func.func @test_fuse_onnx_relu_conv2d(%arg0: tensor<5x3x32x32xf32>, %arg1 : tens
// CHECK-NOT: separator of consecutive DAGs
// CHECK-DAG: [[VAR_3_:%.+]] = "zhigh.Stick"([[VAR_2_]]) {layout = "HWCK"} : (tensor<2x2x3x2xf32>) -> tensor<2x2x3x2xf16, #zhigh.layout<{dataLayout = "HWCK"}>>
// CHECK-DAG: [[VAR_4_:%.+]] = "zhigh.Stick"([[PARAM_2_]]) {layout = "1D"} : (tensor<2xf32>) -> tensor<2xf16, #zhigh.layout<{dataLayout = "1D"}>>
// CHECK: [[VAR_5_:%.+]] = "zhigh.Conv2D"([[VAR_1_]], [[VAR_3_]], [[VAR_4_]]) {act_func = "ACT_RELU", kernel_shape = [2, 2], padding_type = "VALID_PADDING", strides = [1, 1]} : (tensor<5x32x32x3xf16, #zhigh.layout<{dataLayout = "NHWC"}>>, tensor<2x2x3x2xf16, #zhigh.layout<{dataLayout = "HWCK"}>>, tensor<2xf16, #zhigh.layout<{dataLayout = "1D"}>>) -> tensor<*xf16>
// CHECK: [[VAR_6_:%.+]] = "zhigh.Unstick"([[VAR_5_]]) : (tensor<*xf16>) -> tensor<5x2x31x31xf32>
// CHECK: [[VAR_5_:%.+]] = "zhigh.Conv2D"([[VAR_1_]], [[VAR_3_]], [[VAR_4_]]) {act_func = "ACT_RELU", kernel_shape = [2, 2], padding_type = "VALID_PADDING", strides = [1, 1]} : (tensor<5x32x32x3xf16, #zhigh.layout<{dataLayout = "NHWC"}>>, tensor<2x2x3x2xf16, #zhigh.layout<{dataLayout = "HWCK"}>>, tensor<2xf16, #zhigh.layout<{dataLayout = "1D"}>>) -> tensor<5x31x31x2xf16, #zhigh.layout<{dataLayout = "NHWC"}>>
// CHECK: [[VAR_6_:%.+]] = "zhigh.Unstick"([[VAR_5_]]) : (tensor<5x31x31x2xf16, #zhigh.layout<{dataLayout = "NHWC"}>>) -> tensor<5x2x31x31xf32>
// CHECK: return [[VAR_6_]] : tensor<5x2x31x31xf32>
// CHECK: }
}
Expand Down
8 changes: 4 additions & 4 deletions test/mlir/accelerators/nnpa/conversion/onnx-to-zhigh/matmul.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,8 +79,8 @@ func.func @test_onnx_matmul_add_to_zhigh_1D_bias(
// CHECK-DAG: [[VAR_0_:%.+]] = "zhigh.Stick"([[PARAM_0_]]) {layout = "2D"} : (tensor<4x8xf32>) -> tensor<4x8xf16, #zhigh.layout<{dataLayout = "2D"}>>
// CHECK-DAG: [[VAR_1_:%.+]] = "zhigh.Stick"([[PARAM_1_]]) {layout = "2D"} : (tensor<8x16xf32>) -> tensor<8x16xf16, #zhigh.layout<{dataLayout = "2D"}>>
// CHECK-DAG: [[VAR_2_:%.+]] = "zhigh.Stick"([[PARAM_2_]]) {layout = "1D"} : (tensor<16xf32>) -> tensor<16xf16, #zhigh.layout<{dataLayout = "1D"}>>
// CHECK: [[VAR_3_:%.+]] = "zhigh.MatMul"([[VAR_0_]], [[VAR_1_]], [[VAR_2_]]) : (tensor<4x8xf16, #zhigh.layout<{dataLayout = "2D"}>>, tensor<8x16xf16, #zhigh.layout<{dataLayout = "2D"}>>, tensor<16xf16, #zhigh.layout<{dataLayout = "1D"}>>) -> tensor<*xf16>
// CHECK: [[VAR_4_:%.+]] = "zhigh.Unstick"([[VAR_3_]]) : (tensor<*xf16>) -> tensor<4x16xf32>
// CHECK: [[VAR_3_:%.+]] = "zhigh.MatMul"([[VAR_0_]], [[VAR_1_]], [[VAR_2_]]) : (tensor<4x8xf16, #zhigh.layout<{dataLayout = "2D"}>>, tensor<8x16xf16, #zhigh.layout<{dataLayout = "2D"}>>, tensor<16xf16, #zhigh.layout<{dataLayout = "1D"}>>) -> tensor<4x16xf16, #zhigh.layout<{dataLayout = "2D"}>>
// CHECK: [[VAR_4_:%.+]] = "zhigh.Unstick"([[VAR_3_]]) : (tensor<4x16xf16, #zhigh.layout<{dataLayout = "2D"}>>) -> tensor<4x16xf32>
// CHECK: return [[VAR_4_]] : tensor<4x16xf32>
// CHECK: }
// CHECK-NOT: "onnx.Add"
Expand All @@ -105,8 +105,8 @@ func.func @test_onnx_matmul_add_to_zhigh_1D_bias_normalized(
// CHECK-DAG: [[VAR_0_:%.+]] = "zhigh.Stick"([[PARAM_0_]]) {layout = "2D"} : (tensor<4x8xf32>) -> tensor<4x8xf16, #zhigh.layout<{dataLayout = "2D"}>>
// CHECK-DAG: [[VAR_1_:%.+]] = "zhigh.Stick"([[PARAM_1_]]) {layout = "2D"} : (tensor<8x16xf32>) -> tensor<8x16xf16, #zhigh.layout<{dataLayout = "2D"}>>
// CHECK-DAG: [[VAR_2_:%.+]] = "zhigh.Stick"([[PARAM_2_]]) {layout = "1D"} : (tensor<16xf32>) -> tensor<16xf16, #zhigh.layout<{dataLayout = "1D"}>>
// CHECK: [[VAR_3_:%.+]] = "zhigh.MatMul"([[VAR_0_]], [[VAR_1_]], [[VAR_2_]]) : (tensor<4x8xf16, #zhigh.layout<{dataLayout = "2D"}>>, tensor<8x16xf16, #zhigh.layout<{dataLayout = "2D"}>>, tensor<16xf16, #zhigh.layout<{dataLayout = "1D"}>>) -> tensor<*xf16>
// CHECK: [[VAR_4_:%.+]] = "zhigh.Unstick"([[VAR_3_]]) : (tensor<*xf16>) -> tensor<4x16xf32>
// CHECK: [[VAR_3_:%.+]] = "zhigh.MatMul"([[VAR_0_]], [[VAR_1_]], [[VAR_2_]]) : (tensor<4x8xf16, #zhigh.layout<{dataLayout = "2D"}>>, tensor<8x16xf16, #zhigh.layout<{dataLayout = "2D"}>>, tensor<16xf16, #zhigh.layout<{dataLayout = "1D"}>>) -> tensor<4x16xf16, #zhigh.layout<{dataLayout = "2D"}>>
// CHECK: [[VAR_4_:%.+]] = "zhigh.Unstick"([[VAR_3_]]) : (tensor<4x16xf16, #zhigh.layout<{dataLayout = "2D"}>>) -> tensor<4x16xf32>
// CHECK: return [[VAR_4_]] : tensor<4x16xf32>
// CHECK: }
// CHECK-NOT: "onnx.Add"
Expand Down
20 changes: 10 additions & 10 deletions test/mlir/accelerators/nnpa/conversion/onnx-to-zhigh/softmax.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,11 @@ func.func @test_onnx_logsoftmax(%arg0 : tensor<10x10xf32>) -> tensor<*xf32> {

// CHECK-LABEL: func @test_onnx_logsoftmax
// CHECK-SAME: ([[PARAM_0_:%.+]]: tensor<10x10xf32>) -> tensor<10x10xf32> {
// CHECK: [[VAR_0_:%.+]] = "onnx.UnsqueezeV11"([[PARAM_0_]]) {axes = [0]} : (tensor<10x10xf32>) -> tensor<*xf32>
// CHECK: [[VAR_1_:%.+]] = "zhigh.Stick"([[VAR_0_]]) {layout = "3DS"} : (tensor<*xf32>) -> tensor<*xf16>
// CHECK: [[VAR_2_:%.+]] = "zhigh.Softmax"([[VAR_1_]]) {act_func = "ACT_LOG"} : (tensor<*xf16>) -> tensor<*xf16>
// CHECK: [[VAR_3_:%.+]] = "zhigh.Unstick"([[VAR_2_]]) : (tensor<*xf16>) -> tensor<*xf32>
// CHECK: [[VAR_4_:%.+]] = "onnx.SqueezeV11"([[VAR_3_]]) {axes = [0]} : (tensor<*xf32>) -> tensor<10x10xf32>
// CHECK: [[VAR_0_:%.+]] = "onnx.UnsqueezeV11"([[PARAM_0_]]) {axes = [0]} : (tensor<10x10xf32>) -> tensor<1x10x10xf32>
// CHECK: [[VAR_1_:%.+]] = "zhigh.Stick"([[VAR_0_]]) {layout = "3DS"} : (tensor<1x10x10xf32>) -> tensor<1x10x10xf16, #zhigh.layout<{dataLayout = "3DS"}>>
// CHECK: [[VAR_2_:%.+]] = "zhigh.Softmax"([[VAR_1_]]) {act_func = "ACT_LOG"} : (tensor<1x10x10xf16, #zhigh.layout<{dataLayout = "3DS"}>>) -> tensor<1x10x10xf16, #zhigh.layout<{dataLayout = "3DS"}>>
// CHECK: [[VAR_3_:%.+]] = "zhigh.Unstick"([[VAR_2_]]) : (tensor<1x10x10xf16, #zhigh.layout<{dataLayout = "3DS"}>>) -> tensor<1x10x10xf32>
// CHECK: [[VAR_4_:%.+]] = "onnx.SqueezeV11"([[VAR_3_]]) {axes = [0]} : (tensor<1x10x10xf32>) -> tensor<10x10xf32>
// CHECK: return [[VAR_4_]] : tensor<10x10xf32>
// CHECK: }
}
Expand All @@ -57,11 +57,11 @@ func.func @test_onnx_logsoftmax_dyn(%arg0 : tensor<?x?xf32>) -> tensor<*xf32> {

// CHECK-LABEL: func @test_onnx_logsoftmax_dyn
// CHECK-SAME: ([[PARAM_0_:%.+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: [[VAR_0_:%.+]] = "onnx.UnsqueezeV11"([[PARAM_0_]]) {axes = [0]} : (tensor<?x?xf32>) -> tensor<*xf32>
// CHECK: [[VAR_1_:%.+]] = "zhigh.Stick"([[VAR_0_]]) {layout = "3DS"} : (tensor<*xf32>) -> tensor<*xf16>
// CHECK: [[VAR_2_:%.+]] = "zhigh.Softmax"([[VAR_1_]]) {act_func = "ACT_LOG"} : (tensor<*xf16>) -> tensor<*xf16>
// CHECK: [[VAR_3_:%.+]] = "zhigh.Unstick"([[VAR_2_]]) : (tensor<*xf16>) -> tensor<*xf32>
// CHECK: [[VAR_4_:%.+]] = "onnx.SqueezeV11"([[VAR_3_]]) {axes = [0]} : (tensor<*xf32>) -> tensor<?x?xf32>
// CHECK: [[VAR_0_:%.+]] = "onnx.UnsqueezeV11"([[PARAM_0_]]) {axes = [0]} : (tensor<?x?xf32>) -> tensor<1x?x?xf32>
// CHECK: [[VAR_1_:%.+]] = "zhigh.Stick"([[VAR_0_]]) {layout = "3DS"} : (tensor<1x?x?xf32>) -> tensor<1x?x?xf16, #zhigh.layout<{dataLayout = "3DS"}>>
// CHECK: [[VAR_2_:%.+]] = "zhigh.Softmax"([[VAR_1_]]) {act_func = "ACT_LOG"} : (tensor<1x?x?xf16, #zhigh.layout<{dataLayout = "3DS"}>>) -> tensor<1x?x?xf16, #zhigh.layout<{dataLayout = "3DS"}>>
// CHECK: [[VAR_3_:%.+]] = "zhigh.Unstick"([[VAR_2_]]) : (tensor<1x?x?xf16, #zhigh.layout<{dataLayout = "3DS"}>>) -> tensor<1x?x?xf32>
// CHECK: [[VAR_4_:%.+]] = "onnx.SqueezeV11"([[VAR_3_]]) {axes = [0]} : (tensor<1x?x?xf32>) -> tensor<?x?xf32>
// CHECK: return [[VAR_4_]] : tensor<?x?xf32>
// CHECK: }
}
Expand Down

0 comments on commit c585f03

Please sign in to comment.