[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(),
+                                                 PHI->getNumIncomingValues(),
+                                                 PHI->getName()));
+
+          Inst2ColumnMatrix[PHI] = PhiM;
+        }
+      }
+
+      // getMatrix() may insert some instructions for reshaping. The safe place
+      // to insert them is at the end of the parent block, where the register
+      // allocator would have inserted the copies that materialize the PHI.
----------------
fhahn wrote:

Hm, could there be matrix use of MatrixVal between its definition and the end of the block?

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