Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Gemmini] Fuse linalg.matmul and linalg.transpose when lowering to Gemmini DIalect #300

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[Gemmini] Fuse linalg.matmul and linalg.transpose when lowering to Gemmini DIalect #300

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0a991a9
[Gemmini] fuse matmul and transpose
qzylalala May 6, 2024
de0cc06
[Gemmini] add test for fusion
qzylalala May 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions examples/GemminiDialect/makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,18 @@ matmul-os-run:
@riscv64-unknown-linux-gnu-gcc log.o -O2 -static -o a.out
@spike --extension=gemmini pk a.out

matmul-transpose-fuse-run:
@${BUDDY_OPT} ./matmul-transpose-fuse.mlir \
-convert-linalg-to-gemmini \
-convert-linalg-to-loops \
-lower-gemmini | \
${BUDDY_TRANSLATE} --buddy-to-llvmir | \
${BUDDY_LLC} -filetype=obj -mtriple=riscv64 \
-mattr=+buddyext,+D -float-abi=hard \
-o log.o
@riscv64-unknown-linux-gnu-gcc log.o -O2 -static -o a.out
@spike --extension=gemmini pk a.out

compute-accumulated-run:
@${BUDDY_OPT} ./compute-accumulated.mlir -lower-gemmini | \
${BUDDY_TRANSLATE} --buddy-to-llvmir | \
Expand Down
36 changes: 36 additions & 0 deletions examples/GemminiDialect/matmul-transpose-fuse.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
// RUN: buddy-opt %s \
// RUN: --convert-linalg-to-gemmini | \
// RUN: FileCheck %s

memref.global "private" @gv1 : memref<3x4xi8> = dense<[[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12]]>
memref.global "private" @gv2 : memref<4x3xi8> = dense<[[1, 1, 1],
[1, 1, 1],
[1, 1, 1],
[1, 1, 1]]>

func.func @main() -> i8 {
%arrayA = memref.get_global @gv1 : memref<3x4xi8>
%arrayB = memref.get_global @gv2 : memref<4x3xi8>
%arrayC = memref.alloc() : memref<3x3xi8>
%cst0 = arith.constant 0 : i8
gemmini.print %arrayC : memref<3x3xi8>
// Matrix-matrix multiplication
// CHECK: gemmini.tile_matmul %1 %0 %alloc %alloc_0 {aTranspose = true, bTranspose = true} :
// CHECK-SAME: memref<4x3xi8> memref<3x4xi8> memref<3x3xi8> memref<3x4xi32>
linalg.matmul
ins(%arrayA, %arrayB: memref<3x4xi8>, memref<4x3xi8>)
outs(%arrayC: memref<3x3xi8>)

// transpose
linalg.transpose
ins(%arrayC: memref<3x3xi8>)
outs(%arrayC: memref<3x3xi8>)
permutation = [1, 0]

gemmini.print %arrayC : memref<3x3xi8>
memref.dealloc %arrayC : memref<3x3xi8>

return %cst0 : i8
}
40 changes: 36 additions & 4 deletions midend/lib/Conversion/LowerLinalgToGemmini/LowerLinalgToGemmini.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,10 +63,42 @@ class MatmulLowering : public OpRewritePattern<linalg::MatmulOp> {
Value fillOpInputValue =
rewriter.create<arith::ConstantOp>(loc, fillOpInsType, fillOpInputAttr);
rewriter.create<linalg::FillOp>(loc, fillOpInputValue, bias);
rewriter.replaceOpWithNewOp<gemmini::TileMatMulOp>(
matMulOp, input0, input1, output0, bias, /*aScaleFactor = */ scale1,
/*bScaleFactor = */ scale1, /*dScaleFactor = */ scale1, /*act = */ 0,
/*accScale = */ scale1, /*bertScale = */ scale0);

// llvm::outs() << " has "
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clear the comments here.

// << std::distance(output0.getUses().begin(),
// output0.getUses().end())
// << " uses:\n";
// for (Operation *userOp : output0.getUsers()) {
// llvm::outs() << " - " << userOp->getName() << "\n";
// }

// If this matmul operation is followed by a transpose operation, do fusion.
// We should make sure that the result of this matmul op only has one user.
Operation* fuseOp = *output0.user_begin();
int output0Use = 0;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There is a serious error in logic here.I would suggest writing the pattern inside the canonicalizer of the gemmini.tile_matmul op.What you should be determining is whether there is only one user of gemmini.tile_matmul and it is linalg.transpose.bTranspose is a attribute, so let it be the opposite.You shouldn't be judging the number of linalg.transpose, it's completely wrong to do so.A pass is based on a greedy algorithm, so this pass will continue to match against the pattern.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi, thanks for your suggestions. I agree that it is indeed possible to do the match during gemmini.tile_matmul lowering. However, I still don’t understand why there is a problem with greedy matching here. Can you give an example here? Maybe I’m not that familiar with it yet.

Copy link
Member

@linuxlonelyeagle linuxlonelyeagle May 10, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

%1 = gemmini.tile_matmul  {transepose = false}
.....
%4 = linalg.transpose %1 
.....
%7 = linalg.transpose %4
.....
use(%7)

first

%1 = gemmini.tile_matmul  {transepose = true}
.....
%7 = linalg.transpose %1
.....
use(%7)

second

%1 = gemmini.tile_matmul {transepose = false}
.....
use(%1)

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more. 

%1 = gemmini.tile_matmul  {transepose = false}
.....
%4 = linalg.transpose %1 
.....
%7 = linalg.transpose %1

I'm not sure if this happens in the model not sure if it does.
So for this fuse pattern, you should make sure that there is only one op that uses gemmini.tile_matmul and that op is linalg.transpose.

for (auto userOp : output0.getUsers()) {
if (auto transposeOp = dyn_cast<linalg::TransposeOp>(userOp)) {
fuseOp = transposeOp;
output0Use ++;
}
}

if (output0Use) {
// llvm::outs() << "Fuse linalg.matmul and linalg.transpose. \n";
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remove comments

rewriter.replaceOpWithNewOp<gemmini::TileMatMulOp>(
matMulOp, input1, input0, output0, bias, /*aScaleFactor = */ scale1,
/*bScaleFactor = */ scale1, /*dScaleFactor = */ scale1, /*act = */0,
/*accScale = */ scale1, /*bertScale = */ scale0,
/*repeatingBias = */ false, /*aTranspose = */ true,
/*bTranspose = */ true);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

bTranspose is an attribute.Obviously the method to create gemmini.tile_matmul here just write one.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For C = AB, C = C^T, we use this formula : C^T = (AB)^T = B^T A^T.
So, we change the position of input0 and input1, and set aTranspose and bTranspose true.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If so, the reasoning is the same, you just need to make aTranspose the opposite of bTransepose as before.

rewriter.eraseOp(fuseOp);
} else {
rewriter.replaceOpWithNewOp<gemmini::TileMatMulOp>(
matMulOp, input0, input1, output0, bias, /*aScaleFactor = */ scale1,
/*bScaleFactor = */ scale1, /*dScaleFactor = */ scale1, /*act = */ 0,
/*accScale = */ scale1, /*bertScale = */ scale0);
}

rewriter.create<memref::DeallocOp>(loc, bias);
return success();
}
Expand Down
Loading