Page MenuHomePhabricator

[mlir] Add pass to convert elementwise ops to linalg.
ClosedPublic

Authored by silvas on Oct 28 2020, 4:42 PM.

Details

Summary

This patch converts elementwise ops on tensors to linalg.generic ops
with the same elementwise op in the payload (except rewritten to
operate on scalars, obviously). This is a great form for later fusion to
clean up.

E.g.

// Compute: %arg0 + %arg1 - %arg2
func @f(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
  %0 = addf %arg0, %arg1 : tensor<?xf32>
  %1 = subf %0, %arg2 : tensor<?xf32>
  return %1 : tensor<?xf32>
}

Running this through
mlir-opt -convert-std-to-linalg -linalg-fusion-for-tensor-ops we get:

func @f(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
  %0 = linalg.generic {indexing_maps = [#map0, #map0, #map0, #map0], iterator_types = ["parallel"]} ins(%arg0, %arg1, %arg2 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) {
  ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):  // no predecessors
    %1 = addf %arg3, %arg4 : f32
    %2 = subf %1, %arg5 : f32
    linalg.yield %2 : f32 
  } -> tensor<?xf32>
  return %0 : tensor<?xf32>
}

So the elementwise ops on tensors have nicely collapsed into a single
linalg.generic, which is the form we want for further transformations.

Diff Detail

Event Timeline

silvas created this revision.Oct 28 2020, 4:42 PM
silvas requested review of this revision.Oct 28 2020, 4:42 PM
mehdi_amini added inline comments.Oct 28 2020, 8:30 PM
mlir/include/mlir/Conversion/Passes.td
360 ↗(On Diff #301476)

Can you add a more comprehensive "description" field?

mlir/lib/Conversion/PassDetail.h
38 ↗(On Diff #301476)

Closing comment? (clang-tidy)

LG, thanks Sean!
Approving conditioned on the doc improvement.
The elementwise trait is also something we want for vectorization so it will be doubly helpful.

This revision is now accepted and ready to land.Oct 29 2020, 12:13 AM
silvas updated this revision to Diff 301793.Oct 29 2020, 5:21 PM

address comments

silvas marked 2 inline comments as done.Oct 29 2020, 5:21 PM
silvas updated this revision to Diff 302726.Nov 3 2020, 6:25 PM

Update to use new elementwise trait

silvas retitled this revision from [mlir] Add pass to convert std elementwise to linalg. to [mlir] Add pass to convert elementwise ops to linalg..
silvas edited the summary of this revision. (Show Details)
stellaraccident accepted this revision.Nov 8 2020, 4:47 PM
This revision was landed with ongoing or failed builds.Nov 10 2020, 1:46 PM
This revision was automatically updated to reflect the committed changes.