[Matrix] Hoist load/store generation logic, add helpers for tiled access.

[fhahn]

Summary

This patch slightly generalizes the code to emit loads and stores of a
matrix and adds helpers to load/store a tile of a larger matrix.

This will be used in a follow-up patch introducing initial tiling.

[fhahn]

ping

[anemet]

LGTM

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

public:
     ColumnMatrixTy() : Columns() {}
     ColumnMatrixTy(ArrayRef<Value *> Cols)
         : Columns(Cols.begin(), Cols.end()) {}
 
     Value *getColumn(unsigned i) const { return Columns[i]; }
 
     void setColumn(unsigned i, Value *V) { Columns[i] = V; }
 
-    size_t getNumColumns() const { return Columns.size(); }
-    size_t getNumRows() const {
+    unsigned getNumColumns() const { return Columns.size(); }
+    unsigned getNumRows() const {
       assert(Columns.size() > 0 && "Cannot call getNumRows without columns");
       return cast<VectorType>(Columns[0]->getType())->getNumElements();
     }
 
     const SmallVectorImpl<Value *> &getColumnVectors() const { return Columns; }
 
     SmallVectorImpl<Value *> &getColumnVectors() { return Columns; }
 
case Intrinsic::matrix_columnwise_store:
       LowerColumnwiseStore(Inst);
       break;
     default:
       return false;
     }
     return true;
   }
 
-  void LowerLoad(Instruction *Inst, Value *Ptr, Value *Stride,
-                 ShapeInfo Shape) {
-    IRBuilder<> Builder(Inst);
-    auto VType = cast<VectorType>(Inst->getType());
+  /// Load a matrix with \p Shape starting at \p Ptr and using \p Stride between
+  /// columns.
+  ColumnMatrixTy loadMatrix(Type *Ty, Value *Ptr, Value *Stride,
+                            ShapeInfo Shape, IRBuilder<> &Builder) {
+    auto VType = cast<VectorType>(Ty);
     Value *EltPtr = createElementPtr(Ptr, VType->getElementType(), Builder);
     ColumnMatrixTy Result;
     // Distance between start of one column and the start of the next
     for (unsigned C = 0, E = Shape.NumColumns; C < E; ++C) {
       Value *GEP =
           computeColumnAddr(EltPtr, Builder.getInt32(C), Stride, Shape.NumRows,
                             VType->getElementType(), Builder);
       Value *Column = createColumnLoad(GEP, VType->getElementType(), Builder);
       Result.addColumn(Column);
     }
+    return Result.addNumLoads(getNumOps(Result.getColumnTy()) *
+                              Result.getNumColumns());
+  }
+
+  /// Loads a sub-matrix with shape \p ResultShape from a \p R x \p C matrix,
+  /// starting at \p MatrixPtr[I][J].
+  ColumnMatrixTy loadMatrix(Value *MatrixPtr, ShapeInfo MatrixShape, unsigned I,
+                            unsigned J, ShapeInfo ResultShape, Type *EltTy,
+                            IRBuilder<> &Builder) {
+
+    Value *Offset = Builder.CreateAdd(
+        Builder.CreateMul(Builder.getInt32(J),
+                          Builder.getInt32(MatrixShape.NumRows)),
+        Builder.getInt32(I));
+
+    unsigned AS = cast<PointerType>(MatrixPtr->getType())->getAddressSpace();
+    Value *EltPtr =
+        Builder.CreatePointerCast(MatrixPtr, PointerType::get(EltTy, AS));
+    Value *TileStart = Builder.CreateGEP(EltTy, EltPtr, Offset);
+    Type *TileTy =
+        VectorType::get(EltTy, ResultShape.NumRows * ResultShape.NumColumns);
+    Type *TilePtrTy = PointerType::get(TileTy, AS);
+    Value *TilePtr =
+        Builder.CreatePointerCast(TileStart, TilePtrTy, "col.cast");
 
+    return loadMatrix(TileTy, TilePtr, Builder.getInt32(ResultShape.NumRows),
+                      ResultShape, Builder);
+  }
+
+  /// Lower a load instruction with shape information.
+  void LowerLoad(Instruction *Inst, Value *Ptr, Value *Stride,
+                 ShapeInfo Shape) {
+    IRBuilder<> Builder(Inst);
     finalizeLowering(Inst,
-                     Result.addNumLoads(getNumOps(Result.getColumnTy()) *
-                                        Result.getNumColumns()),
+                     loadMatrix(Inst->getType(), Ptr, Stride, Shape, Builder),
                      Builder);
   }
 
   /// Lowers llvm.matrix.columnwise.load.
   ///
   /// The intrinsic loads a matrix from memory using a stride between columns.
   void LowerColumnwiseLoad(CallInst *Inst) {
     Value *Ptr = Inst->getArgOperand(0);
     Value *Stride = Inst->getArgOperand(1);
     LowerLoad(Inst, Ptr, Stride,
               {Inst->getArgOperand(2), Inst->getArgOperand(3)});
   }
 
-  void LowerStore(Instruction *Inst, Value *Matrix, Value *Ptr, Value *Stride,
-                  ShapeInfo Shape) {
-    IRBuilder<> Builder(Inst);
-    auto VType = cast<VectorType>(Matrix->getType());
+  /// Stores a sub-matrix \p StoreVal into the \p R x \p C matrix starting at \p
+  /// MatrixPtr[I][J].
+  void storeMatrix(const ColumnMatrixTy &StoreVal, Value *MatrixPtr,
+                   ShapeInfo MatrixShape, unsigned I, unsigned J, Type *EltTy,
+                   IRBuilder<> &Builder) {
+    Value *Offset = Builder.CreateAdd(
+        Builder.CreateMul(Builder.getInt32(J),
+                          Builder.getInt32(MatrixShape.NumRows)),
+        Builder.getInt32(I));
+
+    unsigned AS = cast<PointerType>(MatrixPtr->getType())->getAddressSpace();
+    Value *EltPtr =
+        Builder.CreatePointerCast(MatrixPtr, PointerType::get(EltTy, AS));
+    Value *TileStart = Builder.CreateGEP(EltTy, EltPtr, Offset);
+    Type *TileTy = VectorType::get(EltTy, StoreVal.getNumRows() *
+                                              StoreVal.getNumColumns());
+    Type *TilePtrTy = PointerType::get(TileTy, AS);
+    Value *TilePtr =
+        Builder.CreatePointerCast(TileStart, TilePtrTy, "col.cast");
+
+    storeMatrix(TileTy, StoreVal, TilePtr,
+                Builder.getInt32(StoreVal.getNumRows()), Builder);
+  }
+
+  /// Store matrix \p StoreVal starting at \p Ptr and using \p Stride between
+  /// columns.
+  ColumnMatrixTy storeMatrix(Type *Ty, ColumnMatrixTy StoreVal, Value *Ptr,
+                             Value *Stride, IRBuilder<> &Builder) {
+    auto VType = cast<VectorType>(Ty);
     Value *EltPtr = createElementPtr(Ptr, VType->getElementType(), Builder);
-    auto LM = getMatrix(Matrix, Shape, Builder);
-    for (auto C : enumerate(LM.columns())) {
-      Value *GEP =
-          computeColumnAddr(EltPtr, Builder.getInt32(C.index()), Stride,
-                            Shape.NumRows, VType->getElementType(), Builder);
+    for (auto C : enumerate(StoreVal.columns())) {
+      Value *GEP = computeColumnAddr(EltPtr, Builder.getInt32(C.index()),
+                                     Stride, StoreVal.getNumRows(),
+                                     VType->getElementType(), Builder);
       createColumnStore(C.value(), GEP, VType->getElementType(), Builder);
     }
-    Inst2ColumnMatrix[Inst] = ColumnMatrixTy().addNumStores(
-        getNumOps(LM.getColumnTy()) * LM.getNumColumns());
+    return ColumnMatrixTy().addNumStores(getNumOps(StoreVal.getColumnTy()) *
+                                         StoreVal.getNumColumns());
+  }
 
-    ToRemove.push_back(Inst);
+  /// Lower a store instruction with shape information.
+  void LowerStore(Instruction *Inst, Value *Matrix, Value *Ptr, Value *Stride,
+                  ShapeInfo Shape) {
+    IRBuilder<> Builder(Inst);
+    auto StoreVal = getMatrix(Matrix, Shape, Builder);
+    finalizeLowering(
+        Inst, storeMatrix(Matrix->getType(), StoreVal, Ptr, Stride, Builder),
+        Builder);
   }
 
   /// Lowers llvm.matrix.columnwise.store.
   ///
   /// The intrinsic store a matrix back memory using a stride between columns.
   void LowerColumnwiseStore(CallInst *Inst) {
     Value *Matrix = Inst->getArgOperand(0);
     Value *Ptr = Inst->getArgOperand(1);