130 lines · plain
1// RUN: mlir-opt %s -affine-expand-index-ops -split-input-file | FileCheck %s2 3// CHECK-LABEL: @delinearize_static_basis4// CHECK-SAME: (%[[IDX:.+]]: index)5// CHECK-DAG: %[[C224:.+]] = arith.constant 224 : index6// CHECK-DAG: %[[C50176:.+]] = arith.constant 50176 : index7// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index8// CHECK: %[[N:.+]] = arith.floordivsi %[[IDX]], %[[C50176]]9// CHECK-DAG: %[[P_REM:.+]] = arith.remsi %[[IDX]], %[[C50176]]10// CHECK-DAG: %[[P_NEG:.+]] = arith.cmpi slt, %[[P_REM]], %[[C0]]11// CHECK-DAG: %[[P_SHIFTED:.+]] = arith.addi %[[P_REM]], %[[C50176]] overflow<nsw>12// CHECK-DAG: %[[P_MOD:.+]] = arith.select %[[P_NEG]], %[[P_SHIFTED]], %[[P_REM]]13// CHECK: %[[P:.+]] = arith.divsi %[[P_MOD]], %[[C224]]14// CHECK-DAG: %[[Q_REM:.+]] = arith.remsi %[[IDX]], %[[C224]]15// CHECK-DAG: %[[Q_NEG:.+]] = arith.cmpi slt, %[[Q_REM]], %[[C0]]16// CHECK-DAG: %[[Q_SHIFTED:.+]] = arith.addi %[[Q_REM]], %[[C224]] overflow<nsw>17// CHECK: %[[Q:.+]] = arith.select %[[Q_NEG]], %[[Q_SHIFTED]], %[[Q_REM]]18// CHECK: return %[[N]], %[[P]], %[[Q]]19func.func @delinearize_static_basis(%linear_index: index) -> (index, index, index) {20 %1:3 = affine.delinearize_index %linear_index into (16, 224, 224) : index, index, index21 return %1#0, %1#1, %1#2 : index, index, index22}23 24// -----25 26// CHECK-LABEL: @delinearize_dynamic_basis27// CHECK-SAME: (%[[IDX:.+]]: index, %[[MEMREF:.+]]: memref28// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index29// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index30// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index31// CHECK: %[[DIM1:.+]] = memref.dim %[[MEMREF]], %[[C1]] :32// CHECK: %[[DIM2:.+]] = memref.dim %[[MEMREF]], %[[C2]] :33// CHECK: %[[STRIDE1:.+]] = arith.muli %[[DIM2]], %[[DIM1]] overflow<nsw, nuw>34// CHECK: %[[N:.+]] = arith.floordivsi %[[IDX]], %[[STRIDE1]]35// CHECK-DAG: %[[P_REM:.+]] = arith.remsi %[[IDX]], %[[STRIDE1]]36// CHECK-DAG: %[[P_NEG:.+]] = arith.cmpi slt, %[[P_REM]], %[[C0]]37// CHECK-DAG: %[[P_SHIFTED:.+]] = arith.addi %[[P_REM]], %[[STRIDE1]] overflow<nsw>38// CHECK-DAG: %[[P_MOD:.+]] = arith.select %[[P_NEG]], %[[P_SHIFTED]], %[[P_REM]]39// CHECK: %[[P:.+]] = arith.divsi %[[P_MOD]], %[[DIM2]]40// CHECK-DAG: %[[Q_REM:.+]] = arith.remsi %[[IDX]], %[[DIM2]]41// CHECK-DAG: %[[Q_NEG:.+]] = arith.cmpi slt, %[[Q_REM]], %[[C0]]42// CHECK-DAG: %[[Q_SHIFTED:.+]] = arith.addi %[[Q_REM]], %[[DIM2]] overflow<nsw>43// CHECK: %[[Q:.+]] = arith.select %[[Q_NEG]], %[[Q_SHIFTED]], %[[Q_REM]]44// CHECK: return %[[N]], %[[P]], %[[Q]]45func.func @delinearize_dynamic_basis(%linear_index: index, %src: memref<?x?x?xf32>) -> (index, index, index) {46 %c1 = arith.constant 1 : index47 %c2 = arith.constant 2 : index48 %b1 = memref.dim %src, %c1 : memref<?x?x?xf32>49 %b2 = memref.dim %src, %c2 : memref<?x?x?xf32>50 // Note: no outer bound.51 %1:3 = affine.delinearize_index %linear_index into (%b1, %b2) : index, index, index52 return %1#0, %1#1, %1#2 : index, index, index53}54 55// -----56 57// CHECK-LABEL: @linearize_static58// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index, %[[arg2:.+]]: index)59// CHECK-DAG: %[[C5:.+]] = arith.constant 5 : index60// CHECK-DAG: %[[C15:.+]] = arith.constant 15 : index61// CHECK: %[[scaled_0:.+]] = arith.muli %[[arg0]], %[[C15]] overflow<nsw>62// CHECK: %[[scaled_1:.+]] = arith.muli %[[arg1]], %[[C5]] overflow<nsw>63// CHECK: %[[val_0:.+]] = arith.addi %[[scaled_0]], %[[scaled_1]] overflow<nsw>64// CHECK: %[[val_1:.+]] = arith.addi %[[val_0]], %[[arg2]] overflow<nsw>65// CHECK: return %[[val_1]]66func.func @linearize_static(%arg0: index, %arg1: index, %arg2: index) -> index {67 %0 = affine.linearize_index [%arg0, %arg1, %arg2] by (2, 3, 5) : index68 func.return %0 : index69}70 71// -----72 73// CHECK-LABEL: @linearize_dynamic74// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index, %[[arg2:.+]]: index, %[[arg3:.+]]: index, %[[arg4:.+]]: index)75// CHECK: %[[stride_0:.+]] = arith.muli %[[arg4]], %[[arg3]] overflow<nsw>76// CHECK: %[[scaled_0:.+]] = arith.muli %[[arg0]], %[[stride_0]] overflow<nsw>77// CHECK: %[[scaled_1:.+]] = arith.muli %[[arg1]], %[[arg4]] overflow<nsw>78// CHECK: %[[val_0:.+]] = arith.addi %[[scaled_0]], %[[scaled_1]] overflow<nsw>79// CHECK: %[[val_1:.+]] = arith.addi %[[val_0]], %[[arg2]] overflow<nsw>80// CHECK: return %[[val_1]]81func.func @linearize_dynamic(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> index {82 // Note: no outer bounds83 %0 = affine.linearize_index [%arg0, %arg1, %arg2] by (%arg3, %arg4) : index84 func.return %0 : index85}86 87// -----88 89// CHECK-LABEL: @linearize_dynamic_disjoint90// CHECK-SAME: (%[[arg0:.+]]: index, %[[arg1:.+]]: index, %[[arg2:.+]]: index, %[[arg3:.+]]: index, %[[arg4:.+]]: index)91// CHECK: %[[stride_0:.+]] = arith.muli %[[arg4]], %[[arg3]] overflow<nsw, nuw>92// CHECK: %[[scaled_0:.+]] = arith.muli %[[arg0]], %[[stride_0]] overflow<nsw>93// CHECK: %[[scaled_1:.+]] = arith.muli %[[arg1]], %[[arg4]] overflow<nsw>94// CHECK: %[[val_0:.+]] = arith.addi %[[scaled_0]], %[[scaled_1]] overflow<nsw>95// CHECK: %[[val_1:.+]] = arith.addi %[[val_0]], %[[arg2]] overflow<nsw>96// CHECK: return %[[val_1]]97func.func @linearize_dynamic_disjoint(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index) -> index {98 // Note: no outer bounds99 %0 = affine.linearize_index disjoint [%arg0, %arg1, %arg2] by (%arg3, %arg4) : index100 func.return %0 : index101}102 103// -----104 105// CHECK-LABEL: @linearize_sort_adds106// CHECK-SAME: (%[[arg0:.+]]: memref<?xi32>, %[[arg1:.+]]: index, %[[arg2:.+]]: index)107// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index108// CHECK: scf.for %[[arg3:.+]] = %{{.*}} to %[[arg2]] step %{{.*}} {109// CHECK: scf.for %[[arg4:.+]] = %{{.*}} to %[[C4]] step %{{.*}} {110// CHECK: %[[stride_0:.+]] = arith.muli %[[arg2]], %[[C4]] overflow<nsw, nuw>111// CHECK: %[[scaled_0:.+]] = arith.muli %[[arg1]], %[[stride_0]] overflow<nsw>112// CHECK: %[[scaled_1:.+]] = arith.muli %[[arg4]], %[[arg2]] overflow<nsw>113// Note: even though %arg3 has a lower stride, we add it first114// CHECK: %[[val_0_2:.+]] = arith.addi %[[scaled_0]], %[[arg3]] overflow<nsw>115// CHECK: %[[val_1:.+]] = arith.addi %[[val_0_2]], %[[scaled_1]] overflow<nsw>116// CHECK: memref.store %{{.*}}, %[[arg0]][%[[val_1]]]117func.func @linearize_sort_adds(%arg0: memref<?xi32>, %arg1: index, %arg2: index) {118 %c0 = arith.constant 0 : index119 %c1 = arith.constant 1 : index120 %c4 = arith.constant 4 : index121 %c0_i32 = arith.constant 0 : i32122 scf.for %arg3 = %c0 to %arg2 step %c1 {123 scf.for %arg4 = %c0 to %c4 step %c1 {124 %idx = affine.linearize_index disjoint [%arg1, %arg4, %arg3] by (4, %arg2) : index125 memref.store %c0_i32, %arg0[%idx] : memref<?xi32>126 }127 }128 return129}130