How to use `tilePerfectlyNested` in a custom pattern?

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!

These are heavier transformation utilities (as opposed to local rewrites) and are thus not suitable for use in rewrite patterns. You can just have a standard IR walk that performs such heavier transformations, and apply any rewriter-based patterns before or after it.

1 Like

Does it mean that I need to create another pass to do the loop transformation? Besides, how to distinguish whether a transform is heavy or not? Thanks!

You can do the walk within the same pass. You can call the rewrite driver any number of times before or after such a walk-based transformation.

Thanks, I’ll try it!