[Matrix] Preserve existing fast-math flags during lowering

[effective-light]

Summary

This patch makes it so, floating-point instructions created in
LowerMatrixIntrinsics retain fast-math flags from instructions that are
higher up the chain.

Fixes https://bugs.llvm.org/show_bug.cgi?id=49738

[fhahn]

Thanks for the patch! Could you add test coverage for lowering of the other FP instructions (fsub, fmul, fneg)? And maybe also with some variations of individual FMFs.

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

1154

I think it would be cleaner to just pass in the FMF object here? We don't need the instruction for anything else.

Also, I think we don't need to separately pass in AllowContraction any longer?

llvm/test/Transforms/LowerMatrixIntrinsics/preserve-existing-fast-math-flags.ll

16

Is there a reason we need the insert element instructions here, rather than passing in the second operand as vector to the function as well?

[effective-light]

Make existing tests conform to the changes.

[anemet]

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

1154

+1

[effective-light]

Pass FastMathFlags directly into emitMatrixMultiply and remove the
AllowContraction operator from emitMatrixMultiply

[effective-light]

Include more tests.

[effective-light]

rebase

[fhahn]

Thanks for the update! A few more suggestions about dealing with the FMF object.

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

407

Instead of returning None, could we just return a FastMathFlags object with all options set the false (which is the default) ?

1174

I think this should be the other way around, if AllowContractEnabled we could set AllowContraction on FMF. This default could be set in getFastMathFlags?

1686

Can we instead use the FMF object in those versions instead? AllowContraction option should probably also add the contract flag to those.

[effective-light]

Restructure as suggested and fix more tests.

[effective-light]

reformat

[fhahn]

LGTM with the inline nits. Thanks!

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

410

nit: drop {}

llvm/test/Transforms/LowerMatrixIntrinsics/preserve-existing-fast-math-flags.ll

179

could you split up the test, possibly into separate ones for each multiply? Otherwise the IR to check for a single function seems unnecessarily big.

[effective-light]

Address nit and break up the test function.

[fhahn]

Thanks, still LGTM

[effective-light]

@fhahn can you land this when you get the chance (I still don't have commit access)?

[fhahn]

In D103233#2795936, @effective-light wrote:

@fhahn can you land this when you get the chance (I still don't have commit access)?

Sure, thanks for the heads-up!

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

class LowerMatrixIntrinsics {
   /// users of a lowered instruction, if shape information is available and
   /// those need to be removed after we finished lowering.
   SmallVector<Instruction *, 16> ToRemove;
 
   /// Map from instructions to their produced column matrix.
   MapVector<Value *, MatrixTy> Inst2ColumnMatrix;
 
 public:
   LowerMatrixIntrinsics(Function &F, TargetTransformInfo &TTI,
                         AliasAnalysis *AA, DominatorTree *DT, LoopInfo *LI,
                         OptimizationRemarkEmitter *ORE)
       : Func(F), DL(F.getParent()->getDataLayout()), TTI(TTI), AA(AA), DT(DT),
         LI(LI), ORE(ORE) {}
 
   unsigned getNumOps(Type *VT) {
     assert(isa<VectorType>(VT) && "Expected vector type");
     return getNumOps(VT->getScalarType(),
                      cast<FixedVectorType>(VT)->getNumElements());
   }
 
public:
   ///
   /// We can fold a transpose into the operand that is used to extract scalars.
   /// This is the first operands with row-major and the second with
   /// column-major.  If \p IsScalarMatrixTransposed we assume the appropriate
   /// operand is transposed.
   void emitMatrixMultiply(MatrixTy &Result, const MatrixTy &A,
                           const MatrixTy &B, bool AllowContraction,
                           IRBuilder<> &Builder, bool IsTiled,
-                          bool IsScalarMatrixTransposed) {
+                          bool IsScalarMatrixTransposed, Instruction *Inst) {
     const unsigned VF = std::max<unsigned>(
         TTI.getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector)
                 .getFixedSize() /
             Result.getElementType()->getPrimitiveSizeInBits().getFixedSize(),
         1U);
     unsigned R = Result.getNumRows();
     unsigned C = Result.getNumColumns();
     unsigned M = A.getNumColumns();
 
     bool IsFP = Result.getElementType()->isFloatingPointTy();
     assert(A.isColumnMajor() == B.isColumnMajor() &&
            Result.isColumnMajor() == A.isColumnMajor() &&
            "operands must agree on matrix layout");
     unsigned NumComputeOps = 0;
+
+    if (auto *FPOp = dyn_cast<FPMathOperator>(Inst)) {
+      Builder.setFastMathFlags(FPOp->getFastMathFlags());
+    }
+
     if (A.isColumnMajor()) {
       // Multiply columns from the first operand with scalars from the second
       // operand. Then move along the K axes and accumulate the columns.  With
       // this the adds can be vectorized without reassociation.
       for (unsigned J = 0; J < C; ++J) {
         unsigned BlockSize = VF;
         // If Result is zero, we don't need to accumulate in the K==0 iteration.
         bool isSumZero = isa<ConstantAggregateZero>(Result.getColumn(J));
 
         for (unsigned I = 0; I < R; I += BlockSize) {
           // Gradually lower the vectorization factor to cover the remainder.
           while (I + BlockSize > R)
             BlockSize /= 2;
 
           Value *Sum = IsTiled ? Result.extractVector(I, J, BlockSize, Builder)
                                : nullptr;
           for (unsigned K = 0; K < M; ++K) {
             Value *L = A.extractVector(I, K, BlockSize, Builder);
             Value *RH = Builder.CreateExtractElement(
                 B.getColumn(IsScalarMatrixTransposed ? K : J),
                 IsScalarMatrixTransposed ? J : K);
             Value *Splat = Builder.CreateVectorSplat(BlockSize, RH, "splat");
             Sum = createMulAdd(isSumZero && K == 0 ? nullptr : Sum, L, Splat,
-                               Result.getElementType()->isFloatingPointTy(),
-                               Builder, AllowContraction, NumComputeOps);
+                               IsFP, Builder, AllowContraction, NumComputeOps);
           }
           Result.setVector(J,
                            insertVector(Result.getVector(J), I, Sum, Builder));
         }
       }
     } else {
       // Multiply rows from the second operand with scalars from the first
       // operand. Then move along the K axes and accumulate the rows.  With this
void createTiledLoops(CallInst *MatMul, Value *LPtr, ShapeInfo LShape,
     // Insert in the inner loop body, which computes
     //   Res += Load(CurrentRow, K) * Load(K, CurrentColumn)
     Builder.SetInsertPoint(InnerBody->getTerminator());
     // Load tiles of the operands.
     MatrixTy A = loadMatrix(LPtr, {}, false, LShape, TI.CurrentRow, TI.CurrentK,
                             {TileSize, TileSize}, EltType, Builder);
     MatrixTy B = loadMatrix(RPtr, {}, false, RShape, TI.CurrentK, TI.CurrentCol,
                             {TileSize, TileSize}, EltType, Builder);
-    emitMatrixMultiply(TileResult, A, B, AllowContract, Builder, true, false);
+    emitMatrixMultiply(TileResult, A, B, AllowContract, Builder, true, false,
+                       MatMul);
     // Store result after the inner loop is done.
     Builder.SetInsertPoint(TI.RowLoopLatch->getTerminator());
     storeMatrix(TileResult, Store->getPointerOperand(), Store->getAlign(),
                 Store->isVolatile(), {LShape.NumRows, RShape.NumColumns},
                 TI.CurrentRow, TI.CurrentCol, EltType, Builder);
 
     for (unsigned I = 0; I < TileResult.getNumVectors(); I++)
       ColumnPhis[I]->addIncoming(TileResult.getVector(I), TI.InnerLoopLatch);
else {
             MatrixTy A =
                 loadMatrix(APtr, LoadOp0->getAlign(), LoadOp0->isVolatile(),
                            LShape, Builder.getInt64(I), Builder.getInt64(K),
                            {TileR, TileM}, EltType, Builder);
             MatrixTy B =
                 loadMatrix(BPtr, LoadOp1->getAlign(), LoadOp1->isVolatile(),
                            RShape, Builder.getInt64(K), Builder.getInt64(J),
                            {TileM, TileC}, EltType, Builder);
-            emitMatrixMultiply(Res, A, B, AllowContract, Builder, true, false);
+            emitMatrixMultiply(Res, A, B, AllowContract, Builder, true, false,
+                               MatMul);
           }
           storeMatrix(Res, CPtr, Store->getAlign(), Store->isVolatile(), {R, M},
                       Builder.getInt64(I), Builder.getInt64(J), EltType,
                       Builder);
         }
     }
 
     // Mark eliminated instructions as fused and remove them.
if (MatrixLayout == MatrixLayoutTy::ColumnMajor
       }
 
       // Initialize the output
       MatrixTy Result(R, C, EltType);
 
       bool AllowContract =
           AllowContractEnabled ||
           (isa<FPMathOperator>(MatMul) && MatMul->hasAllowContract());
-      emitMatrixMultiply(Result, MA, MB, AllowContract, Builder, false, true);
+      emitMatrixMultiply(Result, MA, MB, AllowContract, Builder, false, true,
+                         MatMul);
 
       FusedInsts.insert(MatMul);
       FusedInsts.insert(cast<Instruction>(Transpose));
       if (Transpose->hasOneUse())
         ToRemove.push_back(cast<Instruction>(Transpose));
       finalizeLowering(MatMul, Result, Builder);
       // TODO: add a fake entry for the folded instruction so that this is
       // included in the expression in the remark.
void LowerMultiply(CallInst *MatMul) {
 
     // Initialize the output
     MatrixTy Result(R, C, EltType);
     assert(Lhs.getElementType() == Result.getElementType() &&
            "Matrix multiply result element type does not match arguments.");
 
     bool AllowContract = AllowContractEnabled || (isa<FPMathOperator>(MatMul) &&
                                                   MatMul->hasAllowContract());
-    emitMatrixMultiply(Result, Lhs, Rhs, AllowContract, Builder, false, false);
+    emitMatrixMultiply(Result, Lhs, Rhs, AllowContract, Builder, false, false,
+                       MatMul);
     finalizeLowering(MatMul, Result, Builder);
   }
 
   /// Lowers llvm.matrix.transpose.
   void LowerTranspose(CallInst *Inst) {
     MatrixTy Result;
     IRBuilder<> Builder(Inst);
     Value *InputVal = Inst->getArgOperand(0);
auto BuildVectorOp = [&Builder, Inst](Value *LHS, Value *RHS) {
       switch (Inst->getOpcode()) {
       case Instruction::Add:
         return Builder.CreateAdd(LHS, RHS);
       case Instruction::Mul:
         return Builder.CreateMul(LHS, RHS);
       case Instruction::Sub:
         return Builder.CreateSub(LHS, RHS);
       case Instruction::FAdd:
-        return Builder.CreateFAdd(LHS, RHS);
+        return Builder.CreateFAddFMF(LHS, RHS, Inst);
       case Instruction::FMul:
-        return Builder.CreateFMul(LHS, RHS);
+        return Builder.CreateFMulFMF(LHS, RHS, Inst);
       case Instruction::FSub:
-        return Builder.CreateFSub(LHS, RHS);
+        return Builder.CreateFSubFMF(LHS, RHS, Inst);
       default:
         llvm_unreachable("Unsupported binary operator for matrix");
       }
     };
 
     for (unsigned I = 0; I < Shape.getNumVectors(); ++I)
       Result.addVector(BuildVectorOp(A.getVector(I), B.getVector(I)));
 
bool VisitUnaryOperator(UnaryOperator *Inst) {
 
     MatrixTy Result;
     MatrixTy M = getMatrix(Op, Shape, Builder);
 
     // Helper to perform unary op on vectors.
     auto BuildVectorOp = [&Builder, Inst](Value *Op) {
       switch (Inst->getOpcode()) {
       case Instruction::FNeg:
-        return Builder.CreateFNeg(Op);
+        return Builder.CreateFNegFMF(Op, Inst);
       default:
         llvm_unreachable("Unsupported unary operator for matrix");
       }
     };
 
     for (unsigned I = 0; I < Shape.getNumVectors(); ++I)
       Result.addVector(BuildVectorOp(M.getVector(I)));
 

llvm/test/Transforms/LowerMatrixIntrinsics/preserve-existing-fast-math-flags.ll

+; RUN: opt -lower-matrix-intrinsics -S < %s | FileCheck %s
+
+; Function Attrs: nofree nounwind uwtable willreturn mustprogress
+define <4 x float> @preserve_fmf(<4 x float> %m, float %x, float %y) {
+; CHECK-LABEL: @preserve_fmf(
+; CHECK: fmul fast <1 x float>
+; CHECK: call fast <1 x float> @llvm.fmuladd.v1f32(
+; CHECK: fmul fast <1 x float>
+; CHECK: call fast <1 x float> @llvm.fmuladd.v1f32(
+; CHECK: fmul fast <1 x float>
+; CHECK: call fast <1 x float> @llvm.fmuladd.v1f32(
+; CHECK: fmul fast <1 x float>
+; CHECK: call fast <1 x float> @llvm.fmuladd.v1f32(
+; CHECK: fadd fast <2 x float>
+; CHECK-NEXT: fadd fast <2 x float>
+  %i1 = insertelement <4 x float> <float poison, float 0.000000e+00, float 0.000000e+00, float poison>, float %x, i64 0
+  %i2 = insertelement <4 x float> %i1, float %y, i64 3
+  %res = tail call fast <4 x float> @llvm.matrix.multiply.v4f32.v4f32.v4f32(<4 x float> %m, <4 x float> %i1, i32 2, i32 2, i32 2)
+  %res.2 = fadd fast <4 x float> %res, %m
+  ret <4 x float> %res
+}
+
+declare <4 x float> @llvm.matrix.multiply.v4f32.v4f32.v4f32(<4 x float>, <4 x float>, i32, i32, i32)