mlir/test/Transforms/scf-loop-utils.mlir

   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
   7 func @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
  26 func @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 }