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

[Florian Hahn via llvm-commits]

================
@@ -2228,26 +2262,80 @@ class LowerMatrixIntrinsics {
   }
 
   /// Lower load instructions.
-  MatrixTy VisitLoad(LoadInst *Inst, const ShapeInfo &SI, Value *Ptr) {
-    IRBuilder<> Builder(Inst);
+  MatrixTy VisitLoad(LoadInst *Inst, const ShapeInfo &SI, Value *Ptr,
+                     IRBuilder<> &Builder) {
     return LowerLoad(Inst, Ptr, Inst->getAlign(),
-                     Builder.getInt64(SI.getStride()), Inst->isVolatile(), SI);
+                     Builder.getInt64(SI.getStride()), Inst->isVolatile(), SI,
+                     Builder);
   }
 
   MatrixTy VisitStore(StoreInst *Inst, const ShapeInfo &SI, Value *StoredVal,
-                      Value *Ptr) {
-    IRBuilder<> Builder(Inst);
+                      Value *Ptr, IRBuilder<> &Builder) {
     return LowerStore(Inst, StoredVal, Ptr, Inst->getAlign(),
-                      Builder.getInt64(SI.getStride()), Inst->isVolatile(), SI);
+                      Builder.getInt64(SI.getStride()), Inst->isVolatile(), SI,
+                      Builder);
+  }
+
+  MatrixTy VisitPHI(PHINode *Inst, const ShapeInfo &SI, IRBuilder<> &Builder) {
+    // Shim this->getMatrix to insert split phi's as needed.
+    auto getMatrix = [this, &Builder, SI](Value *MatrixVal) -> MatrixTy {
+      IRBuilder<>::InsertPointGuard IPG(Builder);
+
+      auto I = Inst2ColumnMatrix.find(MatrixVal);
+      if (I == Inst2ColumnMatrix.end()) {
+        if (auto *PHI = dyn_cast<PHINode>(MatrixVal)) {
+          auto *EltTy = cast<VectorType>(PHI->getType())->getElementType();
+          MatrixTy PhiM{SI.NumRows, SI.NumColumns, EltTy};
+
+          Builder.SetInsertPoint(PHI);
+          for (unsigned VI = 0, VE = PhiM.getNumVectors(); VI != VE; ++VI)
+            PhiM.setVector(VI, Builder.CreatePHI(PhiM.getVectorTy(),
----------------
fhahn wrote:

This creates incomplete phis, to be populated later when they are visited, right?

It might be clearer if we pre-process all phis first, creating  incomplete phis for all phis to split, then have `VisitPHI` set the incoming values for all of them during regular traversal?

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