Page MenuHomePhabricator

[mlir][OneShotModuleBufferize] Add a new flag: no-analysis-func-filter
ClosedPublic

Authored by mamrami on Jan 26 2023, 7:35 AM.

Details

Summary

OneShotModuleBufferize fails if the input IR cannot be analyzed.
One can set CopyBeforeWrite=true in order to skip analysis.
In that case, a buffer copy is inserted on every write.
This leads to many copies, also in FuncOps that could be analyzed.

This change aims to copy buffers only when it is a must.
When running OneShotModuleBufferize with CopyBeforeWrite=false,
FuncOps whose names are specified in noAnalysisFuncFilter will not be
analyzed. Ops in these FuncOps will not be analyzed as well.
They will be bufferized with CopyBeforeWrite=true,
while the other ops will be bufferized with CopyBeforeWrite=false.

Diff Detail

Event Timeline

mamrami created this revision.Jan 26 2023, 7:35 AM
Herald added a project: Restricted Project. · View Herald TranscriptJan 26 2023, 7:35 AM
mamrami requested review of this revision.Jan 26 2023, 7:35 AM
mamrami updated this revision to Diff 492448.Jan 26 2023, 7:43 AM

edit title

mamrami retitled this revision from [mlir] OneShotModuleBufferize- Copy buffers in FuncOps that failed during analysis to [mlir] OneShotModuleBufferize- Copy buffers in FuncOps that could not be analyzed.Jan 26 2023, 7:44 AM
mamrami added a reviewer: springerm.
springerm added a comment.EditedJan 26 2023, 8:31 AM

Instead of collecting FuncOps that could not be bufferized, can we pass a list of excluded FuncOps to the bufferization? Does that work in your flow?

We already have some infrastructure for that: BufferizationOptions::isOpAllowed. It returns false for ops that are excluded from the bufferization. I think it currently has no effect for FuncOps. Then you can bufferize with One-Shot Bufferize first, excluding all FuncOps of which you know that they cannot be bufferized. Then run the bufferization with copyBeforeWrite = true on all excluded FuncOps.

I think only small changes to mlir::bufferization::bufferizeModuleOp, mlir::bufferization::analyzeModuleOp and mlir::bufferization::insertTensorCopies should be needed.

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
1007 ↗(On Diff #492448)

Ideally, no changes are needed to this file.

mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
384

Put a check here:

if (!state.getOptions().isOpAllowed(funcOp)) continue;
mamrami updated this revision to Diff 493298.Jan 30 2023, 6:24 AM

Addressing code review

mamrami retitled this revision from [mlir] OneShotModuleBufferize- Copy buffers in FuncOps that could not be analyzed to [mlir][OneShotModuleBufferize] Add a new flag: no-analysis-func-filter.Jan 30 2023, 6:25 AM
mamrami edited the summary of this revision. (Show Details)
mamrami marked 2 inline comments as done.

Then you can bufferize with One-Shot Bufferize first, excluding all FuncOps of which you know that they cannot be bufferized. Then run the bufferization with copyBeforeWrite = true on all excluded FuncOps.

There are cases where a function is called by another function, so bufferizing the caller is impossible when the callee is excluded.
Now the bufferization happens once on the entire IR so I don't have this problem.

springerm accepted this revision.Jan 30 2023, 6:39 AM
This revision is now accepted and ready to land.Jan 30 2023, 6:39 AM