[Matrix] Refactor transpose distribution. NFC

[thegameg]

Summary

Use a function to distribute transposes. Preparation for future patches.

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

388

nit: Should be /// for doc-comment?

Might be good to say this returns the transposed shape?

690

nit: /// for doc-comment

[thegameg]

https://github.com/llvm/llvm-project/commit/bfd3883e83dd61c68bd920d933f8ca679b788ad9

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

unsigned getStride() const {
       return NumColumns;
     }
 
     unsigned getNumVectors() const {
       if (IsColumnMajor)
         return NumColumns;
       return NumRows;
     }
+
+    // Transpose the shape.
+    ShapeInfo t() const { return ShapeInfo(NumColumns, NumRows); }
   };
 
   /// Maps instructions to their shape information. The shape information
   /// describes the shape to be used while lowering. This matches the shape of
   /// the result value of the instruction, with the only exceptions being store
   /// instructions and the matrix_column_major_store intrinsics. For those, the
   /// shape information indicates that those instructions should be lowered
   /// using shape information as well.  A ValueMap is used so that when
while (!WorkList.empty()) {
       for (size_t I = BeforeProcessingV; I != WorkList.size(); I++)
         for (User *U : WorkList[I]->users())
           if (isa<Instruction>(U) && V != U)
             NewWorkList.push_back(cast<Instruction>(U));
     }
     return NewWorkList;
   }
 
+  // (Op0 op Op1)^T -> Op0^T op Op1^T
+  // Transpose \p Op0 and \p Op1 of shape \p Shape0 and \p Shape1, then use them
+  // on both sides of \p Operation.
+  Instruction *distributeTransposes(
+      Value *Op0, ShapeInfo Shape0, Value *Op1, ShapeInfo Shape1,
+      MatrixBuilder &Builder,
+      function_ref<Instruction *(Value *, ShapeInfo, Value *, ShapeInfo)>
+          Operation) {
+    Value *T0 = Builder.CreateMatrixTranspose(
+        Op0, Shape0.NumRows, Shape0.NumColumns, Op0->getName() + "_t");
+    // We are being run after shape prop, add shape for newly created
+    // instructions so that we lower them later.
+    setShapeInfo(T0, Shape0.t());
+    Value *T1 = Builder.CreateMatrixTranspose(
+        Op1, Shape1.NumRows, Shape1.NumColumns, Op1->getName() + "_t");
+    setShapeInfo(T1, Shape1.t());
+    return Operation(T0, Shape0.t(), T1, Shape1.t());
+  }
+
   /// Try moving transposes in order to fold them away or into multiplies.
   void optimizeTransposes() {
     auto ReplaceAllUsesWith = [this](Instruction &Old, Value *New) {
       // We need to remove Old from the ShapeMap otherwise RAUW will replace it
       // with New. We should only add New it it supportsShapeInfo so we insert
       // it conditionally instead.
       auto S = ShapeMap.find(&Old);
       if (S != ShapeMap.end()) {
for (BasicBlock &BB : reverse(Func)) {
             EraseFromParent(TA);
           }
 
           // (A * B)^t -> B^t * A^t
           // RxK KxC      CxK   KxR
           else if (match(TA, m_Intrinsic<Intrinsic::matrix_multiply>(
                                  m_Value(TAMA), m_Value(TAMB), m_ConstantInt(R),
                                  m_ConstantInt(K), m_ConstantInt(C)))) {
-            Value *T0 = Builder.CreateMatrixTranspose(TAMB, K->getZExtValue(),
-                                                      C->getZExtValue(),
-                                                      TAMB->getName() + "_t");
-            // We are being run after shape prop, add shape for newly created
-            // instructions so that we lower them later.
-            setShapeInfo(T0, {C, K});
-            Value *T1 = Builder.CreateMatrixTranspose(TAMA, R->getZExtValue(),
-                                                      K->getZExtValue(),
-                                                      TAMA->getName() + "_t");
-            setShapeInfo(T1, {K, R});
-            NewInst = Builder.CreateMatrixMultiply(T0, T1, C->getZExtValue(),
-                                                   K->getZExtValue(),
-                                                   R->getZExtValue(), "mmul");
+            NewInst = distributeTransposes(
+                TAMB, {K, C}, TAMA, {R, K}, Builder,
+                [&](Value *T0, ShapeInfo Shape0, Value *T1, ShapeInfo Shape1) {
+                  return Builder.CreateMatrixMultiply(
+                      T0, T1, Shape0.NumRows, Shape0.NumColumns,
+                      Shape1.NumColumns, "mmul");
+                });
             ReplaceAllUsesWith(I, NewInst);
             EraseFromParent(&I);
             EraseFromParent(TA);
           }
         }
 
         // If we replaced I with a new instruction, continue from there.
         if (NewInst)