[llvm] [MatrixIntrinsics] Add row-major support for llvm.matrix.transpose

[aartbik]

Summary

Only column-major was supported so far. This adds row-major support as well.
Note that we probably also want very efficient SIMD implementations for the
various target platforms.

Bug:
https://bugs.llvm.org/show_bug.cgi?id=46085

[fhahn]

Thank you very much for the patch! A few suggestions inline.

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

1300

I think it might be possible to share most of the transpose implementation for both row/column major transpose, as we are iterating over the indices of the second dimension while extracting from the vectors of the first dimension in both cases.

Instead of using ArgShape.NumRows, we could just use the second dimension (InputMatrix.isColumnMajor() : ArgShape.NumRows : ArgShape.NumColumns and instead of iterating over columns/rows we could just iterate over the vectors of the leading dimension (InputMatrix.vectors()) and instead of ArgShape.NumColumns/NumRows, we could use ArgShape.getNumVectors().

llvm/test/Transforms/LowerMatrixIntrinsics/pr46085-row.ll

1

for multiply, we just have row/column major test variants in multiply-i32-row-major.ll/multiply-i32.ll and so on. For consistency, it might be slightly preferable to extend the existing transpose-*.ll tests (if the cases in the new test are not covered) and call the row-major equivalent transpose-*-row-major.ll.

[aartbik]

addressed comments, thanks!

[aartbik]

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

1300

That is a good idea. Fits nicely with the provided utilities. Thanks!

llvm/test/Transforms/LowerMatrixIntrinsics/pr46085-row.ll

1

Sounds good. I removed the pr* tests and used the existing transpose tests for RM.

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

public:
   /// Lowers llvm.matrix.transpose.
   void LowerTranspose(CallInst *Inst) {
     MatrixTy Result;
     IRBuilder<> Builder(Inst);
     Value *InputVal = Inst->getArgOperand(0);
     VectorType *VectorTy = cast<VectorType>(InputVal->getType());
     ShapeInfo ArgShape(Inst->getArgOperand(1), Inst->getArgOperand(2));
     MatrixTy InputMatrix = getMatrix(InputVal, ArgShape, Builder);
-    assert(InputMatrix.isColumnMajor() &&
-           "Row-major code-gen not supported yet!");
 
+    if (InputMatrix.isColumnMajor()) {
+      // Column-major tranposition.
       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.
         for (auto C : enumerate(InputMatrix.columns())) {
           Value *Elt = Builder.CreateExtractElement(C.value(), Row);
           // We insert at index Column since that is the row index after the
           // transpose.
           ResultColumn =
               Builder.CreateInsertElement(ResultColumn, Elt, C.index());
         }
         Result.addVector(ResultColumn);
       }
+    } else {
+      // Row-major tranposition.
+      for (unsigned Column = 0; Column < ArgShape.NumColumns; ++Column) {
+        // Build a single row vector for this column. First initialize it.
+        Value *ResultRow = UndefValue::get(
+            VectorType::get(VectorTy->getElementType(), ArgShape.NumRows));
+
+        // Go through the elements of this column and insert it into the
+        // resulting row vector.
+        for (auto R : enumerate(InputMatrix.vectors())) {
+          Value *Elt = Builder.CreateExtractElement(R.value(), Column);
+          // We insert at index Row since that is the column index after the
+          // transpose.
+          ResultRow = Builder.CreateInsertElement(ResultRow, Elt, R.index());
+        }
+        Result.addVector(ResultRow);
+      }
+    }
 
     // TODO: Improve estimate of operations needed for transposes. Currently we
     // just count the insertelement/extractelement instructions, but do not
     // account for later simplifications/combines.
     finalizeLowering(
         Inst,
         Result.addNumComputeOps(2 * ArgShape.NumRows * ArgShape.NumColumns),
         Builder);

llvm/test/Transforms/LowerMatrixIntrinsics/pr46085-col.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -lower-matrix-intrinsics -matrix-default-layout=column-major -S < %s | FileCheck  %s
+
+declare <4 x i32> @llvm.matrix.transpose.v4i32(<4 x i32>, i32, i32)
+declare <6 x i32> @llvm.matrix.transpose.v6i32(<6 x i32>, i32, i32)
+
+define <4 x i32> @transpose_i32_2x2(<4 x i32> %a) {
+; CHECK-LABEL: @transpose_i32_2x2(
+; CHECK-NEXT:    [[SPLIT:%.*]] = shufflevector <4 x i32> [[A:%.*]], <4 x i32> undef, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    [[SPLIT1:%.*]] = shufflevector <4 x i32> [[A]], <4 x i32> undef, <2 x i32> <i32 2, i32 3>
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i32> [[SPLIT]], i64 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <2 x i32> undef, i32 [[TMP1]], i64 0
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <2 x i32> [[SPLIT1]], i64 0
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <2 x i32> [[TMP2]], i32 [[TMP3]], i64 1
+; CHECK-NEXT:    [[TMP5:%.*]] = extractelement <2 x i32> [[SPLIT]], i64 1
+; CHECK-NEXT:    [[TMP6:%.*]] = insertelement <2 x i32> undef, i32 [[TMP5]], i64 0
+; CHECK-NEXT:    [[TMP7:%.*]] = extractelement <2 x i32> [[SPLIT1]], i64 1
+; CHECK-NEXT:    [[TMP8:%.*]] = insertelement <2 x i32> [[TMP6]], i32 [[TMP7]], i64 1
+; CHECK-NEXT:    [[TMP9:%.*]] = shufflevector <2 x i32> [[TMP4]], <2 x i32> [[TMP8]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    ret <4 x i32> [[TMP9]]
+;
+  %t  = call <4 x i32> @llvm.matrix.transpose.v4i32(<4 x i32> %a, i32 2, i32 2)
+  ret <4 x i32> %t
+}
+
+define <6 x i32> @transpose_i32_2x3(<6 x i32> %a) {
+; CHECK-LABEL: @transpose_i32_2x3(
+; CHECK-NEXT:    [[SPLIT:%.*]] = shufflevector <6 x i32> [[A:%.*]], <6 x i32> undef, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    [[SPLIT1:%.*]] = shufflevector <6 x i32> [[A]], <6 x i32> undef, <2 x i32> <i32 2, i32 3>
+; CHECK-NEXT:    [[SPLIT2:%.*]] = shufflevector <6 x i32> [[A]], <6 x i32> undef, <2 x i32> <i32 4, i32 5>
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i32> [[SPLIT]], i64 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <3 x i32> undef, i32 [[TMP1]], i64 0
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <2 x i32> [[SPLIT1]], i64 0
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <3 x i32> [[TMP2]], i32 [[TMP3]], i64 1
+; CHECK-NEXT:    [[TMP5:%.*]] = extractelement <2 x i32> [[SPLIT2]], i64 0
+; CHECK-NEXT:    [[TMP6:%.*]] = insertelement <3 x i32> [[TMP4]], i32 [[TMP5]], i64 2
+; CHECK-NEXT:    [[TMP7:%.*]] = extractelement <2 x i32> [[SPLIT]], i64 1
+; CHECK-NEXT:    [[TMP8:%.*]] = insertelement <3 x i32> undef, i32 [[TMP7]], i64 0
+; CHECK-NEXT:    [[TMP9:%.*]] = extractelement <2 x i32> [[SPLIT1]], i64 1
+; CHECK-NEXT:    [[TMP10:%.*]] = insertelement <3 x i32> [[TMP8]], i32 [[TMP9]], i64 1
+; CHECK-NEXT:    [[TMP11:%.*]] = extractelement <2 x i32> [[SPLIT2]], i64 1
+; CHECK-NEXT:    [[TMP12:%.*]] = insertelement <3 x i32> [[TMP10]], i32 [[TMP11]], i64 2
+; CHECK-NEXT:    [[TMP13:%.*]] = shufflevector <3 x i32> [[TMP6]], <3 x i32> [[TMP12]], <6 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5>
+; CHECK-NEXT:    ret <6 x i32> [[TMP13]]
+;
+  %t  = call <6 x i32> @llvm.matrix.transpose.v6i32(<6 x i32> %a, i32 2, i32 3)
+  ret <6 x i32> %t
+}

llvm/test/Transforms/LowerMatrixIntrinsics/pr46085-row.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -lower-matrix-intrinsics -matrix-default-layout=row-major -S < %s | FileCheck  %s
+
+declare <4 x i32> @llvm.matrix.transpose.v4i32(<4 x i32>, i32, i32)
+declare <6 x i32> @llvm.matrix.transpose.v6i32(<6 x i32>, i32, i32)
+
+define <4 x i32> @transpose_i32_2x2(<4 x i32> %a) {
+; CHECK-LABEL: @transpose_i32_2x2(
+; CHECK-NEXT:    [[SPLIT:%.*]] = shufflevector <4 x i32> [[A:%.*]], <4 x i32> undef, <2 x i32> <i32 0, i32 1>
+; CHECK-NEXT:    [[SPLIT1:%.*]] = shufflevector <4 x i32> [[A]], <4 x i32> undef, <2 x i32> <i32 2, i32 3>
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x i32> [[SPLIT]], i64 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <2 x i32> undef, i32 [[TMP1]], i64 0
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <2 x i32> [[SPLIT1]], i64 0
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <2 x i32> [[TMP2]], i32 [[TMP3]], i64 1
+; CHECK-NEXT:    [[TMP5:%.*]] = extractelement <2 x i32> [[SPLIT]], i64 1
+; CHECK-NEXT:    [[TMP6:%.*]] = insertelement <2 x i32> undef, i32 [[TMP5]], i64 0
+; CHECK-NEXT:    [[TMP7:%.*]] = extractelement <2 x i32> [[SPLIT1]], i64 1
+; CHECK-NEXT:    [[TMP8:%.*]] = insertelement <2 x i32> [[TMP6]], i32 [[TMP7]], i64 1
+; CHECK-NEXT:    [[TMP9:%.*]] = shufflevector <2 x i32> [[TMP4]], <2 x i32> [[TMP8]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    ret <4 x i32> [[TMP9]]
+;
+  %t  = call <4 x i32> @llvm.matrix.transpose.v4i32(<4 x i32> %a, i32 2, i32 2)
+  ret <4 x i32> %t
+}
+
+define <6 x i32> @transpose_i32_2x3(<6 x i32> %a) {
+; CHECK-LABEL: @transpose_i32_2x3(
+; CHECK-NEXT:    [[SPLIT:%.*]] = shufflevector <6 x i32> [[A:%.*]], <6 x i32> undef, <3 x i32> <i32 0, i32 1, i32 2>
+; CHECK-NEXT:    [[SPLIT1:%.*]] = shufflevector <6 x i32> [[A]], <6 x i32> undef, <3 x i32> <i32 3, i32 4, i32 5>
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <3 x i32> [[SPLIT]], i64 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertelement <2 x i32> undef, i32 [[TMP1]], i64 0
+; CHECK-NEXT:    [[TMP3:%.*]] = extractelement <3 x i32> [[SPLIT1]], i64 0
+; CHECK-NEXT:    [[TMP4:%.*]] = insertelement <2 x i32> [[TMP2]], i32 [[TMP3]], i64 1
+; CHECK-NEXT:    [[TMP5:%.*]] = extractelement <3 x i32> [[SPLIT]], i64 1
+; CHECK-NEXT:    [[TMP6:%.*]] = insertelement <2 x i32> undef, i32 [[TMP5]], i64 0
+; CHECK-NEXT:    [[TMP7:%.*]] = extractelement <3 x i32> [[SPLIT1]], i64 1
+; CHECK-NEXT:    [[TMP8:%.*]] = insertelement <2 x i32> [[TMP6]], i32 [[TMP7]], i64 1
+; CHECK-NEXT:    [[TMP9:%.*]] = extractelement <3 x i32> [[SPLIT]], i64 2
+; CHECK-NEXT:    [[TMP10:%.*]] = insertelement <2 x i32> undef, i32 [[TMP9]], i64 0
+; CHECK-NEXT:    [[TMP11:%.*]] = extractelement <3 x i32> [[SPLIT1]], i64 2
+; CHECK-NEXT:    [[TMP12:%.*]] = insertelement <2 x i32> [[TMP10]], i32 [[TMP11]], i64 1
+; CHECK-NEXT:    [[TMP13:%.*]] = shufflevector <2 x i32> [[TMP4]], <2 x i32> [[TMP8]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
+; CHECK-NEXT:    [[TMP14:%.*]] = shufflevector <2 x i32> [[TMP12]], <2 x i32> undef, <4 x i32> <i32 0, i32 1, i32 undef, i32 undef>
+; CHECK-NEXT:    [[TMP15:%.*]] = shufflevector <4 x i32> [[TMP13]], <4 x i32> [[TMP14]], <6 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5>
+; CHECK-NEXT:    ret <6 x i32> [[TMP15]]
+;
+  %t  = call <6 x i32> @llvm.matrix.transpose.v6i32(<6 x i32> %a, i32 2, i32 3)
+  ret <6 x i32> %t
+}