Page MenuHomePhabricator

[mlir][Linalg] Evolve named ops to use assembly form and support linalg on tensors.
Needs ReviewPublic

Authored by nicolasvasilache on Wed, Sep 16, 8:08 AM.

Details

Summary

This revision allows representing a reduction at the level of linalg on tensors for named ops. When a structured op has a reduction and returns tensor(s), new conventions are added and documented.

As an illustration, the syntax for a linalg.matmul writing into a buffer is:

linalg.matmul ins(%a, %b : memref<?x?xf32>, tensor<?x?xf32>)
             outs(%c : memref<?x?xf32>)

, whereas the syntax for a linalg.matmul returning a new tensor is:

%d = linalg.matmul ins(%a, %b : tensor<?x?xf32>, memref<?x?xf32>)
                  init(%c : memref<?x?xf32>)
                    -> tensor<?x?xf32>

Other parts of linalg will be extended accordingly to allow mixed buffer/tensor semantics in the presence of reductions.

Diff Detail

Event Timeline

Herald added a reviewer: aartbik. · View Herald Transcript
Herald added a reviewer: jpienaar. · View Herald Transcript
Herald added a reviewer: aartbik. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
nicolasvasilache requested review of this revision.Wed, Sep 16, 8:08 AM
ftynse accepted this revision.Thu, Sep 17, 6:07 AM
ftynse added inline comments.
mlir/docs/Dialects/Linalg.md
495

Does this support the change in elemental types? Otherwise it's not only the same shape, but the types must match completely.

533

This should be a tensor. If we expect it to be strictly the same type, we can also omit the type here.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td
31

Can this rather be in extraClassDeclaraiton or, even better, a static function in the C++ implementation file? It does not look like this can ever have a non-default implementation so why pay the cost of making it "virtual"?

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

ins, outs and init ?

84

Nit: traits seem to be using singular in their class names, i.e. NamedStructuredOpTrait

109

Nit: init_tensors does not appear in the IR, did you mean init ?

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
318–335

I don't understand why is this necessary.

1272–1276

Nit: now that you take an OpBuidler, I'd advise to use OpBuilder::createBlock instead.

1396–1398

Nit: I'd expect Twine or formatv to be more efficient than stitching std strings

Tmp OpBuilder creation.

nicolasvasilache marked 9 inline comments as done.

Address review

mlir/docs/Dialects/Linalg.md
533

let's remove later when we have some experience using it if it feels too redundant.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
318–335

the goal is to make everyone use the same verifiers but atm Generic and IndexedGeneric have this view property that I need to kill. This will be done in a followup and then all can be unified.

Better building of block.

Adding some comments. Looking through this right now.

mlir/docs/Dialects/Linalg.md
488

This seems to be more complex because the effort is being made to mix tensor and buffer semantics. Is it possible for the time being to just keep them completely separate (at least by convention).

495

Since we are going this route, why not just add a new region to the named op that describes the computation to generate the init tensor. This region has the same semantics as a linalg.generic/linalg.indexed_generic op?

525

This seems to deviate from the existing form of ops in MLIR, i.e.

linalg.matmul ins(%a : memref<?x?xf32>, %b : tensor<?x?xf32>)...

Fix incorrect OpBuilder state with an InsertionGuard.

mravishankar added inline comments.Thu, Sep 17, 11:57 AM
mlir/docs/Dialects/Linalg.md
495

To get a little more specific, we can do

<linalg named-op> (ins ...) (outs ...)
    init (%init : tensor<...f32>) {
    ^bb0(%arg0 : f32) :
         linalg.yield %arg0 : f32
    }

and that could generalize to

<linalg named-op> (ins ...) (outs ...)
    init (%a : tensor<...f32>, %b : tensor<...f32>) {
    ^bb0(%arg0 : f32, %arg1 : f32) :
         %0 = std.addf %arg0, %arg1 : f32
         linalg.yield %0 : f32
    }

where the initialization is done via a computation.

mlir/integration_test/Dialect/Linalg/CPU/test-conv-1d-ncw-call.mlir
34

super nit: align ins and outs ?

nicolasvasilache marked 3 inline comments as done.Thu, Sep 17, 12:07 PM
nicolasvasilache added inline comments.
mlir/docs/Dialects/Linalg.md
488

The bigger underlying achievement allowed by this convention is that we have named ops that are automatically generated from the TC spec and work with either tensors and buffers.
This was the key blocker to scaling these concepts and was introduced by the args_in / args_out that will be removed entirely in a followup commit.
Making them work with mixed tensor / buffer does not brings real additional complexity: if you want buffer only you'd still use ins/outs; if you want tensors only you'd still use ins/result for pointwise and ins/init/result for reductions.

495

That's significantly more work and out of the scope of this CL.

For instance, I do not know yet how to make it work with the TC lang.
Another tricky point is what would the init region look like and expand in higher-D tensors (e.g. imagine a 3-D 2x2x2 tensor, how does a region encode a 2-D tensor broadcast along some dim)? It would seem the indexed_generic would be required.

Still it is an orthogonal improvement that can we can table into a separate discussion once the existing dead-end state is improved.

525

Yes, this is unfortunate and seems to be a byproduct of using the declarative assembly format.

I do not know how to make it generate interleaved types and uses.
If/when it is available we should go to that.

As an illustration, note that ReturnOp uses the declarative assembly but FuncOp has a custom handwritten parser.

nicolasvasilache marked 4 inline comments as done.Thu, Sep 17, 12:13 PM
nicolasvasilache added inline comments.
mlir/docs/Dialects/Linalg.md
495

Looks nice!

I still view it as future improvement though so I'd keep it for a separate PR once the generic ops are up to speed too :)
The args_in / args_out has to be deprecated with fire first.

mlir/integration_test/Dialect/Linalg/CPU/test-conv-1d-ncw-call.mlir
34

I aligned the parenthesis consistently throughout the examples.
I'm somewhat reluctant to change all tests to align the is and the os now :)

nicolasvasilache marked 2 inline comments as done.

InsertionGuard.

Spin a custom parser as declarative assembly extension is blocked for now.

Drop extra dialect name form printing.

Harbormaster completed remote builds in B72154: Diff 292745.
mravishankar accepted this revision.Fri, Sep 18, 6:24 AM

Looks good to me for a first draft.

mlir/docs/Dialects/Linalg.md
495

Thanks! Thinking a bit more about this, I think its fine to go with juts an init tensor you have. I was only trying to handle the case where the initialization is done by a scalar value, but you could use a fill operaiton for this.

burmako accepted this revision.Fri, Sep 18, 9:31 AM