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