[InstCombine] Pattern match expressions for naive complex multiply/divide.

[jcranmer-intel]

Summary

These patterns are currently hidden via a command-line flag, since the complex
intrinsics are still experimental.

Depends on D119287

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

 //===- InstCombineInternal.h - InstCombine pass internals -------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
public:
   Instruction *SimplifyAnyMemSet(AnyMemSetInst *MI);
 
   Value *EvaluateInDifferentType(Value *V, Type *Ty, bool isSigned);
 
   /// Returns a value X such that Val = X * Scale, or null if none.
   ///
   /// If the multiplication is known not to overflow then NoSignedWrap is set.
   Value *Descale(Value *Val, APInt Scale, bool &NoSignedWrap);
+
+  /// Try to match a complex intrinsic that produces the given real/imaginary
+  /// pair. Returns whether or not it was successful.
+  bool createComplexMathInstruction(Value *Real, Value *Imag);
 };
 
 class Negator final {
   /// Top-to-bottom, def-to-use negated instruction tree we produced.
   SmallVector<Instruction *, NegatorMaxNodesSSO> NewInstructions;
 
   using BuilderTy = IRBuilder<TargetFolder, IRBuilderCallbackInserter>;
   BuilderTy Builder;

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

 //===- InstCombineLoadStoreAlloca.cpp -------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
if (StoreInst *PrevSI = dyn_cast<StoreInst>(BBI)) {
         ++NumDeadStore;
         // Manually add back the original store to the worklist now, so it will
         // be processed after the operands of the removed store, as this may
         // expose additional DSE opportunities.
         Worklist.push(&SI);
         eraseInstFromFunction(*PrevSI);
         return nullptr;
       }
+
+      // Is this potentially a complex instruction?
+      auto OurGEP = dyn_cast<GetElementPtrInst>(Ptr);
+      auto TheirGEP = dyn_cast<GetElementPtrInst>(PrevSI->getOperand(1));
+      if (PrevSI->isUnordered() && OurGEP && TheirGEP &&
+          OurGEP->getOperand(0) == TheirGEP->getOperand(0) &&
+          OurGEP->getNumIndices() == TheirGEP->getNumIndices() &&
+          OurGEP->getType() == TheirGEP->getType()) {
+        bool AllMatch = true;
+        unsigned LastIndex = OurGEP->getNumIndices();
+        for (unsigned Index = 1; Index < LastIndex; Index++) {
+          if (OurGEP->getOperand(Index) != TheirGEP->getOperand(Index)) {
+            AllMatch = false;
+            break;
+          }
+        }
+        if (!AllMatch)
+          break;
+        if (match(OurGEP->getOperand(LastIndex), m_ConstantInt<1>()) &&
+            match(TheirGEP->getOperand(LastIndex), m_ConstantInt<0>())) {
+          IRBuilderBase::InsertPointGuard Guard(Builder);
+          Builder.SetInsertPoint(PrevSI);
+          if (createComplexMathInstruction(PrevSI->getOperand(0), Val))
+            return &SI;
+        }
+      }
+
       break;
     }
 
     // If this is a load, we have to stop.  However, if the loaded value is from
     // the pointer we're loading and is producing the pointer we're storing,
     // then *this* store is dead (X = load P; store X -> P).
     if (LoadInst *LI = dyn_cast<LoadInst>(BBI)) {
       if (LI == Val && equivalentAddressValues(LI->getOperand(0), Ptr)) {

llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp

 //===- InstCombineVectorOps.cpp -------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
Instruction *InstCombinerImpl::visitInsertValueInst(InsertValueInst &I) {
   }
 
   if (IsRedundant)
     return replaceInstUsesWith(I, I.getOperand(0));
 
   if (Instruction *NewI = foldAggregateConstructionIntoAggregateReuse(I))
     return NewI;
 
+  // Check if this is potentially a complex instruction that has been manually
+  // expanded.
+  ArrayRef<Type *> Fields = I.getType()->subtypes();
+  if (Fields.size() == 2 && Fields[0] == Fields[1] &&
+      Fields[0]->isFloatingPointTy()) {
+    Value *RealV, *ImgV;
+    if (match(&I, m_InsertValue<1>(m_InsertValue<0>(m_Value(), m_Value(RealV)),
+                                   m_Value(ImgV)))) {
+      IRBuilderBase::InsertPointGuard Guard(Builder);
+      Builder.SetInsertPoint(cast<Instruction>(I.getOperand(0)));
+      if (createComplexMathInstruction(RealV, ImgV))
+        return &I;
+    }
+  }
+
   return nullptr;
 }
 
 static bool isShuffleEquivalentToSelect(ShuffleVectorInst &Shuf) {
   // Can not analyze scalable type, the number of elements is not a compile-time
   // constant.
   if (isa<ScalableVectorType>(Shuf.getOperand(0)->getType()))
     return false;
if (Instruction *Splat = foldInsEltIntoSplat(IE))
     return Splat;
 
   if (Instruction *IdentityShuf = foldInsEltIntoIdentityShuffle(IE))
     return IdentityShuf;
 
   if (Instruction *Ext = narrowInsElt(IE, Builder))
     return Ext;
 
+  // Check for a potential computation of a complex instruction.
+  ElementCount Count = IE.getType()->getElementCount();
+  Value *RealV, *ImagV;
+  if (!Count.isScalable() && Count.getFixedValue() == 2 &&
+      match(&IE, m_InsertElt(
+                     m_InsertElt(m_Value(), m_Value(RealV), m_ConstantInt<0>()),
+                     m_Value(ImagV), m_ConstantInt<1>()))) {
+    if (createComplexMathInstruction(RealV, ImagV))
+      return &IE;
+  }
+
   return nullptr;
 }
 
 /// Return true if we can evaluate the specified expression tree if the vector
 /// elements were shuffled in a different order.
 static bool canEvaluateShuffled(Value *V, ArrayRef<int> Mask,
                                 unsigned Depth = 5) {
   // We can always reorder the elements of a constant.
Instruction *InstCombinerImpl::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
   if (isSplat || newMask == LHSMask || newMask == RHSMask || newMask == Mask) {
     if (!newRHS)
       newRHS = UndefValue::get(newLHS->getType());
     return new ShuffleVectorInst(newLHS, newRHS, newMask);
   }
 
   return MadeChange ? &SVI : nullptr;
 }
+
+static cl::opt<bool> InstCombineComplex(
+    "inst-combine-complex",
+    cl::desc("Enable pattern match to llvm.complex.* intrinsics"));
+
+bool InstCombinerImpl::createComplexMathInstruction(Value *Real, Value *Imag) {
+  if (!InstCombineComplex)
+    return false;
+
+  Instruction *RealI = dyn_cast<Instruction>(Real);
+  Instruction *ImagI = dyn_cast<Instruction>(Imag);
+  if (!RealI || !ImagI)
+    return false;
+
+  // Don't try to handle vector instructions for now.
+  if (RealI->getType()->isVectorTy())
+    return false;
+
+  Value *Op0R, *Op0I, *Op1R, *Op1I, *Scale, *Numerator;
+  // Compute the intersection of all the fast math flags of the entire tree up
+  // to the point that the input complex numbers are specified.
+  auto computeFMF = [&]() {
+    SmallVector<Instruction *, 8> Worklist = {RealI, ImagI};
+    FastMathFlags Flags;
+    Flags.set();
+    while (!Worklist.empty()) {
+      Instruction *I = Worklist.back();
+      Worklist.pop_back();
+      Flags &= I->getFastMathFlags();
+      for (Use &U : I->operands()) {
+        Value *V = U.get();
+        if (V == Op0R || V == Op0I || V == Op1R || V == Op1I)
+          continue;
+        Worklist.push_back(cast<Instruction>(V));
+      }
+    }
+    return Flags;
+  };
+
+  Intrinsic::ID NewIntrinsic = Intrinsic::not_intrinsic;
+  // Check for complex multiply:
+  // real = op0.real * op1.real - op0.imag * op1.imag
+  // imag = op0.real * op1.imag + op1.imag * op0.real
+  if (match(Real, m_FSub(m_OneUse(m_FMul(m_Value(Op0R), m_Value(Op1R))),
+                         m_OneUse(m_FMul(m_Value(Op0I), m_Value(Op1I)))))) {
+    if (match(
+            Imag,
+            m_c_FAdd(m_OneUse(m_c_FMul(m_Specific(Op0R), m_Specific(Op1I))),
+                     m_OneUse(m_c_FMul(m_Specific(Op1R), m_Specific(Op0I)))))) {
+      NewIntrinsic = Intrinsic::experimental_complex_fmul;
+    }
+  }
+  // Check for complex div:
+  // real = (op0.real * op1.real + op0.imag * op1.imag) / scale
+  // imag = (op0.imag * op1.real - op0.real * op1.imag) / scale
+  // where scale = op1.real * op1.real + op1.imag * op1.imag
+  else if (match(Imag, m_FDiv(m_Value(Numerator), m_Value(Scale)))) {
+    if (match(Scale,
+              m_FAdd(m_OneUse(m_FMul(m_Value(Op1R), m_Deferred(Op1R))),
+                     m_OneUse(m_FMul(m_Value(Op1I), m_Deferred(Op1I)))))) {
+      // The matching of Op1R and Op1I are temporary, we may need to reverse the
+      // assignments.
+      auto checkNumerator = [&]() {
+        return match(Numerator,
+                     m_OneUse(m_FSub(
+                         m_OneUse(m_c_FMul(m_Value(Op0I), m_Specific(Op1R))),
+                         m_OneUse(m_c_FMul(m_Value(Op0R), m_Specific(Op1I))))));
+      };
+      bool ImagMatches = checkNumerator();
+      if (!ImagMatches) {
+        std::swap(Op1R, Op1I);
+        ImagMatches = checkNumerator();
+      }
+      if (ImagMatches &&
+          match(Real,
+                m_FDiv(m_OneUse(m_c_FAdd(m_OneUse(m_c_FMul(m_Specific(Op0R),
+                                                           m_Specific(Op1R))),
+                                         m_OneUse(m_c_FMul(m_Specific(Op0I),
+                                                           m_Specific(Op1I))))),
+                       m_Specific(Scale)))) {
+        NewIntrinsic = Intrinsic::experimental_complex_fdiv;
+      }
+    }
+  }
+
+  // Make sure we matched an intrinsic.
+  if (NewIntrinsic == Intrinsic::not_intrinsic)
+    return false;
+
+  // Use the computation tree to capture all of the fast-math flags.
+  IRBuilderBase::FastMathFlagGuard FMFGuard(Builder);
+  Builder.setFastMathFlags(computeFMF());
+
+  Value *Op0 = Builder.CreateComplexValue(Op0R, Op0I);
+  Value *Op1 = Builder.CreateComplexValue(Op1R, Op1I);
+
+  // Create new intrinsics. From our pattern matching of only the direct
+  // arithmetic formulas, we have to create them with complex-range=limited.
+  Value *Result;
+  switch (NewIntrinsic) {
+  case Intrinsic::experimental_complex_fmul:
+    Result = Builder.CreateComplexMul(Op0, Op1, true);
+    break;
+  case Intrinsic::experimental_complex_fdiv:
+    Result = Builder.CreateComplexDiv(Op0, Op1, true);
+    break;
+  default:
+    llvm_unreachable("Unexpected complex intrinsic");
+  }
+
+  replaceInstUsesWith(*RealI,
+                      Builder.CreateExtractElement(Result, uint64_t(0)));
+  replaceInstUsesWith(*ImagI,
+                      Builder.CreateExtractElement(Result, uint64_t(1)));
+
+  return true;
+}