add checks for nontrivial map cases

Author: srcarroll

mlir/test/Dialect/Linalg/fusion-elementwise.mlir

@@ -130,11 +130,14 @@ func.func @elementwise_ops(%in1: tensor<8xf32>, %in2: tensor<8x10xf32>) -> tenso
     return %sqrt : tensor<8x10xf32>
 }
 
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0)>
 // CHECK-LABEL: func @elementwise_ops
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9]+]]: tensor<8xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9]+]]: tensor<8x10xf32>
 //       CHECK:   %[[EMPTY:.+]] = tensor.empty() : tensor<8x10xf32>
 //       CHECK:   %[[FUSED_OP:.+]] = linalg.generic
+//  CHECK-SAME:       indexing_maps = [#[[MAP1]], #[[MAP0]], #[[MAP0]]]
 //  CHECK-SAME:       ins(%[[ARG0]], %[[ARG1]] : {{.*}}) outs(%[[EMPTY]] :
 //  CHECK-NEXT:   ^bb0(%[[IN0:.*]]: f32, %[[IN1:.*]]: f32, %[[OUT:.*]]: f32):
 //  CHECK-NEXT:     %[[ADD:.*]] = arith.addf %[[IN0]], %[[IN1]]
@@ -193,11 +196,14 @@ func.func @map_genric_ops(%arg0: tensor<8xf32>, %arg1: tensor<8x10xf32>) -> tens
     return %sqrt : tensor<8x10xf32>
 }
 
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0)>
 // CHECK-LABEL: func @map_genric_ops
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9]+]]: tensor<8xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9]+]]: tensor<8x10xf32>
 //       CHECK:   %[[EMPTY:.+]] = tensor.empty() : tensor<8x10xf32>
 //       CHECK:   %[[FUSED_OP:.+]] = linalg.generic
+//  CHECK-SAME:       indexing_maps = [#[[MAP1]], #[[MAP0]], #[[MAP0]]]
 //  CHECK-SAME:       ins(%[[ARG0]], %[[ARG1]] : {{.*}}) outs(%[[EMPTY]] :
 //  CHECK-NEXT:   ^bb0(%[[IN0:.*]]: f32, %[[IN1:.*]]: f32, %[[OUT:.*]]: f32):
 //  CHECK-NEXT:     %[[ADD:.*]] = arith.addf %[[IN0]], %[[IN1]]

mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp

@@ -151,9 +151,14 @@ struct TestLinalgElementwiseFusion
     MLIRContext *context = &this->getContext();
     func::FuncOp funcOp = this->getOperation();
 
+    auto controlFn = [](OpOperand *operand) {
+      auto owner = cast<linalg::LinalgOp>(operand->getOwner());
+      auto producer = cast<linalg::LinalgOp>(operand->get().getDefiningOp());
+      return (linalg::isElementwise(owner) && linalg::isElementwise(producer)) && (!isa<linalg::BroadcastOp>(producer) && !isa<linalg::BroadcastOp>(owner));
+    };
     if (fuseGenericOps) {
       RewritePatternSet fusionPatterns(context);
-      auto controlFn = [](OpOperand *operand) { return true; };
+      // auto controlFn = [](OpOperand *operand) { return true; };
       linalg::populateElementwiseOpsFusionPatterns(fusionPatterns, controlFn);
       if (failed(applyPatternsGreedily(funcOp.getBody(),
                                        std::move(fusionPatterns))))