[llvm] 6babae7 - [Matrix] Update load/storeMatrix to take indices as Value* (NFC).

[Florian Hahn via llvm-commits]

Author: Florian Hahn
Date: 2020-04-06T14:48:48+01:00
New Revision: 6babae74c735bb40611cef9aaffd95c19f2c0743

URL: https://github.com/llvm/llvm-project/commit/6babae74c735bb40611cef9aaffd95c19f2c0743
DIFF: https://github.com/llvm/llvm-project/commit/6babae74c735bb40611cef9aaffd95c19f2c0743.diff

LOG: [Matrix] Update load/storeMatrix to take indices as Value* (NFC).

This allows using the functions to be used with loop dependent indices.

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp b/llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
index 81ac2b423c91..ceaa5547a6b3 100644
--- a/llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
+++ b/llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
@@ -795,14 +795,12 @@ class LowerMatrixIntrinsics {
 
   /// Loads a sub-matrix with shape \p ResultShape from a \p R x \p C matrix,
   /// starting at \p MatrixPtr[I][J].
-  MatrixTy loadMatrix(Value *MatrixPtr, ShapeInfo MatrixShape, unsigned I,
-                      unsigned J, ShapeInfo ResultShape, Type *EltTy,
+  MatrixTy loadMatrix(Value *MatrixPtr, ShapeInfo MatrixShape, Value *I,
+                      Value *J, ShapeInfo ResultShape, Type *EltTy,
                       IRBuilder<> &Builder) {
 
     Value *Offset = Builder.CreateAdd(
-        Builder.CreateMul(Builder.getInt32(J),
-                          Builder.getInt32(MatrixShape.getStride())),
-        Builder.getInt32(I));
+        Builder.CreateMul(J, Builder.getInt32(MatrixShape.getStride())), I);
 
     unsigned AS = cast<PointerType>(MatrixPtr->getType())->getAddressSpace();
     Value *EltPtr =
@@ -843,12 +841,10 @@ class LowerMatrixIntrinsics {
   /// Stores a sub-matrix \p StoreVal into the \p R x \p C matrix starting at \p
   /// MatrixPtr[I][J].
   void storeMatrix(const MatrixTy &StoreVal, Value *MatrixPtr,
-                   ShapeInfo MatrixShape, unsigned I, unsigned J, Type *EltTy,
+                   ShapeInfo MatrixShape, Value *I, Value *J, Type *EltTy,
                    IRBuilder<> &Builder) {
     Value *Offset = Builder.CreateAdd(
-        Builder.CreateMul(Builder.getInt32(J),
-                          Builder.getInt32(MatrixShape.getStride())),
-        Builder.getInt32(I));
+        Builder.CreateMul(J, Builder.getInt32(MatrixShape.getStride())), I);
 
     unsigned AS = cast<PointerType>(MatrixPtr->getType())->getAddressSpace();
     Value *EltPtr =
@@ -1214,12 +1210,15 @@ class LowerMatrixIntrinsics {
         for (unsigned K = 0; K < M; K += TileSize) {
           const unsigned TileM = std::min(M - K, unsigned(TileSize));
           MatrixTy A =
-              loadMatrix(APtr, LShape, I, K, {TileR, TileM}, EltType, Builder);
+              loadMatrix(APtr, LShape, Builder.getInt32(I), Builder.getInt32(K),
+                         {TileR, TileM}, EltType, Builder);
           MatrixTy B =
-              loadMatrix(BPtr, RShape, K, J, {TileM, TileC}, EltType, Builder);
+              loadMatrix(BPtr, RShape, Builder.getInt32(K), Builder.getInt32(J),
+                         {TileM, TileC}, EltType, Builder);
           emitMatrixMultiply(Res, A, B, AllowContract, Builder, true);
         }
-        storeMatrix(Res, CPtr, {R, M}, I, J, EltType, Builder);
+        storeMatrix(Res, CPtr, {R, M}, Builder.getInt32(I), Builder.getInt32(J),
+                    EltType, Builder);
       }
 
     // Mark eliminated instructions as fused and remove them.