Add a pass that fuses matmul and transpose operations #567
Conversation
For the transpose-matrix multiplication-transpose pattern in Llama2, perform fusion and vectorization during dialect reduction. The acceleration ratio at the operator level is 1.84 before and after fusion, while it is not visible at the model level.
| Value B = op->getOperand(1); | ||
| Value C = op->getOpResult(0); | ||
|
|
||
| tosa::ReshapeOp reshapeBOp = B.getDefiningOp<tosa::ReshapeOp>(); |
There was a problem hiding this comment.
maybe you can use auto here, getDefiningOptosa::ReshapeOp(); we know the type of the op
| if (!transposeBOp) { | ||
| return failure(); | ||
| } | ||
| Value::user_iterator reshapeCUserIt = C.getUsers().begin(); |
There was a problem hiding this comment.
C.getUsers().empty() is good
| ShapedType newBType = | ||
| cast<ShapedType>(transposeBOp.getOperand(0).getType()); | ||
| ShapedType newCType = | ||
| cast<ShapedType>(transposeCOp->getOpResult(0).getType()); |
There was a problem hiding this comment.
can you use ransposeCOp->getresult->getType()
There was a problem hiding this comment.
or transposeCOp.getType()
| Value vlStep = rewriter.create<arith::ConstantIndexOp>(loc, vecSize); | ||
| Value zero = rewriter.create<arith::ConstantOp>( | ||
| loc, rewriter.getZeroAttr(elementType)); | ||
| const AffineExpr d0 = rewriter.getAffineDimExpr(0); |
|
|
||
| // Create pass through vector. | ||
| Value passThroughVec = rewriter.create<SplatOp>(loc, vectorTy, zero); | ||
| Value newA = rewriter.create<bufferization::ToMemrefOp>( |
There was a problem hiding this comment.
Is it possible to avoid using the bufferize dialect here?This is just a fusion pattern.
| Value aCol = rewriter.create<memref::DimOp>(loc, newA, c2); | ||
| Value bCol = rewriter.create<memref::DimOp>(loc, newB, c3); | ||
|
|
||
| Value upperBoundTmp = rewriter.create<arith::SubIOp>(loc, bCol, vlStep); |
There was a problem hiding this comment.
sub and add we can use affine,Rather than creating add and sub operations.
| // loopBody->addArguments(types, locs); | ||
| Block &loopBody = parOp.getRegion().front(); | ||
| rewriter.setInsertionPointToStart(&loopBody); | ||
| Value ivs0 = loopBody.getArguments()[0]; |
There was a problem hiding this comment.
iv = loopOp.getLoopInductionVar
| newC, | ||
| ValueRange{c0, ivs1, ivs0, iv}); | ||
| Value idx = | ||
| nestedBuilder.create<arith::AddIOp>(nestedLoc, iv, vlStep); |
There was a problem hiding this comment.
Thank you for your feedback. While considering avoiding the bufferize dialect, I discovered this pass can be moved from the TOSA level to be completed at the Linalg level. I will resubmit all modifications after completing this change.
2 participants
For the transpose-matrix multiplication-transpose pattern in Llama2, perform fusion and vectorization during dialect reduction. The acceleration ratio at the operator level is 1.84 before and after fusion, while it is not visible at the model level.
After applying this pass, the computational speed of transpose-matmul-transpose changed from 0.00528216 to 0.00191212.