Skip to content

Commit dbe4143

Browse files
fhahnfhahn
authored
[Matrix] Fix dimensions when hoisting transpose across add. (#81507)
Row and column arguments for matrix_transpose indicate the shape of the operand. When hoisting the transpose to the result of the add, the add operates on the original operand's shape, and so does the hoisted transpose. This patch also adds an assert that the shape for the original add and the transpose match, as well as the shape of the new add matches the cached shape for it. The assert could potentially be moved to updateShapeAndReplaceAllUsesWith.
1 parent 75a1c4e commit dbe4143

File tree

3 files changed

+113
-68
lines changed

3 files changed

+113
-68
lines changed

‎llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -898,20 +898,28 @@ class LowerMatrixIntrinsics {
898898
updateShapeAndReplaceAllUsesWith(I, NewInst);
899899
CleanupBinOp(I, A, B);
900900
}
901-
// A^t + B ^t -> (A + B)^t
901+
// A^t + B ^t -> (A + B)^t. Pick rows and columns from first transpose. If
902+
// the shape of the second transpose is different, there's a shape conflict
903+
// which gets resolved by picking the shape of the first operand.
902904
else if (match(&I, m_FAdd(m_Value(A), m_Value(B))) &&
903905
match(A, m_Intrinsic<Intrinsic::matrix_transpose>(
904906
m_Value(AT), m_ConstantInt(R), m_ConstantInt(C))) &&
905907
match(B, m_Intrinsic<Intrinsic::matrix_transpose>(
906-
m_Value(BT), m_ConstantInt(R), m_ConstantInt(C)))) {
908+
m_Value(BT), m_ConstantInt(), m_ConstantInt()))) {
907909
IRBuilder<> Builder(&I);
908-
Value *Add = cast<Instruction>(Builder.CreateFAdd(AT, BT, "mfadd"));
909-
setShapeInfo(Add, {C, R});
910+
auto *Add = cast<Instruction>(Builder.CreateFAdd(AT, BT, "mfadd"));
911+
setShapeInfo(Add, {R, C});
910912
MatrixBuilder MBuilder(Builder);
911913
Instruction *NewInst = MBuilder.CreateMatrixTranspose(
912-
Add, C->getZExtValue(), R->getZExtValue(), "mfadd_t");
914+
Add, R->getZExtValue(), C->getZExtValue(), "mfadd_t");
913915
updateShapeAndReplaceAllUsesWith(I, NewInst);
916+
assert(computeShapeInfoForInst(NewInst, ShapeMap) ==
917+
computeShapeInfoForInst(&I, ShapeMap) &&
918+
"Shape of new instruction doesn't match original shape.");
914919
CleanupBinOp(I, A, B);
920+
assert(computeShapeInfoForInst(Add, ShapeMap).value_or(ShapeMap[Add]) ==
921+
ShapeMap[Add] &&
922+
"Shape of updated addition doesn't match cached shape.");
915923
}
916924
}
917925

‎llvm/test/Transforms/LowerMatrixIntrinsics/propagate-backward.ll

Lines changed: 60 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -4,31 +4,35 @@
44
define <8 x double> @fadd_transpose(<8 x double> %a, <8 x double> %b) {
55
; CHECK-LABEL: @fadd_transpose(
66
; CHECK-NEXT: entry:
7-
; CHECK-NEXT: [[SPLIT:%.*]] = shufflevector <8 x double> [[A:%.*]], <8 x double> poison, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
8-
; CHECK-NEXT: [[SPLIT2:%.*]] = shufflevector <8 x double> [[A]], <8 x double> poison, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
9-
; CHECK-NEXT: [[SPLIT3:%.*]] = shufflevector <8 x double> [[B:%.*]], <8 x double> poison, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
10-
; CHECK-NEXT: [[SPLIT4:%.*]] = shufflevector <8 x double> [[B]], <8 x double> poison, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
11-
; CHECK-NEXT: [[TMP0:%.*]] = fadd <4 x double> [[SPLIT]], [[SPLIT3]]
12-
; CHECK-NEXT: [[TMP1:%.*]] = fadd <4 x double> [[SPLIT2]], [[SPLIT4]]
13-
; CHECK-NEXT: [[TMP2:%.*]] = extractelement <4 x double> [[TMP0]], i64 0
14-
; CHECK-NEXT: [[TMP3:%.*]] = insertelement <2 x double> poison, double [[TMP2]], i64 0
15-
; CHECK-NEXT: [[TMP4:%.*]] = extractelement <4 x double> [[TMP1]], i64 0
16-
; CHECK-NEXT: [[TMP5:%.*]] = insertelement <2 x double> [[TMP3]], double [[TMP4]], i64 1
17-
; CHECK-NEXT: [[TMP6:%.*]] = extractelement <4 x double> [[TMP0]], i64 1
18-
; CHECK-NEXT: [[TMP7:%.*]] = insertelement <2 x double> poison, double [[TMP6]], i64 0
19-
; CHECK-NEXT: [[TMP8:%.*]] = extractelement <4 x double> [[TMP1]], i64 1
20-
; CHECK-NEXT: [[TMP9:%.*]] = insertelement <2 x double> [[TMP7]], double [[TMP8]], i64 1
21-
; CHECK-NEXT: [[TMP10:%.*]] = extractelement <4 x double> [[TMP0]], i64 2
22-
; CHECK-NEXT: [[TMP11:%.*]] = insertelement <2 x double> poison, double [[TMP10]], i64 0
23-
; CHECK-NEXT: [[TMP12:%.*]] = extractelement <4 x double> [[TMP1]], i64 2
24-
; CHECK-NEXT: [[TMP13:%.*]] = insertelement <2 x double> [[TMP11]], double [[TMP12]], i64 1
25-
; CHECK-NEXT: [[TMP14:%.*]] = extractelement <4 x double> [[TMP0]], i64 3
26-
; CHECK-NEXT: [[TMP15:%.*]] = insertelement <2 x double> poison, double [[TMP14]], i64 0
27-
; CHECK-NEXT: [[TMP16:%.*]] = extractelement <4 x double> [[TMP1]], i64 3
28-
; CHECK-NEXT: [[TMP17:%.*]] = insertelement <2 x double> [[TMP15]], double [[TMP16]], i64 1
29-
; CHECK-NEXT: [[TMP18:%.*]] = shufflevector <2 x double> [[TMP5]], <2 x double> [[TMP9]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
30-
; CHECK-NEXT: [[TMP19:%.*]] = shufflevector <2 x double> [[TMP13]], <2 x double> [[TMP17]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
31-
; CHECK-NEXT: [[TMP20:%.*]] = shufflevector <4 x double> [[TMP18]], <4 x double> [[TMP19]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
7+
; CHECK-NEXT: [[SPLIT:%.*]] = shufflevector <8 x double> [[A:%.*]], <8 x double> poison, <2 x i32> <i32 0, i32 1>
8+
; CHECK-NEXT: [[SPLIT1:%.*]] = shufflevector <8 x double> [[A]], <8 x double> poison, <2 x i32> <i32 2, i32 3>
9+
; CHECK-NEXT: [[SPLIT2:%.*]] = shufflevector <8 x double> [[A]], <8 x double> poison, <2 x i32> <i32 4, i32 5>
10+
; CHECK-NEXT: [[SPLIT3:%.*]] = shufflevector <8 x double> [[A]], <8 x double> poison, <2 x i32> <i32 6, i32 7>
11+
; CHECK-NEXT: [[SPLIT4:%.*]] = shufflevector <8 x double> [[B:%.*]], <8 x double> poison, <2 x i32> <i32 0, i32 1>
12+
; CHECK-NEXT: [[SPLIT5:%.*]] = shufflevector <8 x double> [[B]], <8 x double> poison, <2 x i32> <i32 2, i32 3>
13+
; CHECK-NEXT: [[SPLIT6:%.*]] = shufflevector <8 x double> [[B]], <8 x double> poison, <2 x i32> <i32 4, i32 5>
14+
; CHECK-NEXT: [[SPLIT7:%.*]] = shufflevector <8 x double> [[B]], <8 x double> poison, <2 x i32> <i32 6, i32 7>
15+
; CHECK-NEXT: [[TMP0:%.*]] = fadd <2 x double> [[SPLIT]], [[SPLIT4]]
16+
; CHECK-NEXT: [[TMP1:%.*]] = fadd <2 x double> [[SPLIT1]], [[SPLIT5]]
17+
; CHECK-NEXT: [[TMP2:%.*]] = fadd <2 x double> [[SPLIT2]], [[SPLIT6]]
18+
; CHECK-NEXT: [[TMP3:%.*]] = fadd <2 x double> [[SPLIT3]], [[SPLIT7]]
19+
; CHECK-NEXT: [[TMP4:%.*]] = extractelement <2 x double> [[TMP0]], i64 0
20+
; CHECK-NEXT: [[TMP5:%.*]] = insertelement <4 x double> poison, double [[TMP4]], i64 0
21+
; CHECK-NEXT: [[TMP6:%.*]] = extractelement <2 x double> [[TMP1]], i64 0
22+
; CHECK-NEXT: [[TMP7:%.*]] = insertelement <4 x double> [[TMP5]], double [[TMP6]], i64 1
23+
; CHECK-NEXT: [[TMP8:%.*]] = extractelement <2 x double> [[TMP2]], i64 0
24+
; CHECK-NEXT: [[TMP9:%.*]] = insertelement <4 x double> [[TMP7]], double [[TMP8]], i64 2
25+
; CHECK-NEXT: [[TMP10:%.*]] = extractelement <2 x double> [[TMP3]], i64 0
26+
; CHECK-NEXT: [[TMP11:%.*]] = insertelement <4 x double> [[TMP9]], double [[TMP10]], i64 3
27+
; CHECK-NEXT: [[TMP12:%.*]] = extractelement <2 x double> [[TMP0]], i64 1
28+
; CHECK-NEXT: [[TMP13:%.*]] = insertelement <4 x double> poison, double [[TMP12]], i64 0
29+
; CHECK-NEXT: [[TMP14:%.*]] = extractelement <2 x double> [[TMP1]], i64 1
30+
; CHECK-NEXT: [[TMP15:%.*]] = insertelement <4 x double> [[TMP13]], double [[TMP14]], i64 1
31+
; CHECK-NEXT: [[TMP16:%.*]] = extractelement <2 x double> [[TMP2]], i64 1
32+
; CHECK-NEXT: [[TMP17:%.*]] = insertelement <4 x double> [[TMP15]], double [[TMP16]], i64 2
33+
; CHECK-NEXT: [[TMP18:%.*]] = extractelement <2 x double> [[TMP3]], i64 1
34+
; CHECK-NEXT: [[TMP19:%.*]] = insertelement <4 x double> [[TMP17]], double [[TMP18]], i64 3
35+
; CHECK-NEXT: [[TMP20:%.*]] = shufflevector <4 x double> [[TMP11]], <4 x double> [[TMP19]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
3236
; CHECK-NEXT: ret <8 x double> [[TMP20]]
3337
;
3438
entry:
@@ -42,40 +46,37 @@ define <8 x double> @load_fadd_transpose(ptr %A.Ptr, <8 x double> %b) {
4246
; CHECK-NEXT: entry:
4347
; CHECK-NEXT: [[COL_LOAD:%.*]] = load <2 x double>, ptr [[A_PTR:%.*]], align 8
4448
; CHECK-NEXT: [[VEC_GEP:%.*]] = getelementptr double, ptr [[A_PTR]], i64 2
45-
; CHECK-NEXT: [[COL_LOAD2:%.*]] = load <2 x double>, ptr [[VEC_GEP]], align 8
46-
; CHECK-NEXT: [[VEC_GEP3:%.*]] = getelementptr double, ptr [[A_PTR]], i64 4
47-
; CHECK-NEXT: [[COL_LOAD4:%.*]] = load <2 x double>, ptr [[VEC_GEP3]], align 8
48-
; CHECK-NEXT: [[VEC_GEP5:%.*]] = getelementptr double, ptr [[A_PTR]], i64 6
49-
; CHECK-NEXT: [[COL_LOAD6:%.*]] = load <2 x double>, ptr [[VEC_GEP5]], align 8
50-
; CHECK-NEXT: [[TMP0:%.*]] = shufflevector <2 x double> [[COL_LOAD]], <2 x double> [[COL_LOAD2]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
51-
; CHECK-NEXT: [[TMP1:%.*]] = shufflevector <2 x double> [[COL_LOAD4]], <2 x double> [[COL_LOAD6]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
52-
; CHECK-NEXT: [[TMP2:%.*]] = shufflevector <4 x double> [[TMP0]], <4 x double> [[TMP1]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
53-
; CHECK-NEXT: [[SPLIT:%.*]] = shufflevector <8 x double> [[TMP2]], <8 x double> poison, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
54-
; CHECK-NEXT: [[SPLIT7:%.*]] = shufflevector <8 x double> [[TMP2]], <8 x double> poison, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
55-
; CHECK-NEXT: [[SPLIT8:%.*]] = shufflevector <8 x double> [[B:%.*]], <8 x double> poison, <4 x i32> <i32 0, i32 1, i32 2, i32 3>
56-
; CHECK-NEXT: [[SPLIT9:%.*]] = shufflevector <8 x double> [[B]], <8 x double> poison, <4 x i32> <i32 4, i32 5, i32 6, i32 7>
57-
; CHECK-NEXT: [[TMP3:%.*]] = fadd <4 x double> [[SPLIT]], [[SPLIT8]]
58-
; CHECK-NEXT: [[TMP4:%.*]] = fadd <4 x double> [[SPLIT7]], [[SPLIT9]]
59-
; CHECK-NEXT: [[TMP5:%.*]] = extractelement <4 x double> [[TMP3]], i64 0
60-
; CHECK-NEXT: [[TMP6:%.*]] = insertelement <2 x double> poison, double [[TMP5]], i64 0
61-
; CHECK-NEXT: [[TMP7:%.*]] = extractelement <4 x double> [[TMP4]], i64 0
62-
; CHECK-NEXT: [[TMP8:%.*]] = insertelement <2 x double> [[TMP6]], double [[TMP7]], i64 1
63-
; CHECK-NEXT: [[TMP9:%.*]] = extractelement <4 x double> [[TMP3]], i64 1
64-
; CHECK-NEXT: [[TMP10:%.*]] = insertelement <2 x double> poison, double [[TMP9]], i64 0
65-
; CHECK-NEXT: [[TMP11:%.*]] = extractelement <4 x double> [[TMP4]], i64 1
66-
; CHECK-NEXT: [[TMP12:%.*]] = insertelement <2 x double> [[TMP10]], double [[TMP11]], i64 1
67-
; CHECK-NEXT: [[TMP13:%.*]] = extractelement <4 x double> [[TMP3]], i64 2
68-
; CHECK-NEXT: [[TMP14:%.*]] = insertelement <2 x double> poison, double [[TMP13]], i64 0
69-
; CHECK-NEXT: [[TMP15:%.*]] = extractelement <4 x double> [[TMP4]], i64 2
70-
; CHECK-NEXT: [[TMP16:%.*]] = insertelement <2 x double> [[TMP14]], double [[TMP15]], i64 1
71-
; CHECK-NEXT: [[TMP17:%.*]] = extractelement <4 x double> [[TMP3]], i64 3
72-
; CHECK-NEXT: [[TMP18:%.*]] = insertelement <2 x double> poison, double [[TMP17]], i64 0
73-
; CHECK-NEXT: [[TMP19:%.*]] = extractelement <4 x double> [[TMP4]], i64 3
74-
; CHECK-NEXT: [[TMP20:%.*]] = insertelement <2 x double> [[TMP18]], double [[TMP19]], i64 1
75-
; CHECK-NEXT: [[TMP21:%.*]] = shufflevector <2 x double> [[TMP8]], <2 x double> [[TMP12]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
76-
; CHECK-NEXT: [[TMP22:%.*]] = shufflevector <2 x double> [[TMP16]], <2 x double> [[TMP20]], <4 x i32> <i32 0, i32 1, i32 2, i32 3>
77-
; CHECK-NEXT: [[TMP23:%.*]] = shufflevector <4 x double> [[TMP21]], <4 x double> [[TMP22]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
78-
; CHECK-NEXT: ret <8 x double> [[TMP23]]
49+
; CHECK-NEXT: [[COL_LOAD1:%.*]] = load <2 x double>, ptr [[VEC_GEP]], align 8
50+
; CHECK-NEXT: [[VEC_GEP2:%.*]] = getelementptr double, ptr [[A_PTR]], i64 4
51+
; CHECK-NEXT: [[COL_LOAD3:%.*]] = load <2 x double>, ptr [[VEC_GEP2]], align 8
52+
; CHECK-NEXT: [[VEC_GEP4:%.*]] = getelementptr double, ptr [[A_PTR]], i64 6
53+
; CHECK-NEXT: [[COL_LOAD5:%.*]] = load <2 x double>, ptr [[VEC_GEP4]], align 8
54+
; CHECK-NEXT: [[SPLIT:%.*]] = shufflevector <8 x double> [[B:%.*]], <8 x double> poison, <2 x i32> <i32 0, i32 1>
55+
; CHECK-NEXT: [[SPLIT6:%.*]] = shufflevector <8 x double> [[B]], <8 x double> poison, <2 x i32> <i32 2, i32 3>
56+
; CHECK-NEXT: [[SPLIT7:%.*]] = shufflevector <8 x double> [[B]], <8 x double> poison, <2 x i32> <i32 4, i32 5>
57+
; CHECK-NEXT: [[SPLIT8:%.*]] = shufflevector <8 x double> [[B]], <8 x double> poison, <2 x i32> <i32 6, i32 7>
58+
; CHECK-NEXT: [[TMP0:%.*]] = fadd <2 x double> [[COL_LOAD]], [[SPLIT]]
59+
; CHECK-NEXT: [[TMP1:%.*]] = fadd <2 x double> [[COL_LOAD1]], [[SPLIT6]]
60+
; CHECK-NEXT: [[TMP2:%.*]] = fadd <2 x double> [[COL_LOAD3]], [[SPLIT7]]
61+
; CHECK-NEXT: [[TMP3:%.*]] = fadd <2 x double> [[COL_LOAD5]], [[SPLIT8]]
62+
; CHECK-NEXT: [[TMP4:%.*]] = extractelement <2 x double> [[TMP0]], i64 0
63+
; CHECK-NEXT: [[TMP5:%.*]] = insertelement <4 x double> poison, double [[TMP4]], i64 0
64+
; CHECK-NEXT: [[TMP6:%.*]] = extractelement <2 x double> [[TMP1]], i64 0
65+
; CHECK-NEXT: [[TMP7:%.*]] = insertelement <4 x double> [[TMP5]], double [[TMP6]], i64 1
66+
; CHECK-NEXT: [[TMP8:%.*]] = extractelement <2 x double> [[TMP2]], i64 0
67+
; CHECK-NEXT: [[TMP9:%.*]] = insertelement <4 x double> [[TMP7]], double [[TMP8]], i64 2
68+
; CHECK-NEXT: [[TMP10:%.*]] = extractelement <2 x double> [[TMP3]], i64 0
69+
; CHECK-NEXT: [[TMP11:%.*]] = insertelement <4 x double> [[TMP9]], double [[TMP10]], i64 3
70+
; CHECK-NEXT: [[TMP12:%.*]] = extractelement <2 x double> [[TMP0]], i64 1
71+
; CHECK-NEXT: [[TMP13:%.*]] = insertelement <4 x double> poison, double [[TMP12]], i64 0
72+
; CHECK-NEXT: [[TMP14:%.*]] = extractelement <2 x double> [[TMP1]], i64 1
73+
; CHECK-NEXT: [[TMP15:%.*]] = insertelement <4 x double> [[TMP13]], double [[TMP14]], i64 1
74+
; CHECK-NEXT: [[TMP16:%.*]] = extractelement <2 x double> [[TMP2]], i64 1
75+
; CHECK-NEXT: [[TMP17:%.*]] = insertelement <4 x double> [[TMP15]], double [[TMP16]], i64 2
76+
; CHECK-NEXT: [[TMP18:%.*]] = extractelement <2 x double> [[TMP3]], i64 1
77+
; CHECK-NEXT: [[TMP19:%.*]] = insertelement <4 x double> [[TMP17]], double [[TMP18]], i64 3
78+
; CHECK-NEXT: [[TMP20:%.*]] = shufflevector <4 x double> [[TMP11]], <4 x double> [[TMP19]], <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
79+
; CHECK-NEXT: ret <8 x double> [[TMP20]]
7980
;
8081

8182

‎llvm/test/Transforms/LowerMatrixIntrinsics/transpose-opts-lifting.ll

Lines changed: 40 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ define <6 x double> @lift_through_add_matching_transpose_dimensions(<6 x double>
99
; CHECK-LABEL: define <6 x double> @lift_through_add_matching_transpose_dimensions(<6 x double> %a, <6 x double> %b) {
1010
; CHECK-NEXT: entry:
1111
; CHECK-NEXT: [[A:%.+]] = fadd <6 x double> %a, %b
12-
; CHECK-NEXT: [[T:%.+]] = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> [[A]], i32 2, i32 3)
12+
; CHECK-NEXT: [[T:%.+]] = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> [[A]], i32 3, i32 2)
1313
; CHECK-NEXT: ret <6 x double> [[T]]
1414
;
1515
entry:
@@ -25,7 +25,7 @@ define <6 x double> @lift_through_add_matching_transpose_dimensions_ops_also_hav
2525
; CHECK-NEXT: [[A:%.+]] = load <6 x double>, ptr %a.ptr
2626
; CHECK-NEXT: [[B:%.+]] = load <6 x double>, ptr %b.ptr
2727
; CHECK-NEXT: [[ADD:%.+]] = fadd <6 x double> [[A]], [[B]]
28-
; CHECK-NEXT: [[T:%.+]] = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> [[ADD]], i32 2, i32 3)
28+
; CHECK-NEXT: [[T:%.+]] = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> [[ADD]], i32 3, i32 2)
2929
; CHECK-NEXT: ret <6 x double> [[T]]
3030
;
3131
entry:
@@ -41,10 +41,28 @@ define <6 x double> @lift_through_add_mismatching_dimensions_1(<6 x double> %a,
4141
; CHECK-LABEL: define <6 x double> @lift_through_add_mismatching_dimensions_1(<6 x double> %a, <6 x double> %b) {
4242
; CHECK-NEXT: entry:
4343
; CHECK-NEXT: [[A:%.+]] = fadd <6 x double> %a, %b
44-
; CHECK-NEXT: [[T:%.+]] = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> [[A]], i32 2, i32 3)
44+
; CHECK-NEXT: [[T:%.+]] = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> [[A]], i32 1, i32 6)
45+
; CHECK-NEXT: ret <6 x double> [[T]]
46+
;
47+
entry:
48+
%a.t = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> %a, i32 1, i32 6)
49+
%b.t = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> %b, i32 3, i32 2)
50+
%add = fadd <6 x double> %a.t, %b.t
51+
ret <6 x double> %add
52+
}
53+
54+
define <6 x double> @lift_through_add_mismatching_dimensions_1_transpose_dimensions_ops_also_have_shape_info(ptr %a.ptr, ptr %b.ptr) {
55+
; CHECK-LABEL: define <6 x double> @lift_through_add_mismatching_dimensions_1_transpose_dimensions_ops_also_have_shape_info(ptr %a.ptr, ptr %b.ptr)
56+
; CHECK-NEXT: entry:
57+
; CHECK-NEXT: [[A:%.+]] = load <6 x double>, ptr %a.ptr
58+
; CHECK-NEXT: [[B:%.+]] = load <6 x double>, ptr %b.ptr
59+
; CHECK-NEXT: [[ADD:%.+]] = fadd <6 x double> [[A]], [[B]]
60+
; CHECK-NEXT: [[T:%.+]] = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> [[ADD]], i32 1, i32 6)
4561
; CHECK-NEXT: ret <6 x double> [[T]]
4662
;
4763
entry:
64+
%a = load <6 x double>, ptr %a.ptr
65+
%b = load <6 x double>, ptr %b.ptr
4866
%a.t = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> %a, i32 1, i32 6)
4967
%b.t = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> %b, i32 3, i32 2)
5068
%add = fadd <6 x double> %a.t, %b.t
@@ -55,7 +73,7 @@ define <6 x double> @lift_through_add_mismatching_dimensions_2(<6 x double> %a,
5573
; CHECK-LABEL: define <6 x double> @lift_through_add_mismatching_dimensions_2(<6 x double> %a, <6 x double> %b) {
5674
; CHECK-NEXT: entry:
5775
; CHECK-NEXT: [[A:%.+]] = fadd <6 x double> %a, %b
58-
; CHECK-NEXT: [[T:%.+]] = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> [[A]], i32 1, i32 6)
76+
; CHECK-NEXT: [[T:%.+]] = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> [[A]], i32 3, i32 2)
5977
; CHECK-NEXT: ret <6 x double> [[T]]
6078
;
6179

@@ -66,6 +84,24 @@ entry:
6684
ret <6 x double> %add
6785
}
6886

87+
define <6 x double> @lift_through_add_mismatching_dimensions_2_transpose_dimensions_ops_also_have_shape_info(ptr %a.ptr, ptr %b.ptr) {
88+
; CHECK-LABEL: define <6 x double> @lift_through_add_mismatching_dimensions_2_transpose_dimensions_ops_also_have_shape_info(ptr %a.ptr, ptr %b.ptr)
89+
; CHECK-NEXT: entry:
90+
; CHECK-NEXT: [[A:%.+]] = load <6 x double>, ptr %a.ptr
91+
; CHECK-NEXT: [[B:%.+]] = load <6 x double>, ptr %b.ptr
92+
; CHECK-NEXT: [[ADD:%.+]] = fadd <6 x double> [[A]], [[B]]
93+
; CHECK-NEXT: [[T:%.+]] = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> [[ADD]], i32 3, i32 2)
94+
; CHECK-NEXT: ret <6 x double> [[T]]
95+
;
96+
entry:
97+
%a = load <6 x double>, ptr %a.ptr
98+
%b = load <6 x double>, ptr %b.ptr
99+
%a.t = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> %a, i32 3, i32 2)
100+
%b.t = call <6 x double> @llvm.matrix.transpose.v6f64(<6 x double> %b, i32 6, i32 1)
101+
%add = fadd <6 x double> %a.t, %b.t
102+
ret <6 x double> %add
103+
}
104+
69105
define <9 x double> @lift_through_multiply(<6 x double> %a, <6 x double> %b) {
70106
; CHECK-LABEL: define <9 x double> @lift_through_multiply(<6 x double> %a, <6 x double> %b) {
71107
; CHECK-NEXT: entry:

0 commit comments

Comments
 (0)