You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
41 lines
1.8 KiB
41 lines
1.8 KiB
4 months ago
|
// 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
|
||
|
}
|