71 lines · plain
1// RUN: mlir-opt %s -affine-expand-index-ops-as-affine -split-input-file | FileCheck %s2 3// CHECK-DAG: #[[$map0:.+]] = affine_map<()[s0] -> (s0 floordiv 50176)>4// CHECK-DAG: #[[$map1:.+]] = affine_map<()[s0] -> ((s0 mod 50176) floordiv 224)>5// CHECK-DAG: #[[$map2:.+]] = affine_map<()[s0] -> (s0 mod 224)>6 7// CHECK-LABEL: @static_basis8// CHECK-SAME: (%[[IDX:.+]]: index)9// CHECK: %[[N:.+]] = affine.apply #[[$map0]]()[%[[IDX]]]10// CHECK: %[[P:.+]] = affine.apply #[[$map1]]()[%[[IDX]]]11// CHECK: %[[Q:.+]] = affine.apply #[[$map2]]()[%[[IDX]]]12// CHECK: return %[[N]], %[[P]], %[[Q]]13func.func @static_basis(%linear_index: index) -> (index, index, index) {14 %1:3 = affine.delinearize_index %linear_index into (16, 224, 224) : index, index, index15 return %1#0, %1#1, %1#2 : index, index, index16}17 18// -----19 20// CHECK-DAG: #[[$map0:.+]] = affine_map<()[s0, s1, s2] -> (s2 floordiv (s0 * s1))>21// CHECK-DAG: #[[$map1:.+]] = affine_map<()[s0, s1, s2] -> ((s2 mod (s0 * s1)) floordiv s1)>22// CHECK-DAG: #[[$map2:.+]] = affine_map<()[s0, s1, s2] -> ((s2 mod (s0 * s1)) mod s1)>23 24// CHECK-LABEL: @dynamic_basis25// CHECK-SAME: (%[[IDX:.+]]: index, %[[MEMREF:.+]]: memref26// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index27// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index28// CHECK: %[[DIM1:.+]] = memref.dim %[[MEMREF]], %[[C1]] :29// CHECK: %[[DIM2:.+]] = memref.dim %[[MEMREF]], %[[C2]] :30// CHECK: %[[N:.+]] = affine.apply #[[$map0]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]31// CHECK: %[[P:.+]] = affine.apply #[[$map1]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]32// CHECK: %[[Q:.+]] = affine.apply #[[$map2]]()[%[[DIM1]], %[[DIM2]], %[[IDX]]]33// CHECK: return %[[N]], %[[P]], %[[Q]]34func.func @dynamic_basis(%linear_index: index, %src: memref<?x?x?xf32>) -> (index, index, index) {35 %c0 = arith.constant 0 : index36 %c1 = arith.constant 1 : index37 %c2 = arith.constant 2 : index38 %b1 = memref.dim %src, %c1 : memref<?x?x?xf32>39 %b2 = memref.dim %src, %c2 : memref<?x?x?xf32>40 // Note: no outer bound.41 %1:3 = affine.delinearize_index %linear_index into (%b1, %b2) : index, index, index42 return %1#0, %1#1, %1#2 : index, index, index43}44 45// -----46 47// CHECK-DAG: #[[$map0:.+]] = affine_map<()[s0, s1, s2] -> (s0 * 15 + s1 * 5 + s2)>48 49// CHECK-LABEL: @linearize_static50// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index, %[[arg2:.+]]: index)51// CHECK: %[[val_0:.+]] = affine.apply #[[$map0]]()[%[[arg0]], %[[arg1]], %[[arg2]]]52// CHECK: return %[[val_0]]53func.func @linearize_static(%arg0: index, %arg1: index, %arg2: index) -> index {54 %0 = affine.linearize_index [%arg0, %arg1, %arg2] by (2, 3, 5) : index55 func.return %0 : index56}57 58// -----59 60// CHECK-DAG: #[[$map0:.+]] = affine_map<()[s0, s1, s2, s3, s4] -> (s1 * s2 + s3 + s0 * (s2 * s4))>61 62// CHECK-LABEL: @linearize_dynamic63// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index, %[[arg2:.+]]: index, %[[arg3:.+]]: index, %[[arg4:.+]]: index)64// CHECK: %[[val_0:.+]] = affine.apply #[[$map0]]()[%[[arg0]], %[[arg1]], %[[arg4]], %[[arg2]], %[[arg3]]]65// CHECK: return %[[val_0]]66func.func @linearize_dynamic(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> index {67 // Note: no outer bounds68 %0 = affine.linearize_index [%arg0, %arg1, %arg2] by (%arg3, %arg4) : index69 func.return %0 : index70}71