This is an archive of the discontinued LLVM Phabricator instance.

Differential D120562

[mlir][OpDSL] Rename function to make signedness explicit (NFC).
ClosedPublic

Authored by gysit on Feb 25 2022, 5:49 AM.

Details

Reviewers
nicolasvasilache
aartbik
Commits
rGe9085d0d2558: [mlir][OpDSL] Rename function to make signedness explicit (NFC).
Summary

The revision renames the following OpDSL functions:

TypeFn.cast -> TypeFn.cast_signed
BinaryFn.min -> BinaryFn.min_signed
BinaryFn.max -> BinaryFn.max_signed

The corresponding enum values on the C++ side are renamed accordingly:

#linalg.type_fn<cast> -> #linalg.type_fn<cast_signed>
#linalg.binary_fn<min> -> #linalg.binary_fn<min_signed>
#linalg.binary_fn<max> -> #linalg.binary_fn<max_signed>

Depends On D120110

Diff Detail

Event Timeline

gysit created this revision.Feb 25 2022, 5:49 AM
Herald added subscribers: sdasgup3, wenzhicui, wrengr and 20 others. · View Herald TranscriptFeb 25 2022, 5:49 AM
gysit requested review of this revision.Feb 25 2022, 5:49 AM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptFeb 25 2022, 5:49 AM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: limo1996, stephenneuendorffer, nicolasvasilache. · View Herald Transcript
gysit added a reviewer: aartbik.Feb 25 2022, 5:51 AM

@aartbik I actually the longer names are better in this case!

Harbormaster completed remote builds in B151467: Diff 411394.Feb 25 2022, 6:07 AM
aartbik accepted this revision.Feb 25 2022, 1:34 PM
This revision is now accepted and ready to land.Feb 25 2022, 1:34 PM
gysit updated this revision to Diff 411991.Mar 1 2022, 12:06 AM

Rebase.

This revision was landed with ongoing or failed builds.Mar 1 2022, 12:19 AM
Closed by commit rGe9085d0d2558: [mlir][OpDSL] Rename function to make signedness explicit (NFC). (authored by gysit). · Explain Why
This revision was automatically updated to reflect the committed changes.
gysit added a commit: rGe9085d0d2558: [mlir][OpDSL] Rename function to make signedness explicit (NFC)..
Harbormaster completed remote builds in B151883: Diff 411991.Mar 1 2022, 12:29 AM

Revision Contents

PathSize
mlir/
docs/
Dialects/
Linalg/
44 lines
include/
mlir/
Dialect/
Linalg/
IR/
10 lines
160 lines
lib/
Dialect/
Linalg/
IR/
14 lines
python/
mlir/
dialects/
linalg/
opdsl/
lang/
16 lines
6 lines
ops/
144 lines
test/
mlir-linalg-ods-gen/
12 lines
python/
dialects/
linalg/
opdsl/
4 lines
12 lines
4 lines
2 lines
2 lines
19 lines
8 lines
3 lines
7 lines
4 lines

Diff 411991

mlir/docs/Dialects/Linalg/OpDSL.md

Show First 20 Lines • Show All 50 Lines • ▼ Show 20 Linesdef matmul(A=TensorDef(T1, S.M, S.K),
C=TensorDef(U, S.M, S.N, output=True)): C=TensorDef(U, S.M, S.N, output=True)):
"""Performs a matrix multiplication of two 2D inputs. """Performs a matrix multiplication of two 2D inputs.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
domain(D.m, D.n, D.k) domain(D.m, D.n, D.k)
implements(ContractionOpInterface) implements(ContractionOpInterface)
C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast_signed(
U, A[D.m, D.k]) * TypeFn.cast_signed(U, B[D.k, D.n])
``` ```
Here we have a simple type polymorphic contraction that takes arguments `A` and Here we have a simple type polymorphic contraction that takes arguments `A` and
`B` and outputs `C`. Each is bound to a `TensorDef`, which specifies: `B` and outputs `C`. Each is bound to a `TensorDef`, which specifies:
* The symbolic element type (`T1`, `T2`, `U` above). * The symbolic element type (`T1`, `T2`, `U` above).
* Symbolic shape expressions with symbols that are bound globally for the op ( * Symbolic shape expressions with symbols that are bound globally for the op (
note that in this simple example, the shape expressions are just symbol note that in this simple example, the shape expressions are just symbol
▲ Show 20 Lines • Show All 87 Lines • ▼ Show 20 Lines
```python ```python
@linalg_structured_op @linalg_structured_op
def pooling_poly( def pooling_poly(
I=TensorDef(T1, S.N, S.H, S.W, S.C), I=TensorDef(T1, S.N, S.H, S.W, S.C),
K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
O[D.n, D.oh, D.ow, D.c] += TypeFn.cast(U, O[D.n, D.oh, D.ow, D.c] += TypeFn.cast_signed(U,
I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]) I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])
``` ```
The pooling operation does not access the shape-only tensor `K`. Instead, the The pooling operation does not access the shape-only tensor `K`. Instead, the
shapes `S.KH` and `S.KW` specify the iteration domain for the reduction shapes `S.KH` and `S.KW` specify the iteration domain for the reduction
dimensions `D.kh` and `D.kw`. dimensions `D.kh` and `D.kw`.
## Assignments ## Assignments
The bulk of language consists of assignment expressions of the form above. The The bulk of language consists of assignment expressions of the form above. The
iteration dimension order is determined lexically based on the order encountered iteration dimension order is determined lexically based on the order encountered
in the expression (following operator precedence if math operators are used). in the expression (following operator precedence if math operators are used).
TODO: Introduce a directive to fix the dimension bindings. TODO: Introduce a directive to fix the dimension bindings.
Reduction dimensions are inferred to be any dimensions on the RHS that are not Reduction dimensions are inferred to be any dimensions on the RHS that are not
on the LHS. on the LHS.
A number of unary and binary arithmetic functions are supported: A number of unary and binary arithmetic functions are supported:
* `BinaryFn.add(a, b)` (also via overloading the binary `+` operator) * `BinaryFn.add(a, b)` (also via overloading the binary `+` operator)
* `BinaryFn.mul(a, b)` (also via overloading the binary `*` operator) * `BinaryFn.mul(a, b)` (also via overloading the binary `*` operator)
* `BinaryFn.max(a, b)` * `BinaryFn.max_signed(a, b)`
* `BinaryFn.min(a, b)` * `BinaryFn.min_signed(a, b)`
* `BinaryFn.sub(a, b)` (also via overloading the binary `-` operator) * `BinaryFn.sub(a, b)` (also via overloading the binary `-` operator)
* `BinaryFn.max_unsigned(a, b)` * `BinaryFn.max_unsigned(a, b)`
* `BinaryFn.min_unsigned(a, b)` * `BinaryFn.min_unsigned(a, b)`
* `UnaryFn.exp(a)` * `UnaryFn.exp(a)`
* `UnaryFn.log(a)` * `UnaryFn.log(a)`
As the integer types are signless, signedness is implement by different As the integer types are signless, signedness is implement by different
functions that treat integers as signed or unsigned values. functions that treat integers as signed or unsigned values.
A subset of the arithmetic functions are supported in reductions. These A subset of the arithmetic functions are supported in reductions. These
reduction functions can appear as the outermost function on the RHS: reduction functions can appear as the outermost function on the RHS:
* `ReduceFn.add` (also overloading the inplace `+=` on a LHS) * `ReduceFn.add` (also overloading the inplace `+=` on a LHS)
* `ReduceFn.mul` * `ReduceFn.mul`
* `ReduceFn.max` * `ReduceFn.max_signed`
* `ReduceFn.min` * `ReduceFn.min_signed`
* `ReduceFn.max_unsigned` * `ReduceFn.max_unsigned`
* `ReduceFn.min_unsigned` * `ReduceFn.min_unsigned`
As the integer types are signless, signedness is implement by different As the integer types are signless, signedness is implement by different
functions that treat integers as signed or unsigned values. functions that treat integers as signed or unsigned values.
Additionally, type conversion functions cast an operand to a target type: Additionally, type conversion functions cast an operand to a target type:
* `TypeFn.cast(TypeVar, operand)` * `TypeFn.cast_signed(TypeVar, operand)`
* `TypeFn.cast_unsigned(TypeVar, operand)` * `TypeFn.cast_unsigned(TypeVar, operand)`
As the integer types are signless, signedness is implement by different As the integer types are signless, signedness is implement by different
functions that treat integers as signed (`TypeFn.cast`) or unsigned functions that treat integers as signed (`TypeFn.cast_signed`) or unsigned
(`TypeFn.cast_unsigned`) values. (`TypeFn.cast_unsigned`) values.
There are also special forms: There are also special forms:
* `const(value)` returns a constant value. * `const(value)` returns a constant value.
* `index(dim)` returns the iteration index in the given dimension `dim`. * `index(dim)` returns the iteration index in the given dimension `dim`.
## Function Attributes ## Function Attributes
Function attributes are compile-time constant function parameters. They can be Function attributes are compile-time constant function parameters. They can be
used to parameterize the computation performed by a structured operation, for used to parameterize the computation performed by a structured operation, for
example, to support signed and unsigned computations. example, to support signed and unsigned computations.
The following example demonstrates the use of function attributes: The following example demonstrates the use of function attributes:
```python ```python
@linalg_structured_op @linalg_structured_op
def elemwise_binary( def elemwise_binary(
lhs=TensorDef(T1), lhs=TensorDef(T1),
rhs=TensorDef(T2), rhs=TensorDef(T2),
O=TensorDef(U, output=True), O=TensorDef(U, output=True),
fun=BinaryFnAttrDef(default=BinaryFn.add), fun=BinaryFnAttrDef(default=BinaryFn.add),
cast=TypeFnAttrDef(default=TypeFn.cast)): cast=TypeFnAttrDef(default=TypeFn.cast_signed)):
O[None] = fun(cast(U, lhs[None]), cast(U, rhs[None])) O[None] = fun(cast(U, lhs[None]), cast(U, rhs[None]))
``` ```
The `fun` and `cast` function attributes by default are aliases for their The `fun` and `cast` function attributes by default are aliases for their
default values `BinaryFn.add` and `TypeFn.cast`, respectively. When default values `BinaryFn.add` and `TypeFn.cast_signed`, respectively. When
instantiating the operation, the function attributes may be set to other instantiating the operation, the function attributes may be set to other
functions using optional named arguments: functions using optional named arguments:
```python ```python
elemwise_binary(lhs, rhs, outs=[out_tensor], elemwise_binary(lhs, rhs, outs=[out_tensor],
fun=BinaryFn.mul, cast=TypeFn.cast_unsigned) fun=BinaryFn.mul, cast=TypeFn.cast_unsigned)
``` ```
In the example, the `fun` and `cast` arguments adapt the body of the operation In the example, the `fun` and `cast` arguments adapt the body of the operation
to implement multiplication and unsigned casts instead of addition and signed to implement multiplication and unsigned casts instead of addition and signed
casts. casts.
OpDSL supports unary, binary, and type conversion function attributes. An OpDSL supports unary, binary, and type conversion function attributes. An
operation can take multiple attributes of different kinds placed at the end of operation can take multiple attributes of different kinds placed at the end of
the parameter list. the parameter list.
## Types ## Types
All types in assignment expressions are late bound based on actual input and All types in assignment expressions are late bound based on actual input and
output types of constructed ops. An exception are predefined types such as output types of constructed ops. An exception are predefined types such as
`I32`, `I64`, `F32`, and `F64`. These hardwired types enable intermediate `I32`, `I64`, `F32`, and `F64`. These hardwired types enable intermediate
computations with a type that is independent of the input and output types. For computations with a type that is independent of the input and output types. For
example, parts of floating point computation may require double precision example, parts of floating point computation may require double precision
arithmetic despite all inputs and outputs being single precision values. arithmetic despite all inputs and outputs being single precision values.
Assignment expressions with no `TypeFn.cast` calls will generally require Assignment expressions with no `TypeFn.cast_signed` calls will generally require
uniform types throughout and will fail to verify if violated. The presence of a uniform types throughout and will fail to verify if violated. The presence of a
`TypeFn.cast` or `TypeFn.cast_unsigned` allows for a limited form of numeric `TypeFn.cast_signed` or `TypeFn.cast_unsigned` allows for a limited form of
type conversion between element types that can be derived from inputs and numeric type conversion between element types that can be derived from inputs
outputs (and in the future, attributes). `TypeFn.cast` calls with a `TypeVar` and outputs (and in the future, attributes). `TypeFn.cast_signed` calls with a
first argument are emitted as `type_fn` primitives in the YAML definition. `TypeVar` first argument are emitted as `type_fn` primitives in the YAML
definition.
Casting will perform `int<->float` and `index->int` type conversions and will Casting will perform `int<->float` and `index->int` type conversions and will
perform any necessary extension or truncation within the type family. The perform any necessary extension or truncation within the type family. The
integer types themselves are signless and signedness is implemented by integer types themselves are signless and signedness is implemented by
functions/operations. The `TypeFn.cast` function treats all integers as signed, functions/operations. The `TypeFn.cast_signed` function treats all integers as
while `TypeFn.cast_unsigned` treats them as unsigned. signed, while `TypeFn.cast_unsigned` treats them as unsigned.
The following examples illustrate the lowering of signed and unsigned functions: The following examples illustrate the lowering of signed and unsigned functions:
* cast(I32 -> I64) -> `arith.ExtSIOp` * cast_signed(I32 -> I64) -> `arith.ExtSIOp`
* cast(F32 -> I32) -> `arith.FPToSIOp` * cast_signed(F32 -> I32) -> `arith.FPToSIOp`
* cast_unsigned(I32 -> I64) -> `arith.ExtUIOp` * cast_unsigned(I32 -> I64) -> `arith.ExtUIOp`
* cast_unsigned(F32 -> I32) -> `arith.FPToUIOp` * cast_unsigned(F32 -> I32) -> `arith.FPToUIOp`
* max -> `arith.MaxSIOp` * max_signed -> `arith.MaxSIOp`
* max_unsinged -> `arith.MaxUIOp` * max_unsinged -> `arith.MaxUIOp`
Not all functions are applicable for all numeric types, and on mismatch, op Not all functions are applicable for all numeric types, and on mismatch, op
verification will fail. verification will fail.
## Pointwise Computations ## Pointwise Computations
Pointwise computations are expressible in a rank polymorphic form that supports Pointwise computations are expressible in a rank polymorphic form that supports
arbitrary ranked operands - all of them need to have the same rank - with a arbitrary ranked operands - all of them need to have the same rank - with a
single operation definition. single operation definition.
An example for a rank polymorphic operation is `fill`: An example for a rank polymorphic operation is `fill`:
```python ```python
@linalg_structured_op @linalg_structured_op
def fill(value=ScalarDef(T1), def fill(value=ScalarDef(T1),
O=TensorDef(U, output=True)): O=TensorDef(U, output=True)):
O[None] = TypeFn.cast(U, value) O[None] = TypeFn.cast_signed(U, value)
``` ```
The operation sets the elements of the output tensor `O` to `value`. All The operation sets the elements of the output tensor `O` to `value`. All
operands are either scalars or rank zero tensors that are accessed using the operands are either scalars or rank zero tensors that are accessed using the
index `None`. The operation thus performs a scalar computation that trivially index `None`. The operation thus performs a scalar computation that trivially
extends to a multi-dimensional pointwise computation. As a result, we may use extends to a multi-dimensional pointwise computation. As a result, we may use
`fill` with arbitrary ranked output tensors: `fill` with arbitrary ranked output tensors:
```python ```python
tensor_2d = linalg.InitTensorOp([4, 8], f32) tensor_2d = linalg.InitTensorOp([4, 8], f32)
tensor_3d = linalg.InitTensorOp([4, 8, 16], f32) tensor_3d = linalg.InitTensorOp([4, 8, 16], f32)
fill(value, outs=[tensor_2d]) fill(value, outs=[tensor_2d])
fill(value, outs=[tensor_3d]) fill(value, outs=[tensor_3d])
``` ```

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

Show First 20 Lines • Show All 62 Lines • ▼ Show 20 Linesdef UnaryFn : I32EnumAttr<"UnaryFn", "", [
I32EnumAttrCase<"exp", 0>, I32EnumAttrCase<"exp", 0>,
I32EnumAttrCase<"log", 1> I32EnumAttrCase<"log", 1>
]> { ]> {
let genSpecializedAttr = 0; let genSpecializedAttr = 0;
let cppNamespace = "::mlir::linalg"; let cppNamespace = "::mlir::linalg";
} }
def BinaryFn : I32EnumAttr<"BinaryFn", "", [ def BinaryFn : I32EnumAttr<"BinaryFn", "", [
I32EnumAttrCase<"add", 0>, I32EnumAttrCase<"add", 0>,
I32EnumAttrCase<"mul", 1>, I32EnumAttrCase<"sub", 1>,
I32EnumAttrCase<"max", 2>, I32EnumAttrCase<"mul", 2>,
I32EnumAttrCase<"min", 3>, I32EnumAttrCase<"max_signed", 3>,
I32EnumAttrCase<"sub", 4>, I32EnumAttrCase<"min_signed", 4>,
I32EnumAttrCase<"max_unsigned", 5>, I32EnumAttrCase<"max_unsigned", 5>,
I32EnumAttrCase<"min_unsigned", 6> I32EnumAttrCase<"min_unsigned", 6>
]> { ]> {
let genSpecializedAttr = 0; let genSpecializedAttr = 0;
let cppNamespace = "::mlir::linalg"; let cppNamespace = "::mlir::linalg";
} }
def TypeFn : I32EnumAttr<"TypeFn", "", [ def TypeFn : I32EnumAttr<"TypeFn", "", [
I32EnumAttrCase<"cast", 0>, I32EnumAttrCase<"cast_signed", 0>,
I32EnumAttrCase<"cast_unsigned", 1> I32EnumAttrCase<"cast_unsigned", 1>
]> { ]> {
let genSpecializedAttr = 0; let genSpecializedAttr = 0;
let cppNamespace = "::mlir::linalg"; let cppNamespace = "::mlir::linalg";
} }
def UnaryFnAttr : EnumAttr<Linalg_Dialect, UnaryFn, "unary_fn"> { def UnaryFnAttr : EnumAttr<Linalg_Dialect, UnaryFn, "unary_fn"> {
let assemblyFormat = "`<` $value `>`"; let assemblyFormat = "`<` $value `>`";
Show All 9 Lines

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

Show All 22 Lines - !LinalgOperandDefConfig
shape_map: affine_map<() -> ()> shape_map: affine_map<() -> ()>
- !LinalgOperandDefConfig - !LinalgOperandDefConfig
name: fun name: fun
kind: unary_fn_attr kind: unary_fn_attr
default_fn: exp default_fn: exp
- !LinalgOperandDefConfig - !LinalgOperandDefConfig
name: cast name: cast
kind: type_fn_attr kind: type_fn_attr
default_fn: cast default_fn: cast_signed
indexing_maps: !LinalgIndexingMapsConfig indexing_maps: !LinalgIndexingMapsConfig
static_indexing_maps: static_indexing_maps:
- affine_map<() -> ()> - affine_map<() -> ()>
- affine_map<() -> ()> - affine_map<() -> ()>
iterator_types: [] iterator_types: []
assignments: assignments:
- !ScalarAssign - !ScalarAssign
arg: O arg: O
Show All 38 Lines - !LinalgOperandDefConfig
shape_map: affine_map<() -> ()> shape_map: affine_map<() -> ()>
- !LinalgOperandDefConfig - !LinalgOperandDefConfig
name: fun name: fun
kind: binary_fn_attr kind: binary_fn_attr
default_fn: add default_fn: add
- !LinalgOperandDefConfig - !LinalgOperandDefConfig
name: cast name: cast
kind: type_fn_attr kind: type_fn_attr
default_fn: cast default_fn: cast_signed
indexing_maps: !LinalgIndexingMapsConfig indexing_maps: !LinalgIndexingMapsConfig
static_indexing_maps: static_indexing_maps:
- affine_map<() -> ()> - affine_map<() -> ()>
- affine_map<() -> ()> - affine_map<() -> ()>
- affine_map<() -> ()> - affine_map<() -> ()>
iterator_types: [] iterator_types: []
assignments: assignments:
- !ScalarAssign - !ScalarAssign
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Linesstructured_op: !LinalgStructuredOpConfig
- !LinalgOperandDefConfig - !LinalgOperandDefConfig
name: C name: C
kind: output_tensor kind: output_tensor
type_var: U type_var: U
shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)> shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
- !LinalgOperandDefConfig - !LinalgOperandDefConfig
name: cast name: cast
kind: type_fn_attr kind: type_fn_attr
default_fn: cast default_fn: cast_signed
indexing_maps: !LinalgIndexingMapsConfig indexing_maps: !LinalgIndexingMapsConfig
static_indexing_maps: static_indexing_maps:
- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)> - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)> - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>
- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)> - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>
iterator_types: iterator_types:
- parallel - parallel
- parallel - parallel
▲ Show 20 Lines • Show All 162 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: AZp scalar_arg: AZp
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: BZp scalar_arg: BZp
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: mmt4d name: mmt4d
cpp_class_name: Mmt4DOp cpp_class_name: Mmt4DOp
▲ Show 20 Lines • Show All 54 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: AccumType type_var: AccumType
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: lhs scalar_arg: lhs
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: AccumType type_var: AccumType
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: rhs scalar_arg: rhs
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: batch_matmul name: batch_matmul
cpp_class_name: BatchMatmulOp cpp_class_name: BatchMatmulOp
▲ Show 20 Lines • Show All 44 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: quantized_batch_matmul name: quantized_batch_matmul
cpp_class_name: QuantizedBatchMatmulOp cpp_class_name: QuantizedBatchMatmulOp
▲ Show 20 Lines • Show All 59 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: AZp scalar_arg: AZp
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: BZp scalar_arg: BZp
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: matvec name: matvec
cpp_class_name: MatvecOp cpp_class_name: MatvecOp
▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: y scalar_arg: y
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: vecmat name: vecmat
cpp_class_name: VecmatOp cpp_class_name: VecmatOp
▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: y scalar_arg: y
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: batch_matvec name: batch_matvec
cpp_class_name: BatchMatvecOp cpp_class_name: BatchMatvecOp
▲ Show 20 Lines • Show All 43 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: dot name: dot
cpp_class_name: DotOp cpp_class_name: DotOp
▲ Show 20 Lines • Show All 41 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_1d name: conv_1d
cpp_class_name: Conv1DOp cpp_class_name: Conv1DOp
▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_2d name: conv_2d
cpp_class_name: Conv2DOp cpp_class_name: Conv2DOp
▲ Show 20 Lines • Show All 44 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_3d name: conv_3d
cpp_class_name: Conv3DOp cpp_class_name: Conv3DOp
▲ Show 20 Lines • Show All 47 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_1d_nwc_wcf name: conv_1d_nwc_wcf
cpp_class_name: Conv1DNwcWcfOp cpp_class_name: Conv1DNwcWcfOp
▲ Show 20 Lines • Show All 59 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_2d_nhwc_hwcf name: conv_2d_nhwc_hwcf
cpp_class_name: Conv2DNhwcHwcfOp cpp_class_name: Conv2DNhwcHwcfOp
▲ Show 20 Lines • Show All 73 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_2d_nhwc_hwcf_q name: conv_2d_nhwc_hwcf_q
cpp_class_name: Conv2DNhwcHwcfQOp cpp_class_name: Conv2DNhwcHwcfQOp
▲ Show 20 Lines • Show All 91 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: IZp scalar_arg: IZp
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: KZp scalar_arg: KZp
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_2d_nchw_fchw name: conv_2d_nchw_fchw
cpp_class_name: Conv2DNchwFchwOp cpp_class_name: Conv2DNchwFchwOp
▲ Show 20 Lines • Show All 73 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_3d_ndhwc_dhwcf name: conv_3d_ndhwc_dhwcf
cpp_class_name: Conv3DNdhwcDhwcfOp cpp_class_name: Conv3DNdhwcDhwcfOp
▲ Show 20 Lines • Show All 75 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: depthwise_conv_1d_nwc_wc name: depthwise_conv_1d_nwc_wc
cpp_class_name: DepthwiseConv1DNwcWcOp cpp_class_name: DepthwiseConv1DNwcWcOp
▲ Show 20 Lines • Show All 58 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: depthwise_conv_2d_nhwc_hwc name: depthwise_conv_2d_nhwc_hwc
cpp_class_name: DepthwiseConv2DNhwcHwcOp cpp_class_name: DepthwiseConv2DNhwcHwcOp
▲ Show 20 Lines • Show All 68 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: depthwise_conv_2d_nhwc_hwc_q name: depthwise_conv_2d_nhwc_hwc_q
cpp_class_name: DepthwiseConv2DNhwcHwcQOp cpp_class_name: DepthwiseConv2DNhwcHwcQOp
▲ Show 20 Lines • Show All 84 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: IZp scalar_arg: IZp
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: KZp scalar_arg: KZp
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: depthwise_conv_2d_nhwc_hwcm name: depthwise_conv_2d_nhwc_hwcm
cpp_class_name: DepthwiseConv2DNhwcHwcmOp cpp_class_name: DepthwiseConv2DNhwcHwcmOp
▲ Show 20 Lines • Show All 69 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: depthwise_conv_2d_nhwc_hwcm_q name: depthwise_conv_2d_nhwc_hwcm_q
cpp_class_name: DepthwiseConv2DNhwcHwcmQOp cpp_class_name: DepthwiseConv2DNhwcHwcmQOp
▲ Show 20 Lines • Show All 86 Lines • ▼ Show 20 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: IZp scalar_arg: IZp
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: KZp scalar_arg: KZp
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nhwc_sum name: pooling_nhwc_sum
cpp_class_name: PoolingNhwcSumOp cpp_class_name: PoolingNhwcSumOp
▲ Show 20 Lines • Show All 62 Lines • ▼ Show 20 Lines value: !ScalarExpression
kind: binary kind: binary
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nhwc_max name: pooling_nhwc_max
cpp_class_name: PoolingNhwcMaxOp cpp_class_name: PoolingNhwcMaxOp
▲ Show 20 Lines • Show All 55 Lines • ▼ Show 20 Linesstructured_op: !LinalgStructuredOpConfig
- reduction - reduction
- reduction - reduction
assignments: assignments:
- !ScalarAssign - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: max fn_name: max_signed
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nhwc_max_unsigned name: pooling_nhwc_max_unsigned
cpp_class_name: PoolingNhwcMaxUnsignedOp cpp_class_name: PoolingNhwcMaxUnsignedOp
▲ Show 20 Lines • Show All 134 Lines • ▼ Show 20 Linesstructured_op: !LinalgStructuredOpConfig
- reduction - reduction
- reduction - reduction
assignments: assignments:
- !ScalarAssign - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: max fn_name: max_signed
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nhwc_min name: pooling_nhwc_min
cpp_class_name: PoolingNhwcMinOp cpp_class_name: PoolingNhwcMinOp
▲ Show 20 Lines • Show All 55 Lines • ▼ Show 20 Linesstructured_op: !LinalgStructuredOpConfig
- reduction - reduction
- reduction - reduction
assignments: assignments:
- !ScalarAssign - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: min fn_name: min_signed
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nhwc_min_unsigned name: pooling_nhwc_min_unsigned
cpp_class_name: PoolingNhwcMinUnsignedOp cpp_class_name: PoolingNhwcMinUnsignedOp
▲ Show 20 Lines • Show All 147 Lines • ▼ Show 20 Lines value: !ScalarExpression
kind: binary kind: binary
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_ndhwc_max name: pooling_ndhwc_max
cpp_class_name: PoolingNdhwcMaxOp cpp_class_name: PoolingNdhwcMaxOp
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Linesstructured_op: !LinalgStructuredOpConfig
- reduction - reduction
- reduction - reduction
assignments: assignments:
- !ScalarAssign - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: max fn_name: max_signed
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_ndhwc_min name: pooling_ndhwc_min
cpp_class_name: PoolingNdhwcMinOp cpp_class_name: PoolingNdhwcMinOp
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Linesstructured_op: !LinalgStructuredOpConfig
- reduction - reduction
- reduction - reduction
assignments: assignments:
- !ScalarAssign - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: min fn_name: min_signed
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: fill_tensor name: fill_tensor
cpp_class_name: FillTensorOp cpp_class_name: FillTensorOp
Show All 20 Lines indexing_maps: !LinalgIndexingMapsConfig
- affine_map<() -> ()> - affine_map<() -> ()>
iterator_types: [] iterator_types: []
assignments: assignments:
- !ScalarAssign - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: value scalar_arg: value
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: fill_rng_2d name: fill_rng_2d
cpp_class_name: FillRng2DOp cpp_class_name: FillRng2DOp
Show All 36 Linesstructured_op: !LinalgStructuredOpConfig
- parallel - parallel
- parallel - parallel
assignments: assignments:
- !ScalarAssign - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: T type_var: T
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: F64 type_var: F64
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '2147483647 : i64' scalar_const: '2147483647 : i64'
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: F64 type_var: F64
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_index: 1 scalar_index: 1
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_index: 0 scalar_index: 0
- !ScalarExpression - !ScalarExpression
scalar_arg: seed scalar_arg: seed
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '1103515245 : i64' scalar_const: '1103515245 : i64'
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '12345 : i64' scalar_const: '12345 : i64'
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '1103515245 : i64' scalar_const: '1103515245 : i64'
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '12345 : i64' scalar_const: '12345 : i64'
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: max scalar_arg: max
- !ScalarExpression - !ScalarExpression
scalar_arg: min scalar_arg: min
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: F64 type_var: F64
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '2.3283063999999999E-10 : f64' scalar_const: '2.3283063999999999E-10 : f64'
- !ScalarExpression - !ScalarExpression
scalar_arg: min scalar_arg: min
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
Show All 34 Lines value: !ScalarExpression
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: binary kind: binary
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '1.000000e+00 : f64' scalar_const: '1.000000e+00 : f64'
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: unary kind: unary
fn_name: exp fn_name: exp
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Show First 20 Lines • Show All 154 Lines • ▼ Show 20 Lines Value buildBinaryFn(BinaryFn binaryFn, Value arg0, Value arg1) {
if (!allFloatingPoint && !allInteger) if (!allFloatingPoint && !allInteger)
llvm_unreachable("unsupported non numeric type"); llvm_unreachable("unsupported non numeric type");
OpBuilder builder = getBuilder(); OpBuilder builder = getBuilder();
switch (binaryFn) { switch (binaryFn) {
case BinaryFn::add: case BinaryFn::add:
if (allFloatingPoint) if (allFloatingPoint)
return builder.create<arith::AddFOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::AddFOp>(arg0.getLoc(), arg0, arg1);
return builder.create<arith::AddIOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::AddIOp>(arg0.getLoc(), arg0, arg1);
case BinaryFn::sub:
if (allFloatingPoint)
return builder.create<arith::SubFOp>(arg0.getLoc(), arg0, arg1);
return builder.create<arith::SubIOp>(arg0.getLoc(), arg0, arg1);
case BinaryFn::mul: case BinaryFn::mul:
if (allFloatingPoint) if (allFloatingPoint)
return builder.create<arith::MulFOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::MulFOp>(arg0.getLoc(), arg0, arg1);
return builder.create<arith::MulIOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::MulIOp>(arg0.getLoc(), arg0, arg1);
case BinaryFn::max: case BinaryFn::max_signed:
if (allFloatingPoint) if (allFloatingPoint)
return builder.create<arith::MaxFOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::MaxFOp>(arg0.getLoc(), arg0, arg1);
return builder.create<arith::MaxSIOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::MaxSIOp>(arg0.getLoc(), arg0, arg1);
case BinaryFn::min: case BinaryFn::min_signed:
if (allFloatingPoint) if (allFloatingPoint)
return builder.create<arith::MinFOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::MinFOp>(arg0.getLoc(), arg0, arg1);
return builder.create<arith::MinSIOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::MinSIOp>(arg0.getLoc(), arg0, arg1);
case BinaryFn::sub:
if (allFloatingPoint)
return builder.create<arith::SubFOp>(arg0.getLoc(), arg0, arg1);
return builder.create<arith::SubIOp>(arg0.getLoc(), arg0, arg1);
case BinaryFn::max_unsigned: case BinaryFn::max_unsigned:
if (allFloatingPoint) if (allFloatingPoint)
return builder.create<arith::MaxFOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::MaxFOp>(arg0.getLoc(), arg0, arg1);
return builder.create<arith::MaxUIOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::MaxUIOp>(arg0.getLoc(), arg0, arg1);
case BinaryFn::min_unsigned: case BinaryFn::min_unsigned:
if (allFloatingPoint) if (allFloatingPoint)
return builder.create<arith::MinFOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::MinFOp>(arg0.getLoc(), arg0, arg1);
return builder.create<arith::MinUIOp>(arg0.getLoc(), arg0, arg1); return builder.create<arith::MinUIOp>(arg0.getLoc(), arg0, arg1);
} }
llvm_unreachable("unsupported binary function"); llvm_unreachable("unsupported binary function");
} }
// Build the type functions defined by OpDSL. // Build the type functions defined by OpDSL.
Value buildTypeFn(TypeFn typeFn, Type toType, Value operand) { Value buildTypeFn(TypeFn typeFn, Type toType, Value operand) {
switch (typeFn) { switch (typeFn) {
case TypeFn::cast: case TypeFn::cast_signed:
return cast(toType, operand, false); return cast(toType, operand, false);
case TypeFn::cast_unsigned: case TypeFn::cast_unsigned:
return cast(toType, operand, true); return cast(toType, operand, true);
} }
llvm_unreachable("unsupported type conversion function"); llvm_unreachable("unsupported type conversion function");
} }
void yieldOutputs(ValueRange values) { void yieldOutputs(ValueRange values) {
▲ Show 20 Lines • Show All 1,488 Lines • Show Last 20 Lines

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

Show First 20 Lines • Show All 299 Lines • ▼ Show 20 Linesclass BinaryFn:
As the integer types are signless, signedness is implement by different As the integer types are signless, signedness is implement by different
functions that treat integers as signed or unsigned values. functions that treat integers as signed or unsigned values.
Examples: Examples:
- max -> `arith.MaxSIOp` - max -> `arith.MaxSIOp`
- max_unsinged -> `arith.MaxUIOp` - max_unsinged -> `arith.MaxUIOp`
""" """
add = BinaryFnType("add") add = BinaryFnType("add")
mul = BinaryFnType("mul")
max = BinaryFnType("max")
min = BinaryFnType("min")
sub = BinaryFnType("sub") sub = BinaryFnType("sub")
mul = BinaryFnType("mul")
max_signed = BinaryFnType("max_signed")
min_signed = BinaryFnType("min_signed")
max_unsigned = BinaryFnType("max_unsigned") max_unsigned = BinaryFnType("max_unsigned")
min_unsigned = BinaryFnType("min_unsigned") min_unsigned = BinaryFnType("min_unsigned")
class TypeFnType: class TypeFnType:
"""Type conversion function. """Type conversion function.
A type conversion function takes a target type and a tensor expression and A type conversion function takes a target type and a tensor expression and
Show All 9 Linesclass TypeFnType:
def __repr__(self): def __repr__(self):
return f"{self.fn_name}" return f"{self.fn_name}"
class TypeFn: class TypeFn:
"""Type conversion function namespace. """Type conversion function namespace.
As the integer types are signless, signedness is implement by different cast As the integer types are signless, signedness is implement by different cast
functions that treat integers as signed (`cast`) or unsigned functions that treat integers as signed (`cast_signed`) or unsigned
(`cast_unsigned`) values. (`cast_unsigned`) values.
Examples: Examples:
- cast(I32 -> I64) -> `arith.ExtSIOp` - cast_signed(I32 -> I64) -> `arith.ExtSIOp`
- cast_unsigned(I32 -> I64) -> `arith.ExtUIOp` - cast_unsigned(I32 -> I64) -> `arith.ExtUIOp`
""" """
cast = TypeFnType("cast") cast_signed = TypeFnType("cast_signed")
cast_unsigned = TypeFnType("cast_unsigned") cast_unsigned = TypeFnType("cast_unsigned")
class ReduceFnUse: class ReduceFnUse:
"""Reduction function use. """Reduction function use.
A reduction use specifies the reduction function and dimensions. A reduction use specifies the reduction function and dimensions.
""" """
Show All 31 Linesclass ReduceFnType:
def __repr__(self): def __repr__(self):
return f"reduce_{repr(self.binary_fn)}" return f"reduce_{repr(self.binary_fn)}"
class ReduceFn: class ReduceFn:
add = ReduceFnType(BinaryFn.add) add = ReduceFnType(BinaryFn.add)
mul = ReduceFnType(BinaryFn.mul) mul = ReduceFnType(BinaryFn.mul)
max = ReduceFnType(BinaryFn.max) max_signed = ReduceFnType(BinaryFn.max_signed)
min = ReduceFnType(BinaryFn.min) min_signed = ReduceFnType(BinaryFn.min_signed)
max_unsigned = ReduceFnType(BinaryFn.max_unsigned) max_unsigned = ReduceFnType(BinaryFn.max_unsigned)
min_unsigned = ReduceFnType(BinaryFn.min_unsigned) min_unsigned = ReduceFnType(BinaryFn.min_unsigned)
############################################################################### ###############################################################################
# Operand definitions. # Operand definitions.
############################################################################### ###############################################################################
▲ Show 20 Lines • Show All 360 Lines • Show Last 20 Lines

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

Show First 20 Lines • Show All 364 Lines • ▼ Show 20 Lines def _cast_to_floating_point(self, to_type: Type, operand: Value,
from_width = _get_floating_point_width(operand_type) from_width = _get_floating_point_width(operand_type)
if to_width > from_width: if to_width > from_width:
return arith.ExtFOp(to_type, operand).result return arith.ExtFOp(to_type, operand).result
elif to_width < from_width: elif to_width < from_width:
return arith.TruncFOp(to_type, operand).result return arith.TruncFOp(to_type, operand).result
raise ValueError(f"Unable to cast body expression from {operand_type} to " raise ValueError(f"Unable to cast body expression from {operand_type} to "
f"{to_type}") f"{to_type}")
def _type_cast(self, type_var_name: str, operand: Value) -> Value: def _type_cast_signed(self, type_var_name: str, operand: Value) -> Value:
return self._cast(type_var_name, operand, False) return self._cast(type_var_name, operand, False)
def _type_cast_unsigned(self, type_var_name: str, operand: Value) -> Value: def _type_cast_unsigned(self, type_var_name: str, operand: Value) -> Value:
return self._cast(type_var_name, operand, True) return self._cast(type_var_name, operand, True)
def _unary_exp(self, x: Value) -> Value: def _unary_exp(self, x: Value) -> Value:
if _is_floating_point_type(x.type): if _is_floating_point_type(x.type):
return math.ExpOp(x).result return math.ExpOp(x).result
Show All 20 Linesclass _BodyBuilder:
def _binary_mul(self, lhs: Value, rhs: Value) -> Value: def _binary_mul(self, lhs: Value, rhs: Value) -> Value:
if _is_floating_point_type(lhs.type): if _is_floating_point_type(lhs.type):
return arith.MulFOp(lhs, rhs).result return arith.MulFOp(lhs, rhs).result
if _is_integer_type(lhs.type) or _is_index_type(lhs.type): if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
return arith.MulIOp(lhs, rhs).result return arith.MulIOp(lhs, rhs).result
raise NotImplementedError("Unsupported 'mul' operands: {lhs}, {rhs}") raise NotImplementedError("Unsupported 'mul' operands: {lhs}, {rhs}")
def _binary_max(self, lhs: Value, rhs: Value) -> Value: def _binary_max_signed(self, lhs: Value, rhs: Value) -> Value:
if _is_floating_point_type(lhs.type): if _is_floating_point_type(lhs.type):
return arith.MaxFOp(lhs, rhs).result return arith.MaxFOp(lhs, rhs).result
if _is_integer_type(lhs.type) or _is_index_type(lhs.type): if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
return arith.MaxSIOp(lhs, rhs).result return arith.MaxSIOp(lhs, rhs).result
raise NotImplementedError("Unsupported 'max' operands: {lhs}, {rhs}") raise NotImplementedError("Unsupported 'max' operands: {lhs}, {rhs}")
def _binary_max_unsigned(self, lhs: Value, rhs: Value) -> Value: def _binary_max_unsigned(self, lhs: Value, rhs: Value) -> Value:
if _is_floating_point_type(lhs.type): if _is_floating_point_type(lhs.type):
return arith.MaxFOp(lhs, rhs).result return arith.MaxFOp(lhs, rhs).result
if _is_integer_type(lhs.type) or _is_index_type(lhs.type): if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
return arith.MaxUIOp(lhs, rhs).result return arith.MaxUIOp(lhs, rhs).result
raise NotImplementedError( raise NotImplementedError(
"Unsupported 'max_unsigned' operands: {lhs}, {rhs}") "Unsupported 'max_unsigned' operands: {lhs}, {rhs}")
def _binary_min(self, lhs: Value, rhs: Value) -> Value: def _binary_min_signed(self, lhs: Value, rhs: Value) -> Value:
if _is_floating_point_type(lhs.type): if _is_floating_point_type(lhs.type):
return arith.MinFOp(lhs, rhs).result return arith.MinFOp(lhs, rhs).result
if _is_integer_type(lhs.type) or _is_index_type(lhs.type): if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
return arith.MinSIOp(lhs, rhs).result return arith.MinSIOp(lhs, rhs).result
raise NotImplementedError("Unsupported 'min' operands: {lhs}, {rhs}") raise NotImplementedError("Unsupported 'min' operands: {lhs}, {rhs}")
def _binary_min_unsigned(self, lhs: Value, rhs: Value) -> Value: def _binary_min_unsigned(self, lhs: Value, rhs: Value) -> Value:
if _is_floating_point_type(lhs.type): if _is_floating_point_type(lhs.type):
▲ Show 20 Lines • Show All 84 Lines • Show Last 20 Lines

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

from ..lang import * from ..lang import *
T1 = TV.T1 T1 = TV.T1
T2 = TV.T2 T2 = TV.T2
Batch = S.Batch Batch = S.Batch
@linalg_structured_op @linalg_structured_op
def elemwise_unary( def elemwise_unary(
I=TensorDef(T1), I=TensorDef(T1),
O=TensorDef(U, output=True), O=TensorDef(U, output=True),
fun=UnaryFnAttrDef(default=UnaryFn.exp), fun=UnaryFnAttrDef(default=UnaryFn.exp),
cast=TypeFnAttrDef(default=TypeFn.cast)): cast=TypeFnAttrDef(default=TypeFn.cast_signed)):
"""Applies the unary function fun elementwise. """Applies the unary function fun elementwise.
Numeric casting is performed on the input operand, promoting it to the same Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output. data type as the accumulator/output.
""" """
O[None] = fun(cast(U, I[None])) O[None] = fun(cast(U, I[None]))
@linalg_structured_op @linalg_structured_op
def elemwise_binary( def elemwise_binary(
lhs=TensorDef(T1), lhs=TensorDef(T1),
rhs=TensorDef(T2), rhs=TensorDef(T2),
O=TensorDef(U, output=True), O=TensorDef(U, output=True),
fun=BinaryFnAttrDef(default=BinaryFn.add), fun=BinaryFnAttrDef(default=BinaryFn.add),
cast=TypeFnAttrDef(default=TypeFn.cast)): cast=TypeFnAttrDef(default=TypeFn.cast_signed)):
"""Applies the binary function fun elementwise. """Applies the binary function fun elementwise.
Numeric casting is performed on the input operand, promoting it to the same Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output. data type as the accumulator/output.
""" """
O[None] = fun(cast(U, lhs[None]), cast(U, rhs[None])) O[None] = fun(cast(U, lhs[None]), cast(U, rhs[None]))
@linalg_structured_op @linalg_structured_op
def matmul( def matmul(
A=TensorDef(T1, S.M, S.K), A=TensorDef(T1, S.M, S.K),
B=TensorDef(T2, S.K, S.N), B=TensorDef(T2, S.K, S.N),
C=TensorDef(U, S.M, S.N, output=True), C=TensorDef(U, S.M, S.N, output=True),
cast=TypeFnAttrDef(default=TypeFn.cast)): cast=TypeFnAttrDef(default=TypeFn.cast_signed)):
"""Performs a matrix multiplication of two 2D inputs. """Performs a matrix multiplication of two 2D inputs.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
domain(D.m, D.n, D.k) domain(D.m, D.n, D.k)
implements(ContractionOpInterface) implements(ContractionOpInterface)
C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
Show All 25 Linesdef quantized_matmul(
"""Performs a matrix multiplication of two 2D inputs. """Performs a matrix multiplication of two 2D inputs.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. The quantized variant them to the same data type as the accumulator/output. The quantized variant
includes zero-point adjustments for the left and right operands of the includes zero-point adjustments for the left and right operands of the
matmul. matmul.
""" """
domain(D.m, D.n, D.k) domain(D.m, D.n, D.k)
C[D.m, D.n] += (TypeFn.cast(U, A[D.m, D.k]) - TypeFn.cast(U, AZp)) * ( C[D.m, D.n] += (
TypeFn.cast(U, B[D.k, D.n]) - TypeFn.cast(U, BZp)) TypeFn.cast_signed(U, A[D.m, D.k]) - TypeFn.cast_signed(U, AZp)) * (
TypeFn.cast_signed(U, B[D.k, D.n]) - TypeFn.cast_signed(U, BZp))
@linalg_structured_op @linalg_structured_op
def mmt4d( def mmt4d(
lhs=TensorDef(TV.LhsType, S.M, S.K, S.M0, S.K0), lhs=TensorDef(TV.LhsType, S.M, S.K, S.M0, S.K0),
rhs=TensorDef(TV.RhsType, S.N, S.K, S.N0, S.K0), rhs=TensorDef(TV.RhsType, S.N, S.K, S.N0, S.K0),
accum=TensorDef(TV.AccumType, S.M, S.N, S.M0, S.N0, output=True)): accum=TensorDef(TV.AccumType, S.M, S.N, S.M0, S.N0, output=True)):
"""Performs a matrix-matrix-transpose multiplication of two 4D inputs. """Performs a matrix-matrix-transpose multiplication of two 4D inputs.
Differences from linalg.matmul: Differences from linalg.matmul:
* The right hand side is transposed, whence the 't' in 'mmt'. * The right hand side is transposed, whence the 't' in 'mmt'.
* The input and output tensors have a 4D shape instead of a 2D shape. They * The input and output tensors have a 4D shape instead of a 2D shape. They
are interpreted as 2D matrices with one level of 2D tile subdivision, are interpreted as 2D matrices with one level of 2D tile subdivision,
whence the 2+2=4 dimensions. The inner tile dimensions are identified with whence the 2+2=4 dimensions. The inner tile dimensions are identified with
'0' suffixes below, for instance the LHS matrix shape (M, K, M0, K0) reads '0' suffixes below, for instance the LHS matrix shape (M, K, M0, K0) reads
as: MxK tiles, each of shape M0xK0. as: MxK tiles, each of shape M0xK0.
""" """
domain(D.m, D.n, D.k, D.m0, D.n0, D.k0) domain(D.m, D.n, D.k, D.m0, D.n0, D.k0)
implements(ContractionOpInterface) implements(ContractionOpInterface)
accum[D.m, D.n, D.m0, D.n0] += TypeFn.cast( accum[D.m, D.n, D.m0, D.n0] += TypeFn.cast_signed(
TV.AccumType, lhs[D.m, D.k, D.m0, D.k0]) * TypeFn.cast( TV.AccumType, lhs[D.m, D.k, D.m0, D.k0]) * TypeFn.cast_signed(
TV.AccumType, rhs[D.n, D.k, D.n0, D.k0]) TV.AccumType, rhs[D.n, D.k, D.n0, D.k0])
@linalg_structured_op @linalg_structured_op
def batch_matmul( def batch_matmul(
A=TensorDef(T1, Batch, S.M, S.K), A=TensorDef(T1, Batch, S.M, S.K),
B=TensorDef(T2, Batch, S.K, S.N), B=TensorDef(T2, Batch, S.K, S.N),
C=TensorDef(U, Batch, S.M, S.N, output=True)): C=TensorDef(U, Batch, S.M, S.N, output=True)):
"""Performs a batched matrix multiplication of two 3D inputs. """Performs a batched matrix multiplication of two 3D inputs.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
domain(D.b, D.m, D.n, D.k) domain(D.b, D.m, D.n, D.k)
implements(ContractionOpInterface) implements(ContractionOpInterface)
C[D.b, D.m, C[D.b, D.m,
D.n] += TypeFn.cast(U, A[D.b, D.m, D.k]) * TypeFn.cast(U, B[D.b, D.k, D.n]) D.n] += TypeFn.cast_signed(U, A[D.b, D.m, D.k]) * TypeFn.cast_signed(
U, B[D.b, D.k, D.n])
@linalg_structured_op @linalg_structured_op
def quantized_batch_matmul( def quantized_batch_matmul(
A=TensorDef(T1, Batch, S.M, S.K), A=TensorDef(T1, Batch, S.M, S.K),
B=TensorDef(T2, Batch, S.K, S.N), B=TensorDef(T2, Batch, S.K, S.N),
AZp=ScalarDef(I32), AZp=ScalarDef(I32),
BZp=ScalarDef(I32), BZp=ScalarDef(I32),
C=TensorDef(U, Batch, S.M, S.N, output=True)): C=TensorDef(U, Batch, S.M, S.N, output=True)):
"""Performs a batched matrix multiplication of two 3D inputs. """Performs a batched matrix multiplication of two 3D inputs.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. The quantized variant them to the same data type as the accumulator/output. The quantized variant
includes zero-point adjustments for the left and right operands of the includes zero-point adjustments for the left and right operands of the
matmul. matmul.
""" """
domain(D.b, D.m, D.n, D.k) domain(D.b, D.m, D.n, D.k)
C[D.b, D.m, C[D.b, D.m, D.n] += (
D.n] += (TypeFn.cast(U, A[D.b, D.m, D.k]) - TypeFn.cast(U, AZp)) * ( TypeFn.cast_signed(U, A[D.b, D.m, D.k]) - TypeFn.cast_signed(U, AZp)) * (
TypeFn.cast(U, B[D.b, D.k, D.n]) - TypeFn.cast(U, BZp)) TypeFn.cast_signed(U, B[D.b, D.k, D.n]) - TypeFn.cast_signed(U, BZp))
@linalg_structured_op @linalg_structured_op
def matvec( def matvec(
A=TensorDef(T1, S.M, S.N), A=TensorDef(T1, S.M, S.N),
y=TensorDef(T2, S.N), y=TensorDef(T2, S.N),
x=TensorDef(U, S.M, output=True)): x=TensorDef(U, S.M, output=True)):
"""Performs a matrix-vector multiplication. """Performs a matrix-vector multiplication.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
domain(D.m, D.n) domain(D.m, D.n)
implements(ContractionOpInterface) implements(ContractionOpInterface)
x[D.m] += TypeFn.cast(U, A[D.m, D.n]) * TypeFn.cast(U, y[D.n]) x[D.m] += TypeFn.cast_signed(U, A[D.m, D.n]) * TypeFn.cast_signed(U, y[D.n])
@linalg_structured_op @linalg_structured_op
def vecmat( def vecmat(
y=TensorDef(T1, S.M), y=TensorDef(T1, S.M),
A=TensorDef(T2, S.M, S.N), A=TensorDef(T2, S.M, S.N),
x=TensorDef(U, S.N, output=True)): x=TensorDef(U, S.N, output=True)):
"""Performs a vector-matrix multiplication. """Performs a vector-matrix multiplication.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
domain(D.n, D.m) domain(D.n, D.m)
implements(ContractionOpInterface) implements(ContractionOpInterface)
x[D.n] += TypeFn.cast(U, y[D.m]) * TypeFn.cast(U, A[D.m, D.n]) x[D.n] += TypeFn.cast_signed(U, y[D.m]) * TypeFn.cast_signed(U, A[D.m, D.n])
@linalg_structured_op @linalg_structured_op
def batch_matvec( def batch_matvec(
A=TensorDef(T1, Batch, S.M, S.K), A=TensorDef(T1, Batch, S.M, S.K),
B=TensorDef(T2, Batch, S.K), B=TensorDef(T2, Batch, S.K),
C=TensorDef(U, Batch, S.M, output=True)): C=TensorDef(U, Batch, S.M, output=True)):
"""Performs a batched matrix-vector multiplication. """Performs a batched matrix-vector multiplication.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
domain(D.b, D.m, D.k) domain(D.b, D.m, D.k)
implements(ContractionOpInterface) implements(ContractionOpInterface)
C[D.b, D.m] += TypeFn.cast(U, A[D.b, D.m, D.k]) * TypeFn.cast(U, B[D.b, D.k]) C[D.b, D.m] += TypeFn.cast_signed(U, A[D.b, D.m, D.k]) * TypeFn.cast_signed(
U, B[D.b, D.k])
@linalg_structured_op @linalg_structured_op
def dot( def dot(
A=TensorDef(T1, S.M), B=TensorDef(T2, S.M), C=TensorDef(U, output=True)): A=TensorDef(T1, S.M), B=TensorDef(T2, S.M), C=TensorDef(U, output=True)):
"""Performs a dot product of two vectors to a scalar result. """Performs a dot product of two vectors to a scalar result.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ContractionOpInterface) implements(ContractionOpInterface)
C[None] += TypeFn.cast(U, A[D.m]) * TypeFn.cast(U, B[D.m]) C[None] += TypeFn.cast_signed(U, A[D.m]) * TypeFn.cast_signed(U, B[D.m])
@linalg_structured_op @linalg_structured_op
def conv_1d( def conv_1d(
I=TensorDef(T1, S.OW + S.KW), I=TensorDef(T1, S.OW + S.KW),
K=TensorDef(T2, S.KW), K=TensorDef(T2, S.KW),
O=TensorDef(U, S.OW, output=True)): O=TensorDef(U, S.OW, output=True)):
"""Performs 1-D convolution with no channels. """Performs 1-D convolution with no channels.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.ow, D.kw) domain(D.ow, D.kw)
O[D.ow] += TypeFn.cast(U, I[D.ow + D.kw]) * TypeFn.cast(U, K[D.kw]) O[D.ow] += TypeFn.cast_signed(U, I[D.ow + D.kw]) * TypeFn.cast_signed(
U, K[D.kw])
@linalg_structured_op @linalg_structured_op
def conv_2d( def conv_2d(
I=TensorDef(T1, S.OH + S.KH, S.OW + S.KW), I=TensorDef(T1, S.OH + S.KH, S.OW + S.KW),
K=TensorDef(T2, S.KH, S.KW), K=TensorDef(T2, S.KH, S.KW),
O=TensorDef(U, S.OH, S.OW, output=True)): O=TensorDef(U, S.OH, S.OW, output=True)):
"""Performs 2-D convolution with no channels. """Performs 2-D convolution with no channels.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.oh, D.ow, D.kh, D.kw) domain(D.oh, D.ow, D.kh, D.kw)
O[D.oh, D.ow] += TypeFn.cast(U, I[D.oh + D.kh, D.ow + D.kw]) * TypeFn.cast( O[D.oh, D.ow] += TypeFn.cast_signed(
U, K[D.kh, D.kw]) U, I[D.oh + D.kh, D.ow + D.kw]) * TypeFn.cast_signed(U, K[D.kh, D.kw])
@linalg_structured_op @linalg_structured_op
def conv_3d( def conv_3d(
I=TensorDef(T1, S.OD + S.KD, S.OH + S.KH, S.OW + S.KW), I=TensorDef(T1, S.OD + S.KD, S.OH + S.KH, S.OW + S.KW),
K=TensorDef(T2, S.KD, S.KH, S.KW), K=TensorDef(T2, S.KD, S.KH, S.KW),
O=TensorDef(U, S.OD, S.OH, S.OW, output=True)): O=TensorDef(U, S.OD, S.OH, S.OW, output=True)):
"""Performs 3-D convolution with no channels. """Performs 3-D convolution with no channels.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.od, D.oh, D.ow, D.kd, D.kh, D.kw) domain(D.od, D.oh, D.ow, D.kd, D.kh, D.kw)
O[D.od, D.oh, O[D.od, D.oh, D.ow] += TypeFn.cast_signed(
D.ow] += TypeFn.cast(U, I[D.od + D.kd, D.oh + D.kh, D.ow + U, I[D.od + D.kd, D.oh + D.kh, D.ow + D.kw]) * TypeFn.cast_signed(
D.kw]) * TypeFn.cast(U, K[D.kd, D.kh, D.kw]) U, K[D.kd, D.kh, D.kw])
@linalg_structured_op @linalg_structured_op
def conv_1d_nwc_wcf( def conv_1d_nwc_wcf(
I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.C), I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.C),
K=TensorDef(T2, S.KW, S.C, S.F), K=TensorDef(T2, S.KW, S.C, S.F),
O=TensorDef(U, S.N, S.OW, S.F, output=True), O=TensorDef(U, S.N, S.OW, S.F, output=True),
strides=IndexAttrDef(S.SW, default=[1]), strides=IndexAttrDef(S.SW, default=[1]),
dilations=IndexAttrDef(S.DW, default=[1])): dilations=IndexAttrDef(S.DW, default=[1])):
"""Performs 1-D convolution. """Performs 1-D convolution.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.ow, D.f, D.kw, D.c) domain(D.n, D.ow, D.f, D.kw, D.c)
O[D.n, D.ow, O[D.n, D.ow,
D.f] += TypeFn.cast(U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.f] += TypeFn.cast_signed(U, I[D.n, D.ow * S.SW + D.kw * S.DW,
D.c]) * TypeFn.cast(U, K[D.kw, D.c, D.f]) D.c]) * TypeFn.cast_signed(U, K[D.kw, D.c, D.f])
@linalg_structured_op @linalg_structured_op
def conv_2d_nhwc_hwcf( def conv_2d_nhwc_hwcf(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.C), S.C),
K=TensorDef(T2, S.KH, S.KW, S.C, S.F), K=TensorDef(T2, S.KH, S.KW, S.C, S.F),
O=TensorDef(U, S.N, S.OH, S.OW, S.F, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.F, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs 2-D convolution. """Performs 2-D convolution.
Layout: Layout:
* Input: NHWC. * Input: NHWC.
* Kernel: HWCF. * Kernel: HWCF.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.oh, D.ow, D.f, D.kh, D.kw, D.c) domain(D.n, D.oh, D.ow, D.f, D.kh, D.kw, D.c)
O[D.n, D.oh, D.ow, D.f] += TypeFn.cast( O[D.n, D.oh, D.ow, D.f] += TypeFn.cast_signed(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
D.c]) * TypeFn.cast(U, K[D.kh, D.kw, D.c, D.f]) D.c]) * TypeFn.cast_signed(U, K[D.kh, D.kw, D.c, D.f])
@linalg_structured_op @linalg_structured_op
def conv_2d_nhwc_hwcf_q( def conv_2d_nhwc_hwcf_q(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.C), S.C),
K=TensorDef(T2, S.KH, S.KW, S.C, S.F), K=TensorDef(T2, S.KH, S.KW, S.C, S.F),
IZp=ScalarDef(I32), IZp=ScalarDef(I32),
Show All 9 Linesdef conv_2d_nhwc_hwcf_q(
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. This includes the zero them to the same data type as the accumulator/output. This includes the zero
point offsets common to quantized operations. point offsets common to quantized operations.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.oh, D.ow, D.f, D.kh, D.kw, D.c) domain(D.n, D.oh, D.ow, D.f, D.kh, D.kw, D.c)
O[D.n, D.oh, D.ow, O[D.n, D.oh, D.ow,
D.f] += (TypeFn.cast( D.f] += (TypeFn.cast_signed(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]) - U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]) -
TypeFn.cast(U, IZp)) * ( TypeFn.cast_signed(U, IZp)) * (
TypeFn.cast(U, K[D.kh, D.kw, D.c, D.f]) - TypeFn.cast(U, KZp)) TypeFn.cast_signed(U, K[D.kh, D.kw, D.c, D.f]) -
TypeFn.cast_signed(U, KZp))
@linalg_structured_op @linalg_structured_op
def conv_2d_nchw_fchw( def conv_2d_nchw_fchw(
I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH, I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH,
S.OW * S.SW + S.KW * S.DW), S.OW * S.SW + S.KW * S.DW),
K=TensorDef(T2, S.F, S.C, S.KH, S.KW), K=TensorDef(T2, S.F, S.C, S.KH, S.KW),
O=TensorDef(U, S.N, S.F, S.OH, S.OW, output=True), O=TensorDef(U, S.N, S.F, S.OH, S.OW, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs 2-D convolution. """Performs 2-D convolution.
Layout: Layout:
* Input: NCHW. * Input: NCHW.
* Kernel: FCHW. * Kernel: FCHW.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.f, D.oh, D.ow, D.c, D.kh, D.kw) domain(D.n, D.f, D.oh, D.ow, D.c, D.kh, D.kw)
O[D.n, D.f, D.oh, D.ow] += TypeFn.cast( O[D.n, D.f, D.oh, D.ow] += TypeFn.cast_signed(
U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW +
D.ow * S.SW + D.kw * S.DW]) * TypeFn.cast(U, K[D.f, D.c, D.kh, D.kw]) D.kw * S.DW]) * TypeFn.cast_signed(U, K[D.f, D.c, D.kh, D.kw])
@linalg_structured_op @linalg_structured_op
def conv_3d_ndhwc_dhwcf( def conv_3d_ndhwc_dhwcf(
I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH, I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
S.OW * S.SW + S.KW * S.DW, S.C), S.OW * S.SW + S.KW * S.DW, S.C),
K=TensorDef(T2, S.KD, S.KH, S.KW, S.C, S.F), K=TensorDef(T2, S.KD, S.KH, S.KW, S.C, S.F),
O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.F, output=True), O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.F, output=True),
strides=IndexAttrDef(S.SD, S.SH, S.SW, default=[1, 1, 1]), strides=IndexAttrDef(S.SD, S.SH, S.SW, default=[1, 1, 1]),
dilations=IndexAttrDef(S.DD, S.DH, S.DW, default=[1, 1, 1])): dilations=IndexAttrDef(S.DD, S.DH, S.DW, default=[1, 1, 1])):
"""Performs 3-D convolution. """Performs 3-D convolution.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.od, D.oh, D.ow, D.f, D.kd, D.kh, D.kw, D.c) domain(D.n, D.od, D.oh, D.ow, D.f, D.kd, D.kh, D.kw, D.c)
O[D.n, D.od, D.oh, D.ow, D.f] += TypeFn.cast( O[D.n, D.od, D.oh, D.ow, D.f] += TypeFn.cast_signed(
U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH, U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
D.ow * S.SW + D.kw * S.DW, D.c]) * TypeFn.cast( D.ow * S.SW + D.kw * S.DW, D.c]) * TypeFn.cast_signed(
U, K[D.kd, D.kh, D.kw, D.c, D.f]) U, K[D.kd, D.kh, D.kw, D.c, D.f])
@linalg_structured_op @linalg_structured_op
def depthwise_conv_1d_nwc_wc( def depthwise_conv_1d_nwc_wc(
I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.IC), I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.IC),
K=TensorDef(T2, S.KW, S.IC), K=TensorDef(T2, S.KW, S.IC),
O=TensorDef(U, S.N, S.OW, S.IC, output=True), O=TensorDef(U, S.N, S.OW, S.IC, output=True),
strides=IndexAttrDef(S.SW, default=[1]), strides=IndexAttrDef(S.SW, default=[1]),
dilations=IndexAttrDef(S.DW, default=[1])): dilations=IndexAttrDef(S.DW, default=[1])):
"""Performs depth-wise 1-D convolution. """Performs depth-wise 1-D convolution.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. Multiplier is set to 1 them to the same data type as the accumulator/output. Multiplier is set to 1
which is a special case for most depthwise convolutions. which is a special case for most depthwise convolutions.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.ow, D.ic, D.kw) domain(D.n, D.ow, D.ic, D.kw)
O[D.n, D.ow, D.ic] += \ O[D.n, D.ow, D.ic] += \
TypeFn.cast(U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.ic]) * \ TypeFn.cast_signed(U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.ic]) * \
TypeFn.cast(U, K[D.kw, D.ic]) TypeFn.cast_signed(U, K[D.kw, D.ic])
@linalg_structured_op @linalg_structured_op
def depthwise_conv_2d_nhwc_hwc( def depthwise_conv_2d_nhwc_hwc(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.IC), S.IC),
K=TensorDef(T2, S.KH, S.KW, S.IC), K=TensorDef(T2, S.KH, S.KW, S.IC),
O=TensorDef(U, S.N, S.OH, S.OW, S.IC, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.IC, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs depth-wise 2-D convolution. """Performs depth-wise 2-D convolution.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. Multiplier is set to 1 them to the same data type as the accumulator/output. Multiplier is set to 1
which is a special case for most depthwise convolutions. which is a special case for most depthwise convolutions.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.oh, D.ow, D.ic, D.kh, D.kw) domain(D.n, D.oh, D.ow, D.ic, D.kh, D.kw)
O[D.n, D.oh, D.ow, D.ic] += TypeFn.cast( O[D.n, D.oh, D.ow, D.ic] += TypeFn.cast_signed(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
D.ic]) * TypeFn.cast(U, K[D.kh, D.kw, D.ic]) D.ic]) * TypeFn.cast_signed(U, K[D.kh, D.kw, D.ic])
@linalg_structured_op @linalg_structured_op
def depthwise_conv_2d_nhwc_hwc_q( def depthwise_conv_2d_nhwc_hwc_q(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.IC), S.IC),
K=TensorDef(T2, S.KH, S.KW, S.IC), K=TensorDef(T2, S.KH, S.KW, S.IC),
IZp=ScalarDef(I32), IZp=ScalarDef(I32),
KZp=ScalarDef(I32), KZp=ScalarDef(I32),
O=TensorDef(U, S.N, S.OH, S.OW, S.IC, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.IC, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs depth-wise 2-D convolution. """Performs depth-wise 2-D convolution.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.oh, D.ow, D.ic, D.kh, D.kw) domain(D.n, D.oh, D.ow, D.ic, D.kh, D.kw)
O[D.n, D.oh, D.ow, O[D.n, D.oh, D.ow, D.ic] += ((TypeFn.cast_signed(
D.ic] += ((TypeFn.cast(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.ic]) - U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.ic]) -
TypeFn.cast(U, IZp)) * TypeFn.cast_signed(U, IZp)) *
(TypeFn.cast(U, K[D.kh, D.kw, D.ic]) - TypeFn.cast(U, KZp))) (TypeFn.cast_signed(U, K[D.kh, D.kw, D.ic]) -
TypeFn.cast_signed(U, KZp)))
@linalg_structured_op @linalg_structured_op
def depthwise_conv_2d_nhwc_hwcm( def depthwise_conv_2d_nhwc_hwcm(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.IC), S.IC),
K=TensorDef(T2, S.KH, S.KW, S.IC, S.CM), K=TensorDef(T2, S.KH, S.KW, S.IC, S.CM),
O=TensorDef(U, S.N, S.OH, S.OW, S.IC, S.CM, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.IC, S.CM, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs depth-wise 2-D convolution. """Performs depth-wise 2-D convolution.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.oh, D.ow, D.ic, D.cm, D.kh, D.kw) domain(D.n, D.oh, D.ow, D.ic, D.cm, D.kh, D.kw)
O[D.n, D.oh, D.ow, D.ic, D.cm] += TypeFn.cast( O[D.n, D.oh, D.ow, D.ic, D.cm] += TypeFn.cast_signed(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
D.ic]) * TypeFn.cast(U, K[D.kh, D.kw, D.ic, D.cm]) D.ic]) * TypeFn.cast_signed(U, K[D.kh, D.kw, D.ic, D.cm])
@linalg_structured_op @linalg_structured_op
def depthwise_conv_2d_nhwc_hwcm_q( def depthwise_conv_2d_nhwc_hwcm_q(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.IC), S.IC),
K=TensorDef(T2, S.KH, S.KW, S.IC, S.CM), K=TensorDef(T2, S.KH, S.KW, S.IC, S.CM),
IZp=ScalarDef(I32), IZp=ScalarDef(I32),
KZp=ScalarDef(I32), KZp=ScalarDef(I32),
O=TensorDef(U, S.N, S.OH, S.OW, S.IC, S.CM, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.IC, S.CM, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs depth-wise 2-D convolution. """Performs depth-wise 2-D convolution.
Numeric casting is performed on the operands to the inner multiply, promoting Numeric casting is performed on the operands to the inner multiply, promoting
them to the same data type as the accumulator/output. them to the same data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.oh, D.ow, D.ic, D.cm, D.kh, D.kw) domain(D.n, D.oh, D.ow, D.ic, D.cm, D.kh, D.kw)
O[D.n, D.oh, D.ow, D.ic, O[D.n, D.oh, D.ow, D.ic,
D.cm] += ((TypeFn.cast( D.cm] += ((TypeFn.cast_signed(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.ic]) - U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.ic]) -
TypeFn.cast(U, IZp)) * TypeFn.cast_signed(U, IZp)) *
(TypeFn.cast(U, K[D.kh, D.kw, D.ic, D.cm]) - TypeFn.cast(U, KZp))) (TypeFn.cast_signed(U, K[D.kh, D.kw, D.ic, D.cm]) -
TypeFn.cast_signed(U, KZp)))
@linalg_structured_op @linalg_structured_op
def pooling_nhwc_sum( def pooling_nhwc_sum(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.C), S.C),
K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs sum pooling. """Performs sum pooling.
Numeric casting is performed on the input operand, promoting it to the same Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output. data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw) domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
O[D.n, D.oh, D.ow, D.c] += TypeFn.cast( O[D.n, D.oh, D.ow, D.c] += TypeFn.cast_signed(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]) U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])
@linalg_structured_op @linalg_structured_op
def pooling_nhwc_max( def pooling_nhwc_max(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.C), S.C),
K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs max pooling. """Performs max pooling.
Numeric casting is performed on the input operand, promoting it to the same Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output. data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw) domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
O[D.n, D.oh, D.ow, D.c] = ReduceFn.max[D.kh, D.kw]( O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_signed[D.kh, D.kw](
TypeFn.cast( TypeFn.cast_signed(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])) U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
@linalg_structured_op @linalg_structured_op
def pooling_nhwc_max_unsigned( def pooling_nhwc_max_unsigned(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.C), S.C),
K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
Show All 22 Linesdef pooling_nchw_max(
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs max pooling. """Performs max pooling.
Numeric casting is performed on the input operand, promoting it to the same Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output. data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.c, D.oh, D.ow, D.kh, D.kw) domain(D.n, D.c, D.oh, D.ow, D.kh, D.kw)
O[D.n, D.c, D.oh, D.ow] = ReduceFn.max[D.kh, D.kw]( O[D.n, D.c, D.oh, D.ow] = ReduceFn.max_signed[D.kh, D.kw](
TypeFn.cast( TypeFn.cast_signed(
U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH,
D.ow * S.SW + D.kw * S.DW,])) D.ow * S.SW + D.kw * S.DW,]))
@linalg_structured_op @linalg_structured_op
def pooling_nhwc_min( def pooling_nhwc_min(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.C), S.C),
K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
"""Performs min pooling. """Performs min pooling.
Numeric casting is performed on the input operand, promoting it to the same Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output. data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw) domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
O[D.n, D.oh, D.ow, D.c] = ReduceFn.min[D.kh, D.kw]( O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_signed[D.kh, D.kw](
TypeFn.cast( TypeFn.cast_signed(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])) U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
@linalg_structured_op @linalg_structured_op
def pooling_nhwc_min_unsigned( def pooling_nhwc_min_unsigned(
I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW, I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
S.C), S.C),
K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
Show All 22 Linesdef pooling_ndhwc_sum(
dilations=IndexAttrDef(S.DD, S.DH, S.DW, default=[1, 1, 1])): dilations=IndexAttrDef(S.DD, S.DH, S.DW, default=[1, 1, 1])):
"""Performs 3D sum pooling. """Performs 3D sum pooling.
Numeric casting is performed on the input operand, promoting it to the same Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output. data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw) domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw)
O[D.n, D.od, D.oh, D.ow, D.c] += TypeFn.cast( O[D.n, D.od, D.oh, D.ow, D.c] += TypeFn.cast_signed(
U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH, U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
D.ow * S.SW + D.kw * S.DW, D.c]) D.ow * S.SW + D.kw * S.DW, D.c])
@linalg_structured_op @linalg_structured_op
def pooling_ndhwc_max( def pooling_ndhwc_max(
I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH, I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
S.OW * S.SW + S.KW * S.DW, S.C), S.OW * S.SW + S.KW * S.DW, S.C),
K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]), K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True), O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
strides=IndexAttrDef(S.SD, S.SH, S.SW, default=[1, 1, 1]), strides=IndexAttrDef(S.SD, S.SH, S.SW, default=[1, 1, 1]),
dilations=IndexAttrDef(S.DD, S.DH, S.DW, default=[1, 1, 1])): dilations=IndexAttrDef(S.DD, S.DH, S.DW, default=[1, 1, 1])):
"""Performs 3D max pooling. """Performs 3D max pooling.
Numeric casting is performed on the input operand, promoting it to the same Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output. data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw) domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw)
O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.max[D.kd, D.kh, D.kw]( O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.max_signed[D.kd, D.kh, D.kw](
TypeFn.cast( TypeFn.cast_signed(
U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH, U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
D.ow * S.SW + D.kw * S.DW, D.c])) D.ow * S.SW + D.kw * S.DW, D.c]))
@linalg_structured_op @linalg_structured_op
def pooling_ndhwc_min( def pooling_ndhwc_min(
I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH, I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
S.OW * S.SW + S.KW * S.DW, S.C), S.OW * S.SW + S.KW * S.DW, S.C),
K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]), K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True), O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
strides=IndexAttrDef(S.SD, S.SH, S.SW, default=[1, 1, 1]), strides=IndexAttrDef(S.SD, S.SH, S.SW, default=[1, 1, 1]),
dilations=IndexAttrDef(S.DD, S.DH, S.DW, default=[1, 1, 1])): dilations=IndexAttrDef(S.DD, S.DH, S.DW, default=[1, 1, 1])):
"""Performs 3D min pooling. """Performs 3D min pooling.
Numeric casting is performed on the input operand, promoting it to the same Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output. data type as the accumulator/output.
""" """
implements(ConvolutionOpInterface) implements(ConvolutionOpInterface)
domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw) domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw)
O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.min[D.kd, D.kh, D.kw]( O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.min_signed[D.kd, D.kh, D.kw](
TypeFn.cast( TypeFn.cast_signed(
U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH, U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
D.ow * S.SW + D.kw * S.DW, D.c])) D.ow * S.SW + D.kw * S.DW, D.c]))
@linalg_structured_op @linalg_structured_op
def fill_tensor(value=ScalarDef(T1), O=TensorDef(U, output=True)): def fill_tensor(value=ScalarDef(T1), O=TensorDef(U, output=True)):
"""Fills the output tensor with the given value. """Fills the output tensor with the given value.
Works for arbitrary ranked output tensors since the operation performs scalar Works for arbitrary ranked output tensors since the operation performs scalar
accesses only and is thus rank polymorphic. Numeric casting is performed on accesses only and is thus rank polymorphic. Numeric casting is performed on
the value operand, promoting it to the same data type as the output. the value operand, promoting it to the same data type as the output.
""" """
O[None] = TypeFn.cast(U, value) O[None] = TypeFn.cast_signed(U, value)
@linalg_structured_op @linalg_structured_op
def fill_rng_2d( def fill_rng_2d(
min=ScalarDef(F64), min=ScalarDef(F64),
max=ScalarDef(F64), max=ScalarDef(F64),
seed=ScalarDef(I32), seed=ScalarDef(I32),
O=TensorDef(T, S.M, S.N, output=True)): O=TensorDef(T, S.M, S.N, output=True)):
"""Fills the output tensor with pseudo random numbers. """Fills the output tensor with pseudo random numbers.
The operation generations pseudo random numbers using a linear congruential The operation generations pseudo random numbers using a linear congruential
generator. It provides no guarantees regarding the distribution of the generator. It provides no guarantees regarding the distribution of the
generated random numbers. Instead of generating the random numbers generated random numbers. Instead of generating the random numbers
sequentially, it instantiates one random number generator per data element sequentially, it instantiates one random number generator per data element
and runs them in parallel. The seed operand and the indices of the data and runs them in parallel. The seed operand and the indices of the data
element seed the random number generation. The min and max operands limit element seed the random number generation. The min and max operands limit
the range of the generated random numbers. the range of the generated random numbers.
""" """
domain(D.m, D.n) domain(D.m, D.n)
multiplier = TypeFn.cast(I32, const(1103515245)) multiplier = TypeFn.cast_signed(I32, const(1103515245))
increment = TypeFn.cast(I32, const(12345)) increment = TypeFn.cast_signed(I32, const(12345))
rand1 = (TypeFn.cast(I32, index(D.m)) + seed) * multiplier + increment rand1 = (TypeFn.cast_signed(I32, index(D.m)) + seed) * multiplier + increment
rand2 = (TypeFn.cast(I32, index(D.n)) + rand1) * multiplier + increment rand2 = (TypeFn.cast_signed(I32, index(D.n)) + rand1) * multiplier + increment
inv_range = TypeFn.cast(F64, const(2.3283064e-10)) inv_range = TypeFn.cast_signed(F64, const(2.3283064e-10))
offset = TypeFn.cast(F64, const(2147483647)) offset = TypeFn.cast_signed(F64, const(2147483647))
scaling = (max - min) * inv_range scaling = (max - min) * inv_range
O[D.m, D.n] = TypeFn.cast(T, O[D.m, D.n] = TypeFn.cast_signed(
(offset + TypeFn.cast(F64, rand2)) * scaling + min) T, (offset + TypeFn.cast_signed(F64, rand2)) * scaling + min)
@linalg_structured_op @linalg_structured_op
def soft_plus_2d( def soft_plus_2d(
I=TensorDef(T, S.M, S.N), O=TensorDef(U, S.M, S.N, output=True)): I=TensorDef(T, S.M, S.N), O=TensorDef(U, S.M, S.N, output=True)):
"""Implements the soft plus operator. """Implements the soft plus operator.
Numeric casting is performed on the input operand, promoting it to the same Numeric casting is performed on the input operand, promoting it to the same
data type as the accumulator/output. data type as the accumulator/output.
""" """
domain(D.m, D.n) domain(D.m, D.n)
O[D.m, D.n] = \ O[D.m, D.n] = \
UnaryFn.log(TypeFn.cast(U, const(1.0)) + UnaryFn.exp(TypeFn.cast(U, I[D.m, D.n]))) UnaryFn.log(TypeFn.cast_signed(U, const(1.0)) + UnaryFn.exp(TypeFn.cast_signed(U, I[D.m, D.n])))

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

# RUN: mlir-linalg-ods-yaml-gen %s --o-ods-decl=- | FileCheck %s --check-prefix=ODS # RUN: mlir-linalg-ods-yaml-gen %s --o-ods-decl=- | FileCheck %s --check-prefix=ODS
# RUN: mlir-linalg-ods-yaml-gen %s --o-impl=- | FileCheck %s --check-prefix=IMPL # RUN: mlir-linalg-ods-yaml-gen %s --o-impl=- | FileCheck %s --check-prefix=IMPL
# @linalg_structured_op # @linalg_structured_op
# def test1(O=TensorDef(T, S.M, S.N, output=True), # def test1(O=TensorDef(T, S.M, S.N, output=True),
# cast=TypeFnAttrDef(default=TypeFn.cast)): # cast=TypeFnAttrDef(default=TypeFn.cast_signed)):
# """Title. # """Title.
# Detailed description. # Detailed description.
# """ # """
# O[D.m, D.n] = cast(T, const(42)) + cast(T, index(D.n)) # O[D.m, D.n] = cast(T, const(42)) + cast(T, index(D.n))
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: test1 name: test1
cpp_class_name: Test1Op cpp_class_name: Test1Op
doc: |- doc: |-
Title. Title.
Detailed description. Detailed description.
structured_op: !LinalgStructuredOpConfig structured_op: !LinalgStructuredOpConfig
args: args:
- !LinalgOperandDefConfig - !LinalgOperandDefConfig
name: O name: O
kind: output_tensor kind: output_tensor
type_var: T type_var: T
shape_map: affine_map<()[s0, s1] -> (s0, s1)> shape_map: affine_map<()[s0, s1] -> (s0, s1)>
- !LinalgOperandDefConfig - !LinalgOperandDefConfig
name: cast name: cast
kind: type_fn_attr kind: type_fn_attr
default_fn: cast default_fn: cast_signed
indexing_maps: !LinalgIndexingMapsConfig indexing_maps: !LinalgIndexingMapsConfig
static_indexing_maps: static_indexing_maps:
- affine_map<(d0, d1)[s0, s1] -> (d0, d1)> - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
iterator_types: iterator_types:
- parallel - parallel
- parallel - parallel
assignments: assignments:
- !ScalarAssign - !ScalarAssign
Show All 25 Lines
# ODS: let summary = [{ Title. }]; # ODS: let summary = [{ Title. }];
# ODS-NEXT: let description = [{ # ODS-NEXT: let description = [{
# ODS-NEXT: Detailed description. # ODS-NEXT: Detailed description.
# ODS-NEXT: }]; # ODS-NEXT: }];
# ODS: let arguments = # ODS: let arguments =
# ODS-NEXT: Variadic<AnyType>:$inputs, # ODS-NEXT: Variadic<AnyType>:$inputs,
# ODS-NEXT: Variadic<AnyShaped>:$outputs, # ODS-NEXT: Variadic<AnyShaped>:$outputs,
# ODS-NEXT: DefaultValuedAttr<TypeFnAttr, "TypeFn::cast">:$cast # ODS-NEXT: DefaultValuedAttr<TypeFnAttr, "TypeFn::cast_signed">:$cast
# ODS: let builders = # ODS: let builders =
# ODS: (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs, # ODS: (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
# ODS-NEXT: "ValueRange":$outputs, # ODS-NEXT: "ValueRange":$outputs,
# ODS-NEXT: CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes), # ODS-NEXT: CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
# ODS: (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs, # ODS: (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
# ODS-NEXT: "ValueRange":$outputs, "Attribute":$cast, # ODS-NEXT: "ValueRange":$outputs, "Attribute":$cast,
Show All 12 Lines
# ODS-NEXT: createAndFillStructuredOpRegion<Test1Op>( # ODS-NEXT: createAndFillStructuredOpRegion<Test1Op>(
# ODS-NEXT: $_builder, # ODS-NEXT: $_builder,
# ODS-NEXT: $_state, # ODS-NEXT: $_state,
# ODS-NEXT: TypeRange(inputs), # ODS-NEXT: TypeRange(inputs),
# ODS-NEXT: TypeRange(outputs) # ODS-NEXT: TypeRange(outputs)
# IMPL-LABEL: void Test1Op::regionBuilder(ImplicitLocOpBuilder &b, # IMPL-LABEL: void Test1Op::regionBuilder(ImplicitLocOpBuilder &b,
# IMPL-NEXT: Block &block, ArrayRef<NamedAttribute> attrs) # IMPL-NEXT: Block &block, ArrayRef<NamedAttribute> attrs)
# IMPL: TypeFn castVal = TypeFn::cast; # IMPL: TypeFn castVal = TypeFn::cast_signed;
# IMPL-NEXT: auto castIter = llvm::find_if(attrs, [&](const NamedAttribute &attr) { # IMPL-NEXT: auto castIter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {
# IMPL-NEXT: return attr.getName() == "cast"; }); # IMPL-NEXT: return attr.getName() == "cast"; });
# IMPL-NEXT: if (castIter != attrs.end()) { # IMPL-NEXT: if (castIter != attrs.end()) {
# IMPL-NEXT: if (auto attr = castIter->getValue().dyn_cast<TypeFnAttr>()) # IMPL-NEXT: if (auto attr = castIter->getValue().dyn_cast<TypeFnAttr>())
# IMPL-NEXT: castVal = attr.getValue(); # IMPL-NEXT: castVal = attr.getValue();
# IMPL-NEXT: } # IMPL-NEXT: }
# IMPL: Value [[VAL0:[a-z0-9]+]] = helper.constant("42 : i64"); # IMPL: Value [[VAL0:[a-z0-9]+]] = helper.constant("42 : i64");
▲ Show 20 Lines • Show All 93 Lines • ▼ Show 20 Lines
# @linalg_structured_op # @linalg_structured_op
# def test3(value=ScalarDef(T1), # def test3(value=ScalarDef(T1),
# O=TensorDef(U, output=True)): # O=TensorDef(U, output=True)):
# """Title. # """Title.
# Detailed description. # Detailed description.
# """ # """
# O[None] = TypeFn.cast(U, value) # O[None] = TypeFn.cast_signed(U, value)
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: test3 name: test3
cpp_class_name: Test3Op cpp_class_name: Test3Op
doc: |- doc: |-
Title. Title.
Show All 15 Lines indexing_maps: !LinalgIndexingMapsConfig
- affine_map<() -> ()> - affine_map<() -> ()>
iterator_types: [] iterator_types: []
assignments: assignments:
- !ScalarAssign - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_fn: scalar_fn:
kind: type kind: type
fn_name: cast fn_name: cast_signed
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: value scalar_arg: value
# IMPL: Test3Op::iterator_types() { # IMPL: Test3Op::iterator_types() {
# IMPL-NEXT: int64_t rank = getRank(getOutputOperand(0)); # IMPL-NEXT: int64_t rank = getRank(getOutputOperand(0));
▲ Show 20 Lines • Show All 83 Lines • Show Last 20 Lines

mlir/test/python/dialects/linalg/opdsl/arguments.py

Show All 20 Lines
# CHECK: name: bfn # CHECK: name: bfn
# CHECK: kind: binary_fn_attr # CHECK: kind: binary_fn_attr
# CHECK: default_fn: mul # CHECK: default_fn: mul
# CHECK: name: ufn # CHECK: name: ufn
# CHECK: kind: unary_fn_attr # CHECK: kind: unary_fn_attr
# CHECK: default_fn: exp # CHECK: default_fn: exp
# CHECK: name: cast # CHECK: name: cast
# CHECK: kind: type_fn_attr # CHECK: kind: type_fn_attr
# CHECK: default_fn: cast # CHECK: default_fn: cast_signed
@linalg_structured_op @linalg_structured_op
def matmul( def matmul(
A=TensorDef(T, S.M, S.K), A=TensorDef(T, S.M, S.K),
B=TensorDef(T, S.K, S.N), B=TensorDef(T, S.K, S.N),
C=TensorDef(U, S.M, S.N, output=True), C=TensorDef(U, S.M, S.N, output=True),
bfn=BinaryFnAttrDef(default=BinaryFn.mul), bfn=BinaryFnAttrDef(default=BinaryFn.mul),
ufn=UnaryFnAttrDef(default=UnaryFn.exp), ufn=UnaryFnAttrDef(default=UnaryFn.exp),
cast=TypeFnAttrDef(default=TypeFn.cast)): cast=TypeFnAttrDef(default=TypeFn.cast_signed)):
C[D.m, D.n] += bfn(cast(U, A[D.m, D.k]), cast(U, B[D.k, D.n])) C[D.m, D.n] += bfn(cast(U, A[D.m, D.k]), cast(U, B[D.k, D.n]))
# CHECK: --- # CHECK: ---
# CHECK-LABEL: fill # CHECK-LABEL: fill
# CHECK: args: # CHECK: args:
# CHECK: name: value # CHECK: name: value
# CHECK: kind: scalar # CHECK: kind: scalar
Show All 30 Lines

mlir/test/python/dialects/linalg/opdsl/assignments.py

Show All 29 Lines
# CHECK: operands: # CHECK: operands:
# CHECK: scalar_arg: B # CHECK: scalar_arg: B
@linalg_structured_op @linalg_structured_op
def matmul( def matmul(
A=TensorDef(T, S.M, S.K), A=TensorDef(T, S.M, S.K),
B=TensorDef(T, S.K, S.N), B=TensorDef(T, S.K, S.N),
C=TensorDef(U, S.M, S.N, output=True), C=TensorDef(U, S.M, S.N, output=True),
mul=BinaryFnAttrDef(default=BinaryFn.mul), mul=BinaryFnAttrDef(default=BinaryFn.mul),
cast=TypeFnAttrDef(default=TypeFn.cast)): cast=TypeFnAttrDef(default=TypeFn.cast_signed)):
C[D.m, D.n] += mul(cast(U, A[D.m, D.k]), cast(U, B[D.k, D.n])) C[D.m, D.n] += mul(cast(U, A[D.m, D.k]), cast(U, B[D.k, D.n]))
# CHECK: --- # CHECK: ---
# CHECK-LABEL: constants # CHECK-LABEL: constants
# CHECK: assignments: # CHECK: assignments:
# CHECK: - # CHECK: -
# CHECK: arg: O # CHECK: arg: O
Show All 11 Lines
# CHECK: operands: # CHECK: operands:
# CHECK: scalar_fn: # CHECK: scalar_fn:
# CHECK: kind: type # CHECK: kind: type
# CHECK: type_var: T # CHECK: type_var: T
# CHECK: operands: # CHECK: operands:
# CHECK: scalar_const: '3.1415926535897931 : f64' # CHECK: scalar_const: '3.1415926535897931 : f64'
# CHECK: scalar_fn: # CHECK: scalar_fn:
# CHECK: kind: type # CHECK: kind: type
# CHECK: fn_name: cast # CHECK: fn_name: cast_signed
# CHECK: type_var: T # CHECK: type_var: T
# CHECK: operands: # CHECK: operands:
# CHECK: scalar_const: '42 : i64' # CHECK: scalar_const: '42 : i64'
# CHECK: scalar_fn: # CHECK: scalar_fn:
# CHECK: kind: type # CHECK: kind: type
# CHECK: fn_name: cast # CHECK: fn_name: cast_signed
# CHECK: type_var: T # CHECK: type_var: T
# CHECK: operands: # CHECK: operands:
# CHECK: scalar_fn: # CHECK: scalar_fn:
# CHECK: kind: unary # CHECK: kind: unary
# CHECK: attr_name: exp # CHECK: attr_name: exp
# CHECK: operands: # CHECK: operands:
# CHECK: scalar_const: '1.{{[0]*}}e+03 : f64' # CHECK: scalar_const: '1.{{[0]*}}e+03 : f64'
@linalg_structured_op @linalg_structured_op
def constants( def constants(
O=TensorDef(T, S.M, S.K, output=True), O=TensorDef(T, S.M, S.K, output=True),
exp=UnaryFnAttrDef(default=UnaryFn.exp)): exp=UnaryFnAttrDef(default=UnaryFn.exp)):
pi = TypeFn.cast(T, const(3.1415926535897931)) pi = TypeFn.cast_signed(T, const(3.1415926535897931))
cst42 = TypeFn.cast(T, const(42)) cst42 = TypeFn.cast_signed(T, const(42))
cst1000 = TypeFn.cast(T, exp(const(1e+3))) cst1000 = TypeFn.cast_signed(T, exp(const(1e+3)))
O[D.m, D.n] = UnaryFn.exp(pi) + cst42 - cst1000 O[D.m, D.n] = UnaryFn.exp(pi) + cst42 - cst1000
# CHECK: --- # CHECK: ---
# CHECK-LABEL: indices # CHECK-LABEL: indices
# CHECK: assignments: # CHECK: assignments:
# CHECK: - # CHECK: -
# CHECK: arg: O # CHECK: arg: O
Show All 20 Lines

mlir/test/python/dialects/linalg/opdsl/emit_convolution.py

Show All 13 Lines
@linalg_structured_op @linalg_structured_op
def conv_poly( def conv_poly(
I=TensorDef(T1, S.N, S.IH, S.IW, S.C), I=TensorDef(T1, S.N, S.IH, S.IW, S.C),
K=TensorDef(T2, S.KH, S.KW, S.C), K=TensorDef(T2, S.KH, S.KW, S.C),
O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 2])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 2])):
domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c) domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
O[D.n, D.oh, D.ow, D.c] += TypeFn.cast( O[D.n, D.oh, D.ow, D.c] += TypeFn.cast_signed(
U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
D.c]) * TypeFn.cast(U, K[D.kh, D.kw, D.c]) D.c]) * TypeFn.cast_signed(U, K[D.kh, D.kw, D.c])
with Context() as ctx, Location.unknown(): with Context() as ctx, Location.unknown():
module = Module.create() module = Module.create()
f32 = F32Type.get() f32 = F32Type.get()
i32 = IntegerType.get_signless(32) i32 = IntegerType.get_signless(32)
with InsertionPoint(module.body): with InsertionPoint(module.body):
Show All 27 Lines

mlir/test/python/dialects/linalg/opdsl/emit_fill.py

# RUN: %PYTHON %s | FileCheck %s # RUN: %PYTHON %s | FileCheck %s
from mlir.ir import * from mlir.ir import *
from mlir.dialects import builtin from mlir.dialects import builtin
from mlir.dialects import linalg from mlir.dialects import linalg
from mlir.dialects import std from mlir.dialects import std
from mlir.dialects.linalg.opdsl.lang import * from mlir.dialects.linalg.opdsl.lang import *
T1 = TV.T1 T1 = TV.T1
T2 = TV.T2 T2 = TV.T2
@linalg_structured_op @linalg_structured_op
def fill_poly(value=ScalarDef(T1), O=TensorDef(U, output=True)): def fill_poly(value=ScalarDef(T1), O=TensorDef(U, output=True)):
O[None] = TypeFn.cast(U, value) O[None] = TypeFn.cast_signed(U, value)
with Context() as ctx, Location.unknown(): with Context() as ctx, Location.unknown():
module = Module.create() module = Module.create()
f32 = F32Type.get() f32 = F32Type.get()
with InsertionPoint(module.body): with InsertionPoint(module.body):
# Fill indexing maps. # Fill indexing maps.
Show All 22 Lines

mlir/test/python/dialects/linalg/opdsl/emit_matmul.py

Show All 19 Linesdef matmul_mono(
C[D.m, D.n] += A[D.m, D.k] * B[D.k, D.n] C[D.m, D.n] += A[D.m, D.k] * B[D.k, D.n]
@linalg_structured_op @linalg_structured_op
def matmul_poly( def matmul_poly(
A=TensorDef(T1, S.M, S.K), A=TensorDef(T1, S.M, S.K),
B=TensorDef(T2, S.K, S.N), B=TensorDef(T2, S.K, S.N),
C=TensorDef(U, S.M, S.N, output=True), C=TensorDef(U, S.M, S.N, output=True),
cast=TypeFnAttrDef(default=TypeFn.cast)): cast=TypeFnAttrDef(default=TypeFn.cast_signed)):
domain(D.m, D.n, D.k) domain(D.m, D.n, D.k)
C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
with Context() as ctx, Location.unknown(): with Context() as ctx, Location.unknown():
module = Module.create() module = Module.create()
f16 = F16Type.get() f16 = F16Type.get()
f32 = F32Type.get() f32 = F32Type.get()
▲ Show 20 Lines • Show All 134 Lines • Show Last 20 Lines

mlir/test/python/dialects/linalg/opdsl/emit_misc.py

Show All 15 Lines
@linalg_structured_op @linalg_structured_op
def fill_rng_poly( def fill_rng_poly(
min=ScalarDef(F64), min=ScalarDef(F64),
max=ScalarDef(F64), max=ScalarDef(F64),
seed=ScalarDef(I32), seed=ScalarDef(I32),
O=TensorDef(T, S.M, S.N, output=True)): O=TensorDef(T, S.M, S.N, output=True)):
multiplier = TypeFn.cast(I32, const(1103515245)) multiplier = TypeFn.cast_signed(I32, const(1103515245))
increment = TypeFn.cast(I32, const(12345)) increment = TypeFn.cast_signed(I32, const(12345))
rand1 = (TypeFn.cast(I32, index(D.m)) + seed) * multiplier + increment rand1 = (TypeFn.cast_signed(I32, index(D.m)) + seed) * multiplier + increment
rand2 = (TypeFn.cast(I32, index(D.n)) + rand1) * multiplier + increment rand2 = (TypeFn.cast_signed(I32, index(D.n)) + rand1) * multiplier + increment
inv_range = TypeFn.cast(F64, const(2.3283064e-10)) inv_range = TypeFn.cast_signed(F64, const(2.3283064e-10))
offset = TypeFn.cast(F64, const(2147483647)) offset = TypeFn.cast_signed(F64, const(2147483647))
scaling = (max - min) * inv_range scaling = (max - min) * inv_range
O[D.m, D.n] = TypeFn.cast(T, O[D.m, D.n] = TypeFn.cast_signed(
(offset + TypeFn.cast(F64, rand2)) * scaling + min) T, (offset + TypeFn.cast_signed(F64, rand2)) * scaling + min)
@linalg_structured_op @linalg_structured_op
def soft_plus_poly( def soft_plus_poly(
I=TensorDef(T, S.M, S.N), O=TensorDef(U, S.M, S.N, output=True)): I=TensorDef(T, S.M, S.N), O=TensorDef(U, S.M, S.N, output=True)):
O[D.m, D.n] = UnaryFn.log( O[D.m, D.n] = UnaryFn.log(
TypeFn.cast(U, const(1.0)) + TypeFn.cast(U, UnaryFn.exp(I[D.m, D.n]))) TypeFn.cast_signed(U, const(1.0)) +
TypeFn.cast_signed(U, UnaryFn.exp(I[D.m, D.n])))
@linalg_structured_op(op_name="custom_op_name") @linalg_structured_op(op_name="custom_op_name")
def non_default_op_name(I=TensorDef(T, S.N), O=TensorDef(T, S.N, output=True)): def non_default_op_name(I=TensorDef(T, S.N), O=TensorDef(T, S.N, output=True)):
O[D.n] = I[D.n] O[D.n] = I[D.n]
with Context() as ctx, Location.unknown(): with Context() as ctx, Location.unknown():
▲ Show 20 Lines • Show All 48 Lines • Show Last 20 Lines

mlir/test/python/dialects/linalg/opdsl/emit_pooling.py

Show All 10 Lines
T2 = TV.T2 T2 = TV.T2
@linalg_structured_op @linalg_structured_op
def pooling_poly( def pooling_poly(
I=TensorDef(T1, S.N, S.H, S.W, S.C), I=TensorDef(T1, S.N, S.H, S.W, S.C),
K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]), K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True), O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
reduce=BinaryFnAttrDef(default=BinaryFn.max), reduce=BinaryFnAttrDef(default=BinaryFn.max_signed),
cast=TypeFnAttrDef(default=TypeFn.cast), cast=TypeFnAttrDef(default=TypeFn.cast_signed),
strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]), strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])): dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c) domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
O[D.n, D.oh, D.ow, D.c] = reduce[D.kh, D.kw]( O[D.n, D.oh, D.ow, D.c] = reduce[D.kh, D.kw](
cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW,
D.c])) D.c]))
▲ Show 20 Lines • Show All 65 Lines • ▼ Show 20 Lines @builtin.FuncOp.from_py_func(
RankedTensorType.get((1, 4, 16, 1), f32), RankedTensorType.get((1, 4, 16, 1), f32),
RankedTensorType.get((2, 2), f32), RankedTensorType.get((2, 2), f32),
RankedTensorType.get((1, 2, 4, 1), i32)) RankedTensorType.get((1, 2, 4, 1), i32))
def test_f32i32_min_pooling(input, shape, init_result): def test_f32i32_min_pooling(input, shape, init_result):
return pooling_poly( return pooling_poly(
input, input,
shape, shape,
outs=[init_result], outs=[init_result],
reduce=BinaryFn.min, reduce=BinaryFn.min_signed,
strides=[2, 4], strides=[2, 4],
dilations=[1, 2]) dilations=[1, 2])
# CHECK-LABEL: @test_f32i32_min_unsigned_pooling # CHECK-LABEL: @test_f32i32_min_unsigned_pooling
# CHECK: = arith.fptoui # CHECK: = arith.fptoui
# CHECK: = arith.minui # CHECK: = arith.minui
@builtin.FuncOp.from_py_func( @builtin.FuncOp.from_py_func(
RankedTensorType.get((1, 4, 16, 1), f32), RankedTensorType.get((1, 4, 16, 1), f32),
Show All 15 Lines @builtin.FuncOp.from_py_func(
RankedTensorType.get((1, 4, 16, 1), f32), RankedTensorType.get((1, 4, 16, 1), f32),
RankedTensorType.get((2, 2), f32), RankedTensorType.get((2, 2), f32),
RankedTensorType.get((1, 2, 4, 1), f32)) RankedTensorType.get((1, 2, 4, 1), f32))
def test_f32f32_min_pooling(input, shape, init_result): def test_f32f32_min_pooling(input, shape, init_result):
return pooling_poly( return pooling_poly(
input, input,
shape, shape,
outs=[init_result], outs=[init_result],
reduce=BinaryFn.min, reduce=BinaryFn.min_signed,
strides=[2, 4], strides=[2, 4],
dilations=[1, 2]) dilations=[1, 2])
print(module) print(module)

mlir/test/python/dialects/linalg/opdsl/interfaces.py

# RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s # RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
from mlir.dialects.linalg.opdsl.lang import * from mlir.dialects.linalg.opdsl.lang import *
# CHECK: --- # CHECK: ---
# CHECK-LABEL: matmul # CHECK-LABEL: matmul
# CHECK: implements: # CHECK: implements:
# CHECK-NEXT: - LinalgContractionOpInterface # CHECK-NEXT: - LinalgContractionOpInterface
@linalg_structured_op @linalg_structured_op
def matmul( def matmul(
A=TensorDef(T, S.M, S.K), A=TensorDef(T, S.M, S.K),
B=TensorDef(T, S.K, S.N), B=TensorDef(T, S.K, S.N),
C=TensorDef(U, S.M, S.N, output=True)): C=TensorDef(U, S.M, S.N, output=True)):
implements(ContractionOpInterface) implements(ContractionOpInterface)
C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast_signed(U, A[D.m, D.k]) * TypeFn.cast_signed(
U, B[D.k, D.n])

mlir/test/python/dialects/linalg/opdsl/shape_maps_iteration.py

Show All 18 Lines
# CHECK-NEXT: - parallel # CHECK-NEXT: - parallel
# CHECK-NEXT: - reduction # CHECK-NEXT: - reduction
@linalg_structured_op @linalg_structured_op
def matmul( def matmul(
A=TensorDef(T, S.M, S.K), A=TensorDef(T, S.M, S.K),
B=TensorDef(T, S.K, S.N), B=TensorDef(T, S.K, S.N),
C=TensorDef(U, S.M, S.N, output=True)): C=TensorDef(U, S.M, S.N, output=True)):
domain(D.m, D.n, D.k) domain(D.m, D.n, D.k)
C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast_signed(U, A[D.m, D.k]) * TypeFn.cast_signed(
U, B[D.k, D.n])
# Verifies that assignment to a scalar (represented as [None]) is represented # Verifies that assignment to a scalar (represented as [None]) is represented
# correctly. # correctly.
# CHECK: --- # CHECK: ---
# CHECK-LABEL: dot # CHECK-LABEL: dot
# CHECK: shape_map: affine_map<()[s0] -> (s0)> # CHECK: shape_map: affine_map<()[s0] -> (s0)>
# CHECK: shape_map: affine_map<()[s0] -> (s0)> # CHECK: shape_map: affine_map<()[s0] -> (s0)>
# CHECK: shape_map: affine_map<()[s0] -> ()> # CHECK: shape_map: affine_map<()[s0] -> ()>
# CHECK: static_indexing_maps: # CHECK: static_indexing_maps:
# CHECK-NEXT: - affine_map<(d0)[s0] -> (d0)> # CHECK-NEXT: - affine_map<(d0)[s0] -> (d0)>
# CHECK-NEXT: - affine_map<(d0)[s0] -> (d0)> # CHECK-NEXT: - affine_map<(d0)[s0] -> (d0)>
# CHECK-NEXT: - affine_map<(d0)[s0] -> ()> # CHECK-NEXT: - affine_map<(d0)[s0] -> ()>
# CHECK: iterator_types: # CHECK: iterator_types:
# CHECK-NEXT: - reduction # CHECK-NEXT: - reduction
@linalg_structured_op @linalg_structured_op
def dot(A=TensorDef(T, S.M), B=TensorDef(T, S.M), C=TensorDef(U, output=True)): def dot(A=TensorDef(T, S.M), B=TensorDef(T, S.M), C=TensorDef(U, output=True)):
C[None] += TypeFn.cast(U, A[D.m]) * TypeFn.cast(U, B[D.m]) C[None] += TypeFn.cast_signed(U, A[D.m]) * TypeFn.cast_signed(U, B[D.m])
# Verifies that the index_dims of shape-only operands translate to correct # Verifies that the index_dims of shape-only operands translate to correct
# indexing maps. # indexing maps.
# CHECK: --- # CHECK: ---
# CHECK-LABEL: pool # CHECK-LABEL: pool
# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0)> # CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0)>
# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s1)> # CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s1)>
# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s2)> # CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s2)>
# CHECK: static_indexing_maps: # CHECK: static_indexing_maps:
# CHECK-NEXT: - affine_map<(d0, d1)[s0, s1, s2] -> (d0 * 2 + d1)> # CHECK-NEXT: - affine_map<(d0, d1)[s0, s1, s2] -> (d0 * 2 + d1)>
# CHECK-NEXT: - affine_map<(d0, d1)[s0, s1, s2] -> (d1)> # CHECK-NEXT: - affine_map<(d0, d1)[s0, s1, s2] -> (d1)>
# CHECK-NEXT: - affine_map<(d0, d1)[s0, s1, s2] -> (d0)> # CHECK-NEXT: - affine_map<(d0, d1)[s0, s1, s2] -> (d0)>
# CHECK: iterator_types: # CHECK: iterator_types:
# CHECK-NEXT: - parallel # CHECK-NEXT: - parallel
# CHECK-NEXT: - reduction # CHECK-NEXT: - reduction
@linalg_structured_op @linalg_structured_op
def pool( def pool(
I=TensorDef(T, S.I), I=TensorDef(T, S.I),
K=TensorDef(T, S.K, index_dims=[D.k]), K=TensorDef(T, S.K, index_dims=[D.k]),
O=TensorDef(U, S.O, output=True)): O=TensorDef(U, S.O, output=True)):
domain(D.o, D.k) domain(D.o, D.k)
O[D.o] += TypeFn.cast(U, I[D.o * 2 + D.k]) O[D.o] += TypeFn.cast_signed(U, I[D.o * 2 + D.k])

mlir/test/python/dialects/linalg/ops.py

Show First 20 Lines • Show All 93 Lines • ▼ Show 20 Lines with Context() as ctx, Location.unknown():
with InsertionPoint(module.body): with InsertionPoint(module.body):
@builtin.FuncOp.from_py_func( @builtin.FuncOp.from_py_func(
RankedTensorType.get((4, 8), f32), RankedTensorType.get((4, 8), f32)) RankedTensorType.get((4, 8), f32), RankedTensorType.get((4, 8), f32))
def named_form(lhs, rhs): def named_form(lhs, rhs):
init_result = linalg.InitTensorOp([4, 8], f32) init_result = linalg.InitTensorOp([4, 8], f32)
# Check for the named form with custom format # Check for the named form with custom format
# CHECK: linalg.elemwise_unary # CHECK: linalg.elemwise_unary
# CHECK-SAME: cast = #linalg.type_fn<cast> # CHECK-SAME: cast = #linalg.type_fn<cast_signed>
# CHECK-SAME: fun = #linalg.unary_fn<exp> # CHECK-SAME: fun = #linalg.unary_fn<exp>
# CHECK-SAME: ins(%{{.*}} : tensor<4x8xf32>) outs(%{{.*}} : tensor<4x8xf32>) # CHECK-SAME: ins(%{{.*}} : tensor<4x8xf32>) outs(%{{.*}} : tensor<4x8xf32>)
unary_result = linalg.elemwise_unary(lhs, outs=[init_result.result]) unary_result = linalg.elemwise_unary(lhs, outs=[init_result.result])
# CHECK: linalg.elemwise_binary # CHECK: linalg.elemwise_binary
# CHECK-SAME: cast = #linalg.type_fn<cast_unsigned> # CHECK-SAME: cast = #linalg.type_fn<cast_unsigned>
# CHECK-SAME: fun = #linalg.binary_fn<mul> # CHECK-SAME: fun = #linalg.binary_fn<mul>
# CHECK-SAME: ins(%{{.*}}, %{{.*}} : tensor<4x8xf32>, tensor<4x8xf32>) outs(%{{.*}} : tensor<4x8xf32>) # CHECK-SAME: ins(%{{.*}}, %{{.*}} : tensor<4x8xf32>, tensor<4x8xf32>) outs(%{{.*}} : tensor<4x8xf32>)
# CHECK: return # CHECK: return
Show All 21 Lines with InsertionPoint(module.body):
f32)) f32))
def named_form(lhs, rhs): def named_form(lhs, rhs):
init_result = linalg.InitTensorOp([4, 8], f32) init_result = linalg.InitTensorOp([4, 8], f32)
# CHECK: "linalg.matmul"(%{{.*}}) # CHECK: "linalg.matmul"(%{{.*}})
# CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32): # CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32):
# CHECK-NEXT: arith.mulf{{.*}} (f32, f32) -> f32 # CHECK-NEXT: arith.mulf{{.*}} (f32, f32) -> f32
# CHECK-NEXT: arith.addf{{.*}} (f32, f32) -> f32 # CHECK-NEXT: arith.addf{{.*}} (f32, f32) -> f32
# CHECK-NEXT: linalg.yield{{.*}} (f32) -> () # CHECK-NEXT: linalg.yield{{.*}} (f32) -> ()
# CHECK-NEXT: cast = #linalg.type_fn<cast> # CHECK-NEXT: cast = #linalg.type_fn<cast_signed>
# CHECK-SAME: operand_segment_sizes = dense<[2, 1]> : vector<2xi32> # CHECK-SAME: operand_segment_sizes = dense<[2, 1]> : vector<2xi32>
# CHECK-SAME: (tensor<4x16xf32>, tensor<16x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> # CHECK-SAME: (tensor<4x16xf32>, tensor<16x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32>
return linalg.matmul(lhs, rhs, outs=[init_result.result]) return linalg.matmul(lhs, rhs, outs=[init_result.result])
module.operation.print(print_generic_op_form=True) module.operation.print(print_generic_op_form=True)
# CHECK-LABEL: TEST: testNamedStructuredAsGenericOp # CHECK-LABEL: TEST: testNamedStructuredAsGenericOp
▲ Show 20 Lines • Show All 44 Lines • Show Last 20 Lines