[llvm] [Matrix] Optimize shuffle extracts with ShapeInfo (PR #142276)

[Jon Roelofs via llvm-commits]

================
@@ -0,0 +1,34 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes='lower-matrix-intrinsics' -S < %s | FileCheck %s
+
+define <3 x double> @extract_column(ptr %in, ptr %out) {
+; CHECK-LABEL: @extract_column(
+; CHECK-NEXT:    [[COL_LOAD:%.*]] = load volatile <3 x double>, ptr [[IN:%.*]], align 8
+; CHECK-NEXT:    [[VEC_GEP:%.*]] = getelementptr double, ptr [[IN]], i64 3
+; CHECK-NEXT:    [[COL_LOAD1:%.*]] = load volatile <3 x double>, ptr [[VEC_GEP]], align 8
+; CHECK-NEXT:    [[VEC_GEP2:%.*]] = getelementptr double, ptr [[IN]], i64 6
+; CHECK-NEXT:    [[COL_LOAD3:%.*]] = load volatile <3 x double>, ptr [[VEC_GEP2]], align 8
+; CHECK-NEXT:    ret <3 x double> [[COL_LOAD3]]
+;
+  %inv = call <9 x double> @llvm.matrix.column.major.load(ptr %in, i64 3, i1 1, i32 3, i32 3)
+  %col = shufflevector <9 x double> %inv, <9 x double> poison, <3 x i32> <i32 6, i32 7, i32 8>
+  ret <3 x double> %col
+}
+
+define <3 x double> @extract_row(ptr %in, ptr %out) {
+; CHECK-LABEL: @extract_row(
+; CHECK-NEXT:    [[COL_LOAD:%.*]] = load volatile <3 x double>, ptr [[IN:%.*]], align 8
+; CHECK-NEXT:    [[VEC_GEP:%.*]] = getelementptr double, ptr [[IN]], i64 3
+; CHECK-NEXT:    [[COL_LOAD1:%.*]] = load volatile <3 x double>, ptr [[VEC_GEP]], align 8
+; CHECK-NEXT:    [[VEC_GEP2:%.*]] = getelementptr double, ptr [[IN]], i64 6
+; CHECK-NEXT:    [[COL_LOAD3:%.*]] = load volatile <3 x double>, ptr [[VEC_GEP2]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = shufflevector <3 x double> [[COL_LOAD]], <3 x double> [[COL_LOAD1]], <6 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5>
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <3 x double> [[COL_LOAD3]], <3 x double> poison, <6 x i32> <i32 0, i32 1, i32 2, i32 poison, i32 poison, i32 poison>
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <6 x double> [[TMP1]], <6 x double> [[TMP2]], <9 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8>
+; CHECK-NEXT:    [[ROW:%.*]] = shufflevector <9 x double> [[TMP3]], <9 x double> poison, <3 x i32> <i32 0, i32 3, i32 6>
+; CHECK-NEXT:    ret <3 x double> [[ROW]]
+;
+  %inv = call <9 x double> @llvm.matrix.column.major.load(ptr %in, i64 3, i1 1, i32 3, i32 3)
+  %row = shufflevector <9 x double> %inv, <9 x double> poison, <3 x i32> <i32 0, i32 3, i32 6>
+  ret <3 x double> %row
+}
----------------
jroelofs wrote:

TODO: needs quite a few more tests.

https://github.com/llvm/llvm-project/pull/142276