This is an archive of the discontinued LLVM Phabricator instance.

[mlir][linalg] Support arbitrary element types in named operations via attributes
Needs ReviewPublic

Authored by andidr on Feb 17 2022, 2:12 AM.

Details

Summary

Linalg named operations are currently limited to tensors and memrefs
composed of floating point, integer, bool or complex elements and
using any other element type triggers an assertion.

This change adds support for arbitrary element types through the
specification of the arithmetic operations associated to a type in
specific attributes of a linalg named operation. The attributes' names
correspond to the short form of the arithmetic operator implemented by
the operation (i.e., add, sub, mul, max_signed, max_unsigned,
min_signed, min_unsigned, exp, abs, ceil, negf or log) and receive as
values the name of an operation and optionally a return type given
after a colon.

For example, a linalg.matmul operation multiplying two tensors
composed of a hypothetical type foo with operations foo.add and
foo.mul would be expressed as:

linalg.matmul { add = "foo.add", mul = "foo.mul" }
  ins(%arg1, %arg2 : tensor<?x?x!foo>, tensor<?x?x!foo>)
  outs(%2 : tensor<?x?x!foo>>)

Sensible default values for the attributes are given for float,
integer and complex types, such that the omission of the attributes
results in the original behavior of the named operation before this
change. I.e., the expression:

linalg.matmul { add = "arith.addf", mul = "arith.mulf" }
  ins(%arg1, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)
  outs(%2 : tensor<?x?xf32>)

and:

linalg.matmul
  ins(%arg1, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)
  outs(%2 : tensor<?x?xf32>)

yield identical results.

By default, the result type of an operation implementing an arithmetic
operator is assumed to be identical with the type of the first
argument. If this assumption does not hold for an operation, the
result type must be specified explicitly. E.g., if for the type foo,
the operation foo.mul yields a bar and for the type bar, the
operation bar.add yields a baz, this would have to be specied as:

linalg.matmul { mul = "foo.mul:!bar"
                add = "bar.add:!baz", }
  ins(%arg1, %arg2 : tensor<?x?x!foo>, tensor<?x?x!foo>)
  outs(%2 : tensor<?x?x!bar>)

The extraction of operation names and result types from attributes,
proper instantiation and default values are provided by a set of
operation interfaces (one per operator) in
LinalgFrontendInterfaces.td. The set of operation interfaces
required for a named operation is derived transparently from the
arithmetic expressions in its YAML specification via
mlir-linalg-ods-yaml-gen.

Diff Detail

Event Timeline

andidr created this revision.Feb 17 2022, 2:12 AM
andidr requested review of this revision.Feb 17 2022, 2:12 AM
andidr updated this revision to Diff 438616.Jun 21 2022, 3:02 AM
andidr edited the summary of this revision. (Show Details)

Rebased to the latest HEAD at 3f81841474fefcf10fea0ed7fc21f47af3c7b80a.

Herald added a project: Restricted Project. · View Herald TranscriptJun 21 2022, 3:02 AM
Herald added a subscriber: bzcheeseman. · View Herald Transcript
andidr updated this revision to Diff 438651.Jun 21 2022, 5:20 AM
gysit resigned from this revision.Oct 16 2022, 12:19 AM
andidr updated this revision to Diff 468459.Oct 18 2022, 2:23 AM
andidr edited the summary of this revision. (Show Details)

Rebased to 8ef96d141a7247832986099722907ca44bbb7ec4, included support for booleans, updated description.

Thanks for developing this change. I'm interested in seeing it land although I'm not a committer. These comments address the problems I had when integrating this patch locally.

mlir/include/mlir/Dialect/Linalg/Frontend/LinalgFrontendInterfaces.td
171–173

AFAIU it's the container operation that has this attribute, the one that implements the interface, eg linalg.matmul. However the parameter operation here is the element-wise operation like arith.addi.

224
mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
663–681

LLVM's optional has been replaced by std::optional since then