[mlir][arith] Add pass to globally set fastmath attributes for a module

[jfurtek]

Summary

This diff adds a pass to the arith dialect that will set the fastmath
attribute for all operations that support the ArithFastMathInterface
interface. This can be useful to apply module-level fastmath settings in a way
that is similar to setting LLVM fastmath values for an entire file.

Future pass modifications may include adding an option to perform a bitwise
operation on the fastmath flags, to (for example) combine the pass argument with
an existing fastmath value with a bitwise or.

[tblah]

mlir/lib/Dialect/Arith/Transforms/SetFastMath.cpp

2

nit: "Wide integer operation emulation" - copy+paste error?

[jfurtek]

• Fix cut-and-paste error in comment

[jfurtek]

mlir/lib/Dialect/Arith/Transforms/SetFastMath.cpp

2

Fixed - thank you.

[vzakhari]

LGTM

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

25

Just FYI, there is bitEnumContainsAll auto-generated overload for FastMathFlags in ArithOpsEnums.h.inc:

inline constexpr bool bitEnumContainsAll(FastMathFlags bits, FastMathFlags bit) {
  return (bits & bit) == bit;
}

mlir/lib/Dialect/Arith/Transforms/SetFastMath.cpp

48

nit: redundant braces

[jfurtek]

• Fix cut-and-paste error in comment
• Rebase. Use tblgen-erated bit enum function. Remove unnecessary braces.

[jfurtek]

Addressed comments. Thank you @vzakhari .

mlir/include/mlir/Dialect/Arith/IR/ArithBase.td

def FastMathFlags : I32BitEnumAttr<
   let cppNamespace = "::mlir::arith";
   let genSpecializedAttr = 0;
   let printBitEnumPrimaryGroups = 1;
 }
 
 def Arith_FastMathAttr :
     EnumAttr<Arith_Dialect, FastMathFlags, "fastmath"> {
   let assemblyFormat = "`<` $value `>`";
+  let extraClassDeclaration = [{
+    bool hasFlags(FastMathFlags f) const {
+      return bitEnumContainsAll(getValue(), f);
+    }
+  }];
 }
 
 #endif // ARITH_BASE

mlir/include/mlir/Dialect/Arith/IR/ArithOps.td

 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/VectorInterfaces.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/EnumAttr.td"
 
 // Base class for Arith dialect ops. Ops in this dialect have no memory
 // effects and can be applied element-wise to vectors and tensors.
 class Arith_Op<string mnemonic, list<Trait> traits = []> :
     Op<Arith_Dialect, mnemonic,
        traits #
        [DeclareOpInterfaceMethods<VectorUnrollOpInterface>, NoMemoryEffect] #
        ElementwiseMappable.traits>;
 
 // Base class for integer and floating point arithmetic ops. All ops have one
 // result, require operands and results to be of the same type, and can accept
 // tensors or vectors of integers or floats.
class Arith_TotalIntBinaryOp<string mnemonic, list<Trait> traits = []> :
     Arith_IntBinaryOp<mnemonic, traits # [Pure]>;
 
 // Base class for floating point unary operations.
 class Arith_FloatUnaryOp<string mnemonic, list<Trait> traits = []> :
     Arith_UnaryOp<mnemonic,
       !listconcat([DeclareOpInterfaceMethods<ArithFastMathInterface>],
                   traits)>,
     Arguments<(ins FloatLike:$operand,
-      DefaultValuedAttr<Arith_FastMathAttr, "FastMathFlags::none">:$fastmath)>,
+      DefaultValuedAttr<Arith_FastMathAttr,
+                        "::mlir::arith::FastMathFlags::none">:$fastmath)>,
     Results<(outs FloatLike:$result)> {
   let assemblyFormat = [{ $operand (`fastmath` `` $fastmath^)?
                           attr-dict `:` type($result) }];
 }
 
 // Base class for floating point binary operations.
 class Arith_FloatBinaryOp<string mnemonic, list<Trait> traits = []> :
     Arith_BinaryOp<mnemonic,
       !listconcat([Pure, DeclareOpInterfaceMethods<ArithFastMathInterface>],
                   traits)>,
     Arguments<(ins FloatLike:$lhs, FloatLike:$rhs,
-      DefaultValuedAttr<Arith_FastMathAttr, "FastMathFlags::none">:$fastmath)>,
+      DefaultValuedAttr<Arith_FastMathAttr,
+                        "::mlir::arith::FastMathFlags::none">:$fastmath)>,
     Results<(outs FloatLike:$result)> {
   let assemblyFormat = [{ $lhs `,` $rhs (`fastmath` `` $fastmath^)?
                           attr-dict `:` type($result) }];
 }
 
 // Base class for arithmetic cast operations. Requires a single operand and
 // result. If either is a shaped type, then the other must be of the same shape.
 class Arith_CastOp<string mnemonic, TypeConstraint From, TypeConstraint To,

mlir/include/mlir/Dialect/Arith/IR/ArithOpsInterfaces.td

InterfaceMethod<
       /*methodName=*/  "getFastMathFlagsAttr",
       /*args=*/        (ins),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
         ConcreteOp op = cast<ConcreteOp>(this->getOperation());
         return op.getFastmathAttr();
       }]
       >,
+    InterfaceMethod<
+      /*desc=*/        "Sets the FastMathFlagsAttr attribute for the operation",
+      /*returnType=*/  "void",
+      /*methodName=*/  "setFastMathFlagsAttr",
+      /*args=*/        (ins "FastMathFlagsAttr":$value),
+      /*methodBody=*/  [{}],
+      /*defaultImpl=*/ [{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        op.setFastmathAttr(value);
+      }]
+      >,
     StaticInterfaceMethod<
       /*desc=*/        [{Returns the name of the FastMathFlagsAttr attribute
                          for the operation}],
       /*returnType=*/  "StringRef",
       /*methodName=*/  "getFastMathAttrName",
       /*args=*/        (ins),
       /*methodBody=*/  [{}],
       /*defaultImpl=*/ [{
         return "fastmath";
       }]
       >
-
   ];
 }
 
 #endif // ARITH_OPS_INTERFACES

mlir/include/mlir/Dialect/Arith/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_ARITH_TRANSFORMS_PASSES_H_
 #define MLIR_DIALECT_ARITH_TRANSFORMS_PASSES_H_
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 namespace arith {
 
 #define GEN_PASS_DECL
 #include "mlir/Dialect/Arith/Transforms/Passes.h.inc"
 

 
 /// Create a pass to legalize Arith ops.
 std::unique_ptr<Pass> createArithExpandOpsPass();
 
 /// Create a pass to replace signed ops with unsigned ones where they are proven
 /// equivalent.
 std::unique_ptr<Pass> createArithUnsignedWhenEquivalentPass();
 
+/// Create a pass to set the fastmath value for operations that support the
+/// fastmath interface.
+std::unique_ptr<Pass>
+createArithSetFastMathPass(FastMathFlags fm = FastMathFlags::none);
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Arith/Transforms/Passes.h.inc"
 
 } // namespace arith
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITH_TRANSFORMS_PASSES_H_

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

def ArithEmulateWideInt : Pass<"arith-emulate-wide-int"> {
   }];
   let options = [
     Option<"widestIntSupported", "widest-int-supported", "unsigned",
            /*default=*/"32", "Widest integer type supported by the target">,
   ];
   let dependentDialects = ["vector::VectorDialect"];
 }
 
+def ArithSetFastMath : Pass<"set-fastmath"> {
+  let summary = "Set the fastmath attribute for all fastmath operations";
+  let description = [{
+    This pass sets the arith dialect fastmath attribute for all operations that
+    support the ArithFastMathInterface.
+  }];
+  let options = [
+    ListOption<"flags", "value", "std::string",
+       "Comma separated list of fastmath flag bits">
+  ];
+  let constructor = "mlir::arith::createArithSetFastMathPass()";
+}
+
 #endif // MLIR_DIALECT_ARITH_TRANSFORMS_PASSES

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

 add_mlir_dialect_library(MLIRArithTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   EmulateWideInt.cpp
   ExpandOps.cpp
+  SetFastMath.cpp
   UnsignedWhenEquivalent.cpp
 
   ADDITIONAL_HEADER_DIRS
   {$MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Arith/Transforms
 
   DEPENDS
   MLIRArithTransformsIncGen
 

mlir/lib/Dialect/Arith/Transforms/SetFastMath.cpp

+//===- SetFastMath.cpp - Set fastmath flags globally ------------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
+
+namespace mlir::arith {
+#define GEN_PASS_DEF_ARITHSETFASTMATH
+#include "mlir/Dialect/Arith/Transforms/Passes.h.inc"
+} // namespace mlir::arith
+
+using namespace mlir;
+
+namespace {
+
+// Pass to set the fastmath attribute for ALL operations that implement the
+// arith::ArithFastMathInterface.
+// Currently, this pass overwrites the fastmath attribute with the value
+// provided on the command line (or to the pass constructor).
+// TODO: Provide a way to perform a bitwise operation on the source and
+// destination fastmath values (e.g. 'or').
+struct ArithSetFastMathPass
+    : public arith::impl::ArithSetFastMathBase<ArithSetFastMathPass> {
+  explicit ArithSetFastMathPass(arith::FastMathFlags fm) : fastmathValue(fm) {}
+  LogicalResult initialize(MLIRContext *context) override {
+    // Parse the comma-delimited string and set the corresponding fastmath
+    // bit(s).
+    for (auto &f : flags) {
+      if (auto enumOpt = arith::symbolizeFastMathFlags(f))
+        fastmathValue = (fastmathValue | enumOpt.value());
+      else
+        return failure();
+    }
+    return success();
+  }
+  void runOnOperation() override {
+    // Create an attribute with the desired fastmath value
+    auto fmfAttr = arith::FastMathFlagsAttr::get(getOperation()->getContext(),
+                                                 fastmathValue);
+    getOperation()->walk([&](Operation *op) {
+      auto fmI = dyn_cast<arith::ArithFastMathInterface>(op);
+      if (fmI)
+        fmI.setFastMathFlagsAttr(fmfAttr);
+    });
+  }
+  // FastMath enum value to be applied
+  arith::FastMathFlags fastmathValue;
+};
+
+} // namespace
+
+std::unique_ptr<Pass>
+mlir::arith::createArithSetFastMathPass(mlir::arith::FastMathFlags fm) {
+  return std::make_unique<ArithSetFastMathPass>(fm);
+}

mlir/test/Dialect/Arith/set-fastmath.mlir

+// RUN: mlir-opt %s --set-fastmath="value=nsz,nnan,ninf" | FileCheck %s
+
+// CHECK-LABEL: @fastmath
+func.func @fastmath(%arg0: f32, %arg1: f32, %arg2: i32) {
+// CHECK: {{.*}} = arith.addf %arg0, %arg1 fastmath<nnan,ninf,nsz> : f32
+// CHECK: {{.*}} = arith.subf %arg0, %arg1 fastmath<nnan,ninf,nsz> : f32
+  %0 = arith.addf %arg0, %arg1 : f32
+  %1 = arith.subf %arg0, %arg1 : f32
+  return
+}