mlir/test/Dialect/Affine/access-analysis.mlir

   1 // RUN: mlir-opt %s -split-input-file -test-affine-access-analysis -verify-diagnostics | FileCheck %s
   2
   3 // CHECK-LABEL: func @loop_simple
   4 func.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
  33 func.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
  53 func.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 }