[mlir] [VectorOps] Add the ability to mark FP reductions with "reassociate" attribute

[aartbik]

Summary

Rationale:
In general, passing "fastmath" from MLIR to LLVM backend is not supported, and even just providing such a feature for experimentation is under debate. However, passing fine-grained fastmath related attributes on individual operations is generally accepted. This CL introduces an option to instruct the vector-to-llvm lowering phase to annotate floating-point reductions with the "reassociate" fastmath attribute, which allows the LLVM backend to use SIMD implementations for such constructs. Oher lowering passes can start using this mechanism right away in cases where reassociation is allowed.

Benefit:
For some microbenchmarks on x86-avx2, speedups over 20 were observed for longer vector (due to cleaner, spill-free and SIMD exploiting code).

Usage:
mlir-opt --convert-vector-to-llvm="reassociate-fp-reductions"

[ftynse]

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

226

I'm generally reluctant to having LLVM dialect operations that do don't exist in LLVM IR, as either operations or intrinsics. Could we rather add attributes to the already existing LLVM_VectorReductionV2 for reassociation (and "fast" since there are only two that need to be supported). This should be as simple as

Arguments<(ins LLVM_Type, LLVM_Type, DefaultValuedAttr<BoolAttr, "false">:$reassoc, DefaultValuedAttr<BoolAttr, "false">:$fast)>

and using them in the llvmBuilder. When we have a general mechanism for fast-math flags, we can rely on that instead.

[aartbik]

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

226

Thanks. Agreed an attributed is cleaner. Changed that (note that I just need $reassoc, not $fast ). I had an internal prototype for passing on fast-math flags in general, but that stopped after an internal discussion, since it was deemed to dangerous. Hence the intrinsic focused solution for now.

[aartbik]

used attributed rather than naming for reassoc=true

[ftynse]

Thanks!

[mehdi_amini]

Nice! I agree the attribute is appropriate here.

To be clear: I'm not against a fast-math for experimenting, I wanted to make sure we build what you did here as a "proper" solution and we don't rely on the "big hammer" global fast-math instead.

mlir/include/mlir/Conversion/Passes.td

 //===----------------------------------------------------------------------===//
 // VectorToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToLLVM : Pass<"convert-vector-to-llvm", "ModuleOp"> {
   let summary = "Lower the operations from the vector dialect into the LLVM "
                 "dialect";
   let constructor = "mlir::createConvertVectorToLLVMPass()";
+  let options = [
+    Option<"reassociateFPReductions", "reassociate-fp-reductions",
+           "bool", /*default=*/"false",
+           "Allows llvm to reassociate floating-point reductions for speed">
+  ];
 }
 
 //===----------------------------------------------------------------------===//
 // VectorToROCDL
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToROCDL : Pass<"convert-vector-to-rocdl", "ModuleOp"> {
   let summary = "Lower the operations from the vector dialect into the ROCDL "
                 "dialect";
   let constructor = "mlir::createConvertVectorToROCDLPass()";
 }
 
 #endif // MLIR_CONVERSION_PASSES

mlir/include/mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h

 
 /// Collect a set of patterns to convert from Vector contractions to LLVM Matrix
 /// Intrinsics. To lower to assembly, the LLVM flag -lower-matrix-intrinsics
 /// will be needed when invoking LLVM.
 void populateVectorToLLVMMatrixConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns);
 
 /// Collect a set of patterns to convert from the Vector dialect to LLVM.
-void populateVectorToLLVMConversionPatterns(LLVMTypeConverter &converter,
-                                            OwningRewritePatternList &patterns);
+void populateVectorToLLVMConversionPatterns(
+    LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
+    bool reassociateFPReductions = false);
 
 /// Create a pass to convert vector operations to the LLVMIR dialect.
 std::unique_ptr<OperationPass<ModuleOp>> createConvertVectorToLLVMPass();
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_VECTORTOLLVM_CONVERTVECTORTOLLVM_H_

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

: LLVM_OneResultIntrOp<"experimental.vector.reduce." # mnem, [], [0], []>,
       Arguments<(ins LLVM_Type)>;
 
 // LLVM vector reduction over a single vector, with an initial value.
 class LLVM_VectorReductionV2<string mnem>
     : LLVM_OneResultIntrOp<"experimental.vector.reduce.v2." # mnem,
                            [0], [1], []>,
       Arguments<(ins LLVM_Type, LLVM_Type)>;
 
+// LLVM vector reduction over a single vector, with an initial value,
+// and with permission to reassociate the reduction operations.
+class LLVM_VectorReductionV2R<string mnem>
+    : LLVM_OpBase<LLVM_Dialect, "intr.experimental.vector.reduce.v2r." # mnem, []>,
+      Results<(outs LLVM_Type:$res)>,
+      Arguments<(ins LLVM_Type, LLVM_Type)> {
+  let llvmBuilder = [{
+    llvm::Module *module = builder.GetInsertBlock()->getModule();
+    llvm::Function *fn = llvm::Intrinsic::getDeclaration(
+        module,
+        llvm::Intrinsic::experimental_vector_reduce_v2_}] # mnem # [{,
+        { }] # StrJoin<!listconcat(
+            ListIntSubst<LLVM_IntrPatterns.result, [0]>.lst,
+            ListIntSubst<LLVM_IntrPatterns.operand, [1]>.lst)>.result # [{
+        });
+    auto operands = lookupValues(opInst.getOperands());
+    llvm::FastMathFlags origFM = builder.getFastMathFlags();
+    llvm::FastMathFlags tempFM = origFM;
+    tempFM.setAllowReassoc(true);
+    builder.setFastMathFlags(tempFM);
+    $res = builder.CreateCall(fn, operands);
+    builder.setFastMathFlags(origFM);
+  }];
+}
+
 #endif  // LLVMIR_OP_BASE

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 def LLVM_experimental_vector_reduce_fmin : LLVM_VectorReduction<"fmin">;
 def LLVM_experimental_vector_reduce_or : LLVM_VectorReduction<"or">;
 def LLVM_experimental_vector_reduce_smax : LLVM_VectorReduction<"smax">;
 def LLVM_experimental_vector_reduce_smin : LLVM_VectorReduction<"smin">;
 def LLVM_experimental_vector_reduce_umax : LLVM_VectorReduction<"umax">;
 def LLVM_experimental_vector_reduce_umin : LLVM_VectorReduction<"umin">;
 def LLVM_experimental_vector_reduce_xor : LLVM_VectorReduction<"xor">;
 
+// Floating-point add/mul reduction that may not be reassociated by backend.
 def LLVM_experimental_vector_reduce_v2_fadd : LLVM_VectorReductionV2<"fadd">;
 def LLVM_experimental_vector_reduce_v2_fmul : LLVM_VectorReductionV2<"fmul">;
 
+// Floating-point add/mul reduction that may be reassociated by backend.
+def LLVM_experimental_vector_reduce_v2r_fadd : LLVM_VectorReductionV2R<"fadd">;
+def LLVM_experimental_vector_reduce_v2r_fmul : LLVM_VectorReductionV2R<"fmul">;
+
 //
 // LLVM Matrix operations.
 //
 
 /// Create a column major, strided 2-D matrix load, as specified in the LLVM
 /// MatrixBuilder.
 /// data       - Start address of the matrix read
 /// rows       - Number of rows in matrix (must be a constant)

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

rewriter.replaceOpWithNewOp<LLVM::MatrixTransposeOp>(
         adaptor.matrix(), transOp.rows(), transOp.columns());
     return success();
   }
 };
 
 class VectorReductionOpConversion : public ConvertToLLVMPattern {
 public:
   explicit VectorReductionOpConversion(MLIRContext *context,
-                                       LLVMTypeConverter &typeConverter)
+                                       LLVMTypeConverter &typeConverter,
+                                       bool reassociateFP)
       : ConvertToLLVMPattern(vector::ReductionOp::getOperationName(), context,
-                             typeConverter) {}
+                             typeConverter),
+        reassociateFPReductions(reassociateFP) {}
 
   LogicalResult
   matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     auto reductionOp = cast<vector::ReductionOp>(op);
     auto kind = reductionOp.kind();
     Type eltType = reductionOp.dest().getType();
     Type llvmType = typeConverter.convertType(eltType);
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     } else if (eltType.isF32() || eltType.isF64()) {
       // Floating-point reductions: add/mul/min/max
       if (kind == "add") {
         // Optional accumulator (or zero).
         Value acc = operands.size() > 1 ? operands[1]
                                         : rewriter.create<LLVM::ConstantOp>(
                                               op->getLoc(), llvmType,
                                               rewriter.getZeroAttr(eltType));
+        if (reassociateFPReductions)
+          rewriter
+              .replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2r_fadd>(
+                  op, llvmType, acc, operands[0]);
+        else
           rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fadd>(
               op, llvmType, acc, operands[0]);
       } else if (kind == "mul") {
         // Optional accumulator (or one).
         Value acc = operands.size() > 1
                         ? operands[1]
                         : rewriter.create<LLVM::ConstantOp>(
                               op->getLoc(), llvmType,
                               rewriter.getFloatAttr(eltType, 1.0));
+        if (reassociateFPReductions)
+          rewriter
+              .replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2r_fmul>(
+                  op, llvmType, acc, operands[0]);
+        else
           rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_v2_fmul>(
               op, llvmType, acc, operands[0]);
       } else if (kind == "min")
         rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmin>(
             op, llvmType, operands[0]);
       else if (kind == "max")
         rewriter.replaceOpWithNewOp<LLVM::experimental_vector_reduce_fmax>(
             op, llvmType, operands[0]);
       else
         return failure();
       return success();
     }
     return failure();
   }
+
+private:
+  const bool reassociateFPReductions;
 };
 
 class VectorShuffleOpConversion : public ConvertToLLVMPattern {
 public:
   explicit VectorShuffleOpConversion(MLIRContext *context,
                                      LLVMTypeConverter &typeConverter)
       : ConvertToLLVMPattern(vector::ShuffleOp::getOperationName(), context,
                              typeConverter) {}
public:
   /// bounded as the rank is strictly decreasing.
   bool hasBoundedRewriteRecursion() const final { return true; }
 };
 
 } // namespace
 
 /// Populate the given list with patterns that convert from Vector to LLVM.
 void mlir::populateVectorToLLVMConversionPatterns(
-    LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
+    LLVMTypeConverter &converter, OwningRewritePatternList &patterns,
+    bool reassociateFPReductions) {
   MLIRContext *ctx = converter.getDialect()->getContext();
   // clang-format off
   patterns.insert<VectorFMAOpNDRewritePattern,
                   VectorInsertStridedSliceOpDifferentRankRewritePattern,
                   VectorInsertStridedSliceOpSameRankRewritePattern,
                   VectorStridedSliceOpConversion>(ctx);
+  patterns.insert<VectorReductionOpConversion>(
+      ctx, converter, reassociateFPReductions);
   patterns
-      .insert<VectorReductionOpConversion,
-              VectorShuffleOpConversion,
+      .insert<VectorShuffleOpConversion,
               VectorExtractElementOpConversion,
               VectorExtractOpConversion,
               VectorFMAOp1DConversion,
               VectorInsertElementOpConversion,
               VectorInsertOpConversion,
               VectorPrintOpConversion,
               VectorTransferConversion<TransferReadOp>,
               VectorTransferConversion<TransferWriteOp>,
// all contraction operations. Also applies folding and DCE.
     populateVectorContractLoweringPatterns(patterns, &getContext());
     applyPatternsAndFoldGreedily(getOperation(), patterns);
   }
 
   // Convert to the LLVM IR dialect.
   LLVMTypeConverter converter(&getContext());
   OwningRewritePatternList patterns;
   populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
-  populateVectorToLLVMConversionPatterns(converter, patterns);
+  populateVectorToLLVMConversionPatterns(converter, patterns,
+                                         reassociateFPReductions);
   populateVectorToLLVMMatrixConversionPatterns(converter, patterns);
   populateStdToLLVMConversionPatterns(converter, patterns);
 
   LLVMConversionTarget target(getContext());
   if (failed(applyPartialConversion(getOperation(), target, patterns))) {
     signalPassFailure();
   }
 }
 
 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertVectorToLLVMPass() {
   return std::make_unique<LowerVectorToLLVMPass>();
 }

mlir/test/Conversion/VectorToLLVM/vector-reduction-to-llvm.mlir

+// RUN: mlir-opt %s -convert-vector-to-llvm | FileCheck %s
+// RUN: mlir-opt %s -convert-vector-to-llvm='reassociate-fp-reductions' | FileCheck %s --check-prefix=REASSOC
+
+//
+// CHECK-LABEL: llvm.func @reduce_add_f32(
+// CHECK-SAME: %[[A:.*]]: !llvm<"<16 x float>">)
+//      CHECK: %[[C:.*]] = llvm.mlir.constant(0.000000e+00 : f32) : !llvm.float
+//      CHECK: %[[V:.*]] = "llvm.intr.experimental.vector.reduce.v2.fadd"(%[[C]], %[[A]])
+//      CHECK: llvm.return %[[V]] : !llvm.float
+//
+// REASSOC-LABEL: llvm.func @reduce_add_f32(
+// REASSOC-SAME: %[[A:.*]]: !llvm<"<16 x float>">)
+//      REASSOC: %[[C:.*]] = llvm.mlir.constant(0.000000e+00 : f32) : !llvm.float
+//      REASSOC: %[[V:.*]] = "llvm.intr.experimental.vector.reduce.v2r.fadd"(%[[C]], %[[A]])
+//      REASSOC: llvm.return %[[V]] : !llvm.float
+//
+func @reduce_add_f32(%arg0: vector<16xf32>) -> f32 {
+  %0 = vector.reduction "add", %arg0 : vector<16xf32> into f32
+  return %0 : f32
+}
+
+//
+// CHECK-LABEL: llvm.func @reduce_mul_f32(
+// CHECK-SAME: %[[A:.*]]: !llvm<"<16 x float>">)
+//      CHECK: %[[C:.*]] = llvm.mlir.constant(1.000000e+00 : f32) : !llvm.float
+//      CHECK: %[[V:.*]] = "llvm.intr.experimental.vector.reduce.v2.fmul"(%[[C]], %[[A]])
+//      CHECK: llvm.return %[[V]] : !llvm.float
+//
+// REASSOC-LABEL: llvm.func @reduce_mul_f32(
+// REASSOC-SAME: %[[A:.*]]: !llvm<"<16 x float>">)
+//      REASSOC: %[[C:.*]] = llvm.mlir.constant(1.000000e+00 : f32) : !llvm.float
+//      REASSOC: %[[V:.*]] = "llvm.intr.experimental.vector.reduce.v2r.fmul"(%[[C]], %[[A]])
+//      REASSOC: llvm.return %[[V]] : !llvm.float
+//
+func @reduce_mul_f32(%arg0: vector<16xf32>) -> f32 {
+  %0 = vector.reduction "mul", %arg0 : vector<16xf32> into f32
+  return %0 : f32
+}

mlir/test/Target/llvmir-intrinsics.mlir

llvm.func @vector_reductions(%arg0: !llvm.float, %arg1: !llvm<"<8 x float>">, %arg2: !llvm<"<8 x i32>">) {
   // CHECK: call i32 @llvm.experimental.vector.reduce.umax.v8i32
   "llvm.intr.experimental.vector.reduce.umax"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
   // CHECK: call i32 @llvm.experimental.vector.reduce.umin.v8i32
   "llvm.intr.experimental.vector.reduce.umin"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
   // CHECK: call float @llvm.experimental.vector.reduce.v2.fadd.f32.v8f32
   "llvm.intr.experimental.vector.reduce.v2.fadd"(%arg0, %arg1) : (!llvm.float, !llvm<"<8 x float>">) -> !llvm.float
   // CHECK: call float @llvm.experimental.vector.reduce.v2.fmul.f32.v8f32
   "llvm.intr.experimental.vector.reduce.v2.fmul"(%arg0, %arg1) : (!llvm.float, !llvm<"<8 x float>">) -> !llvm.float
+  // CHECK: call reassoc float @llvm.experimental.vector.reduce.v2.fadd.f32.v8f32
+  "llvm.intr.experimental.vector.reduce.v2r.fadd"(%arg0, %arg1) : (!llvm.float, !llvm<"<8 x float>">) -> !llvm.float
+  // CHECK: call reassoc float @llvm.experimental.vector.reduce.v2.fmul.f32.v8f32
+  "llvm.intr.experimental.vector.reduce.v2r.fmul"(%arg0, %arg1) : (!llvm.float, !llvm<"<8 x float>">) -> !llvm.float
   // CHECK: call i32 @llvm.experimental.vector.reduce.xor.v8i32
   "llvm.intr.experimental.vector.reduce.xor"(%arg2) : (!llvm<"<8 x i32>">) -> !llvm.i32
   llvm.return
 }
 
 // CHECK-LABEL: @matrix_intrinsics
 //                                       4x16                       16x3
 llvm.func @matrix_intrinsics(%A: !llvm<"<64 x float>">, %B: !llvm<"<48 x float>">,