forked from OSchip/llvm-project
Move description from GenericOpBase to linalg.(indexed_)generic.
PiperOrigin-RevId: 279173284
This commit is contained in:
Alexander Belyaev
A. Unique TensorFlower
parent
09e8e7107a
commit
24f306a22b
|
|
@ -198,26 +198,33 @@ def CopyOp : LinalgLibrary_Op<"copy", [NInputsAndOutputs<1, 1>]> {
|
|||
Copies the data in the input view into the output view.
|
||||
|
||||
Usage:
|
||||
```mlir
|
||||
linalg.copy(%arg0, %arg1) : memref<?xf32, stride_specification>,
|
||||
memref<?xf32, stride_specification>
|
||||
```
|
||||
|
||||
One possible lowering to loop form is:
|
||||
```mlir
|
||||
%0 = linalg.dim %arg0, 0 : index
|
||||
loop.for %i0 = %c0 to %0 step %c1 {
|
||||
%1 = linalg.load %arg0[%i0] : memref<?xf32, stride_specification>
|
||||
linalg.store %1, %arg1[%i0] : memref<?xf32, stride_specification>
|
||||
}
|
||||
```
|
||||
|
||||
Optionally, can take `input_permutation` and `output_permutation` attributes
|
||||
to reorder the dimensions of the input and output views.
|
||||
|
||||
Usage:
|
||||
```mlir
|
||||
linalg.copy(%arg0, %arg1) {inputPermutation : (i, j, k) -> (i, k, j),
|
||||
outputPermutation : (i, j, k) -> (k, j, i)} :
|
||||
memref<?x?x?xf32, stride_specification>,
|
||||
memref<?x?x?xf32, stride_specification>
|
||||
```
|
||||
|
||||
One possible lowering to loop form is:
|
||||
```mlir
|
||||
%0 = linalg.dim %arg0, 0
|
||||
%1 = linalg.dim %arg0, 1
|
||||
%2 = linalg.dim %arg0, 2
|
||||
|
|
@ -228,6 +235,7 @@ def CopyOp : LinalgLibrary_Op<"copy", [NInputsAndOutputs<1, 1>]> {
|
|||
memref<?x?x?xf32, stride_specification>
|
||||
linalg.store %3, %arg1[%i2, %i1, %i0] :
|
||||
memref<?x?x?xf32, stride_specification>
|
||||
```
|
||||
|
||||
The views are expected to be compatible for correctness but this is not
|
||||
enforced at the moment.
|
||||
|
|
@ -363,95 +371,6 @@ def ConvOp : LinalgLibrary_Op<"conv", [NInputsAndOutputs<2, 1>]> {
|
|||
}
|
||||
|
||||
class GenericOpBase<string mnemonic> : LinalgLibraryBase_Op<mnemonic, []> {
|
||||
let description = [{
|
||||
Base class for Generic and Indexed Generic Linalg ops. Key properties of
|
||||
the computation are specified as attributes. In pretty form, a
|
||||
linalg.generic op is written as:
|
||||
|
||||
```
|
||||
linalg.generic #trait_attribute %A, %B, %C {other-attributes} :
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>
|
||||
```
|
||||
|
||||
Where #trait_attributes is an alias of a dictionary attribute containing:
|
||||
- doc [optional]: a documentation string
|
||||
- fun: a SymbolRefAttr that must resolve to an existing function symbol.
|
||||
To support inplace updates in a generic fashion, the signature of the
|
||||
function must be:
|
||||
```
|
||||
fun([input views element types], [output views element types])
|
||||
-> ([output views element types])
|
||||
```
|
||||
- indexing_maps: a list of AffineMapAttr, one AffineMapAttr per each input
|
||||
and output view. Such AffineMapAttr specifies the mapping between the
|
||||
loops and the indexing within each view.
|
||||
- library_call [optional]: a StringAttr containing the name of an
|
||||
external library function that the linalg.generic operation maps to.
|
||||
The external library is assumed to be dynamically linked and no strong
|
||||
compile-time guarantees are provided. In the absence of such a library
|
||||
call, linalg.generic will always lower to loops.
|
||||
- n_loops: a triple of I64Attr representing the number of enclosing
|
||||
[parallel, reduction, window] loops respectively.
|
||||
- n_views: a pair of I64Attr representing the number of input (readonly)
|
||||
and output (readwrite) views.
|
||||
|
||||
Example:
|
||||
Defining a #matmul_trait attribute in MLIR can be done as follows:
|
||||
```
|
||||
func @fma(%a: f32, %b: f32, %c: f32) -> f32 {
|
||||
%d = mulf %a, %b: f32
|
||||
%e = addf %c, %d: f32
|
||||
return %e: f32
|
||||
}
|
||||
#matmul_accesses = [
|
||||
(m, n, k) -> (m, k),
|
||||
(m, n, k) -> (k, n),
|
||||
(m, n, k) -> (m, n)
|
||||
]
|
||||
#matmul_trait = {
|
||||
doc = "C(m, n) += A(m, k) * B(k, n)",
|
||||
fun = @fma,
|
||||
indexing_maps = #matmul_accesses,
|
||||
library_call = "linalg_matmul",
|
||||
n_views = [2, 1],
|
||||
n_loop_types = [2, 1, 0]
|
||||
}
|
||||
```
|
||||
|
||||
And can be reused in multiple places as:
|
||||
```
|
||||
linalg.generic #matmul_trait %A, %B, %C [other-attributes] :
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>
|
||||
```
|
||||
|
||||
This may lower to either:
|
||||
```
|
||||
call @linalg_matmul(%A, %B, %C) :
|
||||
(memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>)
|
||||
-> ()
|
||||
```
|
||||
|
||||
or IR resembling:
|
||||
```
|
||||
loop.for %m = %c0 to %M step %c1 {
|
||||
loop.for %n = %c0 to %N step %c1 {
|
||||
loop.for %k = %c0 to %K step %c1 {
|
||||
%a = linalg.load %A[%m, %k] : memref<?x?xf32, stride_specification>
|
||||
%b = linalg.load %B[%k, %n] : memref<?x?xf32, stride_specification>
|
||||
%c = linalg.load %C[%m, %n] : memref<?x?xf32, stride_specification>
|
||||
%d = call @mac(%a, %b, %c) : (f32, f32, f32) -> (f32)
|
||||
linalg.store %d, %C[%m, %n] : memref<?x?x?xf32, stride_specification>
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
}];
|
||||
let arguments = (ins Variadic<AnyStridedMemRef>:$views,
|
||||
AffineMapArrayAttr:$indexing_maps,
|
||||
I64ArrayAttr:$n_loop_types,
|
||||
|
|
@ -532,10 +451,189 @@ class GenericOpBase<string mnemonic> : LinalgLibraryBase_Op<mnemonic, []> {
|
|||
}
|
||||
|
||||
def GenericOp : GenericOpBase<"generic"> {
|
||||
let description = [{
|
||||
Generic Linalg op form where the key properties of the computation are
|
||||
specified as attributes. In pretty form, a linalg.generic op is written as:
|
||||
|
||||
```mlir
|
||||
linalg.generic #trait_attribute %A, %B, %C {other-attributes} :
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>
|
||||
```
|
||||
|
||||
Where #trait_attributes is an alias of a dictionary attribute containing:
|
||||
- doc [optional]: a documentation string
|
||||
- fun: a SymbolRefAttr that must resolve to an existing function symbol.
|
||||
To support inplace updates in a generic fashion, the signature of the
|
||||
function must be:
|
||||
```
|
||||
fun([input views element types], [output views element types])
|
||||
-> ([output views element types])
|
||||
```
|
||||
- indexing_maps: a list of AffineMapAttr, one AffineMapAttr per each input
|
||||
and output view. Such AffineMapAttr specifies the mapping between the
|
||||
loops and the indexing within each view.
|
||||
- library_call [optional]: a StringAttr containing the name of an
|
||||
external library function that the linalg.generic operation maps to.
|
||||
The external library is assumed to be dynamically linked and no strong
|
||||
compile-time guarantees are provided. In the absence of such a library
|
||||
call, linalg.generic will always lower to loops.
|
||||
- n_loops: a triple of I64Attr representing the number of enclosing
|
||||
[parallel, reduction, window] loops respectively.
|
||||
- n_views: a pair of I64Attr representing the number of input (readonly)
|
||||
and output (readwrite) views.
|
||||
|
||||
Example:
|
||||
Defining a #matmul_trait attribute in MLIR can be done as follows:
|
||||
```mlir
|
||||
func @fma(%a: f32, %b: f32, %c: f32) -> f32 {
|
||||
%d = mulf %a, %b: f32
|
||||
%e = addf %c, %d: f32
|
||||
return %e: f32
|
||||
}
|
||||
#matmul_accesses = [
|
||||
(m, n, k) -> (m, k),
|
||||
(m, n, k) -> (k, n),
|
||||
(m, n, k) -> (m, n)
|
||||
]
|
||||
#matmul_trait = {
|
||||
doc = "C(m, n) += A(m, k) * B(k, n)",
|
||||
fun = @fma,
|
||||
indexing_maps = #matmul_accesses,
|
||||
library_call = "linalg_matmul",
|
||||
n_views = [2, 1],
|
||||
n_loop_types = [2, 1, 0]
|
||||
}
|
||||
```
|
||||
|
||||
And can be reused in multiple places as:
|
||||
```mlir
|
||||
linalg.generic #matmul_trait %A, %B, %C [other-attributes] :
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>
|
||||
```
|
||||
|
||||
This may lower to either:
|
||||
```mlir
|
||||
call @linalg_matmul(%A, %B, %C) :
|
||||
(memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>)
|
||||
-> ()
|
||||
```
|
||||
|
||||
or IR resembling:
|
||||
```mlir
|
||||
loop.for %m = %c0 to %M step %c1 {
|
||||
loop.for %n = %c0 to %N step %c1 {
|
||||
loop.for %k = %c0 to %K step %c1 {
|
||||
%a = linalg.load %A[%m, %k] : memref<?x?xf32, stride_specification>
|
||||
%b = linalg.load %B[%k, %n] : memref<?x?xf32, stride_specification>
|
||||
%c = linalg.load %C[%m, %n] : memref<?x?xf32, stride_specification>
|
||||
%d = call @func_of_elements(%a, %b, %c)
|
||||
: (f32, f32, f32) -> (f32)
|
||||
linalg.store %d, %C[%m, %n] : memref<?x?x?xf32, stride_specification>
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
}];
|
||||
let verifier = [{ return ::verify(*this); }];
|
||||
}
|
||||
|
||||
def IndexedGenericOp : GenericOpBase<"indexed_generic"> {
|
||||
let description = [{
|
||||
Indexed Generic Linalg op form where the key properties of the computation
|
||||
are specified as attributes. In pretty form, a linalg.indexed_generic op is
|
||||
written as:
|
||||
|
||||
```mlir
|
||||
linalg.indexed_generic #trait_attribute %A, %B, %C {other-attributes} :
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>
|
||||
```
|
||||
|
||||
Where #trait_attributes is an alias of a dictionary attribute containing:
|
||||
- doc [optional]: a documentation string
|
||||
- fun: a SymbolRefAttr that must resolve to an existing function symbol.
|
||||
To support inplace updates in a generic fashion, the signature of the
|
||||
function must be:
|
||||
```
|
||||
fun([input views element types], [output views element types])
|
||||
-> ([output views element types])
|
||||
```
|
||||
- indexing_maps: a list of AffineMapAttr, one AffineMapAttr per each input
|
||||
and output view. Such AffineMapAttr specifies the mapping between the
|
||||
loops and the indexing within each view.
|
||||
- library_call [optional]: a StringAttr containing the name of an
|
||||
external library function that the linalg.indexed_generic operation
|
||||
maps to. The external library is assumed to be dynamically linked and
|
||||
no strong compile-time guarantees are provided. In the absence of such
|
||||
a library call, linalg.indexed_generic will always lower to loops.
|
||||
- n_loops: a triple of I64Attr representing the number of enclosing
|
||||
[parallel, reduction, window] loops respectively.
|
||||
- n_views: a pair of I64Attr representing the number of input (readonly)
|
||||
and output (readwrite) views.
|
||||
|
||||
Example:
|
||||
Defining a #matmul_trait attribute in MLIR can be done as follows:
|
||||
```mlir
|
||||
func @fma(%a: f32, %b: f32, %c: f32) -> f32 {
|
||||
%d = mulf %a, %b: f32
|
||||
%e = addf %c, %d: f32
|
||||
return %e: f32
|
||||
}
|
||||
#matmul_accesses = [
|
||||
(m, n, k) -> (m, k),
|
||||
(m, n, k) -> (k, n),
|
||||
(m, n, k) -> (m, n)
|
||||
]
|
||||
#matmul_trait = {
|
||||
doc = "C(m, n) += A(m, k) * B(k, n)",
|
||||
fun = @fma,
|
||||
indexing_maps = #matmul_accesses,
|
||||
library_call = "linalg_matmul",
|
||||
n_views = [2, 1],
|
||||
n_loop_types = [2, 1, 0]
|
||||
}
|
||||
```
|
||||
|
||||
And can be reused in multiple places as:
|
||||
```mlir
|
||||
linalg.indexed_generic #matmul_trait %A, %B, %C [other-attributes] :
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>
|
||||
```
|
||||
|
||||
This may lower to either:
|
||||
```mlir
|
||||
call @linalg_matmul(%A, %B, %C) :
|
||||
(memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>,
|
||||
memref<?x?xf32, stride_specification>)
|
||||
-> ()
|
||||
```
|
||||
|
||||
or IR resembling:
|
||||
```mlir
|
||||
loop.for %m = %c0 to %M step %c1 {
|
||||
loop.for %n = %c0 to %N step %c1 {
|
||||
loop.for %k = %c0 to %K step %c1 {
|
||||
%a = linalg.load %A[%m, %k] : memref<?x?xf32, stride_specification>
|
||||
%b = linalg.load %B[%k, %n] : memref<?x?xf32, stride_specification>
|
||||
%c = linalg.load %C[%m, %n] : memref<?x?xf32, stride_specification>
|
||||
%d = call @func_of_elements_and_indices(%m, %n, %k, %a, %b, %c)
|
||||
: (index, index, index, f32, f32, f32) -> (f32)
|
||||
linalg.store %d, %C[%m, %n] : memref<?x?x?xf32, stride_specification>
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
}];
|
||||
let verifier = [{ return ::verify(*this); }];
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue