Use DisposableElementsAttr for ZHigh constant propagation #3013
Conversation
… at file writing (onnx#2917)" This reverts commit 33b466e. Signed-off-by: Tung D. Le <[email protected]>
tungld
force-pushed
the
zhigh-constprop-with-dispose
branch
from
68979b3
to
48cf039
Compare
Signed-off-by: Tung D. Le <[email protected]>
tungld
force-pushed
the
zhigh-constprop-with-dispose
branch
from
48cf039
to
265ff90
Compare
|
@tungld Just to understand the high level and without the class names. You are using Soren's approach of applying "logical" operations to the constants so that for example if we have |
Yes, I extend it for ZHigh operations so the same approach is used for both ONNX and ZHigh until lowering to krnl. We can extend it to cover krnl operations but it needs more work and I didn't do it in this PR. |
Signed-off-by: Tung D. Le <[email protected]>
Signed-off-by: Tung D. Le <[email protected]>
Signed-off-by: Tung D. Le <[email protected]>
Signed-off-by: Tung D. Le <[email protected]>
Signed-off-by: Tung D. Le <[email protected]>
Signed-off-by: Tung D. Le <[email protected]>
Signed-off-by: Tung D. Le <[email protected]>
Signed-off-by: Tung D. Le <[email protected]>
Signed-off-by: Tung D. Le <[email protected]>
Signed-off-by: Tung D. Le <[email protected]>
| struct ConstantStickPattern : public OpRewritePattern<ZHighStickOp> { | ||
| ConstantStickPattern(MLIRContext *context) : OpRewritePattern(context) {} | ||
| LogicalResult matchAndRewrite( | ||
| ZHighStickOp stickOp, PatternRewriter &rewriter) const override { | ||
| Value input = stickOp.getIn(); | ||
| Value output = stickOp.getOut(); | ||
| StringAttr layout = stickOp.getLayoutAttr(); | ||
|
|
||
| // Match | ||
| if (!isDenseONNXConstant(input)) { | ||
| return failure(); | ||
| } | ||
|
|
||
| // Rewrite | ||
| Value stickifiedVal = | ||
| createConstantForStick(rewriter, output, input, layout); | ||
| replaceOpAndGC(rewriter, stickOp, stickifiedVal); | ||
| return success(); | ||
| } | ||
| }; |
There was a problem hiding this comment.
Just to confirm, you replaced tablegen with cpp. Is this because replaceOpAndGC() is difficult to use in tablegen format?
There was a problem hiding this comment.
Yes, I don't know how to do that with tablegen, let me know if you know how to do it. Thanks!
|
Can you post here for ref the improvements you got, just for future reference purpose. Does not need to be super detailed. Thanks |
I put the measurement results of gpt2-large and Mistral-7b. 
In Mistral-7b, the peak memory usage reduced from 33.2 GB to 27.9 GB, and compilation time becomes faster from 17 min 4 sec to 13 min 58 sec. Left graph is current main, and right graph is PR3013. 
|
|
@jenkins-droid test this please |
store-constants-to-fileis off.ONNXConstantOp) toZHighStickifiedConstantOp, so that the zhigh constant prop can benefit from the memory management forDisposableElementsAttr: In particular, it does:DisposableElementsAttrZHighStikifiedConstantOpto acceptDisposableElementsAttr. Its parser and printer are changed to read/write from/toDenseElementsAttrfor lit tests.ZHighConstantPropagationPassso that it reads data directly fromDisposableElementsAttrinstead ofDenseElementsAttr.ZHighDisposableGabageCollectorandZHighScrubDisposablePassto manage buffers used byZHighStickifiedConstantOpQuick experiment: the peak compile memory consumptions of #2917 and this PR when compiling the gpt2-large model for NNPA (744M parameters, the constant file's size is 3.2GB) are quite similar, both are about 9GB.
This patch contains the reverting code so it's no easy to follow. To ease the review, I merge all new changes (not the reverting code) into a single commit: 265ff90. Please look at this commit for review.