mlir/test/Dialect/Affine/constant-fold.mlir

   1 // RUN: mlir-opt -test-constant-fold -split-input-file %s | FileCheck %s
   2
   3 // CHECK-LABEL: func @affine_apply
   4 func.func @affine_apply(%variable : index) -> (index, index, index) {
   5   %c177 = arith.constant 177 : index
   6   %c211 = arith.constant 211 : index
   7   %N = arith.constant 1075 : index
   8
   9   // CHECK:[[C1159:%.+]] = arith.constant 1159 : index
  10   // CHECK:[[C1152:%.+]] = arith.constant 1152 : index
  11   %x0 = affine.apply affine_map<(d0, d1)[S0] -> ( (d0 + 128 * S0) floordiv 128 + d1 mod 128)>
  12            (%c177, %c211)[%N]
  13   %x1 = affine.apply affine_map<(d0, d1)[S0] -> (128 * (S0 ceildiv 128))>
  14            (%c177, %c211)[%N]
  15
  16   // CHECK:[[C42:%.+]] = arith.constant 42 : index
  17   %y = affine.apply affine_map<(d0) -> (42)> (%variable)
  18
  19   // CHECK: return [[C1159]], [[C1152]], [[C42]]
  20   return %x0, %x1, %y : index, index, index
  21 }
  22
  23 // -----
  24
  25 // CHECK: #[[map:.*]] = affine_map<(d0, d1) -> (42, d1)
  26
  27 func.func @affine_min(%variable: index) -> (index, index) {
  28   // CHECK: %[[C42:.*]] = arith.constant 42
  29   %c42 = arith.constant 42 : index
  30   %c44 = arith.constant 44 : index
  31   // Partial folding will use a different map.
  32   // CHECK: %[[r:.*]] = affine.min #[[map]](%[[C42]], %{{.*}})
  33   %0 = affine.min affine_map<(d0, d1) -> (d0, d1)>(%c42, %variable)
  34
  35   // Full folding will remove the operation entirely.
  36   // CHECK-NOT: affine.min
  37   %1 = affine.min affine_map<(d0, d1) -> (d0, d1)>(%c42, %c44)
  38
  39   // CHECK: return %[[r]], %[[C42]]
  40   return %0, %1 : index, index
  41 }
  42
  43 // -----
  44
  45 // CHECK: #[[map:.*]] = affine_map<(d0, d1) -> (42, d1)
  46
  47 func.func @affine_min(%variable: index) -> (index, index) {
  48   // CHECK: %[[C42:.*]] = arith.constant 42
  49   %c42 = arith.constant 42 : index
  50   // CHECK: %[[C44:.*]] = arith.constant 44
  51   %c44 = arith.constant 44 : index
  52   // Partial folding will use a different map.
  53   // CHECK: %[[r:.*]] = affine.max #[[map]](%[[C42]], %{{.*}})
  54   %0 = affine.max affine_map<(d0, d1) -> (d0, d1)>(%c42, %variable)
  55
  56   // Full folding will remove the operation entirely.
  57   // CHECK-NOT: affine.max
  58   %1 = affine.max affine_map<(d0, d1) -> (d0, d1)>(%c42, %c44)
  59
  60   // CHECK: return %[[r]], %[[C44]]
  61   return %0, %1 : index, index
  62 }
  63
  64 // -----
  65
  66 func.func @affine_apply_poison_division_zero() {
  67   // This is just for mlir::context to load ub dialect
  68   %ub = ub.poison : index
  69   %c16 = arith.constant 16 : index
  70   %0 = affine.apply affine_map<(d0)[s0] -> (d0 mod (s0 - s0))>(%c16)[%c16]
  71   %1 = affine.apply affine_map<(d0)[s0] -> (d0 floordiv (s0 - s0))>(%c16)[%c16]
  72   %2 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv (s0 - s0))>(%c16)[%c16]
  73   %alloc = memref.alloc(%0, %1, %2) : memref<?x?x?xi1>
  74   %3 = affine.load %alloc[%0, %1, %2] : memref<?x?x?xi1>
  75   affine.store %3, %alloc[%0, %1, %2] : memref<?x?x?xi1>
  76   return
  77 }
  78
  79 // CHECK-NOT: affine.apply
  80 // CHECK: %[[poison:.*]] = ub.poison : index
  81 // CHECK-NEXT: %[[alloc:.*]] = memref.alloc(%[[poison]], %[[poison]], %[[poison]])
  82 // CHECK-NEXT: %[[load:.*]] = affine.load %[[alloc]][%[[poison]], %[[poison]], %[[poison]]] : memref<?x?x?xi1>
  83 // CHECK-NEXT: affine.store %[[load]], %alloc[%[[poison]], %[[poison]], %[[poison]]] : memref<?x?x?xi1>
  84
  85 // -----
  86
  87 // Check that this doesn't crash because the ub dialect is automatically loaded
  88 func.func @affine_apply_poison_division_zero_2() -> index {
  89   %c16 = arith.constant 16 : index
  90   %0 = affine.apply affine_map<(d0)[s0] -> (d0 mod (s0 - s0))>(%c16)[%c16]
  91   return %0 : index
  92 }