xref: /llvm-project/mlir/test/Dialect/Affine/dma.mlir (revision 227ed2f448e26cfc646e30ac17b03eac9578c297)
1// RUN: mlir-opt %s -split-input-file | FileCheck %s
2
3// -----
4
5// Test with loop IVs.
6func.func @test0(%arg0 : index, %arg1 : index) {
7  %0 = memref.alloc() : memref<100x100xf32>
8  %1 = memref.alloc() : memref<100x100xf32, affine_map<(d0, d1) -> (d0, d1)>, 2>
9  %2 = memref.alloc() : memref<1xi32>
10  %c0 = arith.constant 0 : index
11  %c64 = arith.constant 64 : index
12  affine.for %i0 = 0 to 10 {
13    affine.for %i1 = 0 to 10 {
14      affine.dma_start %0[%i0, %i1], %1[%i0, %i1], %2[%c0], %c64
15        : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
16      affine.dma_wait %2[%c0], %c64 : memref<1xi32>
17// CHECK: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
18// CHECK: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
19    }
20  }
21  return
22}
23
24// -----
25
26// Test with loop IVs and optional stride arguments.
27func.func @test1(%arg0 : index, %arg1 : index) {
28  %0 = memref.alloc() : memref<100x100xf32>
29  %1 = memref.alloc() : memref<100x100xf32, affine_map<(d0, d1) -> (d0, d1)>, 2>
30  %2 = memref.alloc() : memref<1xi32>
31  %c0 = arith.constant 0 : index
32  %c64 = arith.constant 64 : index
33  %c128 = arith.constant 128 : index
34  %c256 = arith.constant 256 : index
35  affine.for %i0 = 0 to 10 {
36    affine.for %i1 = 0 to 10 {
37      affine.dma_start %0[%i0, %i1], %1[%i0, %i1], %2[%c0], %c64, %c128, %c256
38        : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
39      affine.dma_wait %2[%c0], %c64 : memref<1xi32>
40// CHECK: affine.dma_start %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}], %{{.*}}, %{{.*}}, %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
41// CHECK: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
42    }
43  }
44  return
45}
46
47// -----
48
49// Test with loop IVs and symbols (without symbol keyword).
50func.func @test2(%arg0 : index, %arg1 : index) {
51  %0 = memref.alloc() : memref<100x100xf32>
52  %1 = memref.alloc() : memref<100x100xf32, affine_map<(d0, d1) -> (d0, d1)>, 2>
53  %2 = memref.alloc() : memref<1xi32>
54  %c0 = arith.constant 0 : index
55  %c64 = arith.constant 64 : index
56  affine.for %i0 = 0 to 10 {
57    affine.for %i1 = 0 to 10 {
58      affine.dma_start %0[%i0 + %arg0, %i1], %1[%i0, %i1 + %arg1 + 5],
59                       %2[%c0], %c64
60        : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
61      affine.dma_wait %2[%c0], %c64 : memref<1xi32>
62// CHECK: affine.dma_start %{{.*}}[%{{.*}} + %{{.*}}, %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}} + %{{.*}} + 5], %{{.*}}[%{{.*}}], %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
63// CHECK: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
64    }
65  }
66  return
67}
68
69// -----
70
71// Test with loop IVs and symbols (with symbol keyword).
72func.func @test3(%arg0 : index, %arg1 : index) {
73  %0 = memref.alloc() : memref<100x100xf32>
74  %1 = memref.alloc() : memref<100x100xf32, affine_map<(d0, d1) -> (d0, d1)>, 2>
75  %2 = memref.alloc() : memref<1xi32>
76  %c0 = arith.constant 0 : index
77  %c64 = arith.constant 64 : index
78  affine.for %i0 = 0 to 10 {
79    affine.for %i1 = 0 to 10 {
80      affine.dma_start %0[%i0 + symbol(%arg0), %i1],
81                       %1[%i0, %i1 + symbol(%arg1) + 7],
82                       %2[%i0 + %i1 + 11], %c64
83        : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
84      affine.dma_wait %2[%c0], %c64 : memref<1xi32>
85// CHECK: affine.dma_start %{{.*}}[%{{.*}} + symbol(%{{.*}}), %{{.*}}], %{{.*}}[%{{.*}}, %{{.*}} + symbol(%{{.*}}) + 7], %{{.*}}[%{{.*}} + %{{.*}} + 11], %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
86// CHECK: affine.dma_wait %{{.*}}[%{{.*}}], %{{.*}} : memref<1xi32>
87    }
88  }
89  return
90}
91
92// -----
93
94// Test with loop IVs, symbols and constants in nested affine expressions.
95func.func @test4(%arg0 : index, %arg1 : index) {
96  %0 = memref.alloc() : memref<100x100xf32>
97  %1 = memref.alloc() : memref<100x100xf32, 2>
98  %2 = memref.alloc() : memref<1xi32>
99  %c64 = arith.constant 64 : index
100  affine.for %i0 = 0 to 10 {
101    affine.for %i1 = 0 to 10 {
102      affine.dma_start %0[(%i0 + symbol(%arg0)) floordiv 3, %i1],
103                       %1[%i0, (%i1 + symbol(%arg1)) mod 9 + 7],
104                       %2[%i0 + %i1 + 11], %c64
105        : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
106      affine.dma_wait %2[%i0 + %i1 + 11], %c64 : memref<1xi32>
107// CHECK: affine.dma_start %{{.*}}[(%{{.*}} + symbol(%{{.*}})) floordiv 3, %{{.*}}], %{{.*}}[%{{.*}}, (%{{.*}} + symbol(%{{.*}})) mod 9 + 7], %{{.*}}[%{{.*}} + %{{.*}} + 11], %{{.*}} : memref<100x100xf32>, memref<100x100xf32, 2>, memref<1xi32>
108// CHECK: affine.dma_wait %{{.*}}[%{{.*}} + %{{.*}} + 11], %{{.*}} : memref<1xi32>
109    }
110  }
111  return
112}
113