[Matrix] Propagate shape information through PHI insts

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

@@ -32,6 +32,7 @@
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Analysis/VectorUtils.h"
 #include "llvm/IR/CFG.h"
+#include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DerivedTypes.h"
@@ -288,6 +289,7 @@ static bool isUniformShape(Value *V) {
     }
 
   switch (I->getOpcode()) {
+  case Instruction::PHI:
   case Instruction::FNeg:
     return true;
   default:
@@ -1136,7 +1138,27 @@ class LowerMatrixIntrinsics {
 
     Changed |= !FusedInsts.empty();
 
-    // Fourth, lower remaining instructions with shape information.
+    // Fourth, pre-process all the PHINode's. The incoming values will be
+    // assigned later in VisitPHI.
+    for (Instruction *Inst : MatrixInsts) {
+      auto *PHI = dyn_cast<PHINode>(Inst);
+      if (!PHI)
+        continue;
+
+      const ShapeInfo &SI = ShapeMap.at(Inst);
+      auto *EltTy = cast<FixedVectorType>(PHI->getType())->getElementType();
+      MatrixTy PhiM(SI.NumRows, SI.NumColumns, EltTy);
+
+      IRBuilder<> Builder(Inst);
+      for (unsigned VI = 0, VE = PhiM.getNumVectors(); VI != VE; ++VI)
+        PhiM.setVector(VI, Builder.CreatePHI(PhiM.getVectorTy(),
+                                             PHI->getNumIncomingValues(),
+                                             PHI->getName()));
+      assert(!Inst2ColumnMatrix.contains(PHI) && "map already contains phi?");
+      Inst2ColumnMatrix[PHI] = PhiM;
+    }
+
+    // Fifth, lower remaining instructions with shape information.
     for (Instruction *Inst : MatrixInsts) {
       if (FusedInsts.count(Inst))
         continue;
@@ -1161,6 +1183,8 @@ class LowerMatrixIntrinsics {
         Result = VisitLoad(cast<LoadInst>(Inst), SI, Op1, Builder);
       else if (match(Inst, m_Store(m_Value(Op1), m_Value(Op2))))
         Result = VisitStore(cast<StoreInst>(Inst), SI, Op1, Op2, Builder);
+      else if (auto *PHI = dyn_cast<PHINode>(Inst))
+        Result = VisitPHI(PHI, SI, Builder);
       else
         continue;
 
@@ -1458,7 +1482,8 @@ class LowerMatrixIntrinsics {
                         IRBuilder<> &Builder) {
     auto inserted = Inst2ColumnMatrix.insert(std::make_pair(Inst, Matrix));
     (void)inserted;
-    assert(inserted.second && "multiple matrix lowering mapping");
+    assert((inserted.second || isa<PHINode>(Inst)) &&
+           "multiple matrix lowering mapping");
 
     ToRemove.push_back(Inst);
     Value *Flattened = nullptr;
@@ -2239,6 +2264,41 @@ class LowerMatrixIntrinsics {
                       Builder);
   }
 
+  MatrixTy VisitPHI(PHINode *Inst, const ShapeInfo &SI, IRBuilder<> &Builder) {
+    // Shim this->getMatrix to adjust where it creates new instructions, which
+    // it may need to insert for re-shaping.
+    auto GetMatrix = [this, &Builder, SI, Inst](Value *MatrixVal) -> MatrixTy {
+      IRBuilder<>::InsertPointGuard IPG(Builder);
+      if (auto *MatrixInst = dyn_cast<Instruction>(MatrixVal)) {
+        if (auto MaybeIP = MatrixInst->getInsertionPointAfterDef())
+          Builder.SetInsertPoint(*MaybeIP);
+      } else
+        Builder.SetInsertPoint(Inst->getIterator());
+
+      return this->getMatrix(MatrixVal, SI, Builder);
+    };
+
+    MatrixTy PhiM = GetMatrix(Inst);
+
+    for (auto [IncomingV, IncomingB] :
+         llvm::zip_equal(Inst->incoming_values(), Inst->blocks())) {
+      MatrixTy OpM = GetMatrix(IncomingV);
+
+      for (unsigned VI = 0, VE = PhiM.getNumVectors(); VI != VE; ++VI) {
+        PHINode *NewPHI = cast<PHINode>(PhiM.getVector(VI));
+        NewPHI->addIncoming(OpM.getVector(VI), IncomingB);
+      }
+    }
+
+    // finalizeLowering() may also insert instructions in some cases. The safe
+    // place for those is at the end of the initial block of PHIs.
+    auto IP = Inst->getInsertionPointAfterDef();
+    assert(IP.has_value() &&
+           "expected to find a valid insertion point after the phi");
+    Builder.SetInsertPoint(*IP);
+    return PhiM;
+  }
+
   /// Lower binary operators.
   MatrixTy VisitBinaryOperator(BinaryOperator *Inst, const ShapeInfo &SI,
                                IRBuilder<> &Builder) {