Page MenuHomePhabricator

[mlir][Linalg] Update the semantics, verifier and test for Linalg with tensors.
ClosedPublic

Authored by nicolasvasilache on Fri, Jan 10, 11:22 PM.

Details

Summary

This diff fixes issues with the semantics of linalg.generic on tensors that appeared when converting directly from HLO to linalg.generic.
The changes are self-contained within MLIR and can be captured and tested independently of XLA.

The linalg.generic and indexed_generic are updated to:

To allow progressive lowering from the value world (a.k.a tensor values) to
the buffer world (a.k.a memref values), a linalg.generic op accepts
mixing input and output ranked tensor values with input and output memrefs.

%1 = linalg.generic #trait_attribute %A, %B {other-attributes} :
  tensor<?x?xf32>,
  memref<?x?xf32, stride_specification>
  -> (tensor<?x?xf32>)

In this case, the number of outputs (args_out) must match the sum of (1) the
number of output buffer operands and (2) the number of tensor return values.
The semantics is that the linalg.indexed_generic op produces (i.e.
allocates and fills) its return values.

Tensor values must be legalized by a buffer allocation pass before most
transformations can be applied. Such legalization moves tensor return values
into output buffer operands and updates the region argument accordingly.

Transformations that create control-flow around linalg.indexed_generic
operations are not expected to mix with tensors because SSA values do not
escape naturally. Still, transformations and rewrites that take advantage of
tensor SSA values are expected to be useful and will be added in the near
future.

Diff Detail

Event Timeline

Herald added a project: Restricted Project. · View Herald TranscriptFri, Jan 10, 11:22 PM

Unit tests: pass. 61744 tests passed, 0 failed and 780 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Unit tests: pass. 61744 tests passed, 0 failed and 780 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

sanjoy edited reviewers, added: sanjoy.google; removed: sanjoy.Sun, Jan 12, 1:46 PM

I browsed over the code and left some minor comments. The idea and code in verifiers LGTM. I find it much harder to read now with all those similar sounding accessors but that is inevitable I assume if you want to support mixed forms.

It seems that all transformation and code generation still assumes all buffer representations. That is fine (and also makes sense) but I'd prefer to have something emit an error or assert if this is not the case. One way less to shoot myself in the foot.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
693–695

So this example has a single output assuming that there is an order constraint on the arguments (inputs first, then outputs). Could the describe this here?

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h
186

The number of results alone needs to the smaller than nOutputs(). Maybe assert that?

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

This iterators over all operands, including output buffers.

mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp
92–94

This code assumes an all buffer representation? Maybe add an assert?

230

This code assumes an all buffers representation?

nicolasvasilache marked 8 inline comments as done.Mon, Jan 13, 1:28 PM
nicolasvasilache added inline comments.
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
693–695

Updated the example, had forgotten this, thanks!

mlir/lib/Dialect/Linalg/Transforms/LinalgToLoops.cpp
92–94

Sprinkled a bunch in the relevant places, thanks!

230

sprinkled a bunch of asserts in the relevant places, thanks!

nicolasvasilache marked 3 inline comments as done.

Address review comments.

Unit tests: pass. 61794 tests passed, 0 failed and 781 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

ftynse requested changes to this revision.Tue, Jan 14, 7:23 AM
ftynse added inline comments.
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
60

Apologies for being obnoxious, but the review would be significantly faster if NFC code motion changes were separate from functional changes (in the former case I can quickly eyeball it to confirm it's indeed NFC, whereas in the late I need to understand deeply what the code does; without separation, the understanding time also extends on the NFC parts)

132

Why remove the documentation?

573–574

Unclear whether "allocates and fills" extends to both memrefs and tensors, or just tensors since only tensors are "return values".

578

Nit: argument -> arguments

580–582

I cannot understand what "does not mix with tensors" mean in this context. Consider rephrasing or providing an example.

698–702

I wonder if you could write this text only once in a tablegen variable for both linalg.generic and linalg.indexed_generic, and then just concatenate that variable with the rest of the op-specific documentation. This would avoid the duplication and everything that comes with it, but may make it harder for a casual reader of the op definition to follow. Since you are the only person actively supporting this and doing the work twice, your call.

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h
123–124

Nit: values->buffers

136

Typo "depending regardless"

138

Nit: remove this //, I was confused by it to read the block comment above as describing the function below, which is contradictory.

188

Spurious semicolon

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp
142–146

This dependence analysis now only works if the output is a buffer, and cannot be queried if it is a tensor. Consider documenting (and maybe extending later to also work for tensors).

148

Do you need any modifications to support tensor inputs here? From a quick look, it should just work, but better check...

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

Ultra-nit: functions are not necessarily fun, please use "function" in user-visible messages

188–201

Why +1 here? It's inconsistent with block arguments above.

220

The use of 0-based index here is inconsistent with the 1-based index below. Please adopt a convention, document it in the dialect documentation and use everywhere.

225–239

This looks common with then non-indexed version. Can this be factored out into a helper function?

243

Nit: "inputs and buffer operands" sounds like a false dichotomy.

730

Ultra-nit: whitespace before (

This revision now requires changes to proceed.Tue, Jan 14, 7:23 AM
nicolasvasilache marked 25 inline comments as done.

Address review comments.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
60

Apologies for this reordering.
I was contemplating extracting and landing an NFC but please note that the code that moved is clearly marked by phabricator with the yellow vertical bar.
When you hover over it it says Moved from line xxx.

If you think that is still not enough I can rebase the NFC part but I personally have found phabricator to be good for this.

132

Because the proper place for this doc is inside the op interface and it is already there, almost word for word.

698–702

Will consider for a followup, thanks!

mlir/lib/Dialect/Linalg/Analysis/DependenceAnalysis.cpp
142–146

Added an assertion, there is no short term plan to support such transformations and analyses on tensors.
The answer is buffer allocate then it's in the buffer world and these apply.
Mixed mode may appear in the future but would require rethinking things.

148

added assertions above.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
188–201

made it consistent, error messages for function arguments start at 1: e.g. 1st argument etc.

220

Made it consistent, thanks for spotting!

Unit tests: pass. 61794 tests passed, 0 failed and 781 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

ftynse accepted this revision.Tue, Jan 14, 1:55 PM

Something went wrong with formatting in the commit description (mixed double space prefix and backticks?). Please reformat and feel free to land.

This revision is now accepted and ready to land.Tue, Jan 14, 1:55 PM
nicolasvasilache edited the summary of this revision. (Show Details)Tue, Jan 14, 2:22 PM
nicolasvasilache edited the summary of this revision. (Show Details)
nicolasvasilache edited the summary of this revision. (Show Details)Tue, Jan 14, 2:24 PM
nicolasvasilache edited the summary of this revision. (Show Details)
This revision was automatically updated to reflect the committed changes.

Unit tests: pass. 61863 tests passed, 0 failed and 781 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

asaadaldien added inline comments.Tue, Jan 14, 3:24 PM
mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h
204

Do we need to assert(i - getNumInputsAndOutputBuffers() < getOutputTensorTypes().size()) ?