Hello, I’m new to mlir and try to do some loop transformation on gemm in my custom pattern. Currently I create a band of for loops using ConversionPatternRewriter &rewriter.create<AffineForOp>
. Here is what it looks like in my code:
M_loop = rewriter.create<AffineForOp>(
loc, c0_range.vr, rewriter.getDimIdentityMap(), M_range.vr,
rewriter.getDimIdentityMap(), 1, std::nullopt,
[&](OpBuilder &builder, Location loc, Value im, ValueRange iterArgs) {
ValueToRange N_range(N);
N_loop = builder.create<AffineForOp>(
loc, c0_range.vr, builder.getDimIdentityMap(), N_range.vr,
builder.getDimIdentityMap(), 1, std::nullopt,
[&](OpBuilder &builder, Location loc, Value in,
ValueRange iterArgs) {
ValueToRange K_range(K);
K_loop = builder.create<AffineForOp>(
loc, c0_range.vr, builder.getDimIdentityMap(), K_range.vr,
builder.getDimIdentityMap(), 1, std::nullopt,
[&](OpBuilder &builder, Location loc, Value ik,
ValueRange iterArgs) {
... // loop body
builder.create<AffineYieldOp>(loc);
});
Attribute K_Attr = rewriter.getStringAttr("K_loop");
K_loop->setAttr("Dimension", K_Attr);
builder.create<AffineYieldOp>(loc);
});
Attribute N_Attr = rewriter.getStringAttr("N_loop");
N_loop->setAttr("Dimension", N_Attr);
builder.create<AffineYieldOp>(loc);
});
However, when I use tilePerfectlyNested
, it will cause ERROR: AddressSanitizer: heap-use-after-free
. It seems that tilePerfectlyNested
delete the root loop transformation here. However, the rewriter didn’t realize that and try to access it. My current solution to this problem is to copy the body of tilePerfectlyNested
and the functions it calls to my source file, and change that line to rewriter.eraseOp(rootAffineForOp);
, and everything is ok. However, this solution looks quite dumb to me. Is there a better way to achieve this? Or is there a better way to create a for loop nest? I’ve tried buildAffineLoopNest
, however, I can’t get the ForOp that it created, therefore, I can’t do the following loop optimization in the same matchAndRewrite
function.
Thanks in advance!