blob: 2718f9d008058cb1524af310843ed71566901fce [file] [log] [blame]
//===----------------------------------------------------------------------===//
// Utility Methods
//===----------------------------------------------------------------------===//
func.func private @generate_2D_source_f16(%height : index, %width : index) -> tensor<?x?xf16> {
%init_source = tensor.empty(%height, %width) : tensor<?x?xf16>
%source = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]}
outs(%init_source : tensor<?x?xf16>) {
^bb0(%b0 : f16):
%outer = linalg.index 0 : index
%inner = linalg.index 1 : index
%strided = arith.muli %outer, %width : index
%linearized = arith.addi %inner, %strided : index
%linearized_i16 = arith.index_cast %linearized : index to i16
%linearized_f16 = arith.sitofp %linearized_i16 : i16 to f16
linalg.yield %linearized_f16 : f16
} -> tensor<?x?xf16>
// This blocks the fusion for inputs and testing ops.
%0 = util.optimization_barrier %source : tensor<?x?xf16>
%1 = flow.tensor.tie_shape %0 : tensor<?x?xf16>{%height, %width}
return %1 : tensor<?x?xf16>
}
func.func private @generate_2D_source_f32(%height : index, %width : index) -> tensor<?x?xf32> {
%init_source = tensor.empty(%height, %width) : tensor<?x?xf32>
%source = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]}
outs(%init_source : tensor<?x?xf32>) {
^bb0(%b0 : f32):
%outer = linalg.index 0 : index
%inner = linalg.index 1 : index
%strided = arith.muli %outer, %width : index
%linearized = arith.addi %inner, %strided : index
%linearized_i32 = arith.index_cast %linearized : index to i32
%linearized_f32 = arith.sitofp %linearized_i32 : i32 to f32
linalg.yield %linearized_f32 : f32
} -> tensor<?x?xf32>
// This blocks the fusion for inputs and testing ops.
%0 = util.optimization_barrier %source : tensor<?x?xf32>
%1 = flow.tensor.tie_shape %0 : tensor<?x?xf32>{%height, %width}
return %1 : tensor<?x?xf32>
}
func.func private @generate_2D_source_i8(%height : index, %width : index) -> tensor<?x?xi8> {
%c255 = arith.constant 255 : index
%init_source = tensor.empty(%height, %width) : tensor<?x?xi8>
%source = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]}
outs(%init_source : tensor<?x?xi8>) {
^bb0(%b0 : i8):
%outer = linalg.index 0 : index
%inner = linalg.index 1 : index
%strided = arith.muli %outer, %width : index
%linearized = arith.addi %inner, %strided : index
%linearized_rem = arith.remsi %linearized, %c255 : index
%linearized_i8 = arith.index_cast %linearized_rem : index to i8
linalg.yield %linearized_i8 : i8
} -> tensor<?x?xi8>
// This blocks the fusion for inputs and testing ops.
%0 = util.optimization_barrier %source : tensor<?x?xi8>
%1 = flow.tensor.tie_shape %0 : tensor<?x?xi8>{%height, %width}
return %1 : tensor<?x?xi8>
}
func.func private @generate_2D_source_i32(%height : index, %width : index) -> tensor<?x?xi32> {
%init_source = tensor.empty(%height, %width) : tensor<?x?xi32>
%source = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]}
outs(%init_source : tensor<?x?xi32>) {
^bb0(%b0 : i32):
%outer = linalg.index 0 : index
%inner = linalg.index 1 : index
%strided = arith.muli %outer, %width : index
%linearized = arith.addi %inner, %strided : index
%linearized_i32 = arith.index_cast %linearized : index to i32
linalg.yield %linearized_i32 : i32
} -> tensor<?x?xi32>
// This blocks the fusion for inputs and testing ops.
%0 = util.optimization_barrier %source : tensor<?x?xi32>
%1 = flow.tensor.tie_shape %0 : tensor<?x?xi32>{%height, %width}
return %1 : tensor<?x?xi32>
}
//===----------------------------------------------------------------------===//
// f32.f32.f32 variants
//===----------------------------------------------------------------------===//
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
#encoding_f32f32f32_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_f32f32f32_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_f32f32f32_acc = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
func.func @set_encoding_f32f32f32_lhs() {
%height = arith.constant 129 : index
%width = arith.constant 255 : index
%0 = call @generate_2D_source_f32(%height, %width) : (index, index) -> tensor<?x?xf32>
%source = tensor.cast %0 : tensor<?x?xf32> to tensor<129x255xf32>
%1 = iree_encoding.set_encoding %source : tensor<129x255xf32> -> tensor<129x255xf32, #encoding_f32f32f32_lhs>
%barrire = util.optimization_barrier %1 : tensor<129x255xf32, #encoding_f32f32f32_lhs>
%2 = iree_encoding.unset_encoding %1 : tensor<129x255xf32, #encoding_f32f32f32_lhs> -> tensor<129x255xf32>
check.expect_almost_eq(%2, %source) : tensor<129x255xf32>
return
}
func.func @set_encoding_f32f32f32_rhs() {
%height = arith.constant 129 : index
%width = arith.constant 255 : index
%0 = call @generate_2D_source_f32(%height, %width) : (index, index) -> tensor<?x?xf32>
%source = tensor.cast %0 : tensor<?x?xf32> to tensor<129x255xf32>
%1 = iree_encoding.set_encoding %source : tensor<129x255xf32> -> tensor<129x255xf32, #encoding_f32f32f32_rhs>
%barrire = util.optimization_barrier %1 : tensor<129x255xf32, #encoding_f32f32f32_rhs>
%2 = iree_encoding.unset_encoding %1 : tensor<129x255xf32, #encoding_f32f32f32_rhs> -> tensor<129x255xf32>
check.expect_almost_eq(%2, %source) : tensor<129x255xf32>
return
}
func.func @set_encoding_f32f32f32_acc() {
%height = arith.constant 129 : index
%width = arith.constant 255 : index
%0 = call @generate_2D_source_f32(%height, %width) : (index, index) -> tensor<?x?xf32>
%source = tensor.cast %0 : tensor<?x?xf32> to tensor<129x255xf32>
%1 = iree_encoding.set_encoding %source : tensor<129x255xf32> -> tensor<129x255xf32, #encoding_f32f32f32_acc>
%barrire = util.optimization_barrier %1 : tensor<129x255xf32, #encoding_f32f32f32_acc>
%2 = iree_encoding.unset_encoding %1 : tensor<129x255xf32, #encoding_f32f32f32_acc> -> tensor<129x255xf32>
check.expect_almost_eq(%2, %source) : tensor<129x255xf32>
return
}
//===----------------------------------------------------------------------===//
// i8.i8.i32 variants
//===----------------------------------------------------------------------===//
#encoding_i8i8i32_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_i8i8i32_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_i8i8i32_acc = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
func.func @set_encoding_i8i8i32_lhs() {
%height = arith.constant 129 : index
%width = arith.constant 255 : index
%0 = call @generate_2D_source_i8(%height, %width) : (index, index) -> tensor<?x?xi8>
%source = tensor.cast %0 : tensor<?x?xi8> to tensor<129x255xi8>
%1 = iree_encoding.set_encoding %source : tensor<129x255xi8> -> tensor<129x255xi8, #encoding_i8i8i32_lhs>
%barrire = util.optimization_barrier %1 : tensor<129x255xi8, #encoding_i8i8i32_lhs>
%2 = iree_encoding.unset_encoding %1 : tensor<129x255xi8, #encoding_i8i8i32_lhs> -> tensor<129x255xi8>
check.expect_eq(%2, %source) : tensor<129x255xi8>
return
}
func.func @set_encoding_i8i8i32_rhs() {
%height = arith.constant 129 : index
%width = arith.constant 255 : index
%0 = call @generate_2D_source_i8(%height, %width) : (index, index) -> tensor<?x?xi8>
%source = tensor.cast %0 : tensor<?x?xi8> to tensor<129x255xi8>
%1 = iree_encoding.set_encoding %source : tensor<129x255xi8> -> tensor<129x255xi8, #encoding_i8i8i32_rhs>
%barrire = util.optimization_barrier %1 : tensor<129x255xi8, #encoding_i8i8i32_rhs>
%2 = iree_encoding.unset_encoding %1 : tensor<129x255xi8, #encoding_i8i8i32_rhs> -> tensor<129x255xi8>
check.expect_eq(%2, %source) : tensor<129x255xi8>
return
}
func.func @set_encoding_i8i8i32_acc() {
%height = arith.constant 129 : index
%width = arith.constant 255 : index
%0 = call @generate_2D_source_i32(%height, %width) : (index, index) -> tensor<?x?xi32>
%source = tensor.cast %0 : tensor<?x?xi32> to tensor<129x255xi32>
%1 = iree_encoding.set_encoding %source : tensor<129x255xi32> -> tensor<129x255xi32, #encoding_i8i8i32_acc>
%barrire = util.optimization_barrier %1 : tensor<129x255xi32, #encoding_i8i8i32_acc>
%2 = iree_encoding.unset_encoding %1 : tensor<129x255xi32, #encoding_i8i8i32_acc> -> tensor<129x255xi32>
check.expect_eq(%2, %source) : tensor<129x255xi32>
return
}
//===----------------------------------------------------------------------===//
// f16.f16.f32 variants
//===----------------------------------------------------------------------===//
#encoding_f16f16f32_lhs = #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_f16f16f32_rhs = #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
#encoding_f16f16f32_acc = #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f16, f16, f32], user_indexing_maps = [#map, #map1, #map2], round_dims_to = array<i64: 32, 32, 32>>
func.func @set_encoding_f16f16f32_lhs() {
%height = arith.constant 129 : index
%width = arith.constant 255 : index
%0 = call @generate_2D_source_f16(%height, %width) : (index, index) -> tensor<?x?xf16>
%source = tensor.cast %0 : tensor<?x?xf16> to tensor<129x255xf16>
%1 = iree_encoding.set_encoding %source : tensor<129x255xf16> -> tensor<129x255xf16, #encoding_f16f16f32_lhs>
%barrire = util.optimization_barrier %1 : tensor<129x255xf16, #encoding_f16f16f32_lhs>
%2 = iree_encoding.unset_encoding %1 : tensor<129x255xf16, #encoding_f16f16f32_lhs> -> tensor<129x255xf16>
check.expect_eq(%2, %source) : tensor<129x255xf16>
return
}
func.func @set_encoding_f16f16f32_rhs() {
%height = arith.constant 129 : index
%width = arith.constant 255 : index
%0 = call @generate_2D_source_f16(%height, %width) : (index, index) -> tensor<?x?xf16>
%source = tensor.cast %0 : tensor<?x?xf16> to tensor<129x255xf16>
%1 = iree_encoding.set_encoding %source : tensor<129x255xf16> -> tensor<129x255xf16, #encoding_f16f16f32_rhs>
%barrire = util.optimization_barrier %1 : tensor<129x255xf16, #encoding_f16f16f32_rhs>
%2 = iree_encoding.unset_encoding %1 : tensor<129x255xf16, #encoding_f16f16f32_rhs> -> tensor<129x255xf16>
check.expect_eq(%2, %source) : tensor<129x255xf16>
return
}
func.func @set_encoding_f16f16f32_acc() {
%height = arith.constant 129 : index
%width = arith.constant 255 : index
%0 = call @generate_2D_source_f32(%height, %width) : (index, index) -> tensor<?x?xf32>
%source = tensor.cast %0 : tensor<?x?xf32> to tensor<129x255xf32>
%1 = iree_encoding.set_encoding %source : tensor<129x255xf32> -> tensor<129x255xf32, #encoding_f16f16f32_acc>
%barrire = util.optimization_barrier %1 : tensor<129x255xf32, #encoding_f16f16f32_acc>
%2 = iree_encoding.unset_encoding %1 : tensor<129x255xf32, #encoding_f16f16f32_acc> -> tensor<129x255xf32>
check.expect_eq(%2, %source) : tensor<129x255xf32>
return
}