mlir/test/Dialect/Affine/SuperVectorize/compose_maps.mlir

   1 // RUN: mlir-opt -allow-unregistered-dialect %s -affine-super-vectorizer-test=compose-maps -split-input-file 2>&1 |  FileCheck %s
   2
   3 // For all these cases, the test traverses the `test_affine_map` ops and
   4 // composes them in order one-by-one.
   5 // For instance, the pseudo-sequence:
   6 //    "test_affine_map"() { affine_map = f } : () -> ()
   7 //    "test_affine_map"() { affine_map = g } : () -> ()
   8 //    "test_affine_map"() { affine_map = h } : () -> ()
   9 // will produce the sequence of compositions: f, g(f), h(g(f)) and print the
  10 // AffineMap h(g(f)), which is what FileCheck checks against.
  11
  12 func.func @simple1() {
  13   // CHECK: Composed map: (d0) -> (d0)
  14   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 - 1)> } : () -> ()
  15   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 + 1)> } : () -> ()
  16   return
  17 }
  18
  19 // -----
  20
  21 func.func @simple2() {
  22   // CHECK: Composed map: (d0)[s0, s1] -> (d0 - s0 + s1)
  23   "test_affine_map"() { affine_map = affine_map<(d0)[s0] -> (d0 + s0 - 1)> } : () -> ()
  24   "test_affine_map"() { affine_map = affine_map<(d0)[s0] -> (d0 - s0 + 1)> } : () -> ()
  25   return
  26 }
  27
  28 // -----
  29
  30 func.func @simple3a() {
  31   // CHECK: Composed map: (d0, d1)[s0, s1, s2, s3] -> ((d0 ceildiv s2) * s0, (d1 ceildiv s3) * s1)
  32   "test_affine_map"() { affine_map = affine_map<(d0, d1)[s0, s1] -> (d0 ceildiv s0, d1 ceildiv s1)> } : () -> ()
  33   "test_affine_map"() { affine_map = affine_map<(d0, d1)[s0, s1] -> (d0 * s0, d1 * s1)> } : () -> ()
  34   return
  35 }
  36
  37 // -----
  38
  39 func.func @simple3b() {
  40   // CHECK: Composed map: (d0, d1)[s0, s1] -> (d0 mod s0, d1 mod s1)
  41   "test_affine_map"() { affine_map = affine_map<(d0, d1)[s0, s1] -> (d0 mod s0, d1 mod s1)> } : () -> ()
  42   return
  43 }
  44
  45 // -----
  46
  47 func.func @simple3c() {
  48   // CHECK: Composed map: (d0, d1)[s0, s1, s2, s3, s4, s5] -> ((d0 ceildiv s4) * s4 + d0 mod s2, (d1 ceildiv s5) * s5 + d1 mod s3)
  49   "test_affine_map"() { affine_map = affine_map<(d0, d1)[s0, s1] -> ((d0 ceildiv s0) * s0, (d1 ceildiv s1) * s1, d0, d1)> } : () -> ()
  50   "test_affine_map"() { affine_map = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 + d2 mod s2, d1 + d3 mod s3)> } : () -> ()
  51   return
  52 }
  53
  54 // -----
  55
  56 func.func @simple4() {
  57   // CHECK: Composed map: (d0, d1)[s0, s1] -> (d1 * s1, d0 ceildiv s0)
  58   "test_affine_map"() { affine_map = affine_map<(d0, d1) -> (d1, d0)> } : () -> ()
  59   "test_affine_map"() { affine_map = affine_map<(d0, d1)[s0, s1] -> (d0 * s1, d1 ceildiv s0)> } : () -> ()
  60   return
  61 }
  62
  63 // -----
  64
  65 func.func @simple5a() {
  66   // CHECK: Composed map: (d0) -> (d0 * 3 + 18)
  67   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 - 1)> } : () -> ()
  68   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 + 7)> } : () -> ()
  69   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 * 24)> } : () -> ()
  70   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 ceildiv 8)> } : () -> ()
  71   return
  72 }
  73
  74 // -----
  75
  76 func.func @simple5b() {
  77   // CHECK: Composed map: (d0) -> ((d0 + 6) ceildiv 2)
  78   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 - 1)> } : () -> ()
  79   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 + 7)> } : () -> ()
  80   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 * 4)> } : () -> ()
  81   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 ceildiv 8)> } : () -> ()
  82   return
  83 }
  84
  85 // -----
  86
  87 func.func @simple5c() {
  88   // CHECK: Composed map: (d0) -> (d0 * 8 + 48)
  89   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 - 1)> } : () -> ()
  90   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 + 7)> } : () -> ()
  91   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 * 24)> } : () -> ()
  92   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 floordiv 3)> } : () -> ()
  93   return
  94 }
  95
  96 // -----
  97
  98 func.func @simple5d() {
  99   // CHECK: Composed map: (d0) -> ((d0 * 4) floordiv 3 + 8)
 100   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 - 1)> } : () -> ()
 101   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 + 7)> } : () -> ()
 102   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 * 4)> } : () -> ()
 103   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 floordiv 3)> } : () -> ()
 104   return
 105 }
 106
 107 // -----
 108
 109 func.func @simple5e() {
 110   // CHECK: Composed map: (d0) -> ((d0 + 6) ceildiv 8)
 111   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 - 1)> } : () -> ()
 112   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 + 7)> } : () -> ()
 113   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 ceildiv 8)> } : () -> ()
 114   return
 115 }
 116
 117 // -----
 118
 119 func.func @simple5f() {
 120   // CHECK: Composed map: (d0) -> ((d0 * 4 - 4) floordiv 3)
 121   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 - 1)> } : () -> ()
 122   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 * 4)> } : () -> ()
 123   "test_affine_map"() { affine_map = affine_map<(d0) -> (d0 floordiv 3)> } : () -> ()
 124   return
 125 }
 126
 127 // -----
 128
 129 func.func @perm_and_proj() {
 130   // CHECK: Composed map: (d0, d1, d2, d3) -> (d1, d3, d0)
 131   "test_affine_map"() { affine_map = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2, d0)> } : () -> ()
 132   "test_affine_map"() { affine_map = affine_map<(d0, d1, d2, d3) -> (d1, d0, d3)> } : () -> ()
 133   return
 134 }
 135
 136 // -----
 137
 138 func.func @symbols1() {
 139   // CHECK: Composed map: (d0)[s0] -> (d0 + s0 + 1, d0 - s0 - 1)
 140   "test_affine_map"() { affine_map = affine_map<(d0)[s0] -> (d0 + s0, d0 - s0)> } : () -> ()
 141   "test_affine_map"() { affine_map = affine_map<(d0, d1) -> (d0 + 1, d1 - 1)> } : () -> ()
 142   return
 143 }
 144
 145 // -----
 146
 147 func.func @drop() {
 148   // CHECK: Composed map: (d0, d1, d2)[s0, s1] -> (d0 * 2 + d1 + d2 + s1)
 149   "test_affine_map"() { affine_map = affine_map<(d0, d1, d2)[s0, s1] -> (d0 + s1, d1 + s0, d0 + d1 + d2)> } : () -> ()
 150   "test_affine_map"() { affine_map = affine_map<(d0, d1, d2) -> (d0 + d2)> } : () -> ()
 151   return
 152 }
 153
 154 // -----
 155
 156 func.func @multi_symbols() {
 157   // CHECK: Composed map: (d0)[s0, s1, s2] -> (d0 + s1 + s2 + 1, d0 - s0 - s2 - 1)
 158   "test_affine_map"() { affine_map = affine_map<(d0)[s0] -> (d0 + s0, d0 - s0)> } : () -> ()
 159   "test_affine_map"() { affine_map = affine_map<(d0, d1)[s0, s1] -> (d0 + 1 + s1, d1 - 1 - s0)> } : () -> ()
 160   return
 161 }
 162
 163 // -----
 164
 165 // CHECK-LABEL: @no_affine_maps
 166 func.func @no_affine_maps() {
 167   // CHECK: return
 168   return
 169 }