shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity
llvm-project/mlir/test/Dialect/Vector/vector-contract-matvec-transforms.mlir
Praveen Narayanan a65fb1916c Add a "kind" attribute to ContractionOp and OuterProductOp. 
Currently, vector.contract joins the intermediate result and the accumulator
argument (of ranks K) using summation. We desire more joining operations ---
such as max --- to help vector.contract express reductions. This change extends
Vector_ContractionOp to take an optional attribute (called "kind", of enum type
CombiningKind) specifying the joining operation to be add/mul/min/max for int/fp
, and and/or/xor for int only. By default this attribute has value "add".

To implement this we also need to extend vector.outerproduct, since
vector.contract gets transformed to vector.outerproduct (and that to
vector.fma). The extension for vector.outerproduct is also an optional kind
attribute that uses the same enum type and possible values. The default is
"add". In case of max/min we transform vector.outerproduct to a combination of
compare and select.

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D93280
2021-02-12 20:23:59 +00:00

175 lines
8.2 KiB
MLIR
Raw Blame History

// RUN: mlir-opt %s -test-vector-contraction-conversion=vector-outerproduct=1 | FileCheck %s
#matvec_accesses = [
affine_map<(i, j) -> (i, j)>,
affine_map<(i, j) -> (j)>,
affine_map<(i, j) -> (i)>
]
#matvec_trait = {
indexing_maps = #matvec_accesses,
iterator_types = ["parallel", "reduction"]
}
#matvecmax_trait = {
indexing_maps = #matvec_accesses,
iterator_types = ["parallel", "reduction"],
kind = #vector.kind<max>
}
#mattransvec_accesses = [
affine_map<(i, j) -> (j, i)>,
affine_map<(i, j) -> (j)>,
affine_map<(i, j) -> (i)>
]
#mattransvec_trait = {
indexing_maps = #mattransvec_accesses,
iterator_types = ["parallel", "reduction"]
}
#vecmat_accesses = [
affine_map<(i, j) -> (j)>,
affine_map<(i, j) -> (i, j)>,
affine_map<(i, j) -> (i)>
]
#vecmat_trait = {
indexing_maps = #vecmat_accesses,
iterator_types = ["parallel", "reduction"]
}
#vecmattrans_accesses = [
affine_map<(i, j) -> (j)>,
affine_map<(i, j) -> (j, i)>,
affine_map<(i, j) -> (i)>
]
#vecmattrans_trait = {
indexing_maps = #vecmattrans_accesses,
iterator_types = ["parallel", "reduction"]
}
// CHECK-LABEL: func @matvec2x2
// CHECK-SAME: %[[A:.*0]]: memref<vector<2x2xf32>>
// CHECK-SAME: %[[B:.*1]]: memref<vector<2xf32>>
// CHECK-SAME: %[[C:.*2]]: memref<vector<2xf32>>
// CHECK: %[[T0:.*]] = load %[[A]][] : memref<vector<2x2xf32>>
// CHECK: %[[T1:.*]] = load %[[B]][] : memref<vector<2xf32>>
// CHECK: %[[T2:.*]] = load %[[C]][] : memref<vector<2xf32>>
// CHECK: %[[T3:.*]] = vector.transpose %[[T0]], [1, 0] : vector<2x2xf32> to vector<2x2xf32>
// CHECK: %[[T4:.*]] = vector.extract %[[T3]][0] : vector<2x2xf32>
// CHECK: %[[T5:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
// CHECK: %[[T6:.*]] = vector.outerproduct %[[T4]], %[[T5]], %[[T2]] {kind = #vector.kind<add>} : vector<2xf32>, f32
// CHECK: %[[T7:.*]] = vector.extract %[[T3]][1] : vector<2x2xf32>
// CHECK: %[[T8:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
// CHECK: %[[T9:.*]] = vector.outerproduct %[[T7]], %[[T8]], %[[T6]] {kind = #vector.kind<add>} : vector<2xf32>, f32
// CHECK: store %[[T9]], %[[C]][] : memref<vector<2xf32>>
// CHECK: return
func @matvec2x2(%arg0: memref<vector<2x2xf32>>, %arg1: memref<vector<2xf32>>,
%arg2: memref<vector<2xf32>>) {
%A = load %arg0[] : memref<vector<2x2xf32>>
%x = load %arg1[] : memref<vector<2xf32>>
%b = load %arg2[] : memref<vector<2xf32>>
%0 = vector.contract #matvec_trait %A, %x, %b : vector<2x2xf32>, vector<2xf32> into vector<2xf32>
store %0, %arg2[] : memref<vector<2xf32>>
return
}
// CHECK-LABEL: func @matvecmax2x2
// CHECK-SAME: %[[A:.*0]]: memref<vector<2x2xf32>>
// CHECK-SAME: %[[B:.*1]]: memref<vector<2xf32>>
// CHECK-SAME: %[[C:.*2]]: memref<vector<2xf32>>
// CHECK: %[[T0:.*]] = load %[[A]][] : memref<vector<2x2xf32>>
// CHECK: %[[T1:.*]] = load %[[B]][] : memref<vector<2xf32>>
// CHECK: %[[T2:.*]] = load %[[C]][] : memref<vector<2xf32>>
// CHECK: %[[T3:.*]] = vector.transpose %[[T0]], [1, 0] : vector<2x2xf32> to vector<2x2xf32>
// CHECK: %[[T4:.*]] = vector.extract %[[T3]][0] : vector<2x2xf32>
// CHECK: %[[T5:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
// CHECK: %[[T6:.*]] = vector.outerproduct %[[T4]], %[[T5]], %[[T2]] {kind = #vector.kind<max>} : vector<2xf32>, f32
// CHECK: %[[T7:.*]] = vector.extract %[[T3]][1] : vector<2x2xf32>
// CHECK: %[[T8:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
// CHECK: %[[T9:.*]] = vector.outerproduct %[[T7]], %[[T8]], %[[T6]] {kind = #vector.kind<max>} : vector<2xf32>, f32
// CHECK: store %[[T9]], %[[C]][] : memref<vector<2xf32>>
// CHECK: return
func @matvecmax2x2(%arg0: memref<vector<2x2xf32>>, %arg1: memref<vector<2xf32>>,
%arg2: memref<vector<2xf32>>) {
%A = load %arg0[] : memref<vector<2x2xf32>>
%x = load %arg1[] : memref<vector<2xf32>>
%b = load %arg2[] : memref<vector<2xf32>>
%0 = vector.contract #matvecmax_trait %A, %x, %b : vector<2x2xf32>, vector<2xf32> into vector<2xf32>
store %0, %arg2[] : memref<vector<2xf32>>
return
}
// CHECK-LABEL: func @mattransvec2x2
// CHECK-SAME: %[[A:.*0]]: memref<vector<2x2xf32>>
// CHECK-SAME: %[[B:.*1]]: memref<vector<2xf32>>
// CHECK-SAME: %[[C:.*2]]: memref<vector<2xf32>>
// CHECK: %[[T0:.*]] = load %[[A]][] : memref<vector<2x2xf32>>
// CHECK: %[[T1:.*]] = load %[[B]][] : memref<vector<2xf32>>
// CHECK: %[[T2:.*]] = load %[[C]][] : memref<vector<2xf32>>
// CHECK: %[[T3:.*]] = vector.extract %[[T0]][0] : vector<2x2xf32>
// CHECK: %[[T4:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
// CHECK: %[[T5:.*]] = vector.outerproduct %[[T3]], %[[T4]], %[[T2]] {kind = #vector.kind<add>} : vector<2xf32>, f32
// CHECK: %[[T6:.*]] = vector.extract %[[T0]][1] : vector<2x2xf32>
// CHECK: %[[T7:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
// CHECK: %[[T8:.*]] = vector.outerproduct %[[T6]], %[[T7]], %[[T5]] {kind = #vector.kind<add>} : vector<2xf32>, f32
// CHECK: store %[[T8]], %[[C]][] : memref<vector<2xf32>>
// CHECK: return
func @mattransvec2x2(%arg0: memref<vector<2x2xf32>>, %arg1: memref<vector<2xf32>>,
%arg2: memref<vector<2xf32>>) {
%A = load %arg0[] : memref<vector<2x2xf32>>
%x = load %arg1[] : memref<vector<2xf32>>
%b = load %arg2[] : memref<vector<2xf32>>
%0 = vector.contract #mattransvec_trait %A, %x, %b : vector<2x2xf32>, vector<2xf32> into vector<2xf32>
store %0, %arg2[] : memref<vector<2xf32>>
return
}
// CHECK-LABEL: func @vecmat2x2
// CHECK-SAME: %[[A:.*0]]: memref<vector<2x2xf32>>
// CHECK-SAME: %[[B:.*1]]: memref<vector<2xf32>>
// CHECK-SAME: %[[C:.*2]]: memref<vector<2xf32>>
// CHECK: %[[T0:.*]] = load %[[A]][] : memref<vector<2x2xf32>>
// CHECK: %[[T1:.*]] = load %[[B]][] : memref<vector<2xf32>>
// CHECK: %[[T2:.*]] = load %[[C]][] : memref<vector<2xf32>>
// CHECK: %[[T3:.*]] = vector.transpose %[[T0]], [1, 0] : vector<2x2xf32> to vector<2x2xf32>
// CHECK: %[[T4:.*]] = vector.extract %[[T3]][0] : vector<2x2xf32>
// CHECK: %[[T5:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
// CHECK: %[[T6:.*]] = vector.outerproduct %[[T4]], %[[T5]], %[[T2]] {kind = #vector.kind<add>} : vector<2xf32>, f32
// CHECK: %[[T7:.*]] = vector.extract %[[T3]][1] : vector<2x2xf32>
// CHECK: %[[T8:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
// CHECK: %[[T9:.*]] = vector.outerproduct %[[T7]], %[[T8]], %[[T6]] {kind = #vector.kind<add>} : vector<2xf32>, f32
// CHECK: store %[[T9]], %[[C]][] : memref<vector<2xf32>>
// CHECK: return
func @vecmat2x2(%arg0: memref<vector<2x2xf32>>, %arg1: memref<vector<2xf32>>,
%arg2: memref<vector<2xf32>>) {
%A = load %arg0[] : memref<vector<2x2xf32>>
%x = load %arg1[] : memref<vector<2xf32>>
%b = load %arg2[] : memref<vector<2xf32>>
%0 = vector.contract #vecmat_trait %x, %A, %b : vector<2xf32>, vector<2x2xf32> into vector<2xf32>
store %0, %arg2[] : memref<vector<2xf32>>
return
}
// CHECK-LABEL: func @vecmattrans2x2
// CHECK-SAME: %[[A:.*0]]: memref<vector<2x2xf32>>
// CHECK-SAME: %[[B:.*1]]: memref<vector<2xf32>>
// CHECK-SAME: %[[C:.*2]]: memref<vector<2xf32>>
// CHECK: %[[T0:.*]] = load %[[A]][] : memref<vector<2x2xf32>>
// CHECK: %[[T1:.*]] = load %[[B]][] : memref<vector<2xf32>>
// CHECK: %[[T2:.*]] = load %[[C]][] : memref<vector<2xf32>>
// CHECK: %[[T3:.*]] = vector.extract %[[T0]][0] : vector<2x2xf32>
// CHECK: %[[T4:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
// CHECK: %[[T5:.*]] = vector.outerproduct %[[T3]], %[[T4]], %[[T2]] {kind = #vector.kind<add>} : vector<2xf32>, f32
// CHECK: %[[T6:.*]] = vector.extract %[[T0]][1] : vector<2x2xf32>
// CHECK: %[[T7:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
// CHECK: %[[T8:.*]] = vector.outerproduct %[[T6]], %[[T7]], %[[T5]] {kind = #vector.kind<add>} : vector<2xf32>, f32
// CHECK: store %[[T8]], %[[C]][] : memref<vector<2xf32>>
// CHECK: return
func @vecmattrans2x2(%arg0: memref<vector<2x2xf32>>, %arg1: memref<vector<2xf32>>,
%arg2: memref<vector<2xf32>>) {
%A = load %arg0[] : memref<vector<2x2xf32>>
%x = load %arg1[] : memref<vector<2xf32>>
%b = load %arg2[] : memref<vector<2xf32>>
%0 = vector.contract #vecmattrans_trait %x, %A, %b : vector<2xf32>, vector<2x2xf32> into vector<2xf32>
store %0, %arg2[] : memref<vector<2xf32>>
return
}
Reference in New Issue View Git Blame Copy Permalink