Page MenuHomePhabricator

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

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



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.Sep 16 2020, 8:08 AM
ftynse accepted this revision.Sep 17 2020, 6:07 AM
ftynse added inline comments.

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


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


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"?


ins, outs and init ?


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


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


I don't understand why is this necessary.


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


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


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


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.


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).


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?


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.Sep 17 2020, 11:57 AM

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.


super nit: align ins and outs ?

nicolasvasilache marked 3 inline comments as done.Sep 17 2020, 12:07 PM
nicolasvasilache added inline comments.

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.


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.


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.Sep 17 2020, 12:13 PM
nicolasvasilache added inline comments.

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.


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.


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.Sep 18 2020, 6:24 AM

Looks good to me for a first draft.


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.Sep 18 2020, 9:31 AM
pifon2a accepted this revision.Thu, Oct 1, 11:51 AM

Wasn't this submitted?

nicolasvasilache abandoned this revision.Fri, Oct 2, 1:29 AM