xref: /llvm-project/mlir/test/Dialect/Affine/access-analysis.mlir (revision 1e9bfcd9a423765a86b2fc807c3bc3a097823deb)
1// RUN: mlir-opt %s -split-input-file -test-affine-access-analysis -verify-diagnostics | FileCheck %s
2
3// CHECK-LABEL: func @loop_simple
4func.func @loop_simple(%A : memref<?x?xf32>, %B : memref<?x?x?xf32>) {
5   %c0 = arith.constant 0 : index
6   %M = memref.dim %A, %c0 : memref<?x?xf32>
7   affine.for %i = 0 to %M {
8     affine.for %j = 0 to %M {
9       affine.load %A[%c0, %i] : memref<?x?xf32>
10       // expected-remark@above {{contiguous along loop 0}}
11       // expected-remark@above {{invariant along loop 1}}
12       affine.load %A[%c0, 8 * %i + %j] : memref<?x?xf32>
13       // expected-remark@above {{contiguous along loop 1}}
14       // Note/FIXME: access stride isn't being checked.
15       // expected-remark@-3 {{contiguous along loop 0}}
16
17       // These are all non-contiguous along both loops. Nothing is emitted.
18       affine.load %A[%i, %c0] : memref<?x?xf32>
19       // expected-remark@above {{invariant along loop 1}}
20       // Note/FIXME: access stride isn't being checked.
21       affine.load %A[%i, 8 * %j] : memref<?x?xf32>
22       // expected-remark@above {{contiguous along loop 1}}
23       affine.load %A[%j, 4 * %i] : memref<?x?xf32>
24       // expected-remark@above {{contiguous along loop 0}}
25     }
26   }
27   return
28}
29
30// -----
31
32// CHECK-LABEL: func @loop_unsimplified
33func.func @loop_unsimplified(%A : memref<100xf32>) {
34   affine.for %i = 0 to 100 {
35     affine.load %A[2 * %i - %i - %i] : memref<100xf32>
36     // expected-remark@above {{invariant along loop 0}}
37
38     %m = affine.apply affine_map<(d0) -> (-2 * d0)>(%i)
39     %n = affine.apply affine_map<(d0) -> (2 * d0)>(%i)
40     affine.load %A[(%m + %n) floordiv 2] : memref<100xf32>
41     // expected-remark@above {{invariant along loop 0}}
42   }
43   return
44}
45
46// -----
47
48#map = affine_map<(d0) -> (d0 * 16)>
49#map1 = affine_map<(d0) -> (d0 * 16 + 16)>
50#map2 = affine_map<(d0) -> (d0)>
51#map3 = affine_map<(d0) -> (d0 + 1)>
52
53func.func @tiled(%arg0: memref<*xf32>) {
54  %alloc = memref.alloc() {alignment = 64 : i64} : memref<1x224x224x64xf32>
55  %cast = memref.cast %arg0 : memref<*xf32> to memref<64xf32>
56  affine.for %arg1 = 0 to 4 {
57    affine.for %arg2 = 0 to 224 {
58      affine.for %arg3 = 0 to 14 {
59        %alloc_0 = memref.alloc() : memref<1x16x1x16xf32>
60        affine.for %arg4 = #map(%arg1) to #map1(%arg1) {
61          affine.for %arg5 = #map(%arg3) to #map1(%arg3) {
62            // TODO: here and below, the access isn't really invariant
63            // along tile-space IVs where the intra-tile IVs' bounds
64            // depend on them.
65            %0 = affine.load %cast[%arg4] : memref<64xf32>
66            // expected-remark@above {{contiguous along loop 3}}
67            // expected-remark@above {{invariant along loop 0}}
68            // expected-remark@above {{invariant along loop 1}}
69            // expected-remark@above {{invariant along loop 2}}
70            // expected-remark@above {{invariant along loop 4}}
71            affine.store %0, %alloc_0[0, %arg1 * -16 + %arg4, 0, %arg3 * -16 + %arg5] : memref<1x16x1x16xf32>
72            // expected-remark@above {{contiguous along loop 4}}
73            // expected-remark@above {{contiguous along loop 2}}
74            // expected-remark@above {{invariant along loop 1}}
75          }
76        }
77        affine.for %arg4 = #map(%arg1) to #map1(%arg1) {
78          affine.for %arg5 = #map2(%arg2) to #map3(%arg2) {
79            affine.for %arg6 = #map(%arg3) to #map1(%arg3) {
80              %0 = affine.load %alloc_0[0, %arg1 * -16 + %arg4, -%arg2 + %arg5, %arg3 * -16 + %arg6] : memref<1x16x1x16xf32>
81              // expected-remark@above {{contiguous along loop 5}}
82              // expected-remark@above {{contiguous along loop 2}}
83              affine.store %0, %alloc[0, %arg5, %arg6, %arg4] : memref<1x224x224x64xf32>
84              // expected-remark@above {{contiguous along loop 3}}
85              // expected-remark@above {{invariant along loop 0}}
86              // expected-remark@above {{invariant along loop 1}}
87              // expected-remark@above {{invariant along loop 2}}
88            }
89          }
90        }
91        memref.dealloc %alloc_0 : memref<1x16x1x16xf32>
92      }
93    }
94  }
95  return
96}
97