Insertion point should be the dominant op

Author: ashermancinelli

mlir/lib/Dialect/Math/Transforms/SincosFusion.cpp

@@ -9,7 +9,7 @@
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/Math/Transforms/Passes.h"
 #include "mlir/IR/PatternMatch.h"
-#include "mlir/Pass/Pass.h"
+// #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
@@ -39,10 +39,14 @@ struct SincosFusionPattern : OpRewritePattern<math::SinOp> {
     if (!cosOp)
       return failure();
 
+    Operation *firstOp = sinOp->isBeforeInBlock(cosOp) ? sinOp.getOperation()
+                                                       : cosOp.getOperation();
+    rewriter.setInsertionPoint(firstOp);
+
     Type elemType = sinOp.getType();
-    auto sincos = rewriter.create<math::SincosOp>(
-        sinOp.getLoc(), TypeRange{elemType, elemType}, operand,
-        sinOp.getFastmathAttr());
+    auto sincos = math::SincosOp::create(rewriter, firstOp->getLoc(),
+                                         TypeRange{elemType, elemType}, operand,
+                                         sinOp.getFastmathAttr());
 
     rewriter.replaceOp(sinOp, sincos.getSin());
     rewriter.replaceOp(cosOp, sincos.getCos());

mlir/test/Dialect/Math/sincos-fusion.mlir

@@ -17,6 +17,28 @@ func.func @sincos_fusion(%arg0 : f32, %arg1 : f32) -> (f32, f32, f32, f32) {
     func.return %0, %1, %2, %3 : f32, f32, f32, f32
 }
 
+func.func private @sink(%arg0 : f32)
+
+// CHECK:         func.func private @sink(f32)
+// CHECK-LABEL:   func.func @sincos_ensure_ssa_dominance(
+// CHECK-SAME:      %[[ARG0:.*]]: f32,
+// CHECK-SAME:      %[[ARG1:.*]]: f32) -> (f32, f32, f32, f32) {
+// CHECK:           %[[VAL_0:.*]], %[[VAL_1:.*]] = math.sincos %[[ARG0]] : f32
+// CHECK:           call @sink(%[[VAL_0]]) : (f32) -> ()
+// CHECK:           %[[VAL_2:.*]], %[[VAL_3:.*]] = math.sincos %[[ARG1]] : f32
+// CHECK:           call @sink(%[[VAL_3]]) : (f32) -> ()
+// CHECK:           return %[[VAL_0]], %[[VAL_1]], %[[VAL_3]], %[[VAL_2]] : f32, f32, f32, f32
+// CHECK:         }
+func.func @sincos_ensure_ssa_dominance(%arg0 : f32, %arg1 : f32) -> (f32, f32, f32, f32) {
+    %0 = math.sin %arg0 : f32
+    func.call @sink(%0) : (f32) -> ()
+    %1 = math.cos %arg0 : f32
+    %2 = math.cos %arg1 : f32
+    func.call @sink(%2) : (f32) -> ()
+    %3 = math.sin %arg1 : f32
+    func.return %0, %1, %2, %3 : f32, f32, f32, f32
+}
+
 // CHECK-LABEL:   func.func @sincos_fusion_no_match_fmf(
 // CHECK-SAME:      %[[ARG0:.*]]: f32) -> (f32, f32) {
 // CHECK:           %[[VAL_0:.*]] = math.sin %[[ARG0]] fastmath<contract> : f32