mlir/test/Dialect/ControlFlow/one-shot-bufferize.mlir

   1 // RUN: mlir-opt -one-shot-bufferize="bufferize-function-boundaries" -split-input-file %s | FileCheck %s
   2 // RUN: mlir-opt -one-shot-bufferize -split-input-file %s | FileCheck %s --check-prefix=CHECK-NO-FUNC
   3
   4 // CHECK-NO-FUNC-LABEL: func @br(
   5 //  CHECK-NO-FUNC-SAME:     %[[t:.*]]: tensor<5xf32>)
   6 //       CHECK-NO-FUNC:   %[[m:.*]] = bufferization.to_memref %[[t]] : memref<5xf32, strided<[?], offset: ?>>
   7 //       CHECK-NO-FUNC:   %[[r:.*]] = scf.execute_region -> memref<5xf32, strided<[?], offset: ?>> {
   8 //       CHECK-NO-FUNC:     cf.br ^[[block:.*]](%[[m]]
   9 //       CHECK-NO-FUNC:   ^[[block]](%[[arg1:.*]]: memref<5xf32, strided<[?], offset: ?>>):
  10 //       CHECK-NO-FUNC:     scf.yield %[[arg1]]
  11 //       CHECK-NO-FUNC:   }
  12 //       CHECK-NO-FUNC:   return
  13 func.func @br(%t: tensor<5xf32>) {
  14   %0 = scf.execute_region -> tensor<5xf32> {
  15     cf.br ^bb1(%t : tensor<5xf32>)
  16   ^bb1(%arg1 : tensor<5xf32>):
  17     scf.yield %arg1 : tensor<5xf32>
  18   }
  19   return
  20 }
  21
  22 // -----
  23
  24 // CHECK-NO-FUNC-LABEL: func @cond_br(
  25 //  CHECK-NO-FUNC-SAME:     %[[t1:.*]]: tensor<5xf32>,
  26 //       CHECK-NO-FUNC:   %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<5xf32, strided<[?], offset: ?>>
  27 //       CHECK-NO-FUNC:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
  28 //       CHECK-NO-FUNC:   %[[r:.*]] = scf.execute_region -> memref<5xf32, strided<[?], offset: ?>> {
  29 //       CHECK-NO-FUNC:     cf.cond_br %{{.*}}, ^[[block1:.*]](%[[m1]] : {{.*}}), ^[[block2:.*]](%[[alloc]] : {{.*}})
  30 //       CHECK-NO-FUNC:   ^[[block1]](%[[arg1:.*]]: memref<5xf32, strided<[?], offset: ?>>):
  31 //       CHECK-NO-FUNC:     scf.yield %[[arg1]]
  32 //       CHECK-NO-FUNC:   ^[[block2]](%[[arg2:.*]]: memref<5xf32>):
  33 //       CHECK-NO-FUNC:     %[[cast:.*]] = memref.cast %[[arg2]] : memref<5xf32> to memref<5xf32, strided<[?], offset: ?>
  34 //       CHECK-NO-FUNC:     cf.br ^[[block1]](%[[cast]] : {{.*}})
  35 //       CHECK-NO-FUNC:   }
  36 //       CHECK-NO-FUNC:   return
  37 func.func @cond_br(%t1: tensor<5xf32>, %c: i1) {
  38   // Use an alloc for the second block instead of a function block argument.
  39   // A cast must be inserted because the two will have different layout maps.
  40   %t0 = bufferization.alloc_tensor() : tensor<5xf32>
  41   %0 = scf.execute_region -> tensor<5xf32> {
  42     cf.cond_br %c, ^bb1(%t1 : tensor<5xf32>), ^bb2(%t0 : tensor<5xf32>)
  43   ^bb1(%arg1 : tensor<5xf32>):
  44     scf.yield %arg1 : tensor<5xf32>
  45   ^bb2(%arg2 : tensor<5xf32>):
  46     cf.br ^bb1(%arg2 : tensor<5xf32>)
  47   }
  48   return
  49 }
  50
  51 // -----
  52
  53 // CHECK-LABEL: func @looping_branches(
  54 func.func @looping_branches() -> tensor<5xf32> {
  55 // CHECK: %[[alloc:.*]] = memref.alloc
  56   %0 = bufferization.alloc_tensor() : tensor<5xf32>
  57 // CHECK: cf.br {{.*}}(%[[alloc]]
  58   cf.br ^bb1(%0: tensor<5xf32>)
  59 // CHECK: ^{{.*}}(%[[arg1:.*]]: memref<5xf32>):
  60 ^bb1(%arg1: tensor<5xf32>):
  61   %pos = "test.foo"() : () -> (index)
  62   %val = "test.bar"() : () -> (f32)
  63 // CHECK: memref.store %{{.*}}, %[[arg1]]
  64   %inserted = tensor.insert %val into %arg1[%pos] : tensor<5xf32>
  65   %cond = "test.qux"() : () -> (i1)
  66 // CHECK: cf.cond_br {{.*}}(%[[arg1]] {{.*}}(%[[arg1]]
  67   cf.cond_br %cond, ^bb1(%inserted: tensor<5xf32>), ^bb2(%inserted: tensor<5xf32>)
  68 // CHECK: ^{{.*}}(%[[arg2:.*]]: memref<5xf32>):
  69 ^bb2(%arg2: tensor<5xf32>):
  70 // CHECK: return %[[arg2]]
  71   func.return %arg2 : tensor<5xf32>
  72 }