Decompose Hardswish into simpler ONNX ops

[kumarappan-cmyk]

This pass optimizes the ONNX computation graph by decomposing the HardSwish operation into a sequence of simpler ONNX operations. Since HardSwish is not natively supported in ONNX-MLIR, this transformation ensures compatibility by replacing HardSwish with fundamental arithmetic functions.

Y= HardSwish(x) = x * max(0, min(1, alpha * x + beta))

Where alpha = 1/6 and beta = 0.5

[jenkins-droid]

Can one of the admins verify this patch?

[chentong319]

@jenkins-droid test it please.

[chentong319]

This function could be a general utility function. Do you have the interest in putting it into src/Dailect/ONNX/DialectBuilder?

[chentong319]

It is fine to define rewriting pattern in this way. I am just curious whether it is possible to define this rewriting rule with tablegen.

[chentong319]

First of all, you do not have to use .td for this PR.

To do this particular rewriting with table gen, you may follow the code in src/Dialect/ONNX/Transform/Decompose.*, especially the following rule:
def ClipV6Pattern : Pat<
(ONNXClipV6Op $x, $maxAttr, $minAttr),
(ONNXClipV11Op $x, (ONNXConstantOpFromDenseAttr(createScalarDenseAttrRank0 $minAttr)),
(ONNXConstantOpFromDenseAttr(createScalarDenseAttrRank0 $maxAttr)))

;

Here the input of createScalarDenseAttrRank0 is an Attribute, but you need a new function to take a float value. You may create a FloatAttr of the desired constant value for createScalarDenseAttrRank0.

Hopefully, your rule will be
Def HardSwishDecompose: Pat<
(ONNXHardSwishOp:res $x),
(ONNXMulOp $x, (ONNXMaxOp (ONNXConstantOpFromDenseAttr(createScalarDenseAttrRank0 FloatAttr(0.0) )), ONNXMinOp …)

Remind that you may need to provide the output type for intermediate result, such as ONNXMaxOp, ONNXAddOp and etc.

[kumarappan-cmyk]

Thanks will look into it.

[kumarappan-cmyk]

I can create attributes directly
def HardSwishAlpha : NativeCodeCall<
"FloatAttr::get($_builder.getF32Type(), 1.0 / 6.0)">;

But when I have a generic function
def createFloatAttr : NativeCodeCall<
"FloatAttr::get($_builder.getF32Type(), %0)">;

I am unable to pass the desired constant value
Is it that have to have Multiple def functions itself
Or how to pass a constant value to the native code call in tablegen

[chentong319]

In this case, the %0 is not passed to a function call. Instead, it is used by constructor directly. Therefore it is more like a template initialization. You can check the code in src/DIalect/ONNX/ONNXOps/Canonicalize.td:
class HasRankOf<int rank> : Constraint<CPred<"mlir::isa<ShapedType>($0.getType()) && " "mlir::cast<ShapedType>($0.getType()).hasRank() && " "mlir::cast<ShapedType>($0.getType()).getRank() == " # rank>>
Another approach is to pass the value to your createConstantTensor function.
Again, your current implementation is fine. Just let me know if you want just the current implementation.

[kumarappan-cmyk]

Hi, For class in tablegen file I am able to pass the value as int or string but not as float.
And again in src/DIalect/ONNX/ONNXOps/Canonicalize.td
For rewriting Gemm OP there is seperate definition for the constants
def GemmAlpha : NativeCodeCall<"$_builder.getF32FloatAttr(1.0)">;
def GemmBeta : NativeCodeCall<"$_builder.getF32FloatAttr(1.0)">;
I will try what @tungld suggested

[tungld]

I strongly encourage you to lower this operator into KRNL for the best performance so that all primitive computations are inside the innermost loop, instead of decomposing into multiple ONNX ops.

A very similar lowering is for ONNX HardSigmoid y = max(0, min(1, alpha * x + beta)): https://github.com/onnx/onnx-mlir/blob/main/src/Conversion/ONNXToKrnl/Math/Elementwise.cpp#L619. The code would be almost similar.

[jenkins-droid]

Can one of the admins verify this patch?

[kumarappan-cmyk]

@chentong319 @tungld
I have added the decomposition in the krnl dialect
Earlier with the decomposition pattern in the decompose.cpp, The MLIR was

affine.for %arg1 = 0 to 10 {
affine.for %arg2 = 0 to 3 {
affine.for %arg3 = 0 to 512 {
affine.for %arg4 = 0 to 512 {
%4 = affine.load %arg0[%arg1, %arg2, %arg3, %arg4] : memref<10x3x512x512xf32>
%5 = affine.load %0[] : memref
%6 = arith.mulf %4, %5 : f32
%7 = affine.load %1[] : memref
%8 = arith.addf %6, %7 : f32
%9 = affine.load %2[] : memref
%10 = arith.minnumf %8, %9 : f32
%11 = affine.load %3[] : memref
%12 = arith.maxnumf %10, %11 : f32
%13 = affine.load %arg0[%arg1, %arg2, %arg3, %arg4] : memref<10x3x512x512xf32>
%14 = arith.mulf %13, %12 : f32
affine.store %14, %alloc[%arg1, %arg2, %arg3, %arg4] : memref<10x3x512x512xf32>
}
}
}
}

Multiple loads were there, As suggested i moved the conversion to src/Conversion/ONNXToKrnl/Math/Elementwise.cpp
Now the MLIR has fewer loads

affine.for %arg1 = 0 to 10 {
  affine.for %arg2 = 0 to 3 {
    affine.for %arg3 = 0 to 512 {
      affine.for %arg4 = 0 to 512 {
        %0 = affine.load %arg0[%arg1, %arg2, %arg3, %arg4] : memref<10x3x512x512xf32>
        %1 = arith.mulf %0, %cst_0 : f32
        %2 = arith.addf %1, %cst : f32
        %3 = arith.minnumf %2, %cst_1 : f32
        %4 = arith.maxnumf %3, %cst_2 : f32
        %5 = arith.mulf %0, %4 : f32
        affine.store %5, %alloc[%arg1, %arg2, %arg3, %arg4] : memref<10x3x512x512xf32>
      }
    }
  }
}

[chentong319]

It is more efficient to lower to krnl for CPU. I'd like to keep the decomposition for the following reasons:

1. Support Hardswish at onnx level. It may help other user of onnx-mlir with its semantics expressed with other onnx ops.
2. May expose more optimization opportunities.
3. I'd like to know whether the op fusion in krnl lowering will bring the same efficiency.

How about keeping both decomposition and krnl lowering code and choose one of them with an option? You can set lowering to krnl as the default.

[kumarappan-cmyk]

@chentong319 I have added a command line option to specify with decompose-op-in-onnx to decompose in decompose.cpp

[tungld]

Thank you for the changes! They are very clean. Could you add some lit tests for the lowering to Krnl into the following files? Elementwise_with_canonicalize and Elementwise_with_canonicalize_O3.

[tungld]

A small nit: we avoid using (dtype) for type cast but static_cast<int64_t>(value) instead.

[tungld]

This should be OnnxMlirCommonOptions if you want to use it in both onnx-mlir and onnx-mlir-opt commands.

[tungld]

Please remove this llvm::outs() or you can change to using LLVM_DEBUG(llvm::dbgs() << DEBUG_TYPE << "Decomposing ONNX operation: " << op << "\n" instead.

[tungld]

LGTM! Thank you very much for the patch!

Just some nits as commented.

[tungld]

These changes made the clang-format failed. I sometimes get the same problem. The local clang format suggests these changes but the Jenkins clang-format failed. Perhaps, there is a difference in clang-format version.

[tungld]

Thank you for adding this option! For the comment, it should be common for both.

[tungld]

Thanks for adding these lit tests!

[jenkins-droid]

Can one of the admins verify this patch?

[tungld]

@jenkins-droid test this please

[tungld]

@chentong319 does the newly added option meet your requirement? if so we can merge the PR.

[chentong319]

LGTM!

[Arkar-Hema]

Thanks for the review and approval! Some of the tests are still running, once they finish could you please merge the PR?

[jenkins-droid]

Jenkins Linux s390x Build #16480 [push] Decompose Hardswish into... started at 22:12

[jenkins-droid]

Jenkins Linux amd64 Build #16478 [push] Decompose Hardswish into... started at 21:12

[jenkins-droid]

Jenkins Linux ppc64le Build #15461 [push] Decompose Hardswish into... started at 22:12

[jenkins-droid]

Jenkins Linux amd64 Build #16478 [push] Decompose Hardswish into... passed after 1 hr 30 min

[jenkins-droid]

Jenkins Linux s390x Build #16480 [push] Decompose Hardswish into... passed after 1 hr 54 min

[jenkins-droid]

Jenkins Linux ppc64le Build #15461 [push] Decompose Hardswish into... aborted after 2 hr 38 min