xref: /llvm-project/mlir/test/Dialect/MemRef/alloc-to-alloca.mlir (revision e4384149b58f7c3d19c5d38bc46038c660b77ca9)
1// RUN: mlir-opt --transform-interpreter=entry-point=all %s | FileCheck %s --check-prefixes=CHECK,ALL
2// RUN: mlir-opt --transform-interpreter=entry-point=small %s | FileCheck %s --check-prefixes=CHECK,SMALL
3
4func.func private @callee(memref<*xf32>)
5
6// CHECK-LABEL: @large_alloc
7func.func @large_alloc() {
8  // SMALL: memref.alloc()
9  // ALL:   memref.alloca
10  %0 = memref.alloc() : memref<100x100xf32>
11  %1 = memref.cast %0 : memref<100x100xf32> to memref<*xf32>
12  call @callee(%1) : (memref<*xf32>) -> ()
13  // SMALL: memref.dealloc
14  // ALL-NOT: memref.dealloc
15  memref.dealloc %0 : memref<100x100xf32>
16  return
17}
18
19// CHECK-LABEL: @small_alloc
20func.func @small_alloc() {
21  // CHECK: memref.alloca
22  %0 = memref.alloc() : memref<2x2xf32>
23  %1 = memref.cast %0 : memref<2x2xf32> to memref<*xf32>
24  call @callee(%1) : (memref<*xf32>) -> ()
25  // CHECK-NOT: memref.dealloc
26  memref.dealloc %0 : memref<2x2xf32>
27  return
28}
29
30// CHECK-LABEL: @no_dealloc
31func.func @no_dealloc() {
32  // CHECK: memref.alloc()
33  %0 = memref.alloc() : memref<2x2xf32>
34  %1 = memref.cast %0 : memref<2x2xf32> to memref<*xf32>
35  call @callee(%1) : (memref<*xf32>) -> ()
36  return
37}
38
39// CHECK-LABEL: @mismatching_scope
40func.func @mismatching_scope() {
41  // CHECK: memref.alloc()
42  %0 = memref.alloc() : memref<2x2xf32>
43  %1 = memref.cast %0 : memref<2x2xf32> to memref<*xf32>
44  call @callee(%1) : (memref<*xf32>) -> ()
45  scf.execute_region {
46    memref.dealloc %0 : memref<2x2xf32>
47    scf.yield
48  }
49  return
50}
51
52module attributes {transform.with_named_sequence} {
53  transform.named_sequence @all(%arg0: !transform.any_op {transform.readonly}) {
54    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
55    transform.apply_patterns to %0 {
56      transform.apply_patterns.memref.alloc_to_alloca
57    } : !transform.any_op
58    transform.yield
59  }
60}
61
62module attributes {transform.with_named_sequence} {
63  transform.named_sequence @small(%arg0: !transform.any_op {transform.readonly}) {
64    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
65    transform.apply_patterns to %0 {
66      transform.apply_patterns.memref.alloc_to_alloca size_limit(32)
67    } : !transform.any_op
68    transform.yield
69  }
70}
71