Use DimAnalysis in lowering MatMul (#2195)

Signed-off-by: Tung D. Le <[email protected]>
Co-authored-by: Alexandre Eichenberger <[email protected]>

Author: tungld

src/Conversion/ONNXToKrnl/ConvertONNXToKrnl.cpp

@@ -194,7 +194,7 @@ void populateONNXToKrnlConversionPattern(RewritePatternSet &patterns,
   populateLoweringONNXReductionOpPattern(patterns, typeConverter, ctx);
   populateLoweringONNXSoftmaxOpPattern(patterns, typeConverter, ctx);
   populateLoweringONNXTopKOpPattern(patterns, typeConverter, ctx);
-  populateLoweringONNXMatMulOpPattern(patterns, typeConverter, ctx, enableTiling);
+  populateLoweringONNXMatMulOpPattern(patterns, typeConverter, ctx, dimAnalysis, enableTiling);
   populateLoweringONNXMatMulIntegerOpPattern(patterns, typeConverter, ctx);
   populateLoweringONNXRandomNormalOpPattern(patterns, typeConverter, ctx);
   populateLoweringONNXRandomNormalLikeOpPattern(patterns, typeConverter, ctx);

src/Conversion/ONNXToKrnl/Math/MatMul.cpp

@@ -29,9 +29,11 @@ using namespace mlir;
 namespace onnx_mlir {
 
 struct ONNXMatMulOpLowering : public OpConversionPattern<ONNXMatMulOp> {
-  ONNXMatMulOpLowering(
-      TypeConverter &typeConverter, MLIRContext *ctx, bool enableTiling)
-      : OpConversionPattern(typeConverter, ctx), enableTiling(enableTiling) {}
+  ONNXMatMulOpLowering(TypeConverter &typeConverter, MLIRContext *ctx,
+      DimAnalysis *dimAnalysis, bool enableTiling)
+      : OpConversionPattern(typeConverter, ctx), dimAnalysis(dimAnalysis),
+        enableTiling(enableTiling) {}
+  DimAnalysis *dimAnalysis;
   bool enableTiling;
   // Handle the generic cases, including when there are broadcasts.
   void replaceGenericMatmul(ONNXMatMulOpAdaptor &operandAdaptor,
@@ -433,20 +435,20 @@ struct ONNXMatMulOpLowering : public OpConversionPattern<ONNXMatMulOp> {
           /*broadcasting B*/ false,
           /*same static broadcast*/ false, alloc, zero, rewriter, loc);
     } else {
-      // Test if have A and B have identical static broadcast shapes.
-      bool sameStaticBroadcast = (enableTiling && aRank > 2 && aRank == bRank);
-      if (sameStaticBroadcast) {
-        auto aShape = A.getType().cast<MemRefType>().getShape();
-        auto bShape = B.getType().cast<MemRefType>().getShape();
+      // Test if have A and B have identical batch size.
+      bool sameBatchsize = (enableTiling && aRank > 2 && aRank == bRank);
+      if (sameBatchsize) {
         for (int i = 0; i < aRank - 2; ++i)
-          if (aShape[i] == ShapedType::kDynamic || aShape[i] != bShape[i]) {
-            sameStaticBroadcast = false;
+          // Note that using A and B from the operation instead of adaptor.
+          // It's because DimAnalysis has been done on operations.
+          if (!dimAnalysis->sameDim(matMulOp.getA(), i, matMulOp.getB(), i)) {
+            sameBatchsize = false;
             break;
           }
       }
       // While there is technically no broadcasting there, we can use nearly the
       // same logic as in replace2x2Matmul2dBroadcasting. So reuse that code.
-      if (sameStaticBroadcast) {
+      if (sameBatchsize) {
         assert(cRank == aRank && "expected IxK * *xKxJ = *xIxJ result");
         replace2x2Matmul2dBroadcasting(adaptor, elementType, shapeHelper,
             /*broadcasting B*/ true,
@@ -463,8 +465,10 @@ struct ONNXMatMulOpLowering : public OpConversionPattern<ONNXMatMulOp> {
 }; // namespace onnx_mlir
 
 void populateLoweringONNXMatMulOpPattern(RewritePatternSet &patterns,
-    TypeConverter &typeConverter, MLIRContext *ctx, bool enableTiling) {
-  patterns.insert<ONNXMatMulOpLowering>(typeConverter, ctx, enableTiling);
+    TypeConverter &typeConverter, MLIRContext *ctx, DimAnalysis *dimAnalysis,
+    bool enableTiling) {
+  patterns.insert<ONNXMatMulOpLowering>(
+      typeConverter, ctx, dimAnalysis, enableTiling);
 }
 
 } // namespace onnx_mlir

src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp

@@ -308,7 +308,8 @@ void populateLoweringONNXHardmaxOpPattern(
 void populateLoweringONNXLRNOpPattern(
     mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
 void populateLoweringONNXMatMulOpPattern(mlir::RewritePatternSet &,
-    mlir::TypeConverter &, mlir::MLIRContext *, bool enableTiling);
+    mlir::TypeConverter &, mlir::MLIRContext *, DimAnalysis *,
+    bool enableTiling);
 void populateLoweringONNXMatMulIntegerOpPattern(
     mlir::RewritePatternSet &, mlir::TypeConverter &, mlir::MLIRContext *);
 void populateLoweringONNXRandomNormalOpPattern(

test/mlir/onnx/onnx_lowering_with_canonicalize_O3.mlir

@@ -1591,6 +1591,39 @@ func.func private @test_matmul7(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->
 
 // -----
 
+// N-D x N-D
+func.func private @test_matmul8(%arg0 : tensor<?x10x10xf32>) -> tensor<*xf32> {
+  %0 ="onnx.MatMul"(%arg0, %arg0) : (tensor<?x10x10xf32>, tensor<?x10x10xf32>) -> tensor<*xf32>
+  "func.return"(%0) : (tensor<*xf32>) -> ()
+
+// mlir2FileCheck.py -a'["A", "B"]' -n'{"1": "RES"}'
+// CHECK-LABEL:  func.func private @test_matmul8
+// CHECK-SAME:   ([[A_:%.+]]: memref<?x10x10xf32>) -> memref<?x10x10xf32> {
+// CHECK-DAG:       [[CST_0_dot_000000_:%.+]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG:       [[CST_10_:%.+]] = arith.constant 10 : index
+// CHECK-DAG:       [[CST_0_:%.+]] = arith.constant 0 : index
+// CHECK:           [[VAR_dim_:%.+]] = memref.dim [[A_]], [[CST_0_]] : memref<?x10x10xf32>
+// CHECK:           [[RES_:%.+]] = memref.alloc([[VAR_dim_]]) {{.*}}: memref<?x10x10xf32>
+// CHECK:           krnl.memset [[RES_]], [[CST_0_dot_000000_]] : memref<?x10x10xf32>
+// CHECK:           [[LOOP_0_:%.+]] = krnl.define_loops 1
+// CHECK:           krnl.iterate([[LOOP_0_]]) with ([[LOOP_0_]] -> [[B_:%.+]] = [[CST_0_]] to [[VAR_dim_]]){
+// CHECK-DAG:         [[RES_1_:%.+]] = krnl.get_induction_var_value([[LOOP_0_]]) : (!krnl.loop) -> index
+// CHECK-DAG:         [[LOOP_1_:%.+]]:3 = krnl.define_loops 3
+// CHECK:             [[BLOCK_TILE__0_:%.+]], [[BLOCK_IN__0_:%.+]] = krnl.block [[LOOP_1_]]#0 4 : (!krnl.loop) -> (!krnl.loop, !krnl.loop)
+// CHECK:             [[BLOCK_TILE__1_:%.+]], [[BLOCK_IN__1_:%.+]] = krnl.block [[LOOP_1_]]#1 8 : (!krnl.loop) -> (!krnl.loop, !krnl.loop)
+// CHECK:             [[BLOCK_TILE__2_:%.+]], [[BLOCK_IN__2_:%.+]] = krnl.block [[LOOP_1_]]#2 8 : (!krnl.loop) -> (!krnl.loop, !krnl.loop)
+// CHECK:             krnl.permute([[BLOCK_TILE__0_]], [[BLOCK_IN__0_]], [[BLOCK_TILE__0_]]_0, [[BLOCK_IN__0_]]_1, [[BLOCK_TILE__0_]]_2, [[BLOCK_IN__0_]]_3) [0, 3, 1, 4, 2, 5] : !krnl.loop, !krnl.loop, !krnl.loop, !krnl.loop, !krnl.loop, !krnl.loop
+// CHECK:             krnl.iterate([[BLOCK_TILE__0_]], [[BLOCK_TILE__0_]]_0, [[BLOCK_TILE__0_]]_2) with ([[LOOP_1_]]#0 -> [[I_0_:%.+]] = [[CST_0_]] to [[CST_10_]], [[LOOP_1_]]#1 -> [[I_1_:%.+]] = [[CST_0_]] to [[CST_10_]], [[LOOP_1_]]#2 -> [[I_2_:%.+]] = [[CST_0_]] to [[CST_10_]]){
+// CHECK:               [[VAR_3_:%.+]]:3 = krnl.get_induction_var_value([[BLOCK_TILE__0_]], [[BLOCK_TILE__0_]]_0, [[BLOCK_TILE__0_]]_2) : (!krnl.loop, !krnl.loop, !krnl.loop) -> (index, index, index)
+// CHECK:               krnl.matmul [[A_]]{{.}}[[RES_1_]], [[CST_0_]], [[CST_0_]]{{.}}, [[A_]]{{.}}[[RES_1_]], [[CST_0_]], [[CST_0_]]{{.}}, [[RES_]]{{.}}[[RES_1_]], [[CST_0_]], [[CST_0_]]{{.}}, ([[BLOCK_IN__0_]], [[BLOCK_IN__0_]]_1, [[BLOCK_IN__0_]]_3), ([[VAR_3_]]#0, [[VAR_3_]]#1, [[VAR_3_]]#2), ([[CST_10_]], [[CST_10_]], [[CST_10_]]) {aTileSize = [], bTileSize = [], cTileSize = [], computeTileSize = [4, 8, 8]} : memref<?x10x10xf32>, memref<?x10x10xf32>, memref<?x10x10xf32>, (!krnl.loop, !krnl.loop, !krnl.loop)
+// CHECK:             }
+// CHECK:           }
+// CHECK:           return [[RES_]] : memref<?x10x10xf32>
+// CHECK:         }
+}
+
+// -----
+
 func.func private @test_pool_unknown_dimensions(%arg0 : tensor<1x3x?x32xf32>) -> tensor<*xf32> {
   %0 = "onnx.AveragePool"(%arg0) {auto_pad = "NOTSET", kernel_shape = [2, 2]} : (tensor<1x3x?x32xf32>) -> tensor<*xf32>
   "func.return"(%0) : (tensor<*xf32>) -> ()