vectorize_affine_apply.mlir | Explore in Territory

// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=8 test-fastest-varying=0" -split-input-file | FileCheck %s

// CHECK-DAG: #[[$MAP_ID0:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 12)>
// CHECK-DAG: #[[$MAP_ID1:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 16)>

// CHECK-LABEL: vec_affine_apply
// CHECK-SAME:  (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>) {
func.func @vec_affine_apply(%arg0: memref<8x12x16xf32>, %arg1: memref<8x24x48xf32>) {
// CHECK:       affine.for %[[ARG2:.*]] = 0 to 8 {
// CHECK-NEXT:    affine.for %[[ARG3:.*]] = 0 to 24 {
// CHECK-NEXT:      affine.for %[[ARG4:.*]] = 0 to 48 step 8 {
// CHECK-NEXT:        %[[S0:.*]] = affine.apply #[[$MAP_ID0]](%[[ARG3]])
// CHECK-NEXT:        %[[S1:.*]] = affine.apply #[[$MAP_ID1]](%[[ARG4]])
// CHECK-NEXT:        %[[CST:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-NEXT:        %[[S2:.*]] = vector.transfer_read %[[ARG0]][%[[ARG2]], %[[S0]], %[[S1]]], %[[CST]] : memref<8x12x16xf32>, vector<8xf32>
// CHECK-NEXT:        vector.transfer_write %[[S2]], %[[ARG1]][%[[ARG2]], %[[ARG3]], %[[ARG4]]] : vector<8xf32>, memref<8x24x48xf32>
// CHECK-NEXT:      }
// CHECK-NEXT:    }
// CHECK-NEXT:  }
// CHECK-NEXT:  return
  affine.for %arg2 = 0 to 8 {
    affine.for %arg3 = 0 to 24 {
      affine.for %arg4 = 0 to 48 {
        %0 = affine.apply affine_map<(d0) -> (d0 mod 12)>(%arg3)
        %1 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%arg4)
        %2 = affine.load %arg0[%arg2, %0, %1] : memref<8x12x16xf32>
        affine.store %2, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xf32>
      }
    }
  }
  return
}

// -----

// CHECK-DAG: #[[$MAP_ID2:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 16 + 1)>

// CHECK-LABEL: vec_affine_apply_2
// CHECK-SAME:  (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>) {
func.func @vec_affine_apply_2(%arg0: memref<8x12x16xf32>, %arg1: memref<8x24x48xf32>) {
// CHECK:      affine.for %[[ARG2:.*]] = 0 to 8 {
// CHECK-NEXT:   affine.for %[[ARG3:.*]] = 0 to 12 {
// CHECK-NEXT:     affine.for %[[ARG4:.*]] = 0 to 48 step 8 {
// CHECK-NEXT:       %[[S0:.*]] = affine.apply #[[$MAP_ID2]](%[[ARG4]])
// CHECK-NEXT:       %[[CST:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-NEXT:       %[[S1:.*]] = vector.transfer_read %[[ARG0]][%[[ARG2]], %[[ARG3]], %[[S0]]], %[[CST]] : memref<8x12x16xf32>, vector<8xf32>
// CHECK-NEXT:       vector.transfer_write %[[S1]], %[[ARG1]][%[[ARG2]], %[[ARG3]], %[[ARG4]]] : vector<8xf32>, memref<8x24x48xf32>
// CHECK-NEXT:     }
// CHECK-NEXT:   }
// CHECK-NEXT: }
  affine.for %arg2 = 0 to 8 {
    affine.for %arg3 = 0 to 12 {
      affine.for %arg4 = 0 to 48 {
        %1 = affine.apply affine_map<(d0) -> (d0 mod 16 + 1)>(%arg4)
        %2 = affine.load %arg0[%arg2, %arg3, %1] : memref<8x12x16xf32>
        affine.store %2, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xf32>
      }
    }
  }
  return
}

// -----

// CHECK-LABEL: no_vec_affine_apply
// CHECK-SAME:  (%[[ARG0:.*]]: memref<8x12x16xi32>, %[[ARG1:.*]]: memref<8x24x48xi32>) {
func.func @no_vec_affine_apply(%arg0: memref<8x12x16xi32>, %arg1: memref<8x24x48xi32>) {
// CHECK:      affine.for %[[ARG2:.*]] = 0 to 8 {
// CHECK-NEXT:   affine.for %[[ARG3:.*]] = 0 to 24 {
// CHECK-NEXT:     affine.for %[[ARG4:.*]] = 0 to 48 {
// CHECK-NEXT:       %[[S0:.*]] = affine.apply #[[$MAP_ID0]](%[[ARG3]])
// CHECK-NEXT:       %[[S1:.*]] = affine.apply #[[$MAP_ID1]](%[[ARG4]])
// CHECK-NEXT:       %[[S2:.*]] = affine.load %[[ARG0]][%[[ARG2]], %[[S0]], %[[S1]]] : memref<8x12x16xi32>
// CHECK-NEXT:       %[[S3:.*]] = arith.index_cast %[[S2]] : i32 to index
// CHECK-NEXT:       %[[S4:.*]] = affine.apply #[[$MAP_ID1]](%[[S3]])
// CHECK-NEXT:       %[[S5:.*]] = arith.index_cast %[[S4]] : index to i32
// CHECK-NEXT:       affine.store %[[S5]], %[[ARG1]][%[[ARG2]], %[[ARG3]], %[[ARG4]]] : memref<8x24x48xi32>
// CHECK-NEXT:     }
// CHECK-NEXT:   }
// CHECK-NEXT: }
// CHECK-NEXT: return
  affine.for %arg2 = 0 to 8 {
    affine.for %arg3 = 0 to 24 {
      affine.for %arg4 = 0 to 48 {
        %0 = affine.apply affine_map<(d0) -> (d0 mod 12)>(%arg3)
        %1 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%arg4)
        %2 = affine.load %arg0[%arg2, %0, %1] : memref<8x12x16xi32>
        %3 = arith.index_cast %2 : i32 to index
        %4 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%3)
        %5 = arith.index_cast %4 : index to i32
        affine.store %5, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xi32>
      }
    }
  }
  return
}

// -----

// CHECK-DAG: #[[$MAP_ID1:map[0-9a-zA-Z_]*]] = affine_map<(d0) -> (d0 mod 16)>

// CHECK-LABEL: affine_map_with_expr
// CHECK-SAME:  (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>) {
func.func @affine_map_with_expr(%arg0: memref<8x12x16xf32>, %arg1: memref<8x24x48xf32>) {
// CHECK:      affine.for %[[ARG2:.*]] = 0 to 8 {
// CHECK-NEXT:   affine.for %[[ARG3:.*]] = 0 to 12 {
// CHECK-NEXT:     affine.for %[[ARG4:.*]] = 0 to 48 {
// CHECK-NEXT:       %[[S0:.*]] = affine.apply #[[$MAP_ID1]](%[[ARG4]])
// CHECK-NEXT:       %[[S1:.*]] = affine.load %[[ARG0]][%[[ARG2]], %[[ARG3]], %[[S0]] + 1] : memref<8x12x16xf32>
// CHECK-NEXT:       affine.store %[[S1]], %[[ARG1]][%[[ARG2]], %[[ARG3]], %[[ARG4]]] : memref<8x24x48xf32>
// CHECK-NEXT:     }
// CHECK-NEXT:   }
// CHECK-NEXT: }
// CHECK-NEXT: return
  affine.for %arg2 = 0 to 8 {
    affine.for %arg3 = 0 to 12 {
      affine.for %arg4 = 0 to 48 {
        %1 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%arg4)
        %2 = affine.load %arg0[%arg2, %arg3, %1 + 1] : memref<8x12x16xf32>
        affine.store %2, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xf32>
      }
    }
  }
  return
}

// -----

// CHECK-DAG: #[[$MAP_ID3:map[0-9a-zA-Z_]*]] = affine_map<(d0, d1, d2) -> (d0)>
// CHECK-DAG: #[[$MAP_ID4:map[0-9a-zA-Z_]*]] = affine_map<(d0, d1, d2) -> (d1)>
// CHECK-DAG: #[[$MAP_ID5:map[0-9a-zA-Z_]*]] = affine_map<(d0, d1, d2) -> (d2 + 1)>
// CHECK-DAG: #[[$MAP_ID6:map[0-9a-zA-Z_]*]] = affine_map<(d0, d1, d2) -> (0)>

// CHECK-LABEL: affine_map_with_expr_2
// CHECK-SAME:  (%[[ARG0:.*]]: memref<8x12x16xf32>, %[[ARG1:.*]]: memref<8x24x48xf32>, %[[I0:.*]]: index) {
func.func @affine_map_with_expr_2(%arg0: memref<8x12x16xf32>, %arg1: memref<8x24x48xf32>, %i: index) {
// CHECK:      affine.for %[[ARG3:.*]] = 0 to 8 {
// CHECK-NEXT:   affine.for %[[ARG4:.*]] = 0 to 12 {
// CHECK-NEXT:     affine.for %[[ARG5:.*]] = 0 to 48 step 8 {
// CHECK-NEXT:       %[[S0:.*]] = affine.apply #[[$MAP_ID3]](%[[ARG3]], %[[ARG4]], %[[I0]])
// CHECK-NEXT:       %[[S1:.*]] = affine.apply #[[$MAP_ID4]](%[[ARG3]], %[[ARG4]], %[[I0]])
// CHECK-NEXT:       %[[S2:.*]] = affine.apply #[[$MAP_ID5]](%[[ARG3]], %[[ARG4]], %[[I0]])
// CHECK-NEXT:       %[[CST:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-NEXT:       %[[S3:.*]] = vector.transfer_read %[[ARG0]][%[[S0]], %[[S1]], %[[S2]]], %[[CST]] {permutation_map = #[[$MAP_ID6]]} : memref<8x12x16xf32>, vector<8xf32>
// CHECK-NEXT:       vector.transfer_write %[[S3]], %[[ARG1]][%[[ARG3]], %[[ARG4]], %[[ARG5]]] : vector<8xf32>, memref<8x24x48xf32>
// CHECK-NEXT:     }
// CHECK-NEXT:   }
// CHECK-NEXT: }
// CHECK-NEXT: return
  affine.for %arg2 = 0 to 8 {
    affine.for %arg3 = 0 to 12 {
      affine.for %arg4 = 0 to 48 {
        %2 = affine.load %arg0[%arg2, %arg3, %i + 1] : memref<8x12x16xf32>
        affine.store %2, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xf32>
      }
    }
  }
  return
}
llvm/mlir/test/Dialect/Affine/SuperVectorize/vectorize_affine_apply.mlir
