1// RUN: mlir-opt %s | FileCheck %s
2
3// CHECK: #map = affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>
4
5// CHECK-LABEL: func @alloc() {
6func @alloc() {
7^bb0:
8  // Test simple alloc.
9  // CHECK: %0 = alloc() : memref<1024x64xf32, 1>
10  %0 = alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
11
12  %c0 = "std.constant"() {value = 0: index} : () -> index
13  %c1 = "std.constant"() {value = 1: index} : () -> index
14
15  // Test alloc with dynamic dimensions.
16  // CHECK: %1 = alloc(%c0, %c1) : memref<?x?xf32, 1>
17  %1 = alloc(%c0, %c1) : memref<?x?xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
18
19  // Test alloc with no dynamic dimensions and one symbol.
20  // CHECK: %2 = alloc()[%c0] : memref<2x4xf32, #map, 1>
21  %2 = alloc()[%c0] : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>
22
23  // Test alloc with dynamic dimensions and one symbol.
24  // CHECK: %3 = alloc(%c1)[%c0] : memref<2x?xf32, #map, 1>
25  %3 = alloc(%c1)[%c0] : memref<2x?xf32, affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>, 1>
26
27  // Alloc with no mappings.
28  // b/116054838 Parser crash while parsing ill-formed AllocOp
29  // CHECK: %4 = alloc() : memref<2xi32>
30  %4 = alloc() : memref<2 x i32>
31
32  // CHECK:   return
33  return
34}
35
36// CHECK-LABEL: func @alloca() {
37func @alloca() {
38^bb0:
39  // Test simple alloc.
40  // CHECK: %0 = alloca() : memref<1024x64xf32, 1>
41  %0 = alloca() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
42
43  %c0 = "std.constant"() {value = 0: index} : () -> index
44  %c1 = "std.constant"() {value = 1: index} : () -> index
45
46  // Test alloca with dynamic dimensions.
47  // CHECK: %1 = alloca(%c0, %c1) : memref<?x?xf32, 1>
48  %1 = alloca(%c0, %c1) : memref<?x?xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
49
50  // Test alloca with no dynamic dimensions and one symbol.
51  // CHECK: %2 = alloca()[%c0] : memref<2x4xf32, #map, 1>
52  %2 = alloca()[%c0] : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>
53
54  // Test alloca with dynamic dimensions and one symbol.
55  // CHECK: %3 = alloca(%c1)[%c0] : memref<2x?xf32, #map, 1>
56  %3 = alloca(%c1)[%c0] : memref<2x?xf32, affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>, 1>
57
58  // Alloca with no mappings, but with alignment.
59  // CHECK: %4 = alloca() {alignment = 64 : i64} : memref<2xi32>
60  %4 = alloca() {alignment = 64} : memref<2 x i32>
61
62  return
63}
64
65// CHECK-LABEL: func @dealloc() {
66func @dealloc() {
67^bb0:
68  // CHECK: %0 = alloc() : memref<1024x64xf32>
69  %0 = alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 0>
70
71  // CHECK: dealloc %0 : memref<1024x64xf32>
72  dealloc %0 : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 0>
73  return
74}
75
76// CHECK-LABEL: func @load_store
77func @load_store() {
78^bb0:
79  // CHECK: %0 = alloc() : memref<1024x64xf32, 1>
80  %0 = alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
81
82  %1 = constant 0 : index
83  %2 = constant 1 : index
84
85  // CHECK: %1 = load %0[%c0, %c1] : memref<1024x64xf32, 1>
86  %3 = load %0[%1, %2] : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
87
88  // CHECK: store %1, %0[%c0, %c1] : memref<1024x64xf32, 1>
89  store %3, %0[%1, %2] : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
90
91  return
92}
93
94// CHECK-LABEL: func @dma_ops()
95func @dma_ops() {
96  %c0 = constant 0 : index
97  %stride = constant 32 : index
98  %elt_per_stride = constant 16 : index
99
100  %A = alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 0>
101  %Ah = alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 1>
102  %tag = alloc() : memref<1 x f32>
103
104  %num_elements = constant 256 : index
105
106  dma_start %A[%c0], %Ah[%c0], %num_elements, %tag[%c0] : memref<256 x f32>, memref<256 x f32, 1>, memref<1 x f32>
107  dma_wait %tag[%c0], %num_elements : memref<1 x f32>
108  // CHECK: dma_start %0[%c0], %1[%c0], %c256, %2[%c0] : memref<256xf32>, memref<256xf32, 1>, memref<1xf32>
109  // CHECK-NEXT:  dma_wait %2[%c0], %c256 : memref<1xf32>
110
111  // DMA with strides
112  dma_start %A[%c0], %Ah[%c0], %num_elements, %tag[%c0], %stride, %elt_per_stride : memref<256 x f32>, memref<256 x f32, 1>, memref<1 x f32>
113  dma_wait %tag[%c0], %num_elements : memref<1 x f32>
114  // CHECK-NEXT:  dma_start %0[%c0], %1[%c0], %c256, %2[%c0], %c32, %c16 : memref<256xf32>, memref<256xf32, 1>, memref<1xf32>
115  // CHECK-NEXT:  dma_wait %2[%c0], %c256 : memref<1xf32>
116
117  return
118}
119