scf-loop-utils.mlir
1.78 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
// RUN: mlir-opt -allow-unregistered-dialect -test-scf-for-utils -mlir-disable-threading %s | FileCheck %s
// CHECK-LABEL: @hoist
// CHECK-SAME: %[[lb:[a-zA-Z0-9]*]]: index,
// CHECK-SAME: %[[ub:[a-zA-Z0-9]*]]: index,
// CHECK-SAME: %[[step:[a-zA-Z0-9]*]]: index
func @hoist(%lb: index, %ub: index, %step: index) {
// CHECK: %[[A:.*]] = "fake_read"() : () -> index
// CHECK: %[[RES:.*]] = scf.for %[[I:.*]] = %[[lb]] to %[[ub]] step %[[step]] iter_args(%[[VAL:.*]] = %[[A]]) -> (index)
// CHECK: %[[YIELD:.*]] = "fake_compute"(%[[VAL]]) : (index) -> index
// CHECK: scf.yield %[[YIELD]] : index
// CHECK: "fake_write"(%[[RES]]) : (index) -> ()
scf.for %i = %lb to %ub step %step {
%0 = "fake_read"() : () -> (index)
%1 = "fake_compute"(%0) : (index) -> (index)
"fake_write"(%1) : (index) -> ()
}
return
}
// CHECK-LABEL: @hoist2
// CHECK-SAME: %[[lb:[a-zA-Z0-9]*]]: index,
// CHECK-SAME: %[[ub:[a-zA-Z0-9]*]]: index,
// CHECK-SAME: %[[step:[a-zA-Z0-9]*]]: index
// CHECK-SAME: %[[extra_arg:[a-zA-Z0-9]*]]: f32
func @hoist2(%lb: index, %ub: index, %step: index, %extra_arg: f32) -> f32 {
// CHECK: %[[A:.*]] = "fake_read"() : () -> index
// CHECK: %[[RES:.*]]:2 = scf.for %[[I:.*]] = %[[lb]] to %[[ub]] step %[[step]] iter_args(%[[VAL0:.*]] = %[[extra_arg]], %[[VAL1:.*]] = %[[A]]) -> (f32, index)
// CHECK: %[[YIELD:.*]] = "fake_compute"(%[[VAL1]]) : (index) -> index
// CHECK: scf.yield %[[VAL0]], %[[YIELD]] : f32, index
// CHECK: "fake_write"(%[[RES]]#1) : (index) -> ()
// CHECK: return %[[RES]]#0 : f32
%0 = scf.for %i = %lb to %ub step %step iter_args(%iter = %extra_arg) -> (f32) {
%0 = "fake_read"() : () -> (index)
%1 = "fake_compute"(%0) : (index) -> (index)
"fake_write"(%1) : (index) -> ()
scf.yield %iter: f32
}
return %0: f32
}