llvm-project

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() {
6
func.func @alloc() {
7
^bb0:
8
  // Test simple alloc.
9
  // CHECK: %{{.*}} = memref.alloc() : memref<1024x64xf32, 1>
10
  %0 = memref.alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
11

12
  %c0 = "arith.constant"() {value = 0: index} : () -> index
13
  %c1 = "arith.constant"() {value = 1: index} : () -> index
14

15
  // Test alloc with dynamic dimensions.
16
  // CHECK: %{{.*}} = memref.alloc(%{{.*}}, %{{.*}}) : memref<?x?xf32, 1>
17
  %1 = memref.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: %{{.*}} = memref.alloc()[%{{.*}}] : memref<2x4xf32, #[[$MAP]], 1>
21
  %2 = memref.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: %{{.*}} = memref.alloc(%{{.*}})[%{{.*}}] : memref<2x?xf32, #[[$MAP]], 1>
25
  %3 = memref.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: %{{.*}} = memref.alloc() : memref<2xi32>
30
  %4 = memref.alloc() : memref<2 x i32>
31

32
  // CHECK:   return
33
  return
34
}
35

36
// CHECK-LABEL: func @alloca() {
37
func.func @alloca() {
38
^bb0:
39
  // Test simple alloc.
40
  // CHECK: %{{.*}} = memref.alloca() : memref<1024x64xf32, 1>
41
  %0 = memref.alloca() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
42

43
  %c0 = "arith.constant"() {value = 0: index} : () -> index
44
  %c1 = "arith.constant"() {value = 1: index} : () -> index
45

46
  // Test alloca with dynamic dimensions.
47
  // CHECK: %{{.*}} = memref.alloca(%{{.*}}, %{{.*}}) : memref<?x?xf32, 1>
48
  %1 = memref.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: %{{.*}} = memref.alloca()[%{{.*}}] : memref<2x4xf32, #[[$MAP]], 1>
52
  %2 = memref.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: %{{.*}} = memref.alloca(%{{.*}})[%{{.*}}] : memref<2x?xf32, #[[$MAP]], 1>
56
  %3 = memref.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: %{{.*}} = memref.alloca() {alignment = 64 : i64} : memref<2xi32>
60
  %4 = memref.alloca() {alignment = 64} : memref<2 x i32>
61

62
  return
63
}
64

65
// CHECK-LABEL: func @dealloc() {
66
func.func @dealloc() {
67
^bb0:
68
  // CHECK: %{{.*}} = memref.alloc() : memref<1024x64xf32>
69
  %0 = memref.alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 0>
70

71
  // CHECK: memref.dealloc %{{.*}} : memref<1024x64xf32>
72
  memref.dealloc %0 : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 0>
73
  return
74
}
75

76
// CHECK-LABEL: func @load_store
77
func.func @load_store() {
78
^bb0:
79
  // CHECK: %{{.*}} = memref.alloc() : memref<1024x64xf32, 1>
80
  %0 = memref.alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
81

82
  %1 = arith.constant 0 : index
83
  %2 = arith.constant 1 : index
84

85
  // CHECK: %{{.*}} = memref.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x64xf32, 1>
86
  %3 = memref.load %0[%1, %2] : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
87

88
  // CHECK: memref.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<1024x64xf32, 1>
89
  memref.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()
95
func.func @dma_ops() {
96
  %c0 = arith.constant 0 : index
97
  %stride = arith.constant 32 : index
98
  %elt_per_stride = arith.constant 16 : index
99

100
  %A = memref.alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 0>
101
  %Ah = memref.alloc() : memref<256 x f32, affine_map<(d0) -> (d0)>, 1>
102
  %tag = memref.alloc() : memref<1 x f32>
103

104
  %num_elements = arith.constant 256 : index
105

106
  memref.dma_start %A[%c0], %Ah[%c0], %num_elements, %tag[%c0] : memref<256 x f32>, memref<256 x f32, 1>, memref<1 x f32>
107
  memref.dma_wait %tag[%c0], %num_elements : memref<1 x f32>
108
  // CHECK: dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}[%{{.*}}] : memref<256xf32>, memref<256xf32, 1>, memref<1xf32>
109
  // CHECK-NEXT:  dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xf32>
110

111
  // DMA with strides
112
  memref.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
  memref.dma_wait %tag[%c0], %num_elements : memref<1 x f32>
114
  // CHECK-NEXT:  dma_start %{{.*}}[%{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}} : memref<256xf32>, memref<256xf32, 1>, memref<1xf32>
115
  // CHECK-NEXT:  dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xf32>
116

117
  return
118
}
119