[llvm] [Matrix] Propagate shape information through PHI insts (PR #141681)

[Jon Roelofs via llvm-commits]

================
@@ -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();
----------------
jroelofs wrote:

oh, and since we know `Inst` is always a phi, then you're right: we can simplify that to `getFirstInsertionPt()`

https://github.com/llvm/llvm-project/pull/141681