This is an archive of the discontinued LLVM Phabricator instance.

Differential D115239

[mlir][OpDSL] Add `TypeFn` class.
ClosedPublic

Authored by gysit on Dec 7 2021, 4:49 AM.

Details

Reviewers
nicolasvasilache
stellaraccident
Commits
rG15757ea80a41: [mlir][OpDSL] Add `TypeFn` class.
Summary

This revision introduces a the TypeFn class that similar to the PrimFn class contains an extensible set of type conversion functions. Having the same mechanism for both type conversion functions and arithmetic functions improves code consistency. Additionally, having an explicit function class and function name is a prerequisite to specify a conversion or arithmetic function via attribute. In a follow up commits, we will introduce function attributes to make OpDSL operations more generic. In particular, the goal is to handle signed and unsigned computation in one operations. Today, there is a linalg.matmul and a linalg.matmul_unsigned.

The commit implements the following changes:

  • Introduce the class of type conversion functions TypeFn
  • Replace the hardwired cast and cast_unsigned ops by the TypeFn counterparts
  • Adapt the python and C++ code generation paths to support the new cast operations

Example:

cast(U, A[D.m, D.k])

changes to

TypeFn.cast(U, A[D.m, D.k])

Depends On D115237

Diff Detail

Event Timeline

gysit created this revision.Dec 7 2021, 4:49 AM
Herald added subscribers: sdasgup3, wenzhicui, wrengr and 21 others. · View Herald TranscriptDec 7 2021, 4:49 AM
gysit requested review of this revision.Dec 7 2021, 4:49 AM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptDec 7 2021, 4:49 AM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: limo1996, stephenneuendorffer, nicolasvasilache. · View Herald Transcript
gysit added a child revision: D115240: [mlir][OpDSL] Rename `PrimFn` to `ArithFn`..Dec 7 2021, 4:50 AM
gysit updated this revision to Diff 392361.Dec 7 2021, 5:06 AM

Fix formatting of templated code.

gysit added a reviewer: stellaraccident.Dec 7 2021, 5:10 AM
gysit added a comment.Dec 7 2021, 5:16 AM

The name TypeFn and the fact that TypeFn.cast is longer than just cast may be controversial. I am fine with renaming. ConversionFn or TypeConversionFn would be alternatives. Also not that it would be possible to define an alias tf = TypeFn and then write tf.cast. Once OpDSL has support for function attributes we may also redefine the cast function as cast = TypeFnAttrDef(default=TypeFn.cast).

Harbormaster completed remote builds in B137881: Diff 392361.Dec 7 2021, 5:23 AM
gysit updated this revision to Diff 393982.Dec 13 2021, 11:46 AM

Improve doc.

Harbormaster completed remote builds in B139032: Diff 393982.Dec 13 2021, 12:21 PM
stellaraccident accepted this revision.Jan 4 2022, 10:15 AM
This revision is now accepted and ready to land.Jan 4 2022, 10:15 AM
gysit updated this revision to Diff 398069.Jan 7 2022, 12:57 AM

Rebase.

Harbormaster completed remote builds in B142030: Diff 398069.Jan 7 2022, 12:58 AM
This revision was landed with ongoing or failed builds.Jan 7 2022, 4:29 AM
Closed by commit rG15757ea80a41: [mlir][OpDSL] Add `TypeFn` class. (authored by gysit). · Explain Why
This revision was automatically updated to reflect the committed changes.
gysit added a commit: rG15757ea80a41: [mlir][OpDSL] Add `TypeFn` class..

Revision Contents

PathSize
mlir/
docs/
Dialects/
Linalg/
45 lines
include/
mlir/
Dialect/
Linalg/
IR/
308 lines
lib/
Dialect/
Linalg/
IR/
14 lines
python/
mlir/
dialects/
linalg/
opdsl/
lang/
83 lines
55 lines
68 lines
ops/
135 lines
test/
mlir-linalg-ods-gen/
12 lines
python/
dialects/
linalg/
opdsl/
2 lines
18 lines
4 lines
5 lines
20 lines
12 lines
2 lines
6 lines
tools/
mlir-linalg-ods-gen/
70 lines

Diff 398102

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] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(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 86 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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
O[D.n, D.oh, D.ow, D.c] += \ O[D.n, D.oh, D.ow, D.c] += TypeFn.cast(U,
cast(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
Show All 15 Lines
* `PrimFn.sub(a, b)` (also via overloading the binary `-` operator) * `PrimFn.sub(a, b)` (also via overloading the binary `-` operator)
Reduction functions can appear as the outer-most function on the RHS: Reduction functions can appear as the outer-most 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`
Additionally, type conversion functions cast an operand to a target type:
* `TypeFn.cast(TypeVar, operand)`
* `TypeFn.cast_unsigned(TypeVar, operand)`
As the integer types are signless, signedness is implement by different
functions that treat integers as signed (`TypeFn.cast`) or unsigned
(`TypeFn.cast_unsigned`) values.
There are also special forms: There are also special forms:
* `cast(TypeVar, operand)` casts the `operand` to the target type `TypeVar`. * `const(value)` returns a constant value.
* `const(TypeVar, value)` returns a constant value of type `TypeVar`.
* `index(dim)` returns the iteration index in the given dimension `dim`. * `index(dim)` returns the iteration index in the given dimension `dim`.
## 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 `cast` calls will generally require uniform types Assignment expressions with no `TypeFn.cast` calls will generally require
throughout and will fail to verify if violated. The presence of a `cast` allows uniform types throughout and will fail to verify if violated. The presence of a
for a limited form of numeric type conversion between element types that can be `TypeFn.cast` or `TypeFn.cast_unsigned` allows for a limited form of numeric
derived from inputs and outputs (and in the future, attributes). `cast` calls type conversion between element types that can be derived from inputs and
with a `TypeVar` first argument are emitted as `symbolic_cast` primitives in the outputs (and in the future, attributes). `TypeFn.cast` calls with a `TypeVar`
YAML definition. 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 type family. Note that perform any necessary extension or truncation within the type family. The
presently, any integer type is assumed to be signed for the purpose of integer types themselves are signless and signedness is implemented by
determining how to extend or truncate. Supporting unsigned integer types is left functions/operations. The `TypeFn.cast` function treats all integers as signed,
for future work. while `TypeFn.cast_unsigned` treats them as unsigned.
The following examples illustrate the lowering of signed and unsigned functions:
* cast(I32 -> I64) -> `arith.ExtSIOp`
* cast(F32 -> I32) -> `arith.FPToSIOp`
* cast_unsigned(I32 -> I64) -> `arith.ExtUIOp`
* cast_unsigned(F32 -> I32) -> `arith.FPToUIOp`
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.

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

Show First 20 Lines • Show All 45 Lines • ▼ Show 20 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: C scalar_arg: C
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: matmul_unsigned name: matmul_unsigned
cpp_class_name: MatmulUnsignedOp cpp_class_name: MatmulUnsignedOp
doc: |- doc: |-
Performs an unsigned matrix multiplication of two 2D inputs. Performs an unsigned 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
Show All 35 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: C scalar_arg: C
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast_unsigned
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
is_unsigned_cast: true
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast_unsigned
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
is_unsigned_cast: true
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: quantized_matmul name: quantized_matmul
cpp_class_name: QuantizedMatmulOp cpp_class_name: QuantizedMatmulOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: AZp scalar_arg: AZp
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: BZp scalar_arg: BZp
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: mmt4d name: mmt4d
cpp_class_name: Mmt4DOp cpp_class_name: Mmt4DOp
doc: |- doc: |-
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:
▲ Show 20 Lines • Show All 46 Lines • ▼ Show 20 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: accum scalar_arg: accum
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: AccumType type_var: AccumType
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: lhs scalar_arg: lhs
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: AccumType type_var: AccumType
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: rhs scalar_arg: rhs
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: batch_matmul name: batch_matmul
cpp_class_name: BatchMatmulOp cpp_class_name: BatchMatmulOp
doc: |- doc: |-
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
Show All 36 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: C scalar_arg: C
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: quantized_batch_matmul name: quantized_batch_matmul
cpp_class_name: QuantizedBatchMatmulOp cpp_class_name: QuantizedBatchMatmulOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 50 Lines • ▼ Show 20 Lines value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: AZp scalar_arg: AZp
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: BZp scalar_arg: BZp
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: matvec name: matvec
cpp_class_name: MatvecOp cpp_class_name: MatvecOp
doc: |- doc: |-
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
Show All 34 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: x scalar_arg: x
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: y scalar_arg: y
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: vecmat name: vecmat
cpp_class_name: VecmatOp cpp_class_name: VecmatOp
doc: |- doc: |-
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
Show All 34 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: x scalar_arg: x
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: y scalar_arg: y
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: batch_matvec name: batch_matvec
cpp_class_name: BatchMatvecOp cpp_class_name: BatchMatvecOp
doc: |- doc: |-
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
Show All 35 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: C scalar_arg: C
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: dot name: dot
cpp_class_name: DotOp cpp_class_name: DotOp
doc: |- doc: |-
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
Show All 33 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: C scalar_arg: C
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: A scalar_arg: A
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: B scalar_arg: B
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_1d name: conv_1d
cpp_class_name: Conv1DOp cpp_class_name: Conv1DOp
doc: |- doc: |-
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
Show All 34 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_2d name: conv_2d
cpp_class_name: Conv2DOp cpp_class_name: Conv2DOp
doc: |- doc: |-
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
Show All 36 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_3d name: conv_3d
cpp_class_name: Conv3DOp cpp_class_name: Conv3DOp
doc: |- doc: |-
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
Show All 39 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_1d_nwc_wcf name: conv_1d_nwc_wcf
cpp_class_name: Conv1DNwcWcfOp cpp_class_name: Conv1DNwcWcfOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_2d_nhwc_hwcf name: conv_2d_nhwc_hwcf
cpp_class_name: Conv2DNhwcHwcfOp cpp_class_name: Conv2DNhwcHwcfOp
doc: |- doc: |-
Performs 2-D convolution. Performs 2-D convolution.
Layout: Layout:
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
--- !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
doc: |- doc: |-
Performs 2-D convolution with zero point offsets. Performs 2-D convolution with zero point offsets.
Layout: Layout:
▲ Show 20 Lines • Show All 78 Lines • ▼ Show 20 Lines value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: IZp scalar_arg: IZp
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: KZp scalar_arg: KZp
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_2d_nchw_fchw name: conv_2d_nchw_fchw
cpp_class_name: Conv2DNchwFchwOp cpp_class_name: Conv2DNchwFchwOp
doc: |- doc: |-
Performs 2-D convolution. Performs 2-D convolution.
Layout: Layout:
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: conv_3d_ndhwc_dhwcf name: conv_3d_ndhwc_dhwcf
cpp_class_name: Conv3DNdhwcDhwcfOp cpp_class_name: Conv3DNdhwcDhwcfOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
--- !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
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 48 Lines • ▼ Show 20 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
--- !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
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 54 Lines • ▼ Show 20 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
--- !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
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 69 Lines • ▼ Show 20 Lines value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: IZp scalar_arg: IZp
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: KZp scalar_arg: KZp
is_unsigned_cast: false
--- !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
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 57 Lines • ▼ Show 20 Lines value: !ScalarExpression
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
--- !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
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 73 Lines • ▼ Show 20 Lines value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: IZp scalar_arg: IZp
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: sub fn_name: sub
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: K scalar_arg: K
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: KZp scalar_arg: KZp
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nhwc_sum name: pooling_nhwc_sum
cpp_class_name: PoolingNhwcSumOp cpp_class_name: PoolingNhwcSumOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nhwc_max name: pooling_nhwc_max
cpp_class_name: PoolingNhwcMaxOp cpp_class_name: PoolingNhwcMaxOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: max fn_name: max
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nhwc_max_unsigned name: pooling_nhwc_max_unsigned
cpp_class_name: PoolingNhwcMaxUnsignedOp cpp_class_name: PoolingNhwcMaxUnsignedOp
doc: |- doc: |-
Performs unsigned max pooling. Performs unsigned 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
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: max_unsigned fn_name: max_unsigned
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast_unsigned
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: true
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nchw_max name: pooling_nchw_max
cpp_class_name: PoolingNchwMaxOp cpp_class_name: PoolingNchwMaxOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: max fn_name: max
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nhwc_min name: pooling_nhwc_min
cpp_class_name: PoolingNhwcMinOp cpp_class_name: PoolingNhwcMinOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: min fn_name: min
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_nhwc_min_unsigned name: pooling_nhwc_min_unsigned
cpp_class_name: PoolingNhwcMinUnsignedOp cpp_class_name: PoolingNhwcMinUnsignedOp
doc: |- doc: |-
Performs unsigned min pooling. Performs unsigned 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
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: min_unsigned fn_name: min_unsigned
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast_unsigned
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: true
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_ndhwc_sum name: pooling_ndhwc_sum
cpp_class_name: PoolingNdhwcSumOp cpp_class_name: PoolingNdhwcSumOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 55 Lines • ▼ Show 20 Lines - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_ndhwc_max name: pooling_ndhwc_max
cpp_class_name: PoolingNdhwcMaxOp cpp_class_name: PoolingNdhwcMaxOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 55 Lines • ▼ Show 20 Lines - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: max fn_name: max
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: pooling_ndhwc_min name: pooling_ndhwc_min
cpp_class_name: PoolingNdhwcMinOp cpp_class_name: PoolingNdhwcMinOp
doc: |- doc: |-
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
▲ Show 20 Lines • Show All 55 Lines • ▼ Show 20 Lines - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: min fn_name: min
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: O scalar_arg: O
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: fill_rng_2d name: fill_rng_2d
cpp_class_name: FillRng2DOp cpp_class_name: FillRng2DOp
doc: |- doc: |-
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
Show All 30 Lines indexing_maps: !LinalgIndexingMapsConfig
- 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
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: T type_var: T
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: F64 type_var: F64
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '2147483647 : i64' scalar_const: '2147483647 : i64'
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: F64 type_var: F64
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_index: 1 scalar_index: 1
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_index: 0 scalar_index: 0
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
scalar_arg: seed scalar_arg: seed
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '1103515245 : i64' scalar_const: '1103515245 : i64'
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '12345 : i64' scalar_const: '12345 : i64'
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '1103515245 : i64' scalar_const: '1103515245 : i64'
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: I32 type_var: I32
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '12345 : i64' scalar_const: '12345 : i64'
is_unsigned_cast: false
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: mul fn_name: mul
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
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
symbolic_cast: type_fn:
fn_name: cast
type_var: F64 type_var: F64
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '2.3283063999999999E-10 : f64' scalar_const: '2.3283063999999999E-10 : f64'
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
scalar_arg: min scalar_arg: min
is_unsigned_cast: false
--- !LinalgOpConfig --- !LinalgOpConfig
metadata: !LinalgOpMetadata metadata: !LinalgOpMetadata
name: soft_plus_2d name: soft_plus_2d
cpp_class_name: SoftPlus2DOp cpp_class_name: SoftPlus2DOp
doc: |- doc: |-
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
Show All 24 Lines value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: log fn_name: log
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '1.000000e+00 : f64' scalar_const: '1.000000e+00 : f64'
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
scalar_apply: scalar_apply:
fn_name: exp fn_name: exp
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: U type_var: U
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_arg: I scalar_arg: I
is_unsigned_cast: false

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

Show First 20 Lines • Show All 141 Lines • ▼ Show 20 Linesstatic LogicalResult foldMemRefCastInTiledLoopOp(TiledLoopOp op) {
} }
return success(folded); return success(folded);
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Region builder helper. // Region builder helper.
// TODO: Move this to a utility library. // TODO: Move this to a utility library.
// The public methods on this class are referenced directly from generated code // The public methods on this class are referenced directly from generated code
// and bind by name to math functions in the DSL as: // and bind by name to math and type conversion functions in the DSL as:
// `applyfn__{fnName}` // `applyfn__{fnName}`
// `typefn__{fnName}`
// Examples: // Examples:
// `applyfn__add` // `applyfn__add`
// `applyfn__mul` // `applyfn__mul`
// `typefn__cast`
// The naming convention is intentional in order to match snake-cased DSL names. // The naming convention is intentional in order to match snake-cased DSL names.
// See mlir-linalg-ods-yaml-gen.cpp for the code that mates to this class. // See mlir-linalg-ods-yaml-gen.cpp for the code that mates to this class.
// //
// Implementations of the math functions must be polymorphic over numeric types, // Implementations of the math functions must be polymorphic over numeric types,
// internally performing necessary casts. If the function application makes no // internally performing necessary casts. If the function application makes no
// sense, then the only recourse is to assert and return nullptr. This can be // sense, then the only recourse is to assert and return nullptr. This can be
// extended later if it becomes possible to fail construction of the region. The // extended later if it becomes possible to fail construction of the region. The
// invariant should be enforced at a higher level. // invariant should be enforced at a higher level.
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Lines Value cast(Type toType, Value operand, bool isUnsignedCast) {
} }
emitWarning(operand.getLoc()) << "could not cast operand of type " emitWarning(operand.getLoc()) << "could not cast operand of type "
<< operand.getType() << " to " << toType; << operand.getType() << " to " << toType;
return operand; return operand;
} }
// NOLINTNEXTLINE(*-identifier-naming): externally called. // NOLINTNEXTLINE(*-identifier-naming): externally called.
Value typefn__cast(Type toType, Value operand) {
return cast(toType, operand, false);
}
// NOLINTNEXTLINE(*-identifier-naming): externally called.
Value typefn__cast_unsigned(Type toType, Value operand) {
return cast(toType, operand, true);
}
// NOLINTNEXTLINE(*-identifier-naming): externally called.
Value applyfn__add(Value lhs, Value rhs) { Value applyfn__add(Value lhs, Value rhs) {
OpBuilder builder = getBuilder(); OpBuilder builder = getBuilder();
if (isFloatingPoint(lhs)) if (isFloatingPoint(lhs))
return builder.create<arith::AddFOp>(lhs.getLoc(), lhs, rhs); return builder.create<arith::AddFOp>(lhs.getLoc(), lhs, rhs);
if (isInteger(lhs)) if (isInteger(lhs))
return builder.create<arith::AddIOp>(lhs.getLoc(), lhs, rhs); return builder.create<arith::AddIOp>(lhs.getLoc(), lhs, rhs);
llvm_unreachable("unsupported non numeric type"); llvm_unreachable("unsupported non numeric type");
} }
▲ Show 20 Lines • Show All 2,465 Lines • Show Last 20 Lines

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

Show First 20 Lines • Show All 308 Lines • ▼ Show 20 Lines if len(self.definitions) > 1:
values_repr = f"({', '.join(repr(v) for v in self.values)})" values_repr = f"({', '.join(repr(v) for v in self.values)})"
else: else:
defs_repr = f"{repr(self.definitions[0])}" defs_repr = f"{repr(self.definitions[0])}"
values_repr = f"{repr(self.values[0])}" values_repr = f"{repr(self.values[0])}"
return f"{defs_repr} = {values_repr}" return f"{defs_repr} = {values_repr}"
class TypeFnType:
"""Type conversion function.
A type conversion function takes a target type and a tensor expression and
returns the casted tensor expression.
"""
def __init__(self, fn_name: str):
self.fn_name = fn_name
def __call__(self, type_var: TypeVar,
arg: TensorExpression) -> "TensorTypeFn":
return TensorTypeFn(self, type_var, arg)
def __repr__(self):
return f"{self.fn_name}"
class TypeFn:
"""Type conversion function namespace.
As the integer types are signless, signedness is implement by different cast
functions that treat integers as signed (`cast`) or unsigned
(`cast_unsigned`) values.
Examples:
- cast(I32 -> I64) -> `arith.ExtSIOp`
- cast_unsigned(I32 -> I64) -> `arith.ExtUIOp`
"""
cast = TypeFnType("cast")
cast_unsigned = TypeFnType("cast_unsigned")
class PrimFnType: class PrimFnType:
"""Primitive operations.""" """Primitive operations."""
def __init__(self, prim_name: str): def __init__(self, prim_name: str):
self.prim_name = prim_name self.prim_name = prim_name
def __call__(self, *args): def __call__(self, *args):
return PrimApply(self, args) return PrimApply(self, args)
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Lines def visit_tensor_exprs(self, callback):
super().visit_tensor_exprs(callback) super().visit_tensor_exprs(callback)
for arg in self.args: for arg in self.args:
arg.visit_tensor_exprs(callback) arg.visit_tensor_exprs(callback)
def __repr__(self): def __repr__(self):
return f"{repr(self.prim)}({', '.join(repr(a) for a in self.args)})" return f"{repr(self.prim)}({', '.join(repr(a) for a in self.args)})"
class TensorTypeFn(TensorExpression):
"""Application of a type conversion function."""
def __init__(self, type_fn: TypeFn, type_var: TypeVar, arg: TensorExpression):
self.type_fn = type_fn
self.type_var = type_var
self.arg = arg
def to_scalar_expression(self) -> ScalarExpression:
return ScalarTypeFn(self.type_fn.fn_name, self.type_var,
self.arg.to_scalar_expression()).expr()
def visit_tensor_exprs(self, callback):
super().visit_tensor_exprs(callback)
self.arg.visit_tensor_exprs(callback)
def __repr__(self):
return f"{repr(self.type_fn)}({type_var}, {self.arg})"
class const(TensorExpression): class const(TensorExpression):
"""Returns the given constant floating point or integer value.""" """Returns the given constant floating point or integer value."""
def __init__(self, value: Any): def __init__(self, value: Any):
with _ir.Context(): with _ir.Context():
if isinstance(value, float): if isinstance(value, float):
self.value = str(_ir.FloatAttr.get_f64(float(value))) self.value = str(_ir.FloatAttr.get_f64(float(value)))
elif isinstance(value, int): elif isinstance(value, int):
Show All 26 Linesclass index(TensorExpression):
def to_scalar_expression(self) -> ScalarExpression: def to_scalar_expression(self) -> ScalarExpression:
assert self.dim != -1, "Dimension name not resolved" assert self.dim != -1, "Dimension name not resolved"
return ScalarIndex(self.dim).expr() return ScalarIndex(self.dim).expr()
def __repr__(self): def __repr__(self):
return f"index({repr(self.dim)})" return f"index({repr(self.dim)})"
class cast(TensorExpression):
"""Casts the element type to a type (typically symbolic TypeVar)."""
def __init__(self, to_type: TypeVar, operand: TensorExpression):
self.to_type = to_type
self.operand = operand
def to_scalar_expression(self) -> ScalarExpression:
return ScalarSymbolicCast(self.to_type, self.operand.to_scalar_expression(),
False).expr()
def visit_tensor_exprs(self, callback):
super().visit_tensor_exprs(callback)
self.operand.visit_tensor_exprs(callback)
def __repr__(self):
return f"cast({self.to_type}, {repr(self.operand)})"
class cast_unsigned(cast):
"""Casts the element type to an unsigned type (typically symbolic TypeVar)."""
def to_scalar_expression(self) -> ScalarExpression:
return ScalarSymbolicCast(self.to_type, self.operand.to_scalar_expression(),
True).expr()
def __repr__(self):
return f"cast_unsigned({self.to_type}, {repr(self.operand)})"
class ReduceApply(TensorExpression): class ReduceApply(TensorExpression):
"""Application of a reduction. """Application of a reduction.
This captures the lhs separately (initial value) separately from the rhs. This captures the lhs separately (initial value) separately from the rhs.
""" """
def __init__(self, reduce: ReduceFnType, args: Sequence[TensorExpression]): def __init__(self, reduce: ReduceFnType, args: Sequence[TensorExpression]):
self.reduce = reduce self.reduce = reduce
▲ Show 20 Lines • Show All 97 Lines • Show Last 20 Lines

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

# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. # Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information. # See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from typing import Dict, List, Sequence, Tuple, Union from typing import Callable, Dict, List, Sequence, Tuple, Union
from .....ir import * from .....ir import *
from .... import linalg from .... import linalg
from .... import std from .... import std
from .... import math from .... import math
from .... import arith from .... import arith
from ...._ods_common import get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values from ...._ods_common import get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values
from .scalar_expr import * from .scalar_expr import *
from .config import * from .config import *
import numpy as np import numpy as np
__all__ = [ __all__ = [
"emit_generic_structured_op", "emit_generic_structured_op",
"emit_named_structured_op", "emit_named_structured_op",
"ValueList", "ValueList",
] ]
ValueList = Union[Sequence[Value], OpResultList] ValueList = Union[Sequence[Value], OpResultList]
def isa(cls: Type, ty: Type): def isa(cls: Type, ty: Type):
try: try:
cls(ty) cls(ty)
return True return True
except ValueError: except ValueError:
return False return False
▲ Show 20 Lines • Show All 181 Lines • ▼ Show 20 Lines def expression(self, expr: ScalarExpression) -> Value:
elif expr.scalar_const: elif expr.scalar_const:
value_attr = Attribute.parse(expr.scalar_const.value) value_attr = Attribute.parse(expr.scalar_const.value)
return arith.ConstantOp(value_attr.type, value_attr).result return arith.ConstantOp(value_attr.type, value_attr).result
elif expr.scalar_index: elif expr.scalar_index:
dim_attr = IntegerAttr.get( dim_attr = IntegerAttr.get(
IntegerType.get_signless(64), expr.scalar_index.dim) IntegerType.get_signless(64), expr.scalar_index.dim)
return linalg.IndexOp(dim_attr).result return linalg.IndexOp(dim_attr).result
elif expr.scalar_apply: elif expr.scalar_apply:
try: fn = self._get_function(f"_eval_{expr.scalar_apply.fn_name}")
fn = getattr(self, f"_eval_{expr.scalar_apply.fn_name}")
except AttributeError:
raise ValueError(
f"Function '{expr.scalar_apply.fn_name}' is not a known "
"scalar body function")
operand_values = [ operand_values = [
self.expression(operand) for operand in expr.scalar_apply.operands self.expression(operand) for operand in expr.scalar_apply.operands
] ]
return fn(*operand_values) return fn(*operand_values)
elif expr.symbolic_cast: elif expr.type_fn:
operand_value = self.expression(expr.symbolic_cast.operand) fn = self._get_function(f"_typefn_{expr.type_fn.fn_name}")
return self.cast(expr.symbolic_cast.to_type.name, operand_value, operand = self.expression(expr.type_fn.operand)
expr.symbolic_cast.is_unsigned_cast) return fn(expr.type_fn.type_var.name, operand)
raise NotImplementedError(f"Unimplemented scalar body expression: {expr}") raise NotImplementedError(f"Unimplemented scalar body expression: {expr}")
def cast(self, type_var_name: str, operand: Value, def yield_outputs(self, *output_names: str):
is_unsigned_cast: bool) -> Value: output_values = []
for n in output_names:
try:
output_values.append(self.yield_mapping[n])
except KeyError:
raise ValueError(f"Body assignments do not assign all outputs: "
f"missing '{n}'")
linalg.YieldOp(output_values)
def _get_function(self, fn_name: str) -> Callable:
try:
fn = getattr(self, f"{fn_name}")
except AttributeError:
raise ValueError(f"Function '{fn_name}' is not a known function")
return fn
def _cast(self,
type_var_name: str,
operand: Value,
is_unsigned_cast: bool = False) -> Value:
try: try:
to_type = self.type_mapping[type_var_name] to_type = self.type_mapping[type_var_name]
except KeyError: except KeyError:
raise ValueError(f"Unbound type variable '{type_var_name}' (" raise ValueError(f"Unbound type variable '{type_var_name}' ("
f"expected one of {self.type_mapping.keys()}") f"expected one of {self.type_mapping.keys()}")
if operand.type == to_type: if operand.type == to_type:
return operand return operand
if _is_integer_type(to_type): if _is_integer_type(to_type):
Show All 34 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 yield_outputs(self, *output_names: str): def _typefn_cast(self, type_var_name: str, operand: Value) -> Value:
output_values = [] return self._cast(type_var_name, operand, False)
for n in output_names:
try: def _typefn_cast_unsigned(self, type_var_name: str, operand: Value) -> Value:
output_values.append(self.yield_mapping[n]) return self._cast(type_var_name, operand, True)
except KeyError:
raise ValueError(f"Body assignments do not assign all outputs: "
f"missing '{n}'")
linalg.YieldOp(output_values)
def _eval_add(self, lhs: Value, rhs: Value) -> Value: def _eval_add(self, lhs: Value, rhs: Value) -> Value:
if _is_floating_point_type(lhs.type): if _is_floating_point_type(lhs.type):
return arith.AddFOp(lhs, rhs).result return arith.AddFOp(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.AddIOp(lhs, rhs).result return arith.AddIOp(lhs, rhs).result
raise NotImplementedError("Unsupported 'add' operand: {lhs}") raise NotImplementedError("Unsupported 'add' operand: {lhs}")
▲ Show 20 Lines • Show All 130 Lines • Show Last 20 Lines

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

Show All 15 Lines
from typing import Optional, Sequence from typing import Optional, Sequence
from .yaml_helper import * from .yaml_helper import *
from .types import * from .types import *
__all__ = [ __all__ = [
"ScalarAssign", "ScalarAssign",
"ScalarApplyFn", "ScalarApplyFn",
"ScalarTypeFn",
"ScalarArg", "ScalarArg",
"ScalarConst", "ScalarConst",
"ScalarIndex", "ScalarIndex",
"ScalarExpression", "ScalarExpression",
"ScalarSymbolicCast",
] ]
class ScalarApplyFn: class ScalarApplyFn:
"""A type of ScalarExpression that applies a named function to operands.""" """A type of ScalarExpression that applies a named function to operands."""
def __init__(self, fn_name: str, *operands: "ScalarExpression"): def __init__(self, fn_name: str, *operands: "ScalarExpression"):
self.fn_name = fn_name self.fn_name = fn_name
self.operands = operands self.operands = operands
def expr(self) -> "ScalarExpression": def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_apply=self) return ScalarExpression(scalar_apply=self)
def __repr__(self): def __repr__(self):
return f"ScalarApplyFn<{self.fn_name}>({', '.join(self.operands)})" return f"ScalarApplyFn<{self.fn_name}>({', '.join(self.operands)})"
class ScalarTypeFn:
"""A type of ScalarExpression that applies a type conversion function."""
def __init__(self, fn_name: str, type_var: TypeVar,
operand: "ScalarExpression"):
self.fn_name = fn_name
self.type_var = type_var
self.operand = operand
def expr(self) -> "ScalarExpression":
return ScalarExpression(type_fn=self)
def __repr__(self):
return f"ScalarTypeFn<{self.fn_name}>({self.type_var}, {self.operand})"
class ScalarArg: class ScalarArg:
"""A type of ScalarExpression that references a named argument.""" """A type of ScalarExpression that references a named argument."""
def __init__(self, arg: str): def __init__(self, arg: str):
self.arg = arg self.arg = arg
def expr(self) -> "ScalarExpression": def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_arg=self) return ScalarExpression(scalar_arg=self)
Show All 23 Linesclass ScalarIndex:
def expr(self) -> "ScalarExpression": def expr(self) -> "ScalarExpression":
return ScalarExpression(scalar_index=self) return ScalarExpression(scalar_index=self)
def __repr__(self): def __repr__(self):
return f"(ScalarIndex({self.dim})" return f"(ScalarIndex({self.dim})"
class ScalarSymbolicCast:
"""A type of ScalarExpression that symbolically casts an operand to a TypeVar."""
def __init__(self, to_type: TypeVar, operand: "ScalarExpression",
is_unsigned_cast: bool):
self.to_type = to_type
self.operand = operand
self.is_unsigned_cast = is_unsigned_cast
def expr(self) -> "ScalarExpression":
return ScalarExpression(symbolic_cast=self)
def __repr__(self):
return f"ScalarSymbolicCast({self.to_type}, {self.operand}, {self.is_unsigned_cast})"
class ScalarExpression(YAMLObject): class ScalarExpression(YAMLObject):
"""An expression on scalar values. """An expression on scalar values.
Can be one of: Can be one of:
- ScalarApplyFn - ScalarApplyFn
- ScalarTypeFn
- ScalarArg - ScalarArg
- ScalarConst - ScalarConst
- ScalarIndex - ScalarIndex
- ScalarSymbolicCast - ScalarSymbolicCast
""" """
yaml_tag = "!ScalarExpression" yaml_tag = "!ScalarExpression"
def __init__(self, def __init__(self,
scalar_apply: Optional[ScalarApplyFn] = None, scalar_apply: Optional[ScalarApplyFn] = None,
type_fn: Optional[ScalarTypeFn] = None,
scalar_arg: Optional[ScalarArg] = None, scalar_arg: Optional[ScalarArg] = None,
scalar_const: Optional[ScalarConst] = None, scalar_const: Optional[ScalarConst] = None,
scalar_index: Optional[ScalarIndex] = None, scalar_index: Optional[ScalarIndex] = None):
symbolic_cast: Optional[ScalarSymbolicCast] = None): if (bool(scalar_apply) + bool(type_fn) + bool(scalar_arg) +
if (bool(scalar_apply) + bool(scalar_arg) + bool(scalar_const) + bool(scalar_const) + bool(scalar_index)) != 1:
bool(scalar_index) + bool(symbolic_cast)) != 1: raise ValueError("One of 'scalar_apply', 'type_fn', 'scalar_arg', "
raise ValueError("One of 'scalar_apply', 'scalar_arg', 'scalar_const', " "'scalar_const', 'scalar_index', must be specified")
"'scalar_index', 'symbolic_cast' must be specified")
self.scalar_apply = scalar_apply self.scalar_apply = scalar_apply
self.type_fn = type_fn
self.scalar_arg = scalar_arg self.scalar_arg = scalar_arg
self.scalar_const = scalar_const self.scalar_const = scalar_const
self.scalar_index = scalar_index self.scalar_index = scalar_index
self.symbolic_cast = symbolic_cast
def to_yaml_custom_dict(self): def to_yaml_custom_dict(self):
if self.scalar_apply: if self.scalar_apply:
return dict( return dict(
scalar_apply=dict( scalar_apply=dict(
fn_name=self.scalar_apply.fn_name, fn_name=self.scalar_apply.fn_name,
operands=list(self.scalar_apply.operands), operands=list(self.scalar_apply.operands),
)) ))
if self.type_fn:
# Note that even though operands must be arity 1, we write it the
# same way as for apply because it allows handling code to be more
# generic vs having a special form.
return dict(
type_fn=dict(
fn_name=self.type_fn.fn_name,
type_var=self.type_fn.type_var.name,
operands=[self.type_fn.operand],
))
elif self.scalar_arg: elif self.scalar_arg:
return dict(scalar_arg=self.scalar_arg.arg) return dict(scalar_arg=self.scalar_arg.arg)
elif self.scalar_const: elif self.scalar_const:
return dict(scalar_const=self.scalar_const.value) return dict(scalar_const=self.scalar_const.value)
elif self.scalar_index: elif self.scalar_index:
return dict(scalar_index=self.scalar_index.dim) return dict(scalar_index=self.scalar_index.dim)
elif self.symbolic_cast:
# Note that even though operands must be arity 1, we write it the
# same way as for apply because it allows handling code to be more
# generic vs having a special form.
return dict(
symbolic_cast=dict(
type_var=self.symbolic_cast.to_type.name,
operands=[self.symbolic_cast.operand],
is_unsigned_cast=self.symbolic_cast.is_unsigned_cast))
else: else:
raise ValueError(f"Unexpected ScalarExpression type: {self}") raise ValueError(f"Unexpected ScalarExpression type: {self}")
class ScalarAssign(YAMLObject): class ScalarAssign(YAMLObject):
"""An assignment to a named argument (LHS of a comprehension).""" """An assignment to a named argument (LHS of a comprehension)."""
yaml_tag = "!ScalarAssign" yaml_tag = "!ScalarAssign"
Show All 9 Lines

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

Show All 12 Linesdef matmul(
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] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(U, B[D.k, D.n])
@linalg_structured_op @linalg_structured_op
def matmul_unsigned( def matmul_unsigned(
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)):
"""Performs an unsigned matrix multiplication of two 2D inputs. """Performs an unsigned 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_unsigned(U, A[D.m, D.k]) * cast_unsigned(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast_unsigned(U, A[D.m, D.k]) * TypeFn.cast_unsigned(
U, B[D.k, D.n])
@linalg_structured_op @linalg_structured_op
def quantized_matmul( def quantized_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),
AZp=ScalarDef(I32), AZp=ScalarDef(I32),
BZp=ScalarDef(I32), BZp=ScalarDef(I32),
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. 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] += (cast(U, A[D.m, D.k]) - cast(U, AZp)) * ( C[D.m, D.n] += (TypeFn.cast(U, A[D.m, D.k]) - TypeFn.cast(U, AZp)) * (
cast(U, B[D.k, D.n]) - cast(U, BZp)) TypeFn.cast(U, B[D.k, D.n]) - TypeFn.cast(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, accum[D.m, D.n, D.m0, D.n0] += TypeFn.cast(
D.n0] += cast(TV.AccumType, lhs[D.m, D.k, D.m0, D.k0]) * cast( TV.AccumType, lhs[D.m, D.k, D.m0, D.k0]) * TypeFn.cast(
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, D.n] += cast(U, A[D.b, D.m, D.k]) * cast(U, B[D.b, D.k, D.n]) 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])
@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, D.n] += (cast(U, A[D.b, D.m, D.k]) - cast(U, AZp)) * ( C[D.b, D.m,
cast(U, B[D.b, D.k, D.n]) - cast(U, BZp)) D.n] += (TypeFn.cast(U, A[D.b, D.m, D.k]) - TypeFn.cast(U, AZp)) * (
TypeFn.cast(U, B[D.b, D.k, D.n]) - TypeFn.cast(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] += cast(U, A[D.m, D.n]) * cast(U, y[D.n]) x[D.m] += TypeFn.cast(U, A[D.m, D.n]) * TypeFn.cast(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] += cast(U, y[D.m]) * cast(U, A[D.m, D.n]) x[D.n] += TypeFn.cast(U, y[D.m]) * TypeFn.cast(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] += cast(U, A[D.b, D.m, D.k]) * cast(U, B[D.b, D.k]) C[D.b, D.m] += TypeFn.cast(U, A[D.b, D.m, D.k]) * TypeFn.cast(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] += cast(U, A[D.m]) * cast(U, B[D.m]) C[None] += TypeFn.cast(U, A[D.m]) * TypeFn.cast(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] += cast(U, I[D.ow + D.kw]) * cast(U, K[D.kw]) O[D.ow] += TypeFn.cast(U, I[D.ow + D.kw]) * TypeFn.cast(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] += cast(U, I[D.oh + D.kh, D.ow + D.kw]) * cast(U, K[D.kh, D.kw]) O[D.oh, D.ow] += TypeFn.cast(U, I[D.oh + D.kh, D.ow + D.kw]) * TypeFn.cast(
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] += cast(U, I[D.od + D.kd, D.oh + D.kh, D.ow + D.kw]) * cast( D.ow] += TypeFn.cast(U, I[D.od + D.kd, D.oh + D.kh, D.ow +
U, K[D.kd, D.kh, D.kw]) D.kw]) * TypeFn.cast(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), strides=IndexAttrDef(S.SW),
dilations=IndexAttrDef(S.DW)): dilations=IndexAttrDef(S.DW)):
"""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, D.f] += cast(U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.c]) * cast( O[D.n, D.ow,
U, K[D.kw, D.c, D.f]) D.f] += TypeFn.cast(U, I[D.n, D.ow * S.SW + D.kw * S.DW,
D.c]) * TypeFn.cast(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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""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] += cast( O[D.n, D.oh, D.ow, D.f] += TypeFn.cast(
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]) * cast(U, K[D.kh, D.kw, D.c, D.f]) D.c]) * TypeFn.cast(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] += (cast( D.f] += (TypeFn.cast(
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]) -
cast(U, IZp)) * ( TypeFn.cast(U, IZp)) * (
cast(U, K[D.kh, D.kw, D.c, D.f]) - cast(U, KZp)) TypeFn.cast(U, K[D.kh, D.kw, D.c, D.f]) - TypeFn.cast(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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""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] += cast( O[D.n, D.f, D.oh, D.ow] += TypeFn.cast(
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]) * cast(U, K[D.f, D.c, D.kh, D.kw]) D.ow * S.SW + D.kw * S.DW]) * TypeFn.cast(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), strides=IndexAttrDef(S.SD, S.SH, S.SW),
dilations=IndexAttrDef(S.DD, S.DH, S.DW)): dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
"""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] += cast( O[D.n, D.od, D.oh, D.ow, D.f] += TypeFn.cast(
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]) * cast( D.ow * S.SW + D.kw * S.DW, D.c]) * TypeFn.cast(
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), strides=IndexAttrDef(S.SW),
dilations=IndexAttrDef(S.DW)): dilations=IndexAttrDef(S.DW)):
"""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] += \
cast(U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.ic]) * \ TypeFn.cast(U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.ic]) * \
cast(U, K[D.kw, D.ic]) TypeFn.cast(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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""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] += cast( O[D.n, D.oh, D.ow, D.ic] += TypeFn.cast(
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]) * cast(U, K[D.kh, D.kw, D.ic]) D.ic]) * TypeFn.cast(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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""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, D.ic] += ( O[D.n, D.oh, D.ow,
(cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.ic] += ((TypeFn.cast(
D.ic]) - cast(U, IZp)) * U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.ic]) -
(cast(U, K[D.kh, D.kw, D.ic]) - cast(U, KZp))) TypeFn.cast(U, IZp)) *
(TypeFn.cast(U, K[D.kh, D.kw, D.ic]) - TypeFn.cast(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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""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] += cast( O[D.n, D.oh, D.ow, D.ic, D.cm] += TypeFn.cast(
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]) * cast(U, K[D.kh, D.kw, D.ic, D.cm]) D.ic]) * TypeFn.cast(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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""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] += ( O[D.n, D.oh, D.ow, D.ic,
(cast(U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.cm] += ((TypeFn.cast(
D.ic]) - cast(U, IZp)) * U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.ic]) -
(cast(U, K[D.kh, D.kw, D.ic, D.cm]) - cast(U, KZp))) TypeFn.cast(U, IZp)) *
(TypeFn.cast(U, K[D.kh, D.kw, D.ic, D.cm]) - TypeFn.cast(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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""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.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] += cast( O[D.n, D.oh, D.ow, D.c] += TypeFn.cast(
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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""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.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] = ReduceFn.max(D.kh, D.kw)( O[D.n, D.oh, D.ow, D.c] = ReduceFn.max(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, TypeFn.cast(
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]),
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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""Performs unsigned max pooling. """Performs unsigned 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.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] = ReduceFn.max_unsigned(D.kh, D.kw)( O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned(D.kh, D.kw)(
cast_unsigned( TypeFn.cast_unsigned(
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_nchw_max( def pooling_nchw_max(
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.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.C, S.OH, S.OW, output=True), O=TensorDef(U, S.N, S.C, S.OH, S.OW, output=True),
strides=IndexAttrDef(S.SH, S.SW), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""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(D.kh, D.kw)(
cast(U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, TypeFn.cast(
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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""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.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] = ReduceFn.min(D.kh, D.kw)( O[D.n, D.oh, D.ow, D.c] = ReduceFn.min(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, TypeFn.cast(
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]),
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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
"""Performs unsigned min pooling. """Performs unsigned 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.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] = ReduceFn.min_unsigned(D.kh, D.kw)( O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned(D.kh, D.kw)(
cast_unsigned( TypeFn.cast_unsigned(
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_ndhwc_sum( def pooling_ndhwc_sum(
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), strides=IndexAttrDef(S.SD, S.SH, S.SW),
dilations=IndexAttrDef(S.DD, S.DH, S.DW)): dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
"""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.kd, D.kh, D.kw, D.c) domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
O[D.n, D.od, D.oh, D.ow, D.c] += cast( O[D.n, D.od, D.oh, D.ow, D.c] += TypeFn.cast(
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), strides=IndexAttrDef(S.SD, S.SH, S.SW),
dilations=IndexAttrDef(S.DD, S.DH, S.DW)): dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
"""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.kd, D.kh, D.kw, D.c) domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
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(D.kd, D.kh, D.kw)(
cast( TypeFn.cast(
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), strides=IndexAttrDef(S.SD, S.SH, S.SW),
dilations=IndexAttrDef(S.DD, S.DH, S.DW)): dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
"""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.kd, D.kh, D.kw, D.c) domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
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(D.kd, D.kh, D.kw)(
cast( TypeFn.cast(
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_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 = cast(I32, const(1103515245)) multiplier = TypeFn.cast(I32, const(1103515245))
increment = cast(I32, const(12345)) increment = TypeFn.cast(I32, const(12345))
rand1 = (cast(I32, index(D.m)) + seed) * multiplier + increment rand1 = (TypeFn.cast(I32, index(D.m)) + seed) * multiplier + increment
rand2 = (cast(I32, index(D.n)) + rand1) * multiplier + increment rand2 = (TypeFn.cast(I32, index(D.n)) + rand1) * multiplier + increment
inv_range = cast(F64, const(2.3283064e-10)) inv_range = TypeFn.cast(F64, const(2.3283064e-10))
offset = cast(F64, const(2147483647)) offset = TypeFn.cast(F64, const(2147483647))
scaling = (max - min) * inv_range scaling = (max - min) * inv_range
O[D.m, D.n] = cast(T, (offset + cast(F64, rand2)) * scaling + min) O[D.m, D.n] = TypeFn.cast(T,
(offset + TypeFn.cast(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] = \
PrimFn.log(cast(U, const(1.0)) + PrimFn.exp(cast(U, I[D.m, D.n]))) PrimFn.log(TypeFn.cast(U, const(1.0)) + PrimFn.exp(TypeFn.cast(U, I[D.m, D.n])))

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

Show All 32 Linesstructured_op: !LinalgStructuredOpConfig
assignments: assignments:
- !ScalarAssign - !ScalarAssign
arg: O arg: O
value: !ScalarExpression value: !ScalarExpression
scalar_apply: scalar_apply:
fn_name: add fn_name: add
operands: operands:
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast
type_var: T type_var: T
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_const: '42 : i64' scalar_const: '42 : i64'
is_unsigned_cast: false
- !ScalarExpression - !ScalarExpression
symbolic_cast: type_fn:
fn_name: cast_unsigned
type_var: T type_var: T
operands: operands:
- !ScalarExpression - !ScalarExpression
scalar_index: 1 scalar_index: 1
is_unsigned_cast: true
# ODS-LABEL: def Test1Op : LinalgStructuredBase_Op<"test1" # ODS-LABEL: def Test1Op : LinalgStructuredBase_Op<"test1"
# 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: }];
Show All 19 Lines
# 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( # IMPL-LABEL: void Test1Op::regionBuilder(
# IMPL: ImplicitLocOpBuilder &b, Block &block) # IMPL: ImplicitLocOpBuilder &b, Block &block)
# IMPL: Value [[VAL0:[a-z0-9]+]] = helper.constant("42 : i64"); # IMPL: Value [[VAL0:[a-z0-9]+]] = helper.constant("42 : i64");
# IMPL-DAG: Value [[VAL1:[a-z0-9]+]] = helper.cast(block.getArgument(0).getType(), [[VAL0]], false); # IMPL-DAG: Value [[VAL1:[a-z0-9]+]] = helper.typefn__cast(block.getArgument(0).getType(), [[VAL0]]);
# IMPL-DAG: Value [[VAL2:[a-z0-9]+]] = helper.index(1); # IMPL-DAG: Value [[VAL2:[a-z0-9]+]] = helper.index(1);
# IMPL-DAG: Value [[VAL3:[a-z0-9]+]] = helper.cast(block.getArgument(0).getType(), [[VAL2]], true); # IMPL-DAG: Value [[VAL3:[a-z0-9]+]] = helper.typefn__cast_unsigned(block.getArgument(0).getType(), [[VAL2]]);
# IMPL-DAG: Value [[VAL4:[a-z0-9]+]] = helper.applyfn__add([[VAL1]], [[VAL3]]); # IMPL-DAG: Value [[VAL4:[a-z0-9]+]] = helper.applyfn__add([[VAL1]], [[VAL3]]);
# @linalg_structured_op # @linalg_structured_op
# def test2(I=TensorDef(T, S.M, S.N), # def test2(I=TensorDef(T, S.M, S.N),
# O=TensorDef(T, S.M, S.N, output=True), # O=TensorDef(T, S.M, S.N, output=True),
# strides=IndexAttrDef(S.SM, S.SN)): # strides=IndexAttrDef(S.SM, S.SN)):
# """Title. # """Title.
▲ Show 20 Lines • Show All 78 Lines • Show Last 20 Lines

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

Show All 17 Lines
# CHECK: usage: OutputOperand # CHECK: usage: OutputOperand
# CHECK: type_var: U # CHECK: type_var: U
# CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)> # CHECK: shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
@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)):
C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.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: usage: InputOperand # CHECK: usage: InputOperand
# CHECK-NOT: shape_map: # CHECK-NOT: shape_map:
Show All 27 Lines

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

Show All 9 Lines
# CHECK: arg: C # CHECK: arg: C
# CHECK: value: # CHECK: value:
# CHECK: scalar_apply: # CHECK: scalar_apply:
# CHECK: fn_name: add # CHECK: fn_name: add
# CHECK: operands: # CHECK: operands:
# CHECK: scalar_apply: # CHECK: scalar_apply:
# CHECK: fn_name: mul # CHECK: fn_name: mul
# CHECK: operands: # CHECK: operands:
# CHECK: symbolic_cast: # CHECK: type_fn:
# CHECK: type_var: U # CHECK: type_var: U
# CHECK: operands: # CHECK: operands:
# CHECK: scalar_arg: A # CHECK: scalar_arg: A
# CHECK: symbolic_cast: # CHECK: type_fn:
# CHECK: type_var: U # CHECK: type_var: U
# 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)):
C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.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
# CHECK: scalar_apply: # CHECK: scalar_apply:
# CHECK: fn_name: sub # CHECK: fn_name: sub
# CHECK: operands: # CHECK: operands:
# CHECK: scalar_apply: # CHECK: scalar_apply:
# CHECK: fn_name: add # CHECK: fn_name: add
# CHECK: operands: # CHECK: operands:
# CHECK: symbolic_cast: # CHECK: type_fn:
# 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: symbolic_cast: # CHECK: type_fn:
# CHECK: type_var: T # CHECK: type_var: T
# CHECK: operands: # CHECK: operands:
# CHECK: scalar_const: '42 : i64' # CHECK: scalar_const: '42 : i64'
# CHECK: symbolic_cast: # CHECK: type_fn:
# CHECK: type_var: T # CHECK: type_var: T
# 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(O=TensorDef(T, S.M, S.K, output=True)): def constants(O=TensorDef(T, S.M, S.K, output=True)):
pi = cast(T, const(3.1415926535897931)) pi = TypeFn.cast(T, const(3.1415926535897931))
cst42 = cast(T, const(42)) cst42 = TypeFn.cast(T, const(42))
cst1000 = cast(T, const(1e+3)) cst1000 = TypeFn.cast(T, const(1e+3))
O[D.m, D.n] = pi + cst42 - cst1000 O[D.m, D.n] = 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 19 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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
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] += cast( O[D.n, D.oh, D.ow, D.c] += TypeFn.cast(
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]) * cast(U, K[D.kh, D.kw, D.c]) D.c]) * TypeFn.cast(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 26 Lines

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

Show All 20 Lines
@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)):
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] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(U, B[D.k, D.n])
@linalg_structured_op @linalg_structured_op
def matmul_unsigned_poly( def matmul_unsigned_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)):
domain(D.m, D.n, D.k) domain(D.m, D.n, D.k)
C[D.m, D.n] += cast_unsigned(U, A[D.m, D.k]) * cast_unsigned(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast_unsigned(U, A[D.m, D.k]) * TypeFn.cast_unsigned(
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()
f64 = F64Type.get() f64 = F64Type.get()
i8 = IntegerType.get_signless(8) i8 = IntegerType.get_signless(8)
▲ Show 20 Lines • Show All 130 Lines • Show Last 20 Lines

mlir/test/python/dialects/linalg/opdsl/emit_misc.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 *
# This tests miscellaneous features of the emitter that are not tested by the # This tests miscellaneous features of the emitter that are not tested by the
# matmul, convolution, or, pooling tests. The features include: # matmul, convolution, or, pooling tests. The features include:
# - constant defined in the body # - constant defined in the body
# - fix/predefined types # - fix/predefined types
# - exponential functions # - exponential functions
# - custom op names. # - custom op names.
@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 = cast(I32, const(1103515245)) multiplier = TypeFn.cast(I32, const(1103515245))
increment = cast(I32, const(12345)) increment = TypeFn.cast(I32, const(12345))
rand1 = (cast(I32, index(D.m)) + seed) * multiplier + increment rand1 = (TypeFn.cast(I32, index(D.m)) + seed) * multiplier + increment
rand2 = (cast(I32, index(D.n)) + rand1) * multiplier + increment rand2 = (TypeFn.cast(I32, index(D.n)) + rand1) * multiplier + increment
inv_range = cast(F64, const(2.3283064e-10)) inv_range = TypeFn.cast(F64, const(2.3283064e-10))
offset = cast(F64, const(2147483647)) offset = TypeFn.cast(F64, const(2147483647))
scaling = (max - min) * inv_range scaling = (max - min) * inv_range
O[D.m, D.n] = cast(T, (offset + cast(F64, rand2)) * scaling + min) O[D.m, D.n] = TypeFn.cast(T,
(offset + TypeFn.cast(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] = \ O[D.m, D.n] = PrimFn.log(
PrimFn.log(cast(U, const(1.0)) + cast(U, PrimFn.exp(I[D.m, D.n]))) TypeFn.cast(U, const(1.0)) + TypeFn.cast(U, PrimFn.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 14 Lines
def pooling_max_poly( def pooling_max_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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
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] = ReduceFn.max(D.kh, D.kw)( O[D.n, D.oh, D.ow, D.c] = ReduceFn.max(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, TypeFn.cast(
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_max_unsigned_poly( def pooling_max_unsigned_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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
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] = ReduceFn.max_unsigned(D.kh, D.kw)( O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned(D.kh, D.kw)(
cast_unsigned( TypeFn.cast_unsigned(
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_min_poly( def pooling_min_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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
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] = ReduceFn.min(D.kh, D.kw)( O[D.n, D.oh, D.ow, D.c] = ReduceFn.min(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, TypeFn.cast(
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_min_unsigned_poly( def pooling_min_unsigned_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), strides=IndexAttrDef(S.SH, S.SW),
dilations=IndexAttrDef(S.DH, S.DW)): dilations=IndexAttrDef(S.DH, S.DW)):
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] = ReduceFn.min_unsigned(D.kh, D.kw)( O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned(D.kh, D.kw)(
cast_unsigned( TypeFn.cast_unsigned(
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]))
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 20 Lines • Show All 84 Lines • Show Last 20 Lines

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] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(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] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) C[D.m, D.n] += TypeFn.cast(U, A[D.m, D.k]) * TypeFn.cast(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] += cast(U, A[D.m]) * cast(U, B[D.m]) C[None] += TypeFn.cast(U, A[D.m]) * TypeFn.cast(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] += cast(U, I[D.o * 2 + D.k]) O[D.o] += TypeFn.cast(U, I[D.o * 2 + D.k])

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

Show First 20 Lines • Show All 83 Lines • ▼ Show 20 Lines
struct ScalarApply { struct ScalarApply {
std::string fnName; std::string fnName;
// NOTE: Must be pure heap allocated container (not SmallVector) // NOTE: Must be pure heap allocated container (not SmallVector)
// due to recursive data type. // due to recursive data type.
std::vector<ScalarExpression> operands; std::vector<ScalarExpression> operands;
}; };
struct ScalarSymbolicCast { struct ScalarTypeFn {
std::string fnName;
std::string typeVar; std::string typeVar;
// NOTE: This must be of arity 1, but to break the self-referential cycle, // NOTE: This must be of arity 1, but to break the self-referential cycle,
// we use a heap allocated vector. // we use a heap allocated vector.
std::vector<ScalarExpression> operands; std::vector<ScalarExpression> operands;
bool isUnsignedCast;
}; };
struct ScalarExpression { struct ScalarExpression {
Optional<std::string> arg; Optional<std::string> arg;
Optional<std::string> constant; Optional<std::string> constant;
Optional<int64_t> index; Optional<int64_t> index;
Optional<ScalarApply> apply; Optional<ScalarApply> apply;
Optional<ScalarSymbolicCast> symbolicCast; Optional<ScalarTypeFn> typeFn;
}; };
struct ScalarAssign { struct ScalarAssign {
std::string arg; std::string arg;
ScalarExpression value; ScalarExpression value;
}; };
struct LinalgStructuredOpConfig { struct LinalgStructuredOpConfig {
Show All 22 Lines
LLVM_YAML_IS_DOCUMENT_LIST_VECTOR(LinalgOpConfig) LLVM_YAML_IS_DOCUMENT_LIST_VECTOR(LinalgOpConfig)
namespace llvm { namespace llvm {
namespace yaml { namespace yaml {
/// Top-level type containing op metadata and one of a concrete op type. /// Top-level type containing op metadata and one of a concrete op type.
/// Currently, the only defined op type is `structured_op` (maps to /// Currently, the only defined op type is `structured_op` (maps to
/// `LinalgStructuredOpConfig`). /// `LinalgStructuredOpConfig`).
template <> struct MappingTraits<LinalgOpConfig> { template <>
struct MappingTraits<LinalgOpConfig> {
static void mapping(IO &io, LinalgOpConfig &info) { static void mapping(IO &io, LinalgOpConfig &info) {
io.mapOptional("metadata", info.metadata); io.mapOptional("metadata", info.metadata);
io.mapOptional("structured_op", info.structuredOp); io.mapOptional("structured_op", info.structuredOp);
} }
}; };
/// A structured op models (at most) a single contraction by modeling /// A structured op models (at most) a single contraction by modeling
/// - A list of named arguments (`LinalgOperandDef`), which can be inputs, /// - A list of named arguments (`LinalgOperandDef`), which can be inputs,
/// outputs, or index attributes. /// outputs, or index attributes.
/// - List of indexing maps (see `LinalgIndexingMaps`). /// - List of indexing maps (see `LinalgIndexingMaps`).
/// - Iterator types (see `LinalgIteratorTypeDef`). /// - Iterator types (see `LinalgIteratorTypeDef`).
/// - List of scalar level assignment (see `ScalarAssign`). /// - List of scalar level assignment (see `ScalarAssign`).
template <> struct MappingTraits<LinalgStructuredOpConfig> { template <>
struct MappingTraits<LinalgStructuredOpConfig> {
static void mapping(IO &io, LinalgStructuredOpConfig &info) { static void mapping(IO &io, LinalgStructuredOpConfig &info) {
io.mapRequired("args", info.args); io.mapRequired("args", info.args);
io.mapRequired("indexing_maps", info.indexingMaps); io.mapRequired("indexing_maps", info.indexingMaps);
io.mapRequired("iterator_types", info.iteratorTypes); io.mapRequired("iterator_types", info.iteratorTypes);
io.mapRequired("assignments", info.assignments); io.mapRequired("assignments", info.assignments);
} }
}; };
/// Maps a named tensor, scalar or attribute argument to an operation, /// Maps a named tensor, scalar or attribute argument to an operation,
/// consisting of: /// consisting of:
/// - `name`: Must be unique within the operation. /// - `name`: Must be unique within the operation.
/// - `usage`: How the argument is used (input, output, attribute, etc). /// - `usage`: How the argument is used (input, output, attribute, etc).
/// - `type_var`: The symbolic type variable that binds to the element or self /// - `type_var`: The symbolic type variable that binds to the element or self
/// type of the tensor or scalar argument, respectively. /// type of the tensor or scalar argument, respectively.
/// - `shape_map`: An optional AffineMap from all op symbols to the shape of /// - `shape_map`: An optional AffineMap from all op symbols to the shape of
/// the argument. Only tensor arguments have a `shape_map`. Each shape must /// the argument. Only tensor arguments have a `shape_map`. Each shape must
/// be normalized over the same list of symbols and have no dimension /// be normalized over the same list of symbols and have no dimension
/// inputs. /// inputs.
/// - `attribute_map`: An optional AffineMap from all op symbols to the /// - `attribute_map`: An optional AffineMap from all op symbols to the
/// attribute symbols. During op creation these symbols are replaced by the /// attribute symbols. During op creation these symbols are replaced by the
/// corresponding `name` attribute values. Only attribute arguments have /// corresponding `name` attribute values. Only attribute arguments have
/// an `attribute_map`. /// an `attribute_map`.
template <> struct MappingTraits<LinalgOperandDef> { template <>
struct MappingTraits<LinalgOperandDef> {
static void mapping(IO &io, LinalgOperandDef &info) { static void mapping(IO &io, LinalgOperandDef &info) {
io.mapRequired("name", info.name); io.mapRequired("name", info.name);
io.mapRequired("usage", info.usage); io.mapRequired("usage", info.usage);
io.mapRequired("type_var", info.typeVar); io.mapRequired("type_var", info.typeVar);
io.mapOptional("shape_map", info.shapeMap); io.mapOptional("shape_map", info.shapeMap);
io.mapOptional("attribute_map", info.attributeMap); io.mapOptional("attribute_map", info.attributeMap);
} }
}; };
/// Usage enum for a named argument. /// Usage enum for a named argument.
template <> struct ScalarEnumerationTraits<LinalgOperandDefUsage> { template <>
struct ScalarEnumerationTraits<LinalgOperandDefUsage> {
static void enumeration(IO &io, LinalgOperandDefUsage &value) { static void enumeration(IO &io, LinalgOperandDefUsage &value) {
io.enumCase(value, "InputOperand", LinalgOperandDefUsage::input); io.enumCase(value, "InputOperand", LinalgOperandDefUsage::input);
io.enumCase(value, "OutputOperand", LinalgOperandDefUsage::output); io.enumCase(value, "OutputOperand", LinalgOperandDefUsage::output);
io.enumCase(value, "IndexAttribute", LinalgOperandDefUsage::attribute); io.enumCase(value, "IndexAttribute", LinalgOperandDefUsage::attribute);
} }
}; };
/// Iterator type enum. /// Iterator type enum.
template <> struct ScalarEnumerationTraits<LinalgIteratorTypeDef> { template <>
struct ScalarEnumerationTraits<LinalgIteratorTypeDef> {
static void enumeration(IO &io, LinalgIteratorTypeDef &value) { static void enumeration(IO &io, LinalgIteratorTypeDef &value) {
io.enumCase(value, "parallel", LinalgIteratorTypeDef::parallel); io.enumCase(value, "parallel", LinalgIteratorTypeDef::parallel);
io.enumCase(value, "reduction", LinalgIteratorTypeDef::reduction); io.enumCase(value, "reduction", LinalgIteratorTypeDef::reduction);
} }
}; };
/// Metadata about the op (name, C++ name, and documentation). /// Metadata about the op (name, C++ name, and documentation).
template <> struct MappingTraits<LinalgOpMetadata> { template <>
struct MappingTraits<LinalgOpMetadata> {
static void mapping(IO &io, LinalgOpMetadata &info) { static void mapping(IO &io, LinalgOpMetadata &info) {
io.mapRequired("name", info.name); io.mapRequired("name", info.name);
io.mapRequired("cpp_class_name", info.cppClassName); io.mapRequired("cpp_class_name", info.cppClassName);
io.mapOptional("doc", info.doc); io.mapOptional("doc", info.doc);
io.mapOptional("implements", info.implements); io.mapOptional("implements", info.implements);
} }
}; };
/// How the ops indexing maps are produced. Must be one of: /// How the ops indexing maps are produced. Must be one of:
/// - static_indexing_maps: A static list of AffineMaps, possibly with /// - static_indexing_maps: A static list of AffineMaps, possibly with
/// some symbols that bind to attributes of the op. Each indexing map must /// some symbols that bind to attributes of the op. Each indexing map must
/// be normalized over the same list of dimensions, and its symbols must /// be normalized over the same list of dimensions, and its symbols must
/// match the symbols for argument shapes. /// match the symbols for argument shapes.
template <> struct MappingTraits<LinalgIndexingMapsConfig> { template <>
struct MappingTraits<LinalgIndexingMapsConfig> {
static void mapping(IO &io, LinalgIndexingMapsConfig &info) { static void mapping(IO &io, LinalgIndexingMapsConfig &info) {
io.mapOptional("static_indexing_maps", info.staticIndexingMaps); io.mapOptional("static_indexing_maps", info.staticIndexingMaps);
} }
}; };
/// Models an assignment to a named output. /// Models an assignment to a named output.
/// - The `arg` name must match a named output. /// - The `arg` name must match a named output.
/// - The `value` is a scalar expression for computing the value to /// - The `value` is a scalar expression for computing the value to
/// assign (see `ScalarExpression`). /// assign (see `ScalarExpression`).
template <> struct MappingTraits<ScalarAssign> { template <>
struct MappingTraits<ScalarAssign> {
static void mapping(IO &io, ScalarAssign &info) { static void mapping(IO &io, ScalarAssign &info) {
io.mapRequired("arg", info.arg); io.mapRequired("arg", info.arg);
io.mapRequired("value", info.value); io.mapRequired("value", info.value);
} }
}; };
/// A scalar expression (RHS of an assignment). Must be one of: /// A scalar expression (RHS of an assignment). Must be one of:
/// - `scalar_arg`: Name of an argument to the op. /// - `scalar_arg`: Name of an argument to the op.
/// - `scalar_apply`: Result of evaluating a named function (see /// - `scalar_apply`: Result of evaluating a named function (see
/// `ScalarApply`). /// `ScalarApply`).
/// - `symbolic_cast`: Cast to a symbolic TypeVar bound elsewhere. /// - `type_fn`: A named type conversion function (see `ScalarTypeFn`).
template <> struct MappingTraits<ScalarExpression> { template <>
struct MappingTraits<ScalarExpression> {
static void mapping(IO &io, ScalarExpression &info) { static void mapping(IO &io, ScalarExpression &info) {
io.mapOptional("scalar_arg", info.arg); io.mapOptional("scalar_arg", info.arg);
io.mapOptional("scalar_const", info.constant); io.mapOptional("scalar_const", info.constant);
io.mapOptional("scalar_index", info.index); io.mapOptional("scalar_index", info.index);
io.mapOptional("scalar_apply", info.apply); io.mapOptional("scalar_apply", info.apply);
io.mapOptional("symbolic_cast", info.symbolicCast); io.mapOptional("type_fn", info.typeFn);
} }
}; };
/// A scalar expression that evaluates a named function. /// A scalar expression that evaluates a named function.
/// Functions are generally "math" level and type polymorphic. Builtin /// Functions are generally "math" level and type polymorphic. Builtin
/// functions include: /// functions include:
/// - `add(lhs, rhs)` /// - `add(lhs, rhs)`
/// - `mul(lhs, rhs)` /// - `mul(lhs, rhs)`
template <> struct MappingTraits<ScalarApply> { template <>
struct MappingTraits<ScalarApply> {
static void mapping(IO &io, ScalarApply &info) { static void mapping(IO &io, ScalarApply &info) {
io.mapRequired("fn_name", info.fnName); io.mapRequired("fn_name", info.fnName);
io.mapRequired("operands", info.operands); io.mapRequired("operands", info.operands);
} }
}; };
template <> struct MappingTraits<ScalarSymbolicCast> { template <>
static void mapping(IO &io, ScalarSymbolicCast &info) { struct MappingTraits<ScalarTypeFn> {
static void mapping(IO &io, ScalarTypeFn &info) {
io.mapRequired("fn_name", info.fnName);
io.mapRequired("type_var", info.typeVar); io.mapRequired("type_var", info.typeVar);
io.mapRequired("operands", info.operands); io.mapRequired("operands", info.operands);
io.mapRequired("is_unsigned_cast", info.isUnsignedCast);
} }
}; };
/// Helper mapping which accesses an AffineMapAttr as a serialized string of /// Helper mapping which accesses an AffineMapAttr as a serialized string of
/// the same. /// the same.
template <> struct ScalarTraits<SerializedAffineMap> { template <>
struct ScalarTraits<SerializedAffineMap> {
static void output(const SerializedAffineMap &value, void *rawYamlContext, static void output(const SerializedAffineMap &value, void *rawYamlContext,
raw_ostream &out) { raw_ostream &out) {
assert(value.affineMapAttr); assert(value.affineMapAttr);
value.affineMapAttr.print(out); value.affineMapAttr.print(out);
} }
static StringRef input(StringRef scalar, void *rawYamlContext, static StringRef input(StringRef scalar, void *rawYamlContext,
SerializedAffineMap &value) { SerializedAffineMap &value) {
assert(rawYamlContext); assert(rawYamlContext);
▲ Show 20 Lines • Show All 659 Lines • ▼ Show 20 Lines for (LinalgOperandDef &arg : args) {
} }
std::string cppIdent = llvm::formatv("value{0}", ++localCounter); std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
stmts.push_back( stmts.push_back(
llvm::formatv("Value {0} = helper.applyfn__{1}({2});", cppIdent, llvm::formatv("Value {0} = helper.applyfn__{1}({2});", cppIdent,
expression.apply->fnName, expression.apply->fnName,
interleaveToString(operandCppValues, ", "))); interleaveToString(operandCppValues, ", ")));
return cppIdent; return cppIdent;
} }
if (expression.symbolicCast) { if (expression.typeFn) {
// Symbolic cast. // Symbolic cast.
// Operands must be arity 1. // Operands must be arity 1.
if (expression.symbolicCast->operands.size() != 1) { if (expression.typeFn->operands.size() != 1) {
emitError(genContext.getLoc()) emitError(genContext.getLoc())
<< "symbolic_cast operand arity must be 1"; << "type conversion operand arity must be 1";
return None; return None;
} }
Optional<std::string> operandCppValue = Optional<std::string> operandCppValue =
generateExpression(expression.symbolicCast->operands[0]); generateExpression(expression.typeFn->operands[0]);
if (!operandCppValue) if (!operandCppValue)
return None; return None;
Optional<std::string> typeCppValue = Optional<std::string> typeCppValue =
findTypeValue(expression.symbolicCast->typeVar, args); findTypeValue(expression.typeFn->typeVar, args);
if (!typeCppValue) { if (!typeCppValue) {
emitError(genContext.getLoc()) emitError(genContext.getLoc())
<< "type variable " << expression.symbolicCast->typeVar << "type variable " << expression.typeFn->typeVar
<< ", used in a symbolic cast must map to a predefined or " << ", used in a type conversion, must map to a predefined or "
<< "an argument type but it does not"; << "an argument type but it does not";
return None; return None;
} }
std::string cppIdent = llvm::formatv("value{0}", ++localCounter); std::string cppIdent = llvm::formatv("value{0}", ++localCounter);
stmts.push_back( stmts.push_back(
llvm::formatv("Value {0} = helper.cast({1}, {2}, {3});", cppIdent, llvm::formatv("Value {0} = helper.typefn__{1}({2}, {3});",
typeCppValue.getValue(), *operandCppValue, cppIdent, expression.typeFn->fnName,
expression.symbolicCast->isUnsignedCast)); typeCppValue.getValue(), *operandCppValue));
return cppIdent; return cppIdent;
} }
emitError(genContext.getLoc()) << "unknown ScalarExpression type"; emitError(genContext.getLoc()) << "unknown ScalarExpression type";
return None; return None;
}; };
Optional<std::string> cppValue = generateExpression(assignment->value); Optional<std::string> cppValue = generateExpression(assignment->value);
if (!cppValue) if (!cppValue)
return failure(); return failure();
▲ Show 20 Lines • Show All 119 Lines • Show Last 20 Lines