[mlir][arith] Initial support for fastmath flag attributes in the Arithmetic dialect (v2)

This diff adds initial (partial) support for "fastmath" attributes for floating
point operations in the arithmetic dialect. The "fastmath" attributes are
implemented using a default-valued bit enum. The defined flags currently mirror
the fastmath flags in the LLVM dialect (and in LLVM itself). Extending the
set of flags (if necessary) is left as a future task.

In this diff:
- Definition of FastMathAttr as a custom attribute in the Arithmetic dialect
  that inherits from the EnumAttr class.
- Definition of ArithFastMathInterface, which is an interface that is
  implemented by operations that have an arith::fastmath attribute.
- Declaration of a default-valued fastmath attribute for unary and (some) binary
  floating point operations in the Arithmetic dialect.
- Conversion code to lower arithmetic fastmath flags to LLVM fastmath flags

NOT in this diff (but planned or currently in progress):
- Documentation of flag meanings
- Addition of FastMathAttr attributes to other dialects that might lower to the
  Arithmetic dialect (e.g. Math and Complex)
- Folding/rewrite implementations that are enabled by fastmath flags
- Specification of fastmath values from Python bindings (pending other in-
  progress diffs)

Reviewed By: mehdi_amini, vzakhari

Differential Revision: https://reviews.llvm.org/D126305

Author: jfurtek

mlir/include/mlir/Conversion/LLVMCommon/Pattern.h

@@ -22,8 +22,10 @@ namespace detail {
 /// and given operands.
 LogicalResult oneToOneRewrite(Operation *op, StringRef targetOp,
                               ValueRange operands,
+                              ArrayRef<NamedAttribute> targetAttrs,
                               LLVMTypeConverter &typeConverter,
                               ConversionPatternRewriter &rewriter);
+
 } // namespace detail
 } // namespace LLVM
 
@@ -197,7 +199,7 @@ class OneToOneConvertToLLVMPattern : public ConvertOpToLLVMPattern<SourceOp> {
   matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     return LLVM::detail::oneToOneRewrite(op, TargetOp::getOperationName(),
-                                         adaptor.getOperands(),
+                                         adaptor.getOperands(), op->getAttrs(),
                                          *this->getTypeConverter(), rewriter);
   }
 };

mlir/include/mlir/Conversion/LLVMCommon/VectorPattern.h

@@ -56,14 +56,34 @@ LogicalResult handleMultidimensionalVectors(
 
 LogicalResult vectorOneToOneRewrite(Operation *op, StringRef targetOp,
                                     ValueRange operands,
+                                    ArrayRef<NamedAttribute> targetAttrs,
                                     LLVMTypeConverter &typeConverter,
                                     ConversionPatternRewriter &rewriter);
 } // namespace detail
 } // namespace LLVM
 
+// Default attribute conversion class, which passes all source attributes
+// through to the target op, unmodified.
+template <typename SourceOp, typename TargetOp>
+class AttrConvertPassThrough {
+public:
+  AttrConvertPassThrough(SourceOp srcOp) : srcAttrs(srcOp->getAttrs()) {}
+
+  ArrayRef<NamedAttribute> getAttrs() const { return srcAttrs; }
+
+private:
+  ArrayRef<NamedAttribute> srcAttrs;
+};
+
 /// Basic lowering implementation to rewrite Ops with just one result to the
 /// LLVM Dialect. This supports higher-dimensional vector types.
-template <typename SourceOp, typename TargetOp>
+/// The AttrConvert template template parameter should be a template class
+/// with SourceOp and TargetOp type parameters, a constructor that takes
+/// a SourceOp instance, and a getAttrs() method that returns
+/// ArrayRef<NamedAttribute>.
+template <typename SourceOp, typename TargetOp,
+          template <typename, typename> typename AttrConvert =
+              AttrConvertPassThrough>
 class VectorConvertToLLVMPattern : public ConvertOpToLLVMPattern<SourceOp> {
 public:
   using ConvertOpToLLVMPattern<SourceOp>::ConvertOpToLLVMPattern;
@@ -75,9 +95,12 @@ class VectorConvertToLLVMPattern : public ConvertOpToLLVMPattern<SourceOp> {
     static_assert(
         std::is_base_of<OpTrait::OneResult<SourceOp>, SourceOp>::value,
         "expected single result op");
+    // Determine attributes for the target op
+    AttrConvert<SourceOp, TargetOp> attrConvert(op);
+
     return LLVM::detail::vectorOneToOneRewrite(
         op, TargetOp::getOperationName(), adaptor.getOperands(),
-        *this->getTypeConverter(), rewriter);
+        attrConvert.getAttrs(), *this->getTypeConverter(), rewriter);
   }
 };
 } // namespace mlir

mlir/include/mlir/Dialect/Arith/IR/Arith.h

@@ -17,6 +17,7 @@
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
+#include "llvm/ADT/StringExtras.h"
 
 //===----------------------------------------------------------------------===//
 // ArithDialect
@@ -29,6 +30,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/ArithOpsEnums.h.inc"
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/Arith/IR/ArithOpsAttributes.h.inc"
+
+//===----------------------------------------------------------------------===//
+// Arith Interfaces
+//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/Arith/IR/ArithOpsInterfaces.h.inc"
 
 //===----------------------------------------------------------------------===//
 // Arith Dialect Operations

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

@@ -23,6 +23,7 @@ def Arith_Dialect : Dialect {
   }];
 
   let hasConstantMaterializer = 1;
+  let useDefaultAttributePrinterParser = 1;
 }
 
 // The predicate indicates the type of the comparison to perform:
@@ -92,4 +93,32 @@ def AtomicRMWKindAttr : I64EnumAttr<
   let cppNamespace = "::mlir::arith";
 }
 
+def FASTMATH_NONE            : I32BitEnumAttrCaseNone<"none"      >;
+def FASTMATH_REASSOC         : I32BitEnumAttrCaseBit<"reassoc",  0>;
+def FASTMATH_NO_NANS         : I32BitEnumAttrCaseBit<"nnan",     1>;
+def FASTMATH_NO_INFS         : I32BitEnumAttrCaseBit<"ninf",     2>;
+def FASTMATH_NO_SIGNED_ZEROS : I32BitEnumAttrCaseBit<"nsz",      3>;
+def FASTMATH_ALLOW_RECIP     : I32BitEnumAttrCaseBit<"arcp",     4>;
+def FASTMATH_ALLOW_CONTRACT  : I32BitEnumAttrCaseBit<"contract", 5>;
+def FASTMATH_APPROX_FUNC     : I32BitEnumAttrCaseBit<"afn",      6>;
+def FASTMATH_FAST            : I32BitEnumAttrCaseGroup<
+    "fast",
+    [
+      FASTMATH_REASSOC,         FASTMATH_NO_NANS,     FASTMATH_NO_INFS,
+      FASTMATH_NO_SIGNED_ZEROS, FASTMATH_ALLOW_RECIP, FASTMATH_ALLOW_CONTRACT,
+      FASTMATH_APPROX_FUNC]>;
+
+def FastMathFlags : I32BitEnumAttr<
+    "FastMathFlags",
+    "Floating point fast math flags",
+    [
+      FASTMATH_NONE,           FASTMATH_REASSOC,         FASTMATH_NO_NANS,
+      FASTMATH_NO_INFS,        FASTMATH_NO_SIGNED_ZEROS, FASTMATH_ALLOW_RECIP,
+      FASTMATH_ALLOW_CONTRACT, FASTMATH_APPROX_FUNC,     FASTMATH_FAST]> {
+  let separator = ",";
+  let cppNamespace = "::mlir::arith";
+  let genSpecializedAttr = 0;
+  let printBitEnumPrimaryGroups = 1;
+}
+
 #endif // ARITH_BASE

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

@@ -10,13 +10,20 @@
 #define ARITH_OPS
 
 include "mlir/Dialect/Arith/IR/ArithBase.td"
+include "mlir/Dialect/Arith/IR/ArithOpsInterfaces.td"
 include "mlir/Interfaces/CastInterfaces.td"
 include "mlir/Interfaces/InferIntRangeInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.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"
+
+def Arith_FastMathAttr :
+    EnumAttr<Arith_Dialect, FastMathFlags, "fastmath"> {
+  let assemblyFormat = "`<` $value `>`";
+}
 
 // Base class for Arith dialect ops. Ops in this dialect have no side
 // effects and can be applied element-wise to vectors and tensors.
@@ -58,15 +65,27 @@ class Arith_IntBinaryOp<string mnemonic, list<Trait> traits = []> :
 
 // Base class for floating point unary operations.
 class Arith_FloatUnaryOp<string mnemonic, list<Trait> traits = []> :
-    Arith_UnaryOp<mnemonic, traits>,
-    Arguments<(ins FloatLike:$operand)>,
-    Results<(outs FloatLike:$result)>;
+    Arith_UnaryOp<mnemonic,
+      !listconcat([DeclareOpInterfaceMethods<ArithFastMathInterface>],
+                  traits)>,
+    Arguments<(ins FloatLike:$operand,
+      DefaultValuedAttr<Arith_FastMathAttr, "FastMathFlags::none">:$fastmath)>,
+    Results<(outs FloatLike:$result)> {
+  let assemblyFormat = [{ $operand custom<ArithFastMathAttr>($fastmath)
+                          attr-dict `:` type($result) }];
+}
 
 // Base class for floating point binary operations.
 class Arith_FloatBinaryOp<string mnemonic, list<Trait> traits = []> :
-    Arith_BinaryOp<mnemonic, traits>,
-    Arguments<(ins FloatLike:$lhs, FloatLike:$rhs)>,
-    Results<(outs FloatLike:$result)>;
+    Arith_BinaryOp<mnemonic,
+      !listconcat([DeclareOpInterfaceMethods<ArithFastMathInterface>],
+                  traits)>,
+    Arguments<(ins FloatLike:$lhs, FloatLike:$rhs,
+      DefaultValuedAttr<Arith_FastMathAttr, "FastMathFlags::none">:$fastmath)>,
+    Results<(outs FloatLike:$result)> {
+  let assemblyFormat = [{ $lhs `,` $rhs `` custom<ArithFastMathAttr>($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.

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

@@ -0,0 +1,52 @@
+//===-- ArithOpsInterfaces.td - arith op interfaces ---*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This is the Arith interfaces definition file.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef ARITH_OPS_INTERFACES
+#define ARITH_OPS_INTERFACES
+
+include "mlir/IR/OpBase.td"
+
+def ArithFastMathInterface : OpInterface<"ArithFastMathInterface"> {
+  let description = [{
+    Access to operation fastmath flags.
+  }];
+
+  let cppNamespace = "::mlir::arith";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/        "Returns a FastMathFlagsAttr attribute for the operation",
+      /*returnType=*/  "FastMathFlagsAttr",
+      /*methodName=*/  "getFastMathFlagsAttr",
+      /*args=*/        (ins),
+      /*methodBody=*/  [{}],
+      /*defaultImpl=*/ [{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getFastmathAttr();
+      }]
+      >,
+    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/IR/CMakeLists.txt

@@ -1,5 +1,14 @@
 set(LLVM_TARGET_DEFINITIONS ArithOps.td)
 mlir_tablegen(ArithOpsEnums.h.inc -gen-enum-decls)
 mlir_tablegen(ArithOpsEnums.cpp.inc -gen-enum-defs)
+mlir_tablegen(ArithOpsAttributes.h.inc -gen-attrdef-decls
+              -attrdefs-dialect=arith)
+mlir_tablegen(ArithOpsAttributes.cpp.inc -gen-attrdef-defs
+              -attrdefs-dialect=arith)
 add_mlir_dialect(ArithOps arith)
 add_mlir_doc(ArithOps ArithOps Dialects/ -gen-dialect-doc)
+
+set(LLVM_TARGET_DEFINITIONS ArithOpsInterfaces.td)
+mlir_tablegen(ArithOpsInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(ArithOpsInterfaces.cpp.inc -gen-op-interface-defs)
+add_public_tablegen_target(MLIRArithOpsInterfacesIncGen)

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

@@ -48,7 +48,6 @@ class SmartMutex;
 namespace mlir {
 namespace LLVM {
 class LLVMDialect;
-class LoopOptionsAttrBuilder;
 
 namespace detail {
 struct LLVMTypeStorage;

mlir/include/mlir/Dialect/LLVMIR/LLVMOpsInterfaces.td

@@ -23,8 +23,28 @@ def FastmathFlagsInterface : OpInterface<"FastmathFlagsInterface"> {
   let cppNamespace = "::mlir::LLVM";
 
   let methods = [
-    InterfaceMethod<"Get fastmath flags", "::mlir::LLVM::FastmathFlags",
-                    "getFastmathFlags">,
+    InterfaceMethod<
+      /*desc=*/        "Returns a FastmathFlagsAttr attribute for the operation",
+      /*returnType=*/  "FastmathFlagsAttr",
+      /*methodName=*/  "getFastmathAttr",
+      /*args=*/        (ins),
+      /*methodBody=*/  [{}],
+      /*defaultImpl=*/ [{
+        ConcreteOp op = cast<ConcreteOp>(this->getOperation());
+        return op.getFastmathFlagsAttr();
+      }]
+      >,
+    StaticInterfaceMethod<
+      /*desc=*/        [{Returns the name of the FastmathFlagsAttr attribute
+                         for the operation}],
+      /*returnType=*/  "StringRef",
+      /*methodName=*/  "getFastmathAttrName",
+      /*args=*/        (ins),
+      /*methodBody=*/  [{}],
+      /*defaultImpl=*/ [{
+        return "fastmathFlags";
+      }]
+      >
   ];
 }