[mlir][math]Expand Fused math.fmaf to a multiply-add

[bviyer]

Summary

Fused multiply and add are being pushed directly to the libm. This is problematic
for situations where libm is not available. This patch will break down a fused multiply and
add into a multiply followed by an add.

[rsuderman]

mlir/test/Dialect/Math/expand-math.mlir

71

I would not include the CHECK-NOT, if we are matching the replacement we just assume the old op is gone.

72

[[NAME:%.+]] is used to capture a new variable
[[NAME]] is used to reference a previously captured one

Also, we typically avoid capture the % in the matcher. It seems arbitrary but there are certain cases where it makes the test easier to write. Also make sure the capture named makes sense (e.g. MULF). You can also drop types from the verifier and if they are inferred from the parent.

For short, this check should look like

// CHECK-NEXT:     [[MULF:%.+]] = arith.mulf %[[ARG0]], %[[ARG1]]

Repeat for the following lines.

[bviyer]

Added the changes requested by rsuderman

[bviyer]

@rsuderman: I have fixed the changes you have requested.

mlir/include/mlir/Dialect/Math/Transforms/Passes.h

 
 namespace mlir {
 
 class RewritePatternSet;
 
 void populateExpandCtlzPattern(RewritePatternSet &patterns);
 void populateExpandTanPattern(RewritePatternSet &patterns);
 void populateExpandTanhPattern(RewritePatternSet &patterns);
+void populateExpandFmaFPattern(RewritePatternSet &patterns);
 
 void populateMathAlgebraicSimplificationPatterns(RewritePatternSet &patterns);
 
 struct MathPolynomialApproximationOptions {
   // Enables the use of AVX2 intrinsics in some of the approximations.
   bool enableAvx2 = false;
 };
 
 void populateMathPolynomialApproximationPatterns(
     RewritePatternSet &patterns,
     const MathPolynomialApproximationOptions &options = {});
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_

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

static LogicalResult convertTanOp(math::TanOp op, PatternRewriter &rewriter) {
   Type type = operand.getType();
   Value sin = b.create<math::SinOp>(type, operand);
   Value cos = b.create<math::CosOp>(type, operand);
   Value div = b.create<arith::DivFOp>(type, sin, cos);
   rewriter.replaceOp(op, div);
   return success();
 }
 
+static LogicalResult convertFmaFOp(math::FmaOp op, PatternRewriter &rewriter) {
+  ImplicitLocOpBuilder b(op->getLoc(), rewriter);
+  Value operandA = op.getOperand(0);
+  Value operandB = op.getOperand(1);
+  Value operandC = op.getOperand(2);
+  Type type = op.getType();
+  Value mult = b.create<arith::MulFOp>(type, operandA, operandB);
+  Value add = b.create<arith::AddFOp>(type, mult, operandC);
+  rewriter.replaceOp(op, add);
+  return success();
+}
+
 // Converts math.ctlz to scf and arith operations. This is done
 // by performing a binary search on the bits.
 static LogicalResult convertCtlzOp(math::CountLeadingZerosOp op,
                                    PatternRewriter &rewriter) {
   auto operand = op.getOperand();
   auto operandTy = operand.getType();
   auto eTy = getElementTypeOrSelf(operandTy);
   Location loc = op.getLoc();

 
 void mlir::populateExpandTanPattern(RewritePatternSet &patterns) {
   patterns.add(convertTanOp);
 }
 
 void mlir::populateExpandTanhPattern(RewritePatternSet &patterns) {
   patterns.add(convertTanhOp);
 }
+
+void mlir::populateExpandFmaFPattern(RewritePatternSet &patterns) {
+  patterns.add(convertFmaFOp);
+}

mlir/test/Dialect/Math/expand-math.mlir

func.func @ctlz(%arg: i32) -> i32 {
   return %res : i32
 }
 
 // CHECK-LABEL: @ctlz
 // CHECK-SAME: %[[ARG0:.+]]: i32
 // CHECK-DAG: %[[C0:.+]] = arith.constant 0
 // CHECK-DAG: %[[C16:.+]] = arith.constant 16
 // CHECK-DAG: %[[C65535:.+]] = arith.constant 65535
 // CHECK-DAG: %[[C8:.+]] = arith.constant 8
 // CHECK-DAG: %[[C16777215:.+]] = arith.constant 16777215
 // CHECK-DAG: %[[C4:.+]] = arith.constant 4
 // CHECK-DAG: %[[C268435455:.+]] = arith.constant 268435455
 // CHECK-DAG: %[[C2:.+]] = arith.constant 2
 // CHECK-DAG: %[[C1073741823:.+]] = arith.constant 1073741823
 // CHECK-DAG: %[[C1:.+]] = arith.constant 1
 // CHECK-DAG: %[[C2147483647:.+]] = arith.constant 2147483647
 // CHECK-DAG: %[[C32:.+]] = arith.constant 32
 

 
 func.func @ctlz_vector(%arg: vector<4xi32>) -> vector<4xi32> {
   %res = math.ctlz %arg : vector<4xi32>
   return %res : vector<4xi32>
 }
 
 // CHECK-LABEL: @ctlz_vector
 // CHECK-NOT: math.ctlz
+
+// -----
+
+// CHECK-LABEL:    func @fmaf_func
+// CHECK-SAME:     ([[ARG0:%.+]]: f64, [[ARG1:%.+]]: f64, [[ARG2:%.+]]: f64) -> f64
+func.func @fmaf_func(%a: f64, %b: f64, %c: f64) -> f64 {
+  // CHECK-NEXT:     [[MULF:%.+]] = arith.mulf [[ARG0]], [[ARG1]]
+  // CHECK-NEXT:     [[ADDF:%.+]] = arith.addf [[MULF]], [[ARG2]]
+  // CHECK-NEXT:     return [[ADDF]]
+  %ret = math.fma %a, %b, %c : f64
+  return %ret : f64
+}

mlir/test/lib/Dialect/Math/TestExpandMath.cpp

 };
 } // namespace
 
 void TestExpandMathPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
   populateExpandCtlzPattern(patterns);
   populateExpandTanPattern(patterns);
   populateExpandTanhPattern(patterns);
+  populateExpandFmaFPattern(patterns);
   (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
 }
 
 namespace mlir {
 namespace test {
 void registerTestExpandMathPass() { PassRegistration<TestExpandMathPass>(); }
 } // namespace test
 } // namespace mlir