[Matrix] Generalize ColumnMatrixTy to MatrixTy (NFC).

[fhahn]

Summary

This patch sets the stage for supporting both row and column major
layouts for matrixes. It renames ColumnMatrixTy to MatrixTy, adds
booleans indicating the underlying layout to both MatrixTy and ShapeInfo
and generalizes the methods of MatrixTy to support both row and column
major layouts.

[anemet]

LGTM

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

167–168

update comment

[fhahn]

Rebase and update comment.

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

struct OpInfoTy {
 
     OpInfoTy &operator+=(const OpInfoTy &RHS) {
       NumStores += RHS.NumStores;
       NumLoads += RHS.NumLoads;
       NumComputeOps += RHS.NumComputeOps;
       return *this;
     }
   };
 
-  /// Wrapper class representing a matrix as a set of column vectors.
-  /// All column vectors must have the same vector type.
-  class ColumnMatrixTy {
-    SmallVector<Value *, 16> Columns;
+  /// Wrapper class representing a matrix as a set of vectors, either in row or column major layout. All vectors must have the same vector type.
+  class MatrixTy {
+    SmallVector<Value *, 16> Vectors;
 
     OpInfoTy OpInfo;
 
+    bool IsColumnMajor = true;
+
   public:
-    ColumnMatrixTy() : Columns() {}
-    ColumnMatrixTy(ArrayRef<Value *> Cols)
-        : Columns(Cols.begin(), Cols.end()) {}
+    MatrixTy() : Vectors() {}
+    MatrixTy(ArrayRef<Value *> Vectors)
+        : Vectors(Vectors.begin(), Vectors.end()) {}
 
-    Value *getColumn(unsigned i) const { return Columns[i]; }
+    Value *getVector(unsigned i) const { return Vectors[i]; }
+    Value *getColumn(unsigned i) const {
+      assert(isColumnMajor() && "only supported for column-major matrixes");
+      return Vectors[i];
+    }
 
-    void setColumn(unsigned i, Value *V) { Columns[i] = V; }
+    void setColumn(unsigned i, Value *V) { Vectors[i] = V; }
 
-    Type *getElementType() {
-      return cast<VectorType>(Columns[0]->getType())->getElementType();
-    }
+    Type *getElementType() { return getVectorTy()->getElementType(); }
 
-    unsigned getNumColumns() const { return Columns.size(); }
+    unsigned getNumColumns() const {
+      if (isColumnMajor())
+        return Vectors.size();
+      else {
+        assert(Vectors.size() > 0 && "Cannot call getNumRows without columns");
+        return cast<VectorType>(Vectors[0]->getType())->getNumElements();
+      }
+    }
     unsigned getNumRows() const {
-      assert(Columns.size() > 0 && "Cannot call getNumRows without columns");
-      return cast<VectorType>(Columns[0]->getType())->getNumElements();
+      if (isColumnMajor()) {
+        assert(Vectors.size() > 0 && "Cannot call getNumRows without columns");
+        return cast<VectorType>(Vectors[0]->getType())->getNumElements();
+      } else
+        return Vectors.size();
     }
 
-    const SmallVectorImpl<Value *> &getColumnVectors() const { return Columns; }
+    const SmallVectorImpl<Value *> &getColumnVectors() const { return Vectors; }
 
-    SmallVectorImpl<Value *> &getColumnVectors() { return Columns; }
+    SmallVectorImpl<Value *> &getColumnVectors() { return Vectors; }
 
-    void addColumn(Value *V) { Columns.push_back(V); }
+    void addColumn(Value *V) { Vectors.push_back(V); }
 
     VectorType *getColumnTy() {
-      return cast<VectorType>(Columns[0]->getType());
+      assert(isColumnMajor() && "only supported for column-major matrixes");
+      return getVectorTy();
+    }
+
+    VectorType *getVectorTy() {
+      return cast<VectorType>(Vectors[0]->getType());
     }
 
     iterator_range<SmallVector<Value *, 8>::iterator> columns() {
-      return make_range(Columns.begin(), Columns.end());
+      return make_range(Vectors.begin(), Vectors.end());
     }
 
     /// Embed the columns of the matrix into a flat vector by concatenating
     /// them.
     Value *embedInVector(IRBuilder<> &Builder) const {
-      return Columns.size() == 1 ? Columns[0]
-                                 : concatenateVectors(Builder, Columns);
+      return Vectors.size() == 1 ? Vectors[0]
+                                 : concatenateVectors(Builder, Vectors);
     }
 
-    ColumnMatrixTy &addNumLoads(unsigned N) {
+    MatrixTy &addNumLoads(unsigned N) {
       OpInfo.NumLoads += N;
       return *this;
     }
 
     void setNumLoads(unsigned N) { OpInfo.NumLoads = N; }
 
-    ColumnMatrixTy &addNumStores(unsigned N) {
+    MatrixTy &addNumStores(unsigned N) {
       OpInfo.NumStores += N;
       return *this;
     }
 
-    ColumnMatrixTy &addNumComputeOps(unsigned N) {
+    MatrixTy &addNumComputeOps(unsigned N) {
       OpInfo.NumComputeOps += N;
       return *this;
     }
 
     unsigned getNumStores() const { return OpInfo.NumStores; }
     unsigned getNumLoads() const { return OpInfo.NumLoads; }
     unsigned getNumComputeOps() const { return OpInfo.NumComputeOps; }
 
     const OpInfoTy &getOpInfo() const { return OpInfo; }
+
+    bool isColumnMajor() const { return IsColumnMajor; }
   };
 
   struct ShapeInfo {
     unsigned NumRows;
     unsigned NumColumns;
 
     ShapeInfo(unsigned NumRows = 0, unsigned NumColumns = 0)
         : NumRows(NumRows), NumColumns(NumColumns) {}
class LowerMatrixIntrinsics {
   DenseMap<Value *, ShapeInfo> ShapeMap;
 
   /// List of instructions to remove. While lowering, we are not replacing all
   /// users of a lowered instruction, if shape information is available and
   /// those need to be removed after we finished lowering.
   SmallVector<Instruction *, 16> ToRemove;
 
   /// Map from instructions to their produced column matrix.
-  MapVector<Value *, ColumnMatrixTy> Inst2ColumnMatrix;
+  MapVector<Value *, MatrixTy> Inst2ColumnMatrix;
 
 public:
   LowerMatrixIntrinsics(Function &F, TargetTransformInfo &TTI,
                         OptimizationRemarkEmitter &ORE)
       : Func(F), DL(F.getParent()->getDataLayout()), TTI(TTI), ORE(ORE) {}
 
   unsigned getNumOps(Type *VT) {
     assert(isa<VectorType>(VT) && "Expected vector type");
return std::ceil((ST->getPrimitiveSizeInBits() * N).getFixedSize() /
                      double(TTI.getRegisterBitWidth(true)));
   }
 
   /// Return the set of column vectors that a matrix value is lowered to.
   ///
   /// If we lowered \p MatrixVal, just return the cache result column matrix.
   /// Otherwie split the flat vector \p MatrixVal containing a matrix with
   /// shape \p SI into column vectors.
-  ColumnMatrixTy getMatrix(Value *MatrixVal, const ShapeInfo &SI,
-                           IRBuilder<> &Builder) {
+  MatrixTy getMatrix(Value *MatrixVal, const ShapeInfo &SI,
     VectorType *VType = dyn_cast<VectorType>(MatrixVal->getType());
     assert(VType && "MatrixVal must be a vector type");
     assert(VType->getNumElements() == SI.NumRows * SI.NumColumns &&
            "The vector size must match the number of matrix elements");
 
     // Check if we lowered MatrixVal using shape information. In that case,
     // return the existing column matrix, if it matches the requested shape
     // information. If there is a mis-match, embed the result in a flat
     // vector and split it later.
     auto Found = Inst2ColumnMatrix.find(MatrixVal);
     if (Found != Inst2ColumnMatrix.end()) {
-      ColumnMatrixTy &M = Found->second;
+      MatrixTy &M = Found->second;
       // Return the found matrix, if its shape matches the requested shape
       // information
       if (SI.NumRows == M.getNumRows() && SI.NumColumns == M.getNumColumns())
         return M;
 
       MatrixVal = M.embedInVector(Builder);
     }
 
bool VisitCallInst(CallInst *Inst) {
     default:
       return false;
     }
     return true;
   }
 
   /// 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) {
+  MatrixTy 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;
+    MatrixTy 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) {
+  MatrixTy loadMatrix(Value *MatrixPtr, ShapeInfo MatrixShape, unsigned I,
 
     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 =
void LowerColumnwiseLoad(CallInst *Inst) {
     Value *Ptr = Inst->getArgOperand(0);
     Value *Stride = Inst->getArgOperand(1);
     LowerLoad(Inst, Ptr, Stride,
               {Inst->getArgOperand(2), Inst->getArgOperand(3)});
   }
 
   /// 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,
+  void storeMatrix(const MatrixTy &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) {
+  MatrixTy storeMatrix(Type *Ty, MatrixTy StoreVal, Value *Ptr, Value *Stride,
+                       IRBuilder<> &Builder) {
     auto VType = cast<VectorType>(Ty);
     Value *EltPtr = createElementPtr(Ptr, 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);
     }
-    return ColumnMatrixTy().addNumStores(getNumOps(StoreVal.getColumnTy()) *
-                                         StoreVal.getNumColumns());
+    return MatrixTy().addNumStores(getNumOps(StoreVal.getColumnTy()) *
   }
 
   /// 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(
void LowerColumnwiseStore(CallInst *Inst) {
     Value *Ptr = Inst->getArgOperand(1);
     Value *Stride = Inst->getArgOperand(2);
     LowerStore(Inst, Matrix, Ptr, Stride,
                {Inst->getArgOperand(3), Inst->getArgOperand(4)});
   }
 
   /// Extract a column vector of \p NumElts starting at index (\p I, \p J) from
   /// the matrix \p LM represented as a vector of column vectors.
-  Value *extractVector(const ColumnMatrixTy &LM, unsigned I, unsigned J,
+  Value *extractVector(const MatrixTy &LM, unsigned I, unsigned J,
                        unsigned NumElts, IRBuilder<> &Builder) {
     Value *Col = LM.getColumn(J);
     Value *Undef = UndefValue::get(Col->getType());
     Constant *Mask = createSequentialMask(Builder, I, NumElts, 0);
     return Builder.CreateShuffleVector(Col, Undef, Mask, "block");
   }
 
   // Set elements I..I+NumElts-1 to Block
Value *createMulAdd(Value *Sum, Value *A, Value *B, bool UseFPOp,
     return Builder.CreateAdd(Sum, Mul);
   }
 
   /// Cache \p Matrix as result of \p Inst and update the uses of \p Inst. For
   /// users with shape information, there's nothing to do: the will use the
   /// cached value when they are lowered. For other users, \p Matrix is
   /// flattened and the uses are updated to use it. Also marks \p Inst for
   /// deletion.
-  void finalizeLowering(Instruction *Inst, ColumnMatrixTy Matrix,
+  void finalizeLowering(Instruction *Inst, MatrixTy Matrix,
                         IRBuilder<> &Builder) {
     Inst2ColumnMatrix.insert(std::make_pair(Inst, Matrix));
 
     ToRemove.push_back(Inst);
     Value *Flattened = nullptr;
     for (auto I = Inst->use_begin(), E = Inst->use_end(); I != E;) {
       Use &U = *I++;
       if (ShapeMap.find(U.getUser()) == ShapeMap.end()) {
         if (!Flattened)
           Flattened = Matrix.embedInVector(Builder);
         U.set(Flattened);
       }
     }
   }
 
   /// Compute Res += A * B for tile-sized matrices with left-associating
   /// addition.
-  void emitChainedMatrixMultiply(ColumnMatrixTy &Result,
-                                 const ColumnMatrixTy &A,
-                                 const ColumnMatrixTy &B, bool AllowContraction,
+  void emitChainedMatrixMultiply(MatrixTy &Result, const MatrixTy &A,
+                                 const MatrixTy &B, bool AllowContraction,
                                  IRBuilder<> &Builder, bool isTiled) {
     const unsigned VF = std::max<unsigned>(
         TTI.getRegisterBitWidth(true) /
             Result.getElementType()->getPrimitiveSizeInBits().getFixedSize(),
         1U);
     unsigned R = Result.getNumRows();
     unsigned C = Result.getNumColumns();
     unsigned M = A.getNumColumns();
public:
 
   /// Lowers llvm.matrix.multiply.
   void LowerMultiply(CallInst *MatMul) {
     IRBuilder<> Builder(MatMul);
     auto *EltType = cast<VectorType>(MatMul->getType())->getElementType();
     ShapeInfo LShape(MatMul->getArgOperand(2), MatMul->getArgOperand(3));
     ShapeInfo RShape(MatMul->getArgOperand(3), MatMul->getArgOperand(4));
 
-    const ColumnMatrixTy &Lhs =
-        getMatrix(MatMul->getArgOperand(0), LShape, Builder);
-    const ColumnMatrixTy &Rhs =
-        getMatrix(MatMul->getArgOperand(1), RShape, Builder);
+    const MatrixTy &Lhs = getMatrix(MatMul->getArgOperand(0), LShape, Builder);
+    const MatrixTy &Rhs = getMatrix(MatMul->getArgOperand(1), RShape, Builder);
 
     const unsigned R = LShape.NumRows;
     const unsigned C = RShape.NumColumns;
     assert(LShape.NumColumns == RShape.NumRows);
 
     // Initialize the output
-    ColumnMatrixTy Result;
+    MatrixTy Result;
     for (unsigned J = 0; J < C; ++J)
       Result.addColumn(UndefValue::get(VectorType::get(EltType, R)));
 
     bool AllowContract = AllowContractEnabled || (isa<FPMathOperator>(MatMul) &&
                                                   MatMul->hasAllowContract());
     emitChainedMatrixMultiply(Result, Lhs, Rhs, AllowContract, Builder, false);
     finalizeLowering(MatMul, Result, Builder);
   }
 
   /// Lowers llvm.matrix.transpose.
   void LowerTranspose(CallInst *Inst) {
-    ColumnMatrixTy Result;
+    MatrixTy Result;
     IRBuilder<> Builder(Inst);
     Value *InputVal = Inst->getArgOperand(0);
     VectorType *VectorTy = cast<VectorType>(InputVal->getType());
     ShapeInfo ArgShape(Inst->getArgOperand(1), Inst->getArgOperand(2));
-    ColumnMatrixTy InputMatrix = getMatrix(InputVal, ArgShape, Builder);
+    MatrixTy InputMatrix = getMatrix(InputVal, ArgShape, Builder);
 
     for (unsigned Row = 0; Row < ArgShape.NumRows; ++Row) {
       // Build a single column vector for this row. First initialize it.
       Value *ResultColumn = UndefValue::get(
           VectorType::get(VectorTy->getElementType(), ArgShape.NumColumns));
 
       // Go through the elements of this row and insert it into the resulting
       // column vector.
if (I == ShapeMap.end())
       return false;
 
     Value *Lhs = Inst->getOperand(0);
     Value *Rhs = Inst->getOperand(1);
 
     IRBuilder<> Builder(Inst);
     ShapeInfo &Shape = I->second;
 
-    ColumnMatrixTy LoweredLhs = getMatrix(Lhs, Shape, Builder);
-    ColumnMatrixTy LoweredRhs = getMatrix(Rhs, Shape, Builder);
+    MatrixTy LoweredLhs = getMatrix(Lhs, Shape, Builder);
+    MatrixTy LoweredRhs = getMatrix(Rhs, Shape, Builder);
 
     // Add each column and store the result back into the opmapping
-    ColumnMatrixTy Result;
+    MatrixTy Result;
     auto BuildColumnOp = [&Builder, Inst](Value *LHS, Value *RHS) {
       switch (Inst->getOpcode()) {
       case Instruction::Add:
         return Builder.CreateAdd(LHS, RHS);
       case Instruction::Mul:
         return Builder.CreateMul(LHS, RHS);
       case Instruction::Sub:
         return Builder.CreateSub(LHS, RHS);
struct ExprLinearizer {
     unsigned LengthToBreak = 100;
     std::string Str;
     raw_string_ostream Stream;
     unsigned LineLength = 0;
     const DataLayout &DL;
 
     /// Mapping from instructions to column matrixes. It is used to identify
     /// matrix instructions.
-    const MapVector<Value *, ColumnMatrixTy> &Inst2ColumnMatrix;
+    const MapVector<Value *, MatrixTy> &Inst2ColumnMatrix;
 
     /// Mapping from values to the leaves of all expressions that the value is
     /// part of.
     const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared;
 
     /// Set of matrix expressions in the scope of a given DISubprogram.
     const SmallSetVector<Value *, 32> &ExprsInSubprogram;
 
     /// Leaf node of the expression to linearize.
     Value *Leaf;
 
     /// Used to keep track of sub-expressions that get reused while linearizing
     /// the expression. Re-used sub-expressions are marked as (reused).
     SmallPtrSet<Value *, 8> ReusedExprs;
 
     ExprLinearizer(const DataLayout &DL,
-                   const MapVector<Value *, ColumnMatrixTy> &Inst2ColumnMatrix,
+                   const MapVector<Value *, MatrixTy> &Inst2ColumnMatrix,
                    const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared,
                    const SmallSetVector<Value *, 32> &ExprsInSubprogram,
                    Value *Leaf)
         : Str(), Stream(Str), DL(DL), Inst2ColumnMatrix(Inst2ColumnMatrix),
           Shared(Shared), ExprsInSubprogram(ExprsInSubprogram), Leaf(Leaf) {}
 
     void indent(unsigned N) {
       LineLength += N;

   //     mapping.  Leaves are lowered matrix instructions without other matrix
   //     users (like stores) in the current subprogram.
   /// 3. For each leaf, create a remark containing a linearizied version of the
   ///    matrix expression. The expression is linearized by a recursive
   ///    bottom-up traversal of the matrix operands, starting at a leaf. Note
   ///    that multiple leaves can share sub-expressions. Shared subexpressions
   ///    are explicitly marked as shared().
   struct RemarkGenerator {
-    const MapVector<Value *, ColumnMatrixTy> &Inst2ColumnMatrix;
+    const MapVector<Value *, MatrixTy> &Inst2ColumnMatrix;
     OptimizationRemarkEmitter &ORE;
     Function &Func;
     const DataLayout &DL;
 
-    RemarkGenerator(const MapVector<Value *, ColumnMatrixTy> &Inst2ColumnMatrix,
+    RemarkGenerator(const MapVector<Value *, MatrixTy> &Inst2ColumnMatrix,
                     OptimizationRemarkEmitter &ORE, Function &Func)
         : Inst2ColumnMatrix(Inst2ColumnMatrix), ORE(ORE), Func(Func),
           DL(Func.getParent()->getDataLayout()) {}
 
     /// Return all leaves of the expressions in \p ExprsInSubprogram. Those are
     /// instructions in Inst2ColumnMatrix returning void or without any users in
     /// \p ExprsInSubprogram. Currently that should only include stores.
     SmallVector<Value *, 4>