This is an archive of the discontinued LLVM Phabricator instance.

[mlir] Introduce transform.alternatives op
ClosedPublic

Authored by ftynse on Jun 14 2022, 1:19 AM.

Details

Summary

Introduce a transform dialect op that allows one to attempt different
transformation sequences on the same piece of payload IR until one of them
succeeds. This op fundamentally expands the scope of possibilities in the
transform dialect that, until now, could only propagate transformation failure,
at least using in-tree operations. This requires a more detailed specification
of the execution model for the transform dialect that now indicates how failure
is handled and propagated.

Transformations described by transform operations now have tri-state results,
with some errors being fundamentally irrecoverable (e.g., generating malformed
IR) and some others being recoverable by containing ops. Existing transform ops
directly implementing the apply interface method are updated to produce this
directly. Transform ops with the TransformEachTransformOpTrait are currently
considered to produce only irrecoverable failures and will be updated
separately.

Diff Detail

Event Timeline

ftynse created this revision.Jun 14 2022, 1:19 AM
ftynse requested review of this revision.Jun 14 2022, 1:19 AM
Herald added a project: Restricted Project. · View Herald TranscriptJun 14 2022, 1:19 AM
springerm accepted this revision.Jun 14 2022, 3:31 AM
springerm added inline comments.
mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
170

remove

172

immediately

183

register/registry?

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
27

can be

28

maybe put quotes around "reported"

30

either?

47–48

"Typically" or "must"?

737–738

silencable failure

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
30

arbitrary

31

Sounds like this could've been implemented with a single sequence op in each region (thus not having to support multiple ops per region in the alternatives op). Is this design mainly to avoid boilerplate code (multiple nested regions are hard to read)?

33

these

47

Maybe this could be checked to some degree by attaching a listener to the rewriter. If ops are created/deleted outside of the specified scope, we could fail an assertion. But I'm not sure if we can easily attach/detach listeners...

48

arbitrarily

55

is achieved

55

Is it possible that an existing handle to an op is invalidated during rollback? The cloned op will have a different Operation*. I'm thinking of 2 cases:

  1. Handle defined above the alternatives op. First alternative is rolled back, the handle cannot be used in the second alternative. I guess that's not allowed due to "isolated from above"?
  2. Handle defined above the alternatives op. It is not used in the alternatives op, but the op pointed to by the handle is within the specified scope. We have a rollback during the alternatives op. The handle is used after the alternatives op and may now be invalid.
59

lifted

62–64

Maybe add: Each alternative region must return the same number of results.

mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
41

A bit off topic: Looking at this piece of code again. Let's assume that OneShotBufferizeOp has a result (the bufferized op). For most ops (everything except for ModuleOp, maybe FuncOp), we don't have a handle to the bufferized op. In that case, we could just return a silencable error and not populate the transformResults. We could infer from the number of expected return values in the assembly format (one or zero PDL_Operations) if a return handle is expected or not. Then we don't need the targetIsModule attr anymore. Would that work?

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
191–192

You could make DiagnosedSilencableFailure(LogicalResult) constructor public.

mlir/lib/Dialect/Transform/IR/TransformOps.cpp
214–222

I'm not sure what this is doing. deleteClones.release() deactivates the deleteClones function, i.e., it is no longer run at the end of the scope. Don't we want the opposite? If !failed, i.e., if "success", delete the clones. And the for loop restores the original state. I was expecting to see this in an else branch (if failed).

This revision is now accepted and ready to land.Jun 14 2022, 3:31 AM
ftynse updated this revision to Diff 436804.Jun 14 2022, 8:50 AM
ftynse marked 20 inline comments as done.

Review

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
31

Not only harder to read, but also a bit more annoying to reason about.

47

Not all transforms are performed with a rewriter. This is enforced by performing the transformations on a clone of the scope operation that is not inserted into the parent. Thus they are unable to touch anything else.

55

Handle defined above the alternatives op. First alternative is rolled back, the handle cannot be used in the second alternative. I guess that's not allowed due to "isolated from above"?

Indeed, isolated from above requires the handle to be remapped to a new value in each region.

Handle defined above the alternatives op. It is not used in the alternatives op, but the op pointed to by the handle is within the specified scope. We have a rollback during the alternatives op. The handle is used after the alternatives op and may now be invalid.

This is subject to general handle invalidation. AlternativesOp always consumes the scope handle, regardless of it succeeding. This invalidates all handles to any operation nested within the scope per https://mlir.llvm.org/docs/Dialects/Transform/#handle-invalidation. Uses of invalidated handles are caught by the expensive-checks mode.

mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
41

I don't think you can parse that way. The op parser has no access to the parsed LHS values, but still expects the matching number of types to be defined in the result. The attribute indicates whether to add the result type or not.

mlir/lib/Dialect/Transform/IR/TransformOps.cpp
214–222

The transformation is attempted on the clones, not on the original IR, by virtue of mapping the only handle available in the alternative region to the clones. So when the transformation succeeded, we want to keep the clones and erase the original IR, which is done by cancelling the scheduled deletions and calling replaceOp that calls eraseOp internally.

We may be wanting the opposite approach with the transform being applied to the original IR, but doing that will make scoping permeable (e.g., the transformation will be able to inspect and potentially modify the ops outside of the scope, currently the scope operation exists on its own so it is physically impossible for the transformation to do act beyond it). I'd go with the current version and see when we hit the need to inspect the surrounding IR.

Expanded the comment a bit.

This revision was landed with ongoing or failed builds.Jun 14 2022, 8:51 AM
This revision was automatically updated to reflect the committed changes.
mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td
173

typo

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
33

Silenceable in languages without cedillas I guess :) ?

565

Weirdly, I don't see the region being used as != 0 anywhere and I find using the region magic number feels unfortunate.

Can we just drop (or at least use a Region* and interpret nullptr as 0) ?
Seems all the uses are going through the other flavor of mapBlockArguments that take explicit regions?

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td
75

failible -> fallible

mlir/lib/Dialect/Transform/IR/TransformOps.cpp
175

I would add a hard failure here requiring originals to be isolated from above.
In the absence of that, nothing prevents your clone logic below to modify IR it shouldn't and it will likely get surprising.

194

return and leak ?
I always put the guard right after the alloc.

I'd restructure to:

clones = ...
guard = ...

// Map the block argument (the only visible handle) to the cloned scope operations. 
// This effectively prevents the transformation from accessing any IR outside the scope.
if (failed(state.mapBlockArguments(reg.front().getArgument(0), clones)))
  ...
ftynse marked 6 inline comments as done.Jun 15 2022, 9:44 AM
mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.h
565

The number is the positional index, we use them from time to time. I switched to passing a Region & plus an overload without it (null pointer means error most of the time). I also had to assert that the region belongs to the op, which is wasn't necessary with the positional index.