139 lines · plain
1// RUN: mlir-opt %s -allow-unregistered-dialect -affine-raise-from-memref --canonicalize | FileCheck %s2 3// CHECK-LABEL: func @reduce_window_max(4func.func @reduce_window_max() {5 %cst = arith.constant 0.000000e+00 : f326 %0 = memref.alloc() : memref<1x8x8x64xf32>7 %1 = memref.alloc() : memref<1x18x18x64xf32>8 affine.for %arg0 = 0 to 1 {9 affine.for %arg1 = 0 to 8 {10 affine.for %arg2 = 0 to 8 {11 affine.for %arg3 = 0 to 64 {12 memref.store %cst, %0[%arg0, %arg1, %arg2, %arg3] : memref<1x8x8x64xf32>13 }14 }15 }16 }17 affine.for %arg0 = 0 to 1 {18 affine.for %arg1 = 0 to 8 {19 affine.for %arg2 = 0 to 8 {20 affine.for %arg3 = 0 to 64 {21 affine.for %arg4 = 0 to 1 {22 affine.for %arg5 = 0 to 3 {23 affine.for %arg6 = 0 to 3 {24 affine.for %arg7 = 0 to 1 {25 %2 = memref.load %0[%arg0, %arg1, %arg2, %arg3] : memref<1x8x8x64xf32>26 %21 = arith.addi %arg0, %arg4 : index27 %22 = arith.constant 2 : index28 %23 = arith.muli %arg1, %22 : index29 %24 = arith.addi %23, %arg5 : index30 %25 = arith.muli %arg2, %22 : index31 %26 = arith.addi %25, %arg6 : index32 %27 = arith.addi %arg3, %arg7 : index33 %3 = memref.load %1[%21, %24, %26, %27] : memref<1x18x18x64xf32>34 %4 = arith.cmpf ogt, %2, %3 : f3235 %5 = arith.select %4, %2, %3 : f3236 memref.store %5, %0[%arg0, %arg1, %arg2, %arg3] : memref<1x8x8x64xf32>37 }38 }39 }40 }41 }42 }43 }44 }45 return46}47 48// CHECK: %[[cst:.*]] = arith.constant 049// CHECK: %[[v0:.*]] = memref.alloc() : memref<1x8x8x64xf32>50// CHECK: %[[v1:.*]] = memref.alloc() : memref<1x18x18x64xf32>51// CHECK: affine.for %[[arg0:.*]] =52// CHECK: affine.for %[[arg1:.*]] =53// CHECK: affine.for %[[arg2:.*]] =54// CHECK: affine.for %[[arg3:.*]] =55// CHECK: affine.store %[[cst]], %[[v0]][%[[arg0]], %[[arg1]], %[[arg2]], %[[arg3]]] :56// CHECK: affine.for %[[a0:.*]] =57// CHECK: affine.for %[[a1:.*]] =58// CHECK: affine.for %[[a2:.*]] =59// CHECK: affine.for %[[a3:.*]] =60// CHECK: affine.for %[[a4:.*]] =61// CHECK: affine.for %[[a5:.*]] =62// CHECK: affine.for %[[a6:.*]] =63// CHECK: affine.for %[[a7:.*]] =64// CHECK: %[[lhs:.*]] = affine.load %[[v0]][%[[a0]], %[[a1]], %[[a2]], %[[a3]]] :65// CHECK: %[[rhs:.*]] = affine.load %[[v1]][%[[a0]] + %[[a4]], %[[a1]] * 2 + %[[a5]], %[[a2]] * 2 + %[[a6]], %[[a3]] + %[[a7]]] :66// CHECK: %[[res:.*]] = arith.cmpf ogt, %[[lhs]], %[[rhs]] : f3267// CHECK: %[[sel:.*]] = arith.select %[[res]], %[[lhs]], %[[rhs]] : f3268// CHECK: affine.store %[[sel]], %[[v0]][%[[a0]], %[[a1]], %[[a2]], %[[a3]]] :69 70// CHECK-LABEL: func @symbols(71func.func @symbols(%N : index) {72 %0 = memref.alloc() : memref<1024x1024xf32>73 %1 = memref.alloc() : memref<1024x1024xf32>74 %2 = memref.alloc() : memref<1024x1024xf32>75 %cst1 = arith.constant 1 : index76 %cst2 = arith.constant 2 : index77 affine.for %i = 0 to %N {78 affine.for %j = 0 to %N {79 %7 = memref.load %2[%i, %j] : memref<1024x1024xf32>80 %10 = affine.for %k = 0 to %N iter_args(%ax = %cst1) -> index {81 %12 = arith.muli %N, %cst2 : index82 %13 = arith.addi %12, %cst1 : index83 %14 = arith.addi %13, %j : index84 %5 = memref.load %0[%i, %12] : memref<1024x1024xf32>85 %6 = memref.load %1[%14, %j] : memref<1024x1024xf32>86 %8 = arith.mulf %5, %6 : f3287 %9 = arith.addf %7, %8 : f3288 %4 = arith.addi %N, %cst1 : index89 %11 = arith.addi %ax, %cst1 : index90 memref.store %9, %2[%i, %4] : memref<1024x1024xf32> // this uses an expression of the symbol91 memref.store %9, %2[%i, %11] : memref<1024x1024xf32> // this uses an iter_args and cannot be raised92 %something = "ab.v"() : () -> index93 memref.store %9, %2[%i, %something] : memref<1024x1024xf32> // this cannot be raised94 affine.yield %11 : index95 }96 }97 }98 return99}100 101// CHECK: %[[cst1:.*]] = arith.constant 1 : index102// CHECK: %[[v0:.*]] = memref.alloc() : memref<103// CHECK: %[[v1:.*]] = memref.alloc() : memref<104// CHECK: %[[v2:.*]] = memref.alloc() : memref<105// CHECK: affine.for %[[a1:.*]] = 0 to %arg0 {106// CHECK: affine.for %[[a2:.*]] = 0 to %arg0 {107// CHECK: %[[lhs:.*]] = affine.load %{{.*}}[%[[a1]], %[[a2]]] : memref<1024x1024xf32>108// CHECK: affine.for %[[a3:.*]] = 0 to %arg0 iter_args(%[[a4:.*]] = %[[cst1]]) -> (index) {109// CHECK: %[[lhs2:.*]] = affine.load %{{.*}}[%[[a1]], symbol(%arg0) * 2] :110// CHECK: %[[lhs3:.*]] = affine.load %{{.*}}[%[[a2]] + symbol(%arg0) * 2 + 1, %[[a2]]] :111// CHECK: %[[lhs4:.*]] = arith.mulf %[[lhs2]], %[[lhs3]]112// CHECK: %[[lhs5:.*]] = arith.addf %[[lhs]], %[[lhs4]]113// CHECK: %[[lhs6:.*]] = arith.addi %[[a4]], %[[cst1]]114// CHECK: affine.store %[[lhs5]], %{{.*}}[%[[a1]], symbol(%arg0) + 1] :115// CHECK: memref.store %[[lhs5]], %{{.*}}[%[[a1]], %[[lhs6]]] :116// CHECK: %[[lhs7:.*]] = "ab.v"117// CHECK: memref.store %[[lhs5]], %{{.*}}[%[[a1]], %[[lhs7]]] :118// CHECK: affine.yield %[[lhs6]]119 120 121// CHECK-LABEL: func @non_affine(122func.func @non_affine(%N : index) {123 %2 = memref.alloc() : memref<1024x1024xf32>124 affine.for %i = 0 to %N {125 affine.for %j = 0 to %N {126 %ij = arith.muli %i, %j : index127 %7 = memref.load %2[%i, %ij] : memref<1024x1024xf32>128 memref.store %7, %2[%ij, %ij] : memref<1024x1024xf32>129 }130 }131 return132}133 134// CHECK: affine.for %[[i:.*]] =135// CHECK: affine.for %[[j:.*]] =136// CHECK: %[[ij:.*]] = arith.muli %[[i]], %[[j]]137// CHECK: %[[v:.*]] = memref.load %{{.*}}[%[[i]], %[[ij]]]138// CHECK: memref.store %[[v]], %{{.*}}[%[[ij]], %[[ij]]]139