[mlir][math] Uplift from arith to math.fma

[Hardcode84]

Summary

Add pass to uplift from arith mulf + addf ops to math.fma if fastmath flags allow it.

[kuhar]

Just some nits

mlir/lib/Dialect/Arith/Utils/Utils.cpp

175–182

Could we use llvm::enum_seq to iterate over these?

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

34

nit: public is the default for structs. Instead, consider adding final:

struct UpliftFma final : OpRewrite...

57

nit: The type of fmt it not obvious here -- we should spell it out instead of using auto

63–64

nit: also here

mlir/test/Dialect/Math/uplift-to-fma.mlir

1

nit: this mixes short (-xyz) and long options (--xyz)

[Hardcode84]

review comments

[Hardcode84]

mlir/lib/Dialect/Arith/Utils/Utils.cpp

175–182

They are bitfields and not consecutive values, I don't this it will work.

[kuhar]

The implementation looks good to me, but you may want to wait for a second approval from someone with more context

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

20

I don't think we need arith here if the pass doesn't create any arith ops.

mlir/lib/Dialect/Arith/Utils/Utils.cpp

175–182

Ack

[tschuett]

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

47

Do you really need the braces?

49

Why else if you return anyway?

[tschuett]

mlir/lib/Dialect/Arith/Utils/Utils.cpp

183

I always check (a & flag) == flag.

[tschuett]

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

49

Please ignore that comment. Apologies.

[Hardcode84]

mlir/lib/Dialect/Arith/Utils/Utils.cpp

183

Actually, this entire function can be replaced by just a & b, don't know why I decided to this in a such convoluted way ))

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

47

Braces need to silence == vs = in condition warnings.

[Hardcode84]

remove arith from dependentDialect, remove FMFIntersect fuction

[Hardcode84]

rebase

[Hardcode84]

In D152633#4411719, @kuhar wrote:

The implementation looks good to me, but you may want to wait for a second approval from someone with more context

@kuhar @tschuett I'm going to merge this, as there are no more comments.

[jpienaar]

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

53

Nit: prefer rewriter.notifyMatchFailure to be able to give context during debugging as to match failures.

[Hardcode84]

rebase, notifyMatchFailure

mlir/include/mlir/Dialect/Arith/Utils/Utils.h

struct ArithBuilder {
   Value select(Value cmp, Value lhs, Value rhs);
   Value sgt(Value lhs, Value rhs);
   Value slt(Value lhs, Value rhs);
 
 private:
   OpBuilder &b;
   Location loc;
 };
+
+/// Computes intersection between two fastmath flags.
+arith::FastMathFlags arithFMFIntersect(arith::FastMathFlags a,
+                                       arith::FastMathFlags b);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITH_UTILS_UTILS_H

mlir/include/mlir/Dialect/Math/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(Transforms)

mlir/include/mlir/Dialect/Math/Transforms/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name Math)
+add_public_tablegen_target(MLIRMathTransformsIncGen)
+
+add_mlir_doc(Passes MathPasses ./ -gen-pass-doc)

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

 //===- Passes.h - Pass Entrypoints ------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_
 
+#include "mlir/Pass/Pass.h"
+
 namespace mlir {
+namespace math {
+#define GEN_PASS_DECL
+#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
+#define GEN_PASS_DECL_MATHUPLIFTTOFMA
+#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
+} // namespace math
 
 class RewritePatternSet;
 
 void populateExpandCtlzPattern(RewritePatternSet &patterns);
 void populateExpandTanPattern(RewritePatternSet &patterns);
 void populateExpandTanhPattern(RewritePatternSet &patterns);
 void populateExpandFmaFPattern(RewritePatternSet &patterns);
 void populateExpandFloorFPattern(RewritePatternSet &patterns);
 void populateExpandCeilFPattern(RewritePatternSet &patterns);
 void populateExpandExp2FPattern(RewritePatternSet &patterns);
 void populateExpandPowFPattern(RewritePatternSet &patterns);
 void populateExpandRoundFPattern(RewritePatternSet &patterns);
 void populateExpandRoundEvenPattern(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 = {});
 
+void populateUpliftToFMAPatterns(RewritePatternSet &patterns);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MATH_TRANSFORMS_PASSES_H_

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

+//===-- Passes.td - Math pass definition file --------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_MATH_TRANSFORMS_PASSES
+#define MLIR_DIALECT_MATH_TRANSFORMS_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def MathUpliftToFMA : Pass<"math-uplift-to-fma"> {
+  let summary = "Uplift arith ops to math.fma.";
+  let description = [{
+    Uplift sequence of addf and mulf ops to math.fma if fastmath flags allows it.
+  }];
+  let dependentDialects = ["arith::ArithDialect", "math::MathDialect"];
+}
+
+#endif // MLIR_DIALECT_MATH_TRANSFORMS_PASSES

mlir/include/mlir/InitAllPasses.h

 #include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/ArmSME/Transforms/Passes.h"
 #include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Func/Transforms/Passes.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
 #include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/NVGPU/Passes.h"
 #include "mlir/Dialect/SCF/Transforms/Passes.h"
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
 #include "mlir/Dialect/Shape/Transforms/Passes.h"
 #include "mlir/Dialect/SparseTensor/Pipelines/Passes.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/Transforms/Passes.h"
inline void registerAllPasses() {
   func::registerFuncPasses();
   registerGPUPasses();
   registerGpuSerializeToCubinPass();
   registerGpuSerializeToHsacoPass();
   registerLinalgPasses();
   registerNVGPUPasses();
   registerSparseTensorPasses();
   LLVM::registerLLVMPasses();
+  math::registerMathPasses();
   memref::registerMemRefPasses();
   registerSCFPasses();
   registerShapePasses();
   spirv::registerSPIRVPasses();
   tensor::registerTensorPasses();
   tosa::registerTosaOptPasses();
   transform::registerTransformPasses();
   vector::registerVectorPasses();

mlir/lib/Dialect/Arith/Utils/Utils.cpp

 Value ArithBuilder::slt(Value lhs, Value rhs) {
   if (isa<FloatType>(lhs.getType()))
     return b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::OLT, lhs, rhs);
   return b.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, lhs, rhs);
 }
 Value ArithBuilder::select(Value cmp, Value lhs, Value rhs) {
   return b.create<arith::SelectOp>(loc, cmp, lhs, rhs);
 }
+
+arith::FastMathFlags mlir::arithFMFIntersect(arith::FastMathFlags a,
+                                             arith::FastMathFlags b) {
+  auto res = arith::FastMathFlags::none;
+  const arith::FastMathFlags flags[] = {
+      arith::FastMathFlags::reassoc, arith::FastMathFlags::nnan,
+      arith::FastMathFlags::ninf,    arith::FastMathFlags::nsz,
+      arith::FastMathFlags::arcp,    arith::FastMathFlags::contract,
+      arith::FastMathFlags::afn,
+  };
+
+  for (auto flag : flags)
+    if (static_cast<bool>(a & flag) && static_cast<bool>(b & flag))
+      res = res | flag;
+
+  return res;
+}

mlir/lib/Dialect/Math/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRMathTransforms
   AlgebraicSimplification.cpp
   ExpandPatterns.cpp
   PolynomialApproximation.cpp
+  UpliftToFMA.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Math/Transforms
 
+  DEPENDS
+  MLIRMathTransformsIncGen
+
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRDialectUtils
   MLIRIR
   MLIRMathDialect
   MLIRSCFDialect
   MLIRPass
   MLIRTransforms
   MLIRX86VectorDialect
   MLIRVectorDialect
   )

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

+//===- UpliftToFMA.cpp - Arith to FMA uplifting ---------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements uplifting from arith ops to math.fma.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir::math {
+#define GEN_PASS_DEF_MATHUPLIFTTOFMA
+#include "mlir/Dialect/Math/Transforms/Passes.h.inc"
+} // namespace mlir::math
+
+using namespace mlir;
+
+template <typename Op>
+static bool isValidForFMA(Op op) {
+  return static_cast<bool>(op.getFastmath() & arith::FastMathFlags::contract);
+}
+
+namespace {
+
+struct UpliftFma : public OpRewritePattern<arith::AddFOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(arith::AddFOp op,
+                                PatternRewriter &rewriter) const override {
+    if (!isValidForFMA(op))
+      return failure();
+
+    Value c;
+    arith::MulFOp ab;
+    if ((ab = op.getLhs().getDefiningOp<arith::MulFOp>())) {
+      c = op.getRhs();
+    } else if ((ab = op.getRhs().getDefiningOp<arith::MulFOp>())) {
+      c = op.getLhs();
+    } else {
+      return failure();
+    }
+
+    if (!isValidForFMA(ab))
+      return failure();
+
+    Value a = ab.getLhs();
+    Value b = ab.getRhs();
+    auto fmf = arithFMFIntersect(op.getFastmath(), ab.getFastmath());
+    rewriter.replaceOpWithNewOp<math::FmaOp>(op, a, b, c, fmf);
+    return success();
+  }
+};
+
+struct MathUpliftToFMA
+    : public math::impl::MathUpliftToFMABase<MathUpliftToFMA> {
+  using MathUpliftToFMABase::MathUpliftToFMABase;
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    populateUpliftToFMAPatterns(patterns);
+    if (failed(
+            applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
+      return signalPassFailure();
+  }
+};
+
+} // namespace
+
+void mlir::populateUpliftToFMAPatterns(RewritePatternSet &patterns) {
+  patterns.insert<UpliftFma>(patterns.getContext());
+}

mlir/test/Dialect/Math/uplift-to-fma.mlir

+// RUN: mlir-opt %s --split-input-file -math-uplift-to-fma | FileCheck %s
+
+// No uplifting without fastmath flags
+// CHECK-LABEL: func @test
+//  CHECK-SAME: (%[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: f32)
+//       CHECK: %[[V1:.*]] = arith.mulf %[[ARG1]], %[[ARG2]]
+//       CHECK: %[[V2:.*]] = arith.addf %[[V1]], %[[ARG3]]
+//       CHECK: return %[[V2]]
+func.func @test(%arg1: f32, %arg2: f32, %arg3: f32) -> f32 {
+  %1 = arith.mulf %arg1, %arg2 : f32
+  %2 = arith.addf %1, %arg3 : f32
+  return %2 : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @test
+//  CHECK-SAME: (%[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: f32)
+//       CHECK: %[[RES:.*]] = math.fma %[[ARG1]], %[[ARG2]], %[[ARG3]] fastmath<fast> : f32
+//       CHECK: return %[[RES]]
+func.func @test(%arg1: f32, %arg2: f32, %arg3: f32) -> f32 {
+  %1 = arith.mulf %arg1, %arg2 fastmath<fast> : f32
+  %2 = arith.addf %1, %arg3 fastmath<fast> : f32
+  return %2 : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @test
+//  CHECK-SAME: (%[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: f32)
+//       CHECK: %[[RES:.*]] = math.fma %[[ARG1]], %[[ARG2]], %[[ARG3]] fastmath<contract> : f32
+//       CHECK: return %[[RES]]
+func.func @test(%arg1: f32, %arg2: f32, %arg3: f32) -> f32 {
+  %1 = arith.mulf %arg1, %arg2 fastmath<fast> : f32
+  %2 = arith.addf %arg3, %1 fastmath<contract> : f32
+  return %2 : f32
+}