Page MenuHomePhabricator

[MLIR] Support for ReturnOps in memref map layout normalization
ClosedPublic

Authored by avarmapml on Aug 4 2020, 10:30 AM.

Details

Summary
  • This commit handles the returnOp in memref map layout normalization.
  • An initial filter is applied on FuncOps which helps us know which functions can be a suitable candidate for memref normalization which doesn't lead to invalid IR.
  • Handles memref map normalization for external function assuming the external function is normalizable.

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

Diff Detail

Event Timeline

avarmapml created this revision.Aug 4 2020, 10:30 AM
Herald added a project: Restricted Project. · View Herald TranscriptAug 4 2020, 10:30 AM
avarmapml requested review of this revision.Aug 4 2020, 10:30 AM
avarmapml edited the summary of this revision. (Show Details)Aug 4 2020, 11:29 AM
bondhugula requested changes to this revision.Aug 4 2020, 12:43 PM
bondhugula added inline comments.
mlir/lib/Transforms/NormalizeMemRefs.cpp
131–174

Nit: caller -> calling

158–159

MLIR doesn't have any overloading. The function names are unique.

mlir/lib/Transforms/Utils/Utils.cpp
282

You don't need the extra parentheses around isa.

This revision now requires changes to proceed.Aug 4 2020, 12:43 PM
avarmapml updated this revision to Diff 283117.Aug 4 2020, 8:41 PM
avarmapml marked 2 inline comments as done.

Minor review changes

avarmapml marked an inline comment as done.Aug 4 2020, 8:54 PM

@AlexEichenberger Would you be able to review this - esp. for the part that's connected to what you wanted?

The code looks good, and I like the addition of skipping the analysis for "extern" functions. Because dialect ops can in some ways also represent "future" external functions, a natural way of handling "not yet expanded" external functions would to add an interface query that ask if a given op is "external" also. Default would be false, and dialects that needs dome/all of their ops to be external would be able to refine the interface to provide the necessary functionality.

Second, the presence of "dim" currently prevent normalization when the dim act on a memref wit maps. It is my firm belief that "dim" should return the logical iterations, not the mapped ones. Would it be possible that when looking at a dim as below:

%d = dim %a, %c0 : memref<?x?x?xf32,  map>

the proper value is directly returned?

Its fine if this functionality is not part of this patch, I believe they will be important for many practical cases.

The code looks good, and I like the addition of skipping the analysis for "extern" functions. Because dialect ops can in some ways also represent "future" external functions, a natural way of handling "not yet expanded" external functions would to add an interface query that ask if a given op is "external" also. Default would be false, and dialects that needs dome/all of their ops to be external would be able to refine the interface to provide the necessary functionality.

Second, the presence of "dim" currently prevent normalization when the dim act on a memref wit maps. It is my firm belief that "dim" should return the logical iterations, not the mapped ones. Would it be possible that when looking at a dim as below:

Hi Alex, dim does already return the extents of the logical dimensions, whatever you see in the shape signature of the memref (... x ... x ...), and not the physical ones. So, the dim is oblivious to the layout map.

%d = dim %a, %c0 : memref<?x?x?xf32,  map>

the proper value is directly returned?

Its fine if this functionality is not part of this patch, I believe they will be important for many practical cases.

Looks great. Some initial comments.

mlir/lib/Transforms/NormalizeMemRefs.cpp
133

Should this be renamed setCalleesCallersNonNormalizable - most accurate (plural) and as long.

157

Terminate with a full stop.

334

Comment here.

341–345

Drop the obvious "This function".

mlir/lib/Transforms/Utils/Utils.cpp
279–282

types -> ops

I tried the patch and this simple code appears to disable memref normalization

#map0 = affine_map<(d0, d1) -> (d1, d0)>
module {
  func @test_simplification(%arg : memref<5x10xf32>) {
    %a = alloc() : memref<10x5xf32, #map0>
    %c1 = constant 1 :index
    %d = dim %arg, %c1 : memref<5x10xf32>
    dealloc %a : memref<10x5xf32, #map0>
    return 
  }
}

The dim is used on a memref that does not have a map, the presence of the dim for arg variable (no maps) seems to disable the normalization of the a variable (with map)

bondhugula added a comment.EditedAug 11 2020, 1:21 PM

I tried the patch and this simple code appears to disable memref normalization

#map0 = affine_map<(d0, d1) -> (d1, d0)>
module {
  func @test_simplification(%arg : memref<5x10xf32>) {
    %a = alloc() : memref<10x5xf32, #map0>
    %c1 = constant 1 :index
    %d = dim %arg, %c1 : memref<5x10xf32>
    dealloc %a : memref<10x5xf32, #map0>
    return 
  }
}

The dim is used on a memref that does not have a map, the presence of the dim for arg variable (no maps) seems to disable the normalization of the a variable (with map)

@AlexEichenberger The patch doesn't handle/treat dim specially at all - irrespective of its map; it's being treated as a non-dereferencing op and so is bailing out on such ops. For static shapes, the approach would be to just fold the dim op at the start or expect the input to have been canonicalized, which will fold it away. For dynamic shapes, it's quite involved and will have to be definitely dealt with separately.

@bondhugula Static is easy. We are working on making it work for the dynamic shapes, currently in our code base but we could migrate it to the MLIR as part of another patch. Basically, we go back to the alloc responsible for the array, and parse the memref to pick either a constant or the right alloc parameter corresponding to the n th '?' in the memref of the alloc.

avarmapml updated this revision to Diff 284964.Aug 11 2020, 9:59 PM

Addressed review changes.

avarmapml marked 5 inline comments as done.Aug 11 2020, 10:00 PM

@bondhugula Static is easy. We are working on making it work for the dynamic shapes, currently in our code base but we could migrate it to the MLIR as part of another patch. Basically, we go back to the alloc responsible for the array, and parse the memref to pick either a constant or the right alloc parameter corresponding to the n th '?' in the memref of the alloc.

That would be great, thanks!

bondhugula accepted this revision.Aug 12 2020, 2:57 AM

Looks good - thanks.

This revision is now accepted and ready to land.Aug 12 2020, 2:57 AM

Please add a TODO somewhere on the dim op handling part.

avarmapml updated this revision to Diff 285039.Aug 12 2020, 4:40 AM

Added TODO section for DimOp.

bondhugula accepted this revision.Aug 13 2020, 6:38 AM

Looking into extending this approach to handle dialect ops that will become calls to external functions. Specifically, I am looking for a way to express that an operation guarantees that (when it will be lowered to an external function call), the external function be aware of the data layout of the maps and thus it is safe to normalize the op memref parameters.

MLIR experts out there, what is the preferred way to do this? Adding a interface on the dialect? Adding an interface on the operation (like AffineReadOpInterface, e.g. ExpectNormalizedMemRefParameters)?

Once this patch is in, I have a follow on patch that introduces a new trait, MemRefNormalizable to be used by operations (such as DeallocOp, CallOp, ReturnOp) as well as any ops, for example, such as operations of an dialect that 1) will be expanded to external calls at some later time and 2) are ok with having their memrefs normalized to get rid of their maps.

This trait can also be used to special case behavior for std.dim or other such operations

Looking into extending this approach to handle dialect ops that will become calls to external functions. Specifically, I am looking for a way to express that an operation guarantees that (when it will be lowered to an external function call), the external function be aware of the data layout of the maps and thus it is safe to normalize the op memref parameters.

MLIR experts out there, what is the preferred way to do this? Adding a interface on the dialect? Adding an interface on the operation (like AffineReadOpInterface, e.g. ExpectNormalizedMemRefParameters)?

I missed these comments. This discussion is probably better taken on the discourse thread we already have? Having a trait like this sounds useful.

Once this patch is in, I have a follow on patch that introduces a new trait, MemRefNormalizable to be used by operations (such as DeallocOp, CallOp, ReturnOp) as well as any ops, for example, such as operations of an dialect that 1) will be expanded to external calls at some later time and 2) are ok with having their memrefs normalized to get rid of their maps.

This trait can also be used to special case behavior for std.dim or other such operations

This patch is already in. We can take this discussion on the discourse thread.