Page MenuHomePhabricator

Add a "kind" attribute to ContractionOp and OuterProductOp.
ClosedPublic

Authored by pravnar on Dec 14 2020, 10:37 PM.

Details

Summary

Currently, vector.contract joins the intermediate result and the accumulator argument (of ranks K) using summation. We desire more joining operations --- such as max --- to help vector.contract express reductions. This change extends Vector_ContractionOp to take an optional attribute (called "kind", of enum type CombiningKind) specifying the joining operation to be add/mul/min/max for int/fp , and and/or/xor for int only. By default this attribute has value "add".

To implement this we also need to extend vector.outerproduct, since vector.contract gets transformed to vector.outerproduct (and that to vector.fma). The extension for vector.outerproduct is also an optional kind attribute that uses the same enum type and possible values. The default is "add". In case of max/min we transform vector.outerproduct to a combination of compare and select.

Diff Detail

Event Timeline

pravnar created this revision.Dec 14 2020, 10:37 PM
pravnar requested review of this revision.Dec 14 2020, 10:37 PM
pravnar edited the summary of this revision. (Show Details)Dec 14 2020, 10:39 PM
pravnar added a reviewer: bkramer.
pravnar added a subscriber: burmako.
mehdi_amini added inline comments.Dec 16 2020, 12:02 PM
mlir/include/mlir/Dialect/Vector/VectorOps.td
124

Can we use an enum instead of string since the number of values is predefined?

aartbik added inline comments.Dec 17 2020, 9:43 AM
mlir/include/mlir/Dialect/Vector/VectorOps.td
124

If you are extending this, can we please extend this to all operations that currently are supported by reductions (add/mul/min/max for int/fp and and/or/xor for int only). Also, that way we can use one struct for both vector.contract and vector.reduction, which will address Mehdi's request.

pravnar updated this revision to Diff 318284.Jan 21 2021, 12:18 PM

(Test) Changing the revision message.

pravnar updated this revision to Diff 319666.Jan 27 2021, 1:56 PM

Add and use an enum to express combining functions

pravnar edited the summary of this revision. (Show Details)Jan 27 2021, 2:03 PM
pravnar added inline comments.
mlir/include/mlir/Dialect/Vector/VectorOps.td
124

Thanks for the suggestions :). I've added a new enum for expressing the combining functions, and used this enum for defining vector.contract and vector.outerproduct. We can update vector.reduction to use this enum as a separate change.

pravnar updated this revision to Diff 319675.Jan 27 2021, 2:17 PM

Update documentation for contract and outerproduct.

mehdi_amini added inline comments.
mlir/include/mlir/Dialect/Vector/VectorOps.td
246–247

Can you just write kind()? There should be a generated accessor here.

mlir/lib/Dialect/Vector/VectorOps.cpp
310

It may be best to change the build declaration to add a defaulted extra argument for kind to align with the usual generated TableGen builder?
@ftynse may have an opinion here.

mlir/lib/Dialect/Vector/VectorTransforms.cpp
1367

Nit: remove trivial braces (here and below / elsewhere)

pravnar updated this revision to Diff 319866.Jan 28 2021, 7:41 AM

Use kind() and other edits.

pravnar marked 2 inline comments as done.Jan 28 2021, 7:43 AM
pravnar added inline comments.
mlir/include/mlir/Dialect/Vector/VectorOps.td
246–247

That works, thanks.

pravnar marked an inline comment as done.Jan 28 2021, 7:43 AM
mehdi_amini added a comment.EditedJan 28 2021, 9:58 AM

@aartbik / @nicolasvasilache / @ftynse any more comment from the vector dialect point of view? Are you OK with this addition conceptually?

mlir/include/mlir/Dialect/Vector/VectorOps.td
124

Actually it seems that you're still using a StringAttr, if we don't expect this to grow over 64 elements, could we use a BitEnumAttr instead?

pravnar added inline comments.Jan 28 2021, 10:31 AM
mlir/include/mlir/Dialect/Vector/VectorOps.td
124

Ah, I guess BitEnumAttr would provide better performance?

I see that BitEnumAttr is stored as an IntegerAttr. One reason I chose StrEnumAttr over something like I64EnumAttr was because the op attributes were printed like kind = "max" rather than kind = 4 : i64. In the latter case one has to look up what the "4th kind" is defined to be. How much do we care about that? IIUC, in the BitEnumAttr case it would look something like kind = 0x8 : i32.

mehdi_amini added inline comments.Jan 28 2021, 12:26 PM
mlir/include/mlir/Dialect/Vector/VectorOps.td
124

No it can be printed with textual support, take a look into FastmathFlags in LLVM dialect.

pravnar marked an inline comment as not done.Jan 29 2021, 12:15 PM
pravnar added inline comments.
mlir/include/mlir/Dialect/Vector/VectorOps.td
124

I see! That looks great, thanks. I can change this to BitEnumAttr.

thanks for leading this change!

mlir/include/mlir/Dialect/Vector/VectorOps.td
124

+1 on using bits (and symbolics) over strings after this change!

mlir/lib/Dialect/Vector/VectorOps.cpp
556

period at end of comment

1713

can we combine this with the similar code above in a helper?

mlir/lib/Dialect/Vector/VectorTransforms.cpp
1363

genMultX is probably a bit of a misnomer with the generalization

nicolasvasilache requested changes to this revision.Tue, Feb 9, 12:15 AM
nicolasvasilache added inline comments.
mlir/include/mlir/Dialect/Vector/VectorOps.td
40

This is just expressing half of the semantics, we need 2 ops (see below).
Since some combinations are going to be invalid I'd recommend listing the supported cases explicitly rather than pair of ops.

AddMul,
MaxAdd,
etc..
122

the combining operator is only 1/2 of the op semantics.

910

This seems incorrect to me: it assumes multiply is always one of the correct operations but I don't think it is (e.g. {max, +}, {abs, +}, ...).

mlir/lib/Dialect/Vector/VectorTransforms.cpp
1401

If we could use llvm::TypeSwitch or an equivalent, things will look much nicer.
I am unclear it works on int value switching though, maybe @mehdi_amini has a suggestion?

This revision now requires changes to proceed.Tue, Feb 9, 12:15 AM
pravnar updated this revision to Diff 322567.Tue, Feb 9, 7:20 PM

Implement CombiningKind using BitEnumAttr.
Add a wrapper CKAttr and define custom parsing/printing of the attribute.
Refactor duplicate verification functions into a helper.
Add a test for matvec transform with max.

pravnar marked 3 inline comments as done.Tue, Feb 9, 7:22 PM
mehdi_amini added inline comments.Tue, Feb 9, 7:45 PM
mlir/include/mlir/Dialect/Vector/VectorOps.h
78

I'm not fond of the acronym usage, can we spell it out entirely?

Ok, let me not hold you up here: this is already progress and we can separate the generalization of the mul to a followup CL, this is fine for now.

This revision is now accepted and ready to land.Wed, Feb 10, 12:11 PM
pravnar updated this revision to Diff 322808.Wed, Feb 10, 1:23 PM

s/CKAttr/CombiningKindAttr

pravnar marked an inline comment as done.Wed, Feb 10, 1:25 PM
pravnar updated this revision to Diff 322826.Wed, Feb 10, 2:25 PM

Add add_public_tablegen_target.

pravnar updated this revision to Diff 322884.Wed, Feb 10, 5:57 PM

Add new target MLIRVectorOpsEnumsIncGen to address build errors.

pravnar updated this revision to Diff 323130.Thu, Feb 11, 1:27 PM
pravnar marked 3 inline comments as done.

Move kind attribute args to the end to make gen'ed build methods use default vals.

pravnar marked an inline comment as done.Thu, Feb 11, 1:30 PM
pravnar updated this revision to Diff 323203.Thu, Feb 11, 7:09 PM

Sync to latest repo version.