mirror of
https://github.com/intel/llvm.git
synced 2026-01-31 16:25:01 +08:00
[MLIR][Linalg] Named op 'add' element-wise
This adds the first strict element-wise named op to Linalg.
The semantics here is to not allow auto-cast, broadcast semantics and to
restrict the operations only to identical types. The remaining semantics
must come in the form of surrounding operations on operands, to avoid
ambiguity.
Examples:
```
// Cast int-to-fp
%0 = linalg.copy ins(%in: tensor<32x32xi32>)
outs(%out: tensor<32x32xf32>)
%1 = linalg.add ins(%arg, %0: tensor<32x32xf32>, tensor<32x32xf32>)
outs(%0: tensor<32x32xf32>)
// This can be lowered to
%1 = linalg.generic {...}
ins(%arg, %in: tensor<32x32xf32>, tensor<32x32xi32>)
outs(%0: tensor<32x32xf32>) {
^bb0(%a: f32, %i: i32, %out: f32):
%f = arith.uitofp %i : f32
%0 = arith.addf %a, %f : f32
linalg.yield %0 : f32
}
// Broadcast
%0 = linalg.broadcast ins(%in: tensor<32xf32>)
init(%out: tensor<32x32xf32>)
%1 = linalg.add ins(%arg, %0: tensor<32x32xf32>, tensor<32x32xf32>)
outs(%0: tensor<32x32xf32>)
// This can be lowered to
#bcast_map = affine_map<(d0, d1) -> (d0)>
%1 = linalg.generic {... #bcast_map] }
ins(%arg, %in: tensor<32x32xf32>, tensor<32xf32>)
outs(%0: tensor<32x32xf32>) {
^bb0(%a: f32, %b: f32, %out: f32):
%0 = arith.addf %a, %b : f32
linalg.yield %0 : f32
}
```
Once this gets accepted, other arithmetic and maths operations will be
added accordingly, with the same semantics.
Differential Revision: https://reviews.llvm.org/D154500
This commit is contained in:
@@ -51,6 +51,25 @@ def elemwise_binary(
|
||||
O[None] = fun(cast(U, lhs[None]), cast(U, rhs[None]))
|
||||
|
||||
|
||||
@linalg_structured_op
|
||||
def add(
|
||||
lhs=TensorDef(T1),
|
||||
rhs=TensorDef(T1),
|
||||
O=TensorDef(T1, output=True),
|
||||
):
|
||||
""" Adds two tensors elementwise.
|
||||
|
||||
The shapes and element types must be identical. The appropriate casts,
|
||||
broadcasts and reductions should be done previously to calling this op.
|
||||
|
||||
This means reduction/broadcast/element cast semantics is explicit. Further
|
||||
passes can take that into account when lowering this code. For example,
|
||||
a `linalg.broadcast` + `linalg.add` sequence can be lowered to a
|
||||
`linalg.generic` with different affine maps for the two operands.
|
||||
"""
|
||||
O[None] = lhs[None] + rhs[None]
|
||||
|
||||
|
||||
@linalg_structured_op
|
||||
def matmul(
|
||||
A=TensorDef(T1, S.M, S.K),
|
||||
|
||||
Reference in New Issue
Block a user