// RUN: mlir-opt %s --transform-interpreter | FileCheck %s
func.func @transfer_read_rank_reducing(
%arg : memref<1x1x3x2xi8, strided<[6, 6, 2, 1], offset: ?>>) -> vector<3x2xi8> {
%c0 = arith.constant 0 : index
%cst = arith.constant 0 : i8
%v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst :
memref<1x1x3x2xi8, strided<[6, 6, 2, 1], offset: ?>>, vector<3x2xi8>
return %v : vector<3x2xi8>
}
// CHECK-LABEL: func @transfer_read_rank_reducing
// CHECK-SAME: %[[ARG:.+]]: memref<1x1x3x2xi8
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0] [1, 1, 3, 2] [1, 1, 1, 1]
// CHECK-SAME: memref<1x1x3x2xi8, {{.*}}> to memref<3x2xi8, {{.*}}>
// CHECK: vector.transfer_read %[[SUBVIEW]]
func.func @transfer_write_rank_reducing(%arg : memref<1x1x3x2xi8, strided<[6, 6, 2, 1], offset: ?>>, %vec : vector<3x2xi8>) {
%c0 = arith.constant 0 : index
vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0] :
vector<3x2xi8>, memref<1x1x3x2xi8, strided<[6, 6, 2, 1], offset: ?>>
return
}
// CHECK-LABEL: func @transfer_write_rank_reducing
// CHECK-SAME: %[[ARG:.+]]: memref<1x1x3x2xi8
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0] [1, 1, 3, 2] [1, 1, 1, 1]
// CHECK-SAME: memref<1x1x3x2xi8, {{.*}}> to memref<3x2xi8, {{.*}}>
// CHECK: vector.transfer_write %{{.*}}, %[[SUBVIEW]]
func.func @transfer_read_and_vector_rank_reducing(
%arg : memref<1x1x3x2x1xf32>) -> vector<3x2x1xf32> {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.0 : f32
%v = vector.transfer_read %arg[%c0, %c0, %c0, %c0, %c0], %cst :
memref<1x1x3x2x1xf32>, vector<3x2x1xf32>
return %v : vector<3x2x1xf32>
}
// CHECK-LABEL: func @transfer_read_and_vector_rank_reducing
// CHECK-SAME: %[[ARG:.+]]: memref<1x1x3x2x1xf32>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0, 0] [1, 1, 3, 2, 1] [1, 1, 1, 1, 1]
// CHECK-SAME: memref<1x1x3x2x1xf32> to memref<3x2xf32>
// CHECK: vector.transfer_read %[[SUBVIEW]]{{.*}} {in_bounds = [true, true]} : memref<3x2xf32>, vector<3x2xf32>
func.func @transfer_write_and_vector_rank_reducing(
%arg : memref<1x1x3x2x1xf32>,
%vec : vector<3x2x1xf32>) {
%c0 = arith.constant 0 : index
vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0, %c0] :
vector<3x2x1xf32>, memref<1x1x3x2x1xf32>
return
}
// CHECK-LABEL: func @transfer_write_and_vector_rank_reducing
// CHECK-SAME: %[[ARG:.+]]: memref<1x1x3x2x1xf32>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0, 0] [1, 1, 3, 2, 1] [1, 1, 1, 1, 1]
// CHECK-SAME: memref<1x1x3x2x1xf32> to memref<3x2xf32>
// CHECK: vector.transfer_write %{{.*}}, %[[SUBVIEW]]{{.*}} {in_bounds = [true, true]} : vector<3x2xf32>, memref<3x2xf32>
func.func @transfer_read_and_vector_rank_reducing_to_0d(
%arg : memref<1x1x1x1x1xf32>) -> vector<1x1x1xf32> {
%c0 = arith.constant 0 : index
%cst = arith.constant 0.0 : f32
%v = vector.transfer_read %arg[%c0, %c0, %c0, %c0, %c0], %cst :
memref<1x1x1x1x1xf32>, vector<1x1x1xf32>
return %v : vector<1x1x1xf32>
}
// CHECK-LABEL: func @transfer_read_and_vector_rank_reducing_to_0d
// CHECK-SAME: %[[MEMREF:.+]]: memref<1x1x1x1x1xf32>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[MEMREF]][0, 0, 0, 0, 0] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] : memref<1x1x1x1x1xf32> to memref<f32>
// CHECK: %[[READ:.+]] = vector.transfer_read %[[SUBVIEW]]{{.*}} : memref<f32>, vector<f32>
// CHECK: vector.shape_cast %[[READ]] : vector<f32> to vector<1x1x1xf32>
func.func @transfer_write_and_vector_rank_reducing_to_0d(
%arg : memref<1x1x1x1x1xf32>,
%vec : vector<1x1x1xf32>) {
%c0 = arith.constant 0 : index
vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0, %c0] :
vector<1x1x1xf32>, memref<1x1x1x1x1xf32>
return
}
// CHECK-LABEL: func @transfer_write_and_vector_rank_reducing_to_0d
// CHECK-SAME: %[[MEMREF:.+]]: memref<1x1x1x1x1xf32>, %[[VECTOR:.+]]: vector<1x1x1xf32>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[MEMREF]][0, 0, 0, 0, 0] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] : memref<1x1x1x1x1xf32> to memref<f32>
// CHECK: %[[SHCAST:.+]] = vector.shape_cast %[[VECTOR]] : vector<1x1x1xf32> to vector<f32>
// CHECK: vector.transfer_write %[[SHCAST]], %[[SUBVIEW]]{{.*}} : vector<f32>, memref<f32>
func.func @transfer_read_dynamic_rank_reducing(
%arg : memref<?x1xi8, strided<[?, ?], offset: ?>>) -> vector<[16]x1xi8> {
%c0 = arith.constant 0 : index
%pad = arith.constant 0 : i8
%v = vector.transfer_read %arg[%c0, %c0], %pad {in_bounds = [true, true]} :
memref<?x1xi8, strided<[?, ?], offset: ?>>, vector<[16]x1xi8>
return %v : vector<[16]x1xi8>
}
// CHECK-LABEL: func @transfer_read_dynamic_rank_reducing
// CHECK-SAME: %[[ARG:.+]]: memref<?x1xi8
// CHECK: %[[C0:.+]] = arith.constant 0 : index
// CHECK: %[[DIM0:.+]] = memref.dim %[[ARG]], %[[C0]] : memref<?x1xi8, strided<[?, ?], offset: ?>>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0] [%[[DIM0]], 1] [1, 1] : memref<?x1xi8, {{.*}}> to memref<?xi8, {{.*}}>
// CHECK: vector.transfer_read %[[SUBVIEW]]{{.*}} : memref<?xi8, {{.*}}>, vector<[16]xi8>
func.func @masked_transfer_read_dynamic_rank_reducing_1(
%arg : memref<?x1xi8, strided<[?, ?], offset: ?>>,
%mask_dim0 : index) -> vector<[16]x1xi8> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%pad = arith.constant 0 : i8
%mask = vector.create_mask %mask_dim0, %c1 : vector<[16]x1xi1>
%v = vector.transfer_read %arg[%c0, %c0], %pad, %mask {in_bounds = [true, true]} :
memref<?x1xi8, strided<[?, ?], offset: ?>>, vector<[16]x1xi8>
return %v : vector<[16]x1xi8>
}
// CHECK-LABEL: func @masked_transfer_read_dynamic_rank_reducing_1
// CHECK-SAME: %[[ARG:.+]]: memref<?x1xi8
// CHECK-SAME: %[[MASK_DIM0:.+]]: index
// CHECK: %[[C0:.+]] = arith.constant 0 : index
// CHECK: %[[PAD:.+]] = arith.constant 0 : i8
// CHECK: %[[MASK:.+]] = vector.create_mask %[[MASK_DIM0]] : vector<[16]xi1>
// CHECK: %[[DIM0:.+]] = memref.dim %[[ARG]], %[[C0]] : memref<?x1xi8, strided<[?, ?], offset: ?>>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0] [%[[DIM0]], 1] [1, 1] : memref<?x1xi8, {{.*}}> to memref<?xi8, {{.*}}>
// CHECK: vector.transfer_read %[[SUBVIEW]][{{.*}}], %[[PAD]], %[[MASK]] {in_bounds = [true]} : memref<?xi8, {{.*}}>, vector<[16]xi8>
func.func @masked_transfer_read_dynamic_rank_reducing_2(
%arg : memref<1x?x3x1x?x1xi8, strided<[?, ?, ?, ?, ?, ?], offset: ?>>,
%mask_dim1 : index, %mask_dim4 : index) -> vector<1x[1]x3x1x[16]x1xi8> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%pad = arith.constant 0 : i8
%mask = vector.create_mask %c1, %mask_dim1, %c2, %c1, %mask_dim4, %c1 : vector<1x[1]x3x1x[16]x1xi1>
%v = vector.transfer_read %arg[%c0, %c0, %c0, %c0, %c0, %c0], %pad, %mask {in_bounds = [true, true, true, true, true, true]} :
memref<1x?x3x1x?x1xi8, strided<[?, ?, ?, ?, ?, ?], offset: ?>>, vector<1x[1]x3x1x[16]x1xi8>
return %v : vector<1x[1]x3x1x[16]x1xi8>
}
// CHECK-LABEL: func @masked_transfer_read_dynamic_rank_reducing_2
// CHECK-SAME: %[[ARG:.+]]: memref<1x?x3x1x?x1xi8
// CHECK-SAME: %[[MASK_DIM1:.+]]: index, %[[MASK_DIM4:.+]]: index
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
// CHECK-DAG: %[[PAD:.+]] = arith.constant 0 : i8
// CHECK: %[[MASK:.+]] = vector.create_mask %[[MASK_DIM1]], %[[C2]], %[[MASK_DIM4]] : vector<[1]x3x[16]xi1>
// CHECK: %[[DIM1:.+]] = memref.dim %[[ARG]], %[[C1]] : memref<1x?x3x1x?x1xi8, strided<[?, ?, ?, ?, ?, ?], offset: ?>>
// CHECK: %[[DIM4:.+]] = memref.dim %[[ARG]], %[[C4]] : memref<1x?x3x1x?x1xi8, strided<[?, ?, ?, ?, ?, ?], offset: ?>>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0, 0, 0] [1, %[[DIM1]], 3, 1, %[[DIM4]], 1] [1, 1, 1, 1, 1, 1] : memref<1x?x3x1x?x1xi8, {{.*}}> to memref<?x3x?xi8, {{.*}}>
// CHECK: vector.transfer_read %[[SUBVIEW]][{{.*}}], %[[PAD]], %[[MASK]] {in_bounds = [true, true, true]} : memref<?x3x?xi8, {{.*}}>, vector<[1]x3x[16]xi8>
func.func @masked_transfer_write_and_vector_rank_reducing(
%arg : memref<1x1x3x1x16x1xf32>,
%vec : vector<1x3x1x16x1xf32>,
%mask_dim1 : index,
%mask_dim2 : index) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%mask = vector.create_mask %c1, %mask_dim1, %c1, %mask_dim2, %c1 : vector<1x3x1x16x1xi1>
vector.transfer_write %vec, %arg[%c0, %c0, %c0, %c0, %c0, %c0], %mask :
vector<1x3x1x16x1xf32>, memref<1x1x3x1x16x1xf32>
return
}
// CHECK-LABEL: func @masked_transfer_write_and_vector_rank_reducing
// CHECK-SAME: %[[ARG:.+]]: memref<1x1x3x1x16x1xf32>
// CHECK-SAME: {{.*}}: vector<1x3x1x16x1xf32>,
// CHECK-SAME: %[[MASKDIM1:.+]]: index,
// CHECK-SAME: %[[MASKDIM2:.+]]: index
// CHECK: %[[MASK:.+]] = vector.create_mask %[[MASKDIM1]], %[[MASKDIM2]] : vector<3x16xi1>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0, 0, 0] [1, 1, 3, 1, 16, 1] [1, 1, 1, 1, 1, 1]
// CHECK-SAME: memref<1x1x3x1x16x1xf32> to memref<3x16xf32>
// CHECK: vector.transfer_write %{{.*}}, %[[SUBVIEW]]{{.*}}, %[[MASK]] {in_bounds = [true, true]} : vector<3x16xf32>, memref<3x16xf32>
func.func @masked_transfer_write_dynamic_rank_reducing(
%arg : memref<?x1xi8, strided<[?, ?], offset: ?>>,
%vec : vector<[16]x1xi8>,
%mask_dim0 : index) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%pad = arith.constant 0 : i8
%mask = vector.create_mask %mask_dim0, %c1 : vector<[16]x1xi1>
vector.transfer_write %vec, %arg[%c0, %c0], %mask {in_bounds = [true, true]} :
vector<[16]x1xi8>, memref<?x1xi8, strided<[?, ?], offset: ?>>
return
}
// CHECK-LABEL: func @masked_transfer_write_dynamic_rank_reducing
// CHECK-SAME: %[[ARG:.+]]: memref<?x1xi8
// CHECK-SAME: %{{.*}}: vector<[16]x1xi8>,
// CHECK-SAME: %[[MASK_DIM0:.+]]: index
// CHECK: %[[C0:.+]] = arith.constant 0 : index
// CHECK: %[[MASK:.+]] = vector.create_mask %[[MASK_DIM0]] : vector<[16]xi1>
// CHECK: %[[DIM0:.+]] = memref.dim %[[ARG]], %[[C0]] : memref<?x1xi8, strided<[?, ?], offset: ?>>
// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0] [%[[DIM0]], 1] [1, 1] : memref<?x1xi8, {{.*}}> to memref<?xi8, {{.*}}>
// CHECK: vector.transfer_write {{.*}}, %[[SUBVIEW]][%[[C0]]], %[[MASK]] {in_bounds = [true]} : vector<[16]xi8>, memref<?xi8, {{.*}}>
/// Only masks operands of vector.create_mask are currently supported.
func.func @unsupported_masked_transfer_read_dynamic_rank_reducing_1(
%arg : memref<?x1xi8, strided<[?, ?], offset: ?>>,
%mask : vector<[16]x1xi1>) -> vector<[16]x1xi8> {
%c0 = arith.constant 0 : index
%pad = arith.constant 0 : i8
%v = vector.transfer_read %arg[%c0, %c0], %pad, %mask {in_bounds = [true, true]} :
memref<?x1xi8, strided<[?, ?], offset: ?>>, vector<[16]x1xi8>
return %v : vector<[16]x1xi8>
}
// CHECK-LABEL: func @unsupported_masked_transfer_read_dynamic_rank_reducing_1
// CHECK-SAME: %[[ARG:.+]]: memref<?x1xi8
// CHECK-NOT: vector.create_mask
// CHECK-NOT: memref.subview
// CHECK: vector.transfer_read %[[ARG]]
/// Unit dim mask must be constant of 1.
func.func @unsupported_masked_transfer_read_dynamic_rank_reducing_2(
%arg : memref<?x1xi8, strided<[?, ?], offset: ?>>,
%mask_dim0 : index, %mask_dim1 : index) -> vector<[16]x1xi8> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%pad = arith.constant 0 : i8
%mask = vector.create_mask %mask_dim0, %mask_dim1 : vector<[16]x1xi1>
%v = vector.transfer_read %arg[%c0, %c0], %pad, %mask {in_bounds = [true, true]} :
memref<?x1xi8, strided<[?, ?], offset: ?>>, vector<[16]x1xi8>
return %v : vector<[16]x1xi8>
}
// CHECK-LABEL: func @unsupported_masked_transfer_read_dynamic_rank_reducing_2
// CHECK-SAME: %[[ARG:.+]]: memref<?x1xi8
// CHECK-NOT: memref.subview
// CHECK: vector.transfer_read {{.*}} vector<[16]x1xi8>
/// Unit dim must be non-scalable.
func.func @masked_transfer_read_dynamic_rank_reducing_scalable_unit_dim(
%arg : memref<?x1xi8, strided<[?, ?], offset: ?>>,
%mask_dim0 : index) -> vector<[16]x[1]xi8> {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%pad = arith.constant 0 : i8
%mask = vector.create_mask %mask_dim0, %c1 : vector<[16]x[1]xi1>
%v = vector.transfer_read %arg[%c0, %c0], %pad, %mask {in_bounds = [true, true]} :
memref<?x1xi8, strided<[?, ?], offset: ?>>, vector<[16]x[1]xi8>
return %v : vector<[16]x[1]xi8>
}
// CHECK-LABEL: func @masked_transfer_read_dynamic_rank_reducing_scalable_unit_dim
// CHECK-SAME: %[[ARG:.+]]: memref<?x1xi8
// CHECK-NOT: memref.subview
// CHECK: vector.transfer_read {{.*}} vector<[16]x[1]xi8>
module attributes {transform.with_named_sequence} {
transform.named_sequence @__transform_main(%root : !transform.any_op {transform.readonly}) {
%func_op = transform.structured.match ops{["func.func"]} in %root : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op {
transform.apply_patterns.vector.rank_reducing_subview_patterns
} : !transform.op<"func.func">
transform.yield
}
}
Codebase Browser
llvm