Page MenuHomePhabricator

[MLIR] Introduce inter-procedural memref layout normalization
ClosedPublic

Authored by avarmapml on Jul 23 2020, 11:55 PM.

Details

Summary
  • Introduces a pass that normalizes the affine layout maps to the identity layout map both within and across functions by rewriting function arguments and call operands where necessary.
  • Memref normalization is now implemented entirely in the module pass '-normalize-memrefs' and the limited intra-procedural version has been removed from '-simplify-affine-structures'.
  • Run using -normalize-memrefs.
  • Return ops are not handled and would be handled in the subsequent revisions.

Signed-off-by: Abhishek Varma <abhishek.varma@polymagelabs.com>

Diff Detail

Event Timeline

There are a very large number of changes, so older changes are hidden. Show Older Changes
rriddle added inline comments.Jul 26 2020, 9:18 PM
mlir/include/mlir/Transforms/Passes.td
413

Please keep these sorted alphabetically.

mlir/lib/Transforms/NormalizeMemRefs.cpp
2

This line should be 80 lines long.

20

Can you trim these includes? Some shouldn't be necessary, e.g. Operation.h and FunctionSupport shouldn't be necessary.

70

nit: Please add variable names to declarations as well.

80

nit: I don't think this comment is necessary/useful.

86

nit: Please remove the user name from the TODO.

90

This is a very large function, can you split it up?

96

nit: Drop the trivial braces here and below.

104

nit: Can you use a longer name than ft?

110

nit: This could iterate over ft.getResults(), or use something like llvm::to_vector<4>(ft.getResults()).

116

Can you just use llvm::enumerate(funcOp.getArguments()) instead? It provides you with a value+index pair.

121

nit: Can you merge this comment with the one before the if?

140

insertArgument returns the inserted argument.

154

You aren't adding anything to argTypes, is that intended?

161

nit: Remove the user name here and other places.

163

nit: Can you iterate resultTypes directly instead?

for (Type &resultType : resultTypes)
  if (resultType == memrefType)
    resultType = newMemRef.getType();
181

nit: Remove the extra * before context.

185

funcOp.getSymbolUses(moduleOp)?

195

Please do not use OperationState directly, use builder.create<CallOp> instead.

231

This looks incorrect, if any of the argument types changed you can't keep the same call op.

mlir/lib/Transforms/Utils/Utils.cpp
90

Same here, please avoid using OperationState.

417

nit: Remove the wrapping ().

It would be best to handle return ops (where they return memref values with non-identity layouts to be normalized) separately in the next revision - since they add significant complexity.

avarmapml updated this revision to Diff 281033.EditedJul 27 2020, 1:00 PM
avarmapml marked 17 inline comments as done.

Memref map layout normalization for function argument and AllocOps

  • Support for ReturnOp will be handled in the next revision. So now we only deal with function arguments and allocOps whose uses are :-
    1. Dereferencing types
    2. Non-dereferencing types : DeallocOp, CallOp
  • Code review changes.
  • Added test cases.
avarmapml marked 8 inline comments as done.Jul 27 2020, 1:07 PM
avarmapml added inline comments.
mlir/lib/Transforms/NormalizeMemRefs.cpp
195

So, the ReturnOps will be handled in the next revision and I'll use this there.
Thank you!

231

I understand. The next revision will take care of this.

mlir/lib/Transforms/Utils/Utils.cpp
90

Unlike the previous use of OperationState where it was known to be CallOp here it can be of any dereferencing type, so OpTy for builder.create<OpTy> is unknown.

avarmapml updated this revision to Diff 281108.Jul 27 2020, 9:39 PM
avarmapml marked 3 inline comments as done.

Reformatted test cases

Iterating over function arguments

  • llvm::enumerate didn't work on MutableArrayRef in the previous build.

Could you update the commit summary on this revision page to update the return op handing part? (It doesn't pick it up from arc - you'll have to manually hit "Edit Revision" to edit it.

mlir/lib/Transforms/NormalizeMemRefs.cpp
179–181

Drop trivial braces.

bondhugula added inline comments.Jul 28 2020, 5:32 AM
mlir/include/mlir/Transforms/Passes.td
313

Normalize memrefs
(since there is no separate intra-procedural version).

mlir/include/mlir/Transforms/Utils.h
52

a non-dereferencing type -> a non-dereferencing user

74

handleDeallocOp -> replaceInDeallocOp ?

mlir/lib/Transforms/NormalizeMemRefs.cpp
84

Move the "TODO" to before the sentence.

94

-> isMemRefNormalizable?

Also, what's this mlir::detail::inLineOpResult ? Could you use the right type and not anything from mlir::detail. ArrayRef<Operation *>?

108

-> areMemRefsNormalizable?

130

Instead of getKind(), use isa.

mlir/test/Transforms/normalize-memrefs.mlir
201

Nit: ... because the function is returning the memref.

bondhugula requested changes to this revision.Jul 28 2020, 5:32 AM
This revision now requires changes to proceed.Jul 28 2020, 5:32 AM
bondhugula added inline comments.Jul 28 2020, 5:35 AM
mlir/lib/Transforms/NormalizeMemRefs.cpp
110

Looks like you really don't need to collect/store the alloc ops in a vector. You can just run the check inline in the lambda.

AlexEichenberger requested changes to this revision.Jul 28 2020, 7:33 AM
This comment was removed by AlexEichenberger.
AlexEichenberger accepted this revision.EditedJul 28 2020, 7:40 AM

sorry, requested a revision by mistake. Repost as a comment below

This patch does the advertised functionality well, as shown with the little example below:

#map0 = affine_map<(d0, d1) -> (d0, d1 floordiv 32, d1 mod 32)>
module {
  func @test(%in : memref<5x10xf32>, %out : memref<5x10xf32, #map0>) {
      affine.for %i = 0 to 5 {
          affine.for %j = 0 to 10 {
              %v = affine.load %in[%i, %j] :  memref<5x10xf32>
              affine.store %v, %out[%i, %j] : memref<5x10xf32, #map0>
          }
      }
      return
  }
  func @test_simplification() {
    %0 = alloc() : memref<5x10xf32>
    %1 = alloc() : memref<5x10xf32, #map0>
    //"test.test"(%0, %1) : (memref<5x10xf32>, memref<5x10xf32, #map0>) -> ()
    call @test(%0, %1) : (memref<5x10xf32>, memref<5x10xf32, #map0>) -> ()
    dealloc %1 : memref<5x10xf32, #map0>
    dealloc %0 : memref<5x10xf32>
    return
  }
}

transformed into

module {
  func @test(%arg0: memref<5x10xf32>, %arg1: memref<5x1x32xf32>) {
    affine.for %arg2 = 0 to 5 {
      affine.for %arg3 = 0 to 10 {
        %0 = affine.load %arg0[%arg2, %arg3] : memref<5x10xf32>
        affine.store %0, %arg1[%arg2, %arg3 floordiv 32, %arg3 mod 32] : memref<5x1x32xf32>
      }
    }
    return
  }
  func @test_simplification() {
    %0 = alloc() : memref<5x10xf32>
    %1 = alloc() : memref<5x1x32xf32>
    call @test(%0, %1) : (memref<5x10xf32>, memref<5x1x32xf32>) -> ()
    dealloc %1 : memref<5x1x32xf32>
    dealloc %0 : memref<5x10xf32>
    return
  }
}

where all the maps are eliminated from the load/store and alloc/dealloc.

However, a pattern that we see often is the lowering to external implementations not expressed in MLIR (think CUDNN calls or the like). In the above example, if we comment the "@test" lines and uncomment the line with "test.test", then the optimization will not remove any of the simplifications, as shown below.

func @test_simplification() {
  %0 = alloc() : memref<5x10xf32>
  %1 = alloc() : memref<5x10xf32, #map1>
  "test.test"(%0, %1) : (memref<5x10xf32>, memref<5x10xf32, #map1>) -> ()
  dealloc %1 : memref<5x10xf32, #map1>
  dealloc %0 : memref<5x10xf32>
  return
}

Is there a way to force the optimization through dialect operations? In our case, these maps were introduced especially to satisfy expected layout by the dialect "test." So we can take full responsibility that accesses within test.test will be fine. Happy with either a declarative approach or with a flag for a given dialect.

Thanks

avarmapml marked 10 inline comments as done.

Review changes and code formatting

mlir/lib/Transforms/NormalizeMemRefs.cpp
94

.getUsers() is returning Value::user_range and there is no known conversion from that to `ArrayRef<Operation *>.
I'm using Value::user_range as the type instead.

@AlexEichenberger : Such operation handling can be a part of subsequent revision.
Thank you.

avarmapml edited the summary of this revision. (Show Details)Jul 28 2020, 10:43 AM
bondhugula accepted this revision.Jul 28 2020, 10:48 AM
bondhugula added inline comments.
mlir/lib/Transforms/NormalizeMemRefs.cpp
145–147

Can drop trivial braces.

179–181

Can drop trivial braces.

mlir/lib/Transforms/Utils/Utils.cpp
278

Nit: double dots at the end. Drop space before colon.

This revision is now accepted and ready to land.Jul 28 2020, 10:48 AM
avarmapml marked 3 inline comments as done.

Minor formatting of code comment

bondhugula accepted this revision.Jul 29 2020, 7:37 AM

Just a few minor comments.

mlir/include/mlir/Transforms/Utils.h
78

and we provide replacement -> where we replace

mlir/lib/Transforms/NormalizeMemRefs.cpp
26

"interprocedural memrefs" isn't meaningful - instead, memrefs passed across functions.

44

Nit: You can drop the module { } wrapping around.

avarmapml updated this revision to Diff 281589.Jul 29 2020, 7:55 AM

Handled minor comments

avarmapml marked 3 inline comments as done.Jul 29 2020, 7:56 AM

@AlexEichenberger : Such operation handling can be a part of subsequent revision.
Thank you.

Thanks, I think this would help quite a few folks, thanks. Maybe adding an interface that enables operations to trigger this behavior?

Can you please rebase and fix merge conflicts?

avarmapml updated this revision to Diff 281845.EditedJul 30 2020, 2:20 AM

Resolved merge conflicts and involved another test case

avarmapml updated this revision to Diff 281880.Jul 30 2020, 4:52 AM

Rebased to latest master branch

avarmapml updated this revision to Diff 281886.Jul 30 2020, 5:15 AM

Clang-format error fixed

This revision was automatically updated to reflect the committed changes.