Page MenuHomePhabricator

Add a shape function library op
ClosedPublic

Authored by jpienaar on Nov 17 2020, 5:46 PM.

Details

Summary

Op with mapping from ops to corresponding shape functions for those op in the library and mechanism to associate shape functions to functions. The mapping of operand to shape function is kept separate from the shape functions themselves as the operation is associated to the shape function and not vice versa, and one could have a common library of shape functions that can be used in different contexts.

Diff Detail

Event Timeline

jpienaar created this revision.Nov 17 2020, 5:46 PM
jpienaar requested review of this revision.Nov 17 2020, 5:46 PM
rriddle added inline comments.Nov 17 2020, 6:12 PM
mlir/include/mlir/Dialect/Shape/IR/Shape.h
17

Where is the dependency on standard coming from?

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
820

Operation& -> Operation &

Also I'd probably use Operation * here instead given that this is the overwhelming common way of passing an operation, and the one that has the most nice implicit conversions.

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp
26

nit: Move this out of the anonymous namespace.

36

nit: Drop all of the mlir::.

jpienaar updated this revision to Diff 306140.Nov 18 2020, 9:41 AM
jpienaar marked 3 inline comments as done.

Switching to Operation* and moving out anon namespace.

mlir/include/mlir/Dialect/Shape/IR/Shape.h
17

Removed, I needed BuiltinDialect instead for FuncOp.

mehdi_amini added inline comments.Nov 18 2020, 10:29 AM
mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp
32

Walk is recursive, if the assumption is that the library is at the top-level this would be more efficient:

for (auto libraryOp : module.getOps<shape::FunctionLibraryOp>()) {
   ...
}
mehdi_amini added inline comments.Nov 18 2020, 10:35 AM
mlir/test/Analysis/test-shape-fn-report.mlir
6

This attribute seems redundant with the mapping of the library?

More importantly, this SymbolRef does not seem correct to me: it'll lookup in the current symbol table (current module) and not in the function library.

jpienaar updated this revision to Diff 306268.Nov 18 2020, 5:29 PM
jpienaar marked an inline comment as done.

Don't do recursive walk when finding shape fn lib in test pass.

mlir/test/Analysis/test-shape-fn-report.mlir
6

Good question.

The mapping in the library is on ops rather than functions. This is associated with the function and hence attached to the function. The ones below is mapping between operation and associated shape function. I see them connected differently/non-redundantly:

  • shape function library function is a shape transfer function/utility function, they are neither op nor computation specific
  • mapping in shape.function_library is between an op type and a shape function library function, but associated with the op
  • a function can have a shape function attached and that shape function refers to a shape function in the shape function library

But the shape function library I see as independent of the computations, it could be in a separate module even (and that may be more common from other feedback I got :))

What do you mean wrong lookup? It is up to the user to do the lookup. In this case the consumer of the op/attribute would have to know that the shape.function is in the shape.function_library. So it is a convention. Now that would fail if the library was in a different context than this one I believe, but beyond that I think it would work. Now, of course if something blindly were to check symbolrefs then yes this would fail, and I'd need to switch to a string. But I don't think we've generally enforced that all symbolrefattr have to be defined in the same module or even in the same symbol table that the op is defined in.

tpopp accepted this revision.Nov 18 2020, 6:40 PM

Only nits.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
813

Maybe use the OneRegion trait instead to catch any code expecting that trait instead.

828

I think this should be without shape in the op name as that's redundant (shape.shape_fn_lib_terminator)

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp
12–13

Does automatic formatting consider actual ascii values? IR vs Interfaces

mlir/test/lib/Dialect/Test/TestOps.td
139

operand without the 's'

This revision is now accepted and ready to land.Nov 18 2020, 6:40 PM
rriddle added inline comments.Nov 18 2020, 6:41 PM
mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
813

You should never be adding that yourself, it's automatic.

mehdi_amini added inline comments.Nov 18 2020, 7:37 PM
mlir/test/Analysis/test-shape-fn-report.mlir
6

What do you mean wrong lookup? It is up to the user to do the lookup. In this case the consumer of the op/attribute would have to know that the shape.function is in the shape.function_library. So it is a convention. Now that would fail if the library was in a different context than this one I believe, but beyond that I think it would work. Now, of course if something blindly were to check symbolrefs then yes this would fail, and I'd need to switch to a string. But I don't think we've generally enforced that all symbolrefattr have to be defined in the same module or even in the same symbol table that the op is defined in.

I don't think this is correct, basically my take is that this is not allowed / breaking MLIR symbol resolution assumptions I believe. We could have a symbol with the same name in the module and this would count as a use of this symbol, this would be walked by the symbol use/def chains, etc.
You shouldn't use a SymbolRef but a StringAttr if it does not refer to a symbol in the current module.

mehdi_amini added inline comments.Nov 18 2020, 7:38 PM
mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
813

SizedRegion<1> is to express that the region has a single block, not that there is a single region.
(it is also redundant with the SingleBlockImplicitTerminator trait, I wonder if this trait shouldn't be instead placed on the region declaration here)

jpienaar updated this revision to Diff 307220.Nov 23 2020, 5:15 PM
jpienaar marked 8 inline comments as done.

Switched to stringattr & AnyRegion

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp
12–13

I think so, I just rely on clang-format for these.

mehdi_amini accepted this revision.Nov 23 2020, 5:30 PM
herhut accepted this revision.Nov 24 2020, 2:15 AM

I am fine either way, just a suggestion.

mlir/test/Analysis/test-shape-fn-report.mlir
6

Now that the shape function library is a symbol table itself and has an (optional) name, you could use fully qualified symbol names here instead of strings. The advantage would be that this composes nicely (multiple shape fun libraries, etc.) and one can use standard means to drop unused shape functions when they are no longer referenced. Could come in handy when model sizes become an issue.

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp
52

You would not need to search for the single shapeFnLib with qualified symbols here but instead could just do the lookup at the scope of the function.

mehdi_amini added inline comments.Nov 24 2020, 12:50 PM
mlir/test/Analysis/test-shape-fn-report.mlir
6

That's a good point: you marked the shape function library as "Symbol" but the name is optional? I'm even surprised this is allowed.

mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp
52

You can't inspect siblings operations though, this can't be used in function passes.

herhut added inline comments.Nov 25 2020, 8:23 AM
mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp
52

That is an interesting restriction and I wonder whether it should be lifted. If I have an operation pass for OpA, I should be allowed to inspect OpB in the parent, because those are not impacted by the operation pass on OpA. With restrictions lifted that way, it would be feasible to lookup shape functions and associations even from a function pass.

mehdi_amini added inline comments.Nov 25 2020, 9:22 AM
mlir/test/lib/Dialect/Shape/TestShapeFunctions.cpp
52

It isn't clear to me that function pass and a "FunctionLibraryOp" pass can't run in parallel?

jpienaar updated this revision to Diff 308190.Nov 28 2020, 3:52 PM
jpienaar marked an inline comment as done.

Switched to fully qualified names

mlir/test/Analysis/test-shape-fn-report.mlir
6

Module also follows this form, by setting the flag that it optionally defines a symbol. I've removed that for now as Stephan made a good point about allowing multiple.

This revision was landed with ongoing or failed builds.Nov 28 2020, 3:54 PM
This revision was automatically updated to reflect the committed changes.