Page MenuHomePhabricator

[MLIR][SCF] Combine adjacent scf.if with same condition
ClosedPublic

Authored by wsmoses on Mon, May 3, 4:24 PM.

Diff Detail

Event Timeline

wsmoses created this revision.Mon, May 3, 4:24 PM
wsmoses requested review of this revision.Mon, May 3, 4:24 PM
Herald added a project: Restricted Project. · View Herald TranscriptMon, May 3, 4:24 PM
mehdi_amini added inline comments.Mon, May 3, 6:13 PM
mlir/lib/Dialect/SCF/SCF.cpp
1289

The restriction on the first if isn't intrinsic to the fusion right? We could fuse the two even if they both return results, as long as the else branch of the second does not use a result of the first if I think.

1303
wsmoses updated this revision to Diff 342619.Mon, May 3, 6:26 PM
wsmoses marked 2 inline comments as done.

Address comment

mlir/lib/Dialect/SCF/SCF.cpp
1289

I think the condition you mention isn't quite sufficient, but ensuring the second if doesn't use the result of the first if should be sufficient. Adding a note that this can be extended to support that.

mehdi_amini added inline comments.Mon, May 3, 6:59 PM
mlir/lib/Dialect/SCF/SCF.cpp
1289

Actually neither the then or the else branch should use the result indeed.
Can we implement the general fusion directly instead of implementing this partial one?

wsmoses updated this revision to Diff 342627.Mon, May 3, 8:17 PM

Permit yielded values in first scf.if

mehdi_amini added inline comments.Mon, May 3, 8:57 PM
mlir/lib/Dialect/SCF/SCF.cpp
1308–1310

You have a getUsers() on Operation which allows to directly iterates on this nested loop (see inline edit)

1313

Nit: no need for braces here.

1318

(spurious braces)

and you can likely just assign() here:

mergedTypes.assign(nextIf.getResultTypes().begin(), nextIf.getResultTypes().end());
1326

rewriter.mergeBlocks will do the splicing, but at the end.
Can you merge into the prevIf and avoid manual splicing here?

1329

Nit: .front() would avoid derefencing an iterator (may just be more readable)

1338–1340
1363

Just like the other revision: there are far too much back() (and even back().back()!) here, I'd like to see some refactoring into better accessors.

1366–1368
1385

Can you use llvm::enumerate instead of a side indexing?

wsmoses updated this revision to Diff 342633.Mon, May 3, 9:28 PM

Refactor API usage

wsmoses marked 9 inline comments as done.Mon, May 3, 9:36 PM
ftynse added inline comments.Tue, May 4, 2:42 AM
mlir/lib/Dialect/SCF/SCF.cpp
1317

Ultra-nit: we tend to use the /*hasElse=*/false style.

1330–1331

setInsertionPointToEnd(combinedIf.thenBlock())

1334–1335

Nit: there is llvm::append_range.

wsmoses updated this revision to Diff 342749.Tue, May 4, 8:21 AM
wsmoses marked an inline comment as done.

Simplify splicing

mehdi_amini added inline comments.Tue, May 4, 5:49 PM
mlir/lib/Dialect/SCF/SCF.cpp
1342
1343

Why do you still have some "splice" above instead of rewriter.mergeBlocks()?

It seems that these splice are actually adding blocks to the region instead of merging the blocks? I'm not sure I read the code right...

wsmoses added inline comments.Tue, May 4, 5:53 PM
mlir/lib/Dialect/SCF/SCF.cpp
1342

I don't think that's correct since the splice moves the (optional) block from the nextIf to the new if, whereas your change would move instructions (rather than blocks).

1343

The splice is to conditionally add the blocks of the else region (if the block exists).

mehdi_amini accepted this revision.Tue, May 4, 6:44 PM
mehdi_amini added inline comments.
mlir/lib/Dialect/SCF/SCF.cpp
1343

Ah OK I see the logic, it is actually moving and preserving the block itself. And we need this because the else region can be empty but not the then one.
It didn't seem obvious to me on first read somehow, either I'm too tired to review right now, or some high-level comments could help along the way.

This revision is now accepted and ready to land.Tue, May 4, 6:44 PM
This revision was automatically updated to reflect the committed changes.