Skip to content

Commit 4d629fd

Browse files
jumerckxjumerckx
committed
filecheck test
1 parent 6df5ab3 commit 4d629fd

File tree

1 file changed

+140
-0
lines changed

1 file changed

+140
-0
lines changed

‎test/lit_tests/reduce_unused_multi_slice.mlir‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
// RUN: enzymexlamlir-opt --enzyme-hlo-generate-td="patterns=reduce_unused_multislice" --transform-interpreter --enzyme-hlo-remove-transform %s | FileCheck %s
2+
3+
// Test 1: Only center result used - should become a regular slice
4+
func.func @multi_slice_only_center(%arg0: tensor<20x24x80xf64>) -> tensor<1x8x72xf64> {
5+
%0, %1, %2, %3, %4, %5 = "enzymexla.multi_slice"(%arg0) <{
6+
start_indices = array<i64: 1, 0, 3>,
7+
limit_indices = array<i64: 2, 8, 75>,
8+
strides = array<i64: 1, 1, 1>,
9+
dimension = 2 : si32,
10+
left_amount = 2 : si32,
11+
right_amount = 3 : si32
12+
}> : (tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>)
13+
return %2 : tensor<1x8x72xf64>
14+
}
15+
16+
// CHECK-LABEL: func.func @multi_slice_only_center(
17+
// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: tensor<20x24x80xf64>) -> tensor<1x8x72xf64> {
18+
// CHECK: %[[VAL_1:.*]] = stablehlo.slice %[[VAL_0]] [1:2, 0:8, 3:75] : (tensor<20x24x80xf64>) -> tensor<1x8x72xf64>
19+
// CHECK: return %[[VAL_1]] : tensor<1x8x72xf64>
20+
// CHECK: }
21+
22+
23+
// Test 2: Only left-most result used - should become a regular slice
24+
func.func @multi_slice_only_left(%arg0: tensor<20x24x80xf64>) -> tensor<1x8x72xf64> {
25+
%0, %1, %2, %3, %4, %5 = "enzymexla.multi_slice"(%arg0) <{
26+
start_indices = array<i64: 1, 0, 3>,
27+
limit_indices = array<i64: 2, 8, 75>,
28+
strides = array<i64: 1, 1, 1>,
29+
dimension = 2 : si32,
30+
left_amount = 2 : si32,
31+
right_amount = 3 : si32
32+
}> : (tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>)
33+
return %0 : tensor<1x8x72xf64>
34+
}
35+
36+
// CHECK-LABEL: func.func @multi_slice_only_left(
37+
// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: tensor<20x24x80xf64>) -> tensor<1x8x72xf64> {
38+
// CHECK: %[[VAL_1:.*]] = stablehlo.slice %[[VAL_0]] [1:2, 0:8, 1:73] : (tensor<20x24x80xf64>) -> tensor<1x8x72xf64>
39+
// CHECK: return %[[VAL_1]] : tensor<1x8x72xf64>
40+
// CHECK: }
41+
42+
43+
// Test 3: Only right-most result used - should become a regular slice
44+
func.func @multi_slice_only_right(%arg0: tensor<20x24x80xf64>) -> tensor<1x8x72xf64> {
45+
%0, %1, %2, %3, %4, %5 = "enzymexla.multi_slice"(%arg0) <{
46+
start_indices = array<i64: 1, 0, 3>,
47+
limit_indices = array<i64: 2, 8, 75>,
48+
strides = array<i64: 1, 1, 1>,
49+
dimension = 2 : si32,
50+
left_amount = 2 : si32,
51+
right_amount = 3 : si32
52+
}> : (tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>)
53+
return %5 : tensor<1x8x72xf64>
54+
}
55+
56+
// CHECK-LABEL: func.func @multi_slice_only_right(
57+
// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: tensor<20x24x80xf64>) -> tensor<1x8x72xf64> {
58+
// CHECK: %[[VAL_1:.*]] = stablehlo.slice %[[VAL_0]] [1:2, 0:8, 6:78] : (tensor<20x24x80xf64>) -> tensor<1x8x72xf64>
59+
// CHECK: return %[[VAL_1]] : tensor<1x8x72xf64>
60+
// CHECK: }
61+
62+
63+
// Test 4: Two consecutive results used - should become smaller multi_slice
64+
func.func @multi_slice_consecutive(%arg0: tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>) {
65+
%0, %1, %2, %3, %4, %5 = "enzymexla.multi_slice"(%arg0) <{
66+
start_indices = array<i64: 1, 0, 3>,
67+
limit_indices = array<i64: 2, 8, 75>,
68+
strides = array<i64: 1, 1, 1>,
69+
dimension = 2 : si32,
70+
left_amount = 2 : si32,
71+
right_amount = 3 : si32
72+
}> : (tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>)
73+
return %2, %3 : tensor<1x8x72xf64>, tensor<1x8x72xf64>
74+
}
75+
76+
// CHECK-LABEL: func.func @multi_slice_consecutive(
77+
// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>) {
78+
// CHECK: %[[VAL_1:.*]]:2 = "enzymexla.multi_slice"(%[[VAL_0]]) <{dimension = 2 : si32, left_amount = 0 : si32, limit_indices = array<i64: 2, 8, 75>, right_amount = 1 : si32, start_indices = array<i64: 1, 0, 3>, strides = array<i64: 1, 1, 1>}> : (tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>)
79+
// CHECK: return %[[VAL_1]]#0, %[[VAL_1]]#1 : tensor<1x8x72xf64>, tensor<1x8x72xf64>
80+
// CHECK: }
81+
82+
83+
// Test 5: Non-contiguous results used - should keep range between first and last used
84+
func.func @multi_slice_non_contiguous(%arg0: tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>) {
85+
%0, %1, %2, %3, %4, %5 = "enzymexla.multi_slice"(%arg0) <{
86+
start_indices = array<i64: 1, 0, 3>,
87+
limit_indices = array<i64: 2, 8, 75>,
88+
strides = array<i64: 1, 1, 1>,
89+
dimension = 2 : si32,
90+
left_amount = 2 : si32,
91+
right_amount = 3 : si32
92+
}> : (tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>)
93+
return %2, %5 : tensor<1x8x72xf64>, tensor<1x8x72xf64>
94+
}
95+
96+
// CHECK-LABEL: func.func @multi_slice_non_contiguous(
97+
// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>) {
98+
// CHECK: %[[VAL_1:.*]]:4 = "enzymexla.multi_slice"(%[[VAL_0]]) <{dimension = 2 : si32, left_amount = 0 : si32, limit_indices = array<i64: 2, 8, 75>, right_amount = 3 : si32, start_indices = array<i64: 1, 0, 3>, strides = array<i64: 1, 1, 1>}> : (tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>)
99+
// CHECK: return %[[VAL_1]]#0, %[[VAL_1]]#3 : tensor<1x8x72xf64>, tensor<1x8x72xf64>
100+
// CHECK: }
101+
102+
103+
// Test 6: All results used - should not change
104+
func.func @multi_slice_all_used(%arg0: tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>) {
105+
%0, %1, %2, %3, %4, %5 = "enzymexla.multi_slice"(%arg0) <{
106+
start_indices = array<i64: 1, 0, 3>,
107+
limit_indices = array<i64: 2, 8, 75>,
108+
strides = array<i64: 1, 1, 1>,
109+
dimension = 2 : si32,
110+
left_amount = 2 : si32,
111+
right_amount = 3 : si32
112+
}> : (tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>)
113+
return %0, %1, %2, %3, %4, %5 : tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>
114+
}
115+
116+
// CHECK-LABEL: func.func @multi_slice_all_used(
117+
// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>) {
118+
// CHECK: %[[VAL_1:.*]]:6 = "enzymexla.multi_slice"(%[[VAL_0]]) <{dimension = 2 : si32, left_amount = 2 : si32, limit_indices = array<i64: 2, 8, 75>, right_amount = 3 : si32, start_indices = array<i64: 1, 0, 3>, strides = array<i64: 1, 1, 1>}> : (tensor<20x24x80xf64>) -> (tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>)
119+
// CHECK: return %[[VAL_1]]#0, %[[VAL_1]]#1, %[[VAL_1]]#2, %[[VAL_1]]#3, %[[VAL_1]]#4, %[[VAL_1]]#5 : tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>, tensor<1x8x72xf64>
120+
// CHECK: }
121+
122+
123+
// Test 7: Different dimension - test on dimension 0
124+
func.func @multi_slice_dim0(%arg0: tensor<20x24x80xf64>) -> tensor<4x24x80xf64> {
125+
%0, %1, %2, %3, %4 = "enzymexla.multi_slice"(%arg0) <{
126+
start_indices = array<i64: 8, 0, 0>,
127+
limit_indices = array<i64: 12, 24, 80>,
128+
strides = array<i64: 1, 1, 1>,
129+
dimension = 0 : si32,
130+
left_amount = 2 : si32,
131+
right_amount = 2 : si32
132+
}> : (tensor<20x24x80xf64>) -> (tensor<4x24x80xf64>, tensor<4x24x80xf64>, tensor<4x24x80xf64>, tensor<4x24x80xf64>, tensor<4x24x80xf64>)
133+
return %2 : tensor<4x24x80xf64>
134+
}
135+
136+
// CHECK-LABEL: func.func @multi_slice_dim0(
137+
// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: tensor<20x24x80xf64>) -> tensor<4x24x80xf64> {
138+
// CHECK: %[[VAL_1:.*]] = stablehlo.slice %[[VAL_0]] [8:12, 0:24, 0:80] : (tensor<20x24x80xf64>) -> tensor<4x24x80xf64>
139+
// CHECK: return %[[VAL_1]] : tensor<4x24x80xf64>
140+
// CHECK: }

0 commit comments

Comments
 (0)