add a couple more tests

Author: jroelofs

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

@@ -596,11 +596,11 @@ class LowerMatrixIntrinsics {
         // that embedInVector created.
         dbgs() << "matrix reshape from " << Found->second.shape() << " to "
                << SI << " using at least " << SplitVecs.size()
-               << " shuffles on behalf of " << *Inst << '\n';
+               << " shuffles on behalf of:\n" << *Inst << '\n';
       } else if (!ShapeMap.contains(MatrixVal)) {
         dbgs() << "splitting a " << SI << " matrix with " << SplitVecs.size()
                << " shuffles beacuse we do not have a shape-aware lowering for "
-                  "its def: "
+                  "its def:\n"
                << *Inst << '\n';
       } else {
         // The ShapeMap has it, so it's a case where we're being lowered
@@ -1424,9 +1424,9 @@ class LowerMatrixIntrinsics {
         Flattened = Matrix.embedInVector(Builder);
         LLVM_DEBUG(
             if (Instruction *User = dyn_cast<Instruction>(U.getUser())) dbgs()
-                << "flattening a " << Matrix.shape() << " matrix " << *Inst
-                << " because we do not have a shape-aware lowering for its "
-                   "user: "
+                << "flattening a " << Matrix.shape() << " matrix:\n" << *Inst
+                << "\nbecause we do not have a shape-aware lowering for its "
+                   "user:\n"
                 << *User << '\n';);
       }
       U.set(Flattened);

llvm/test/Transforms/LowerMatrixIntrinsics/flatten.ll

@@ -1,26 +1,49 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -passes='lower-matrix-intrinsics' -S < %s | FileCheck %s
-; RUN: opt -passes=lower-matrix-intrinsics -debug-only=lower-matrix-intrinsics -disable-output < %s 2>&1 | FileCheck %s --check-prefix=PASS-DEBUG
+; RUN: opt -passes=lower-matrix-intrinsics -debug-only=lower-matrix-intrinsics -disable-output < %s 2>&1 | FileCheck %s --check-prefix=CHECK
 ; REQUIRES: asserts
 
 define void @diag_3x3(ptr %in, ptr %out) {
-; CHECK-LABEL: @diag_3x3(
-; CHECK-NEXT:    [[COL_LOAD:%.*]] = load <3 x float>, ptr [[IN:%.*]], align 4
-; CHECK-NEXT:    [[VEC_GEP:%.*]] = getelementptr float, ptr [[IN]], i64 3
-; CHECK-NEXT:    [[COL_LOAD1:%.*]] = load <3 x float>, ptr [[VEC_GEP]], align 4
-; CHECK-NEXT:    [[VEC_GEP2:%.*]] = getelementptr float, ptr [[IN]], i64 6
-; CHECK-NEXT:    [[COL_LOAD3:%.*]] = load <3 x float>, ptr [[VEC_GEP2]], align 4
-; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <3 x float> [[COL_LOAD]], <3 x float> [[COL_LOAD1]], <6 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5>
-; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <3 x float> [[COL_LOAD3]], <3 x float> poison, <6 x i32> <i32 0, i32 1, i32 2, i32 poison, i32 poison, i32 poison>
-; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <6 x float> [[TMP1]], <6 x float> [[TMP2]], <9 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8>
-; CHECK-NEXT:    [[DIAG:%.*]] = shufflevector <9 x float> [[TMP3]], <9 x float> poison, <3 x i32> <i32 0, i32 4, i32 8>
-; CHECK-NEXT:    store <3 x float> [[DIAG]], ptr [[OUT:%.*]], align 16
-; CHECK-NEXT:    ret void
-;
   %inv = call <9 x float> @llvm.matrix.column.major.load(ptr %in, i64 3, i1 false, i32 3, i32 3)
   %diag = shufflevector <9 x float> %inv, <9 x float> poison, <3 x i32> <i32 0, i32 4, i32 8>
   store <3 x float> %diag, ptr %out
   ret void
 }
-;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
-; PASS-DEBUG: flattening a 3x3 matrix %{{.*}} = call <9 x float> @llvm.matrix.column.major.load.v9f32.i64(ptr %{{.*}}, i64 3, i1 false, i32 3, i32 3) because we do not have a shape-aware lowering for its user: %{{.*}} = shufflevector <9 x float> %{{.*}}, <9 x float> poison, <3 x i32> <i32 0, i32 4, i32 8>
+; CHECK-LABEL: flattening a 3x3 matrix:
+; CHECK-NEXT: %{{.*}} = call <9 x float> @llvm.matrix.column.major.load.v9f32.i64(ptr %{{.*}}, i64 3, i1 false, i32 3, i32 3)
+; CHECK-NEXT: because we do not have a shape-aware lowering for its user:
+; CHECK-NEXT: %{{.*}} = shufflevector <9 x float> %{{.*}}, <9 x float> poison, <3 x i32> <i32 0, i32 4, i32 8>
+
+define void @reshape(ptr %in, ptr %out) {
+entry:
+  %0 = load <4 x double>, ptr %in, align 8
+  %1 = tail call <4 x double> @llvm.matrix.transpose.v4f64(<4 x double> %0, i32 4, i32 1)
+  %2 = tail call <4 x double> @llvm.matrix.transpose.v4f64(<4 x double> %1, i32 1, i32 4)
+  %3 = tail call <4 x double> @llvm.matrix.transpose.v4f64(<4 x double> %2, i32 2, i32 2)
+  %4 = tail call <4 x double> @llvm.matrix.transpose.v4f64(<4 x double> %3, i32 2, i32 2)
+  %5 = tail call <4 x double> @llvm.matrix.transpose.v4f64(<4 x double> %4, i32 2, i32 2)
+  store <4 x double> %5, ptr %out, align 8
+  ret void
+}
+; CHECK-LABEL: matrix reshape from 4x1 to 2x2 using at least 2 shuffles on behalf of:
+; CHECK-NEXT: %{{.*}} = load <4 x double>, ptr %{{.*}}, align 8
+
+define void @multiply_ntt(ptr %A, ptr %B, ptr %C, ptr %R) {
+entry:
+  %a = load <6 x double>, ptr %A, align 16
+  %b = load <6 x double>, ptr %B, align 16
+  %c = load <8 x double>, ptr %C, align 16
+  %b_t = call <6 x double> @llvm.matrix.transpose.v6f64.v6f64(<6 x double> %b, i32 2, i32 3)
+  %c_t = call <8 x double> @llvm.matrix.transpose.v8f64.v8f64(<8 x double> %c, i32 4, i32 2)
+  %m1 = call <12 x double> @llvm.matrix.multiply.v12f64.v6f64.v8f64(<6 x double> %b_t, <8 x double> %c_t, i32 3, i32 2, i32 4)
+  %m2 = call <8 x double> @llvm.matrix.multiply.v8f64.v6f64.v12f64(<6 x double> %a, <12 x double> %m1, i32 2, i32 3, i32 4)
+  store <8 x double> %m2, ptr %R, align 16
+  ret void
+}
+; CHECK-LABEL: flattening a 2x3 matrix:
+; CHECK-NEXT: %{{.*}} = load <6 x double>, ptr %{{.*}}, align 16
+; CHECK-NEXT: because we do not have a shape-aware lowering for its user:
+; CHECK-NEXT: %{{.*}} = shufflevector <6 x double> %{{.*}}, <6 x double> poison, <2 x i32> <i32 4, i32 5>
+
+; CHECK: flattening a 4x3 matrix:
+; CHECK-NEXT: %{{.*}} = call <12 x double> @llvm.matrix.multiply.v12f64.v8f64.v6f64(<8 x double> %{{.*}}, <6 x double> %{{.*}}, i32 4, i32 2, i32 3)
+; CHECK-NEXT: because we do not have a shape-aware lowering for its user:
+; CHECK-NEXT: %{{.*}} = shufflevector <12 x double> %{{.*}}, <12 x double> poison, <4 x i32> <i32 8, i32 9, i32 10, i32 11>