[SCF] add debug label for MergeNestedParallelLoops pattern

[rscottmanley]

Adding debug labels to rewrite patterns gives the canonicalizer the ability to disable specific patterns if undesired.

The motivation here is that since there is currently no way to handle attributes on ops that are changed when these patterns are applied, this at least provides a workaround.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-scf

Author: Scott Manley (rscottmanley)

Changes

Adding debug labels to rewrite patterns gives the canonicalizer the ability to disable specific patterns if undesired.

The motivation here is that since there is currently no way to handle attributes on ops that are changed when these patterns are applied, this at least provides a workaround.

---

Full diff: https://github.com/llvm/llvm-project/pull/135980.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/SCF/IR/SCF.cpp (+3-3)

diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index c454f342b4204..04e89105c056a 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -3202,9 +3202,9 @@ struct MergeNestedParallelLoops : public OpRewritePattern<ParallelOp> {
 
 void ParallelOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
-  results
-      .add<ParallelOpSingleOrZeroIterationDimsFolder, MergeNestedParallelLoops>(
-          context);
+  results.add<ParallelOpSingleOrZeroIterationDimsFolder>(context);
+  results.addWithLabel<MergeNestedParallelLoops>({"MergeNestedParallelLoops"},
+                                                 context);
 }
 
 /// Given the region at `index`, or the parent operation if `index` is None,

[joker-eph]

Can you add a name on the pattern instead of a label here?
A label to identify a single pattern does not seem the intended purpose of labels.

[rscottmanley]

Okay I agree after I dug around some more. I took my guidance from what I read in include/mlir/Transforms/Passes.h which seems to say you can do it both ways:

/// ... Debug labels must be set explicitly on patterns or when adding
/// them with `RewritePatternSet::addWithLabel`. Debug names may be empty, but
/// patterns created with `RewritePattern::create` have their default debug name
/// set to their type name.

Does this suggest you could have individual pattern names and also a label (in this case, like SCFParallelPatterns) at the same time?

[rscottmanley]

updated with initialize() in the pattern itself

[Hardcode84]

RewritePatternSet.add<MergeNestedParallelLoops> should already set debug name to "MergeNestedParallelLoops" automatically

[rscottmanley]

RewritePatternSet.add<MergeNestedParallelLoops> should already set debug name to "MergeNestedParallelLoops" automatically

No. That appears to happen with RewritePatternSet.insert<> but to quote PatternMatch.h // TODO: These are soft deprecated in favor of the 'add' methods above.

I could add it to the add though

[rscottmanley]

RewritePatternSet.add<MergeNestedParallelLoops> should already set debug name to "MergeNestedParallelLoops" automatically

No. That appears to happen with RewritePatternSet.insert<> but to quote PatternMatch.h // TODO: These are soft deprecated in favor of the 'add' methods above.

I could add it to the add though

Actually it doesnt look like works in that case either.

[Hardcode84]

Both add and insert should set debug name. If they doesn't it's probably a bug.

[rscottmanley]

Both add and insert should set debug name. If they doesn't it's probably a bug.

Yeah, I see what you mean, through RewritePattern::create.. but it's definitely not working. Ill see if I can track it down. If that's the expected behaviour I would much rather have that anyway.

[rscottmanley]

I think I know what's happening.. llvm::getTypeName<T>() is trying to return the demangled name and it's not going to be consistent. I just dumped all of them when they are added and here are some examples:

(anonymous namespace)::MergeNestedParallelLoops
FoldConstantCase
CanonicalizeContractAdd<mlir::arith::AddFOp>
UndoComplexPattern<mlir::arith::AddFOp, fir::AddcOp>
mlir::ComposeReassociativeReshapeOps<mlir::memref::CollapseShapeOp, mlir::ReshapeOpKind::kCollapse>

[rscottmanley]

I'm not sure there's a simple solution to resolve this universally.. it seems like explicitly specifying the debug name is probably the best way to get what you would reasonably "expect"

[rscottmanley]

even the getTypeName() documentation says:

We provide a function which tries to compute the (demangled) name of a type statically.

This routine may fail on some platforms or for particularly unusual types. Do not use it for anything other than logging and debugging aids. It isn't portable or dependendable in any real sense.

I would like to proceed with this change if there is no significant objection.

[Hardcode84]

Well, it's called setDebugName and was added (by myself) explicitly for debugging (for the --debug-only=pattern-application specifically) so it is to be expected. Filtering patterns by debug name was added kind of ad-hoc later. No one guaranteed stable or predictable names (they will most likely differ between os/compilers too).

[rscottmanley]

No one guaranteed stable or predictable names (they will most likely differ between os/compilers too).

I understand and there's no criticism intended - I am just asking that this one be made predictable. There's probably a worthwhile larger conversation/RFC on perhaps requiring a name for patterns so they can be selectively disabled but this feels like a reasonable fix for now?

[Hardcode84]

I have no objections to the patch itself, but it's not very robust. If you really want to prevent nested loops merging I would suggest instead of using custom attributes on parallel op, introduce a custom region op to encode info and separate the ops:

scf.parallel ... {
    my_region_op {
        scf.parallel ... {
            ...
        }
    }
}

This worked good enough for us downstream

[rscottmanley]

I have no objections to the patch itself, but it's not very robust.

Perhaps - I am probably not going to sign up to add a name to every Pattern that exists. I would maybe say it provides an example for others should they want to do something similar.

If you really want to prevent nested loops merging I would suggest instead of using custom attributes on parallel op, introduce a custom region op to encode info

Yes, that's exactly what I was trying to avoid doing. To me this just introduces its own set of challenges to workaround. E.g. if you were hoist an operation out of the inner loop, does it/should it go into the region op or into the body of the outer loop? Some existing xforms don't work unless they have the same parent region, etc.

[joker-eph]

The problem I see here is that you want to attach semantics through the debug name, however here is how they are documented:

  /// Return a readable name for this pattern. This name should only be used for
  /// debugging purposes, and may be empty.
  StringRef getDebugName() const { return debugName; }

  /// Set the human readable debug name used for this pattern. This name will
  /// only be used for debugging purposes.
  void setDebugName(StringRef name) { debugName = name; }

That means we should be able to strip these out of a release build for example without breaking anything, which you may object to in the future if you depends on this behavior.

[rscottmanley]

That means we should be able to strip these out of a release build for example without breaking anything, which you may object to in the future if you depends on this behavior.

That's a fair point. Like I said earlier, an RFC to require a name is probably the larger fix. I mean, can we change "debugName" to "patternName" and then my commit is acceptable? It seems to me like this is a neat feature for customizing what you want or not want to be canonicalized depending on various needs of your downstream implementation.

[joker-eph]

That's a fair point. Like I said earlier, an RFC to require a name is probably the larger fix. I mean, can we change "debugName" to "patternName" and then my commit is acceptable?

Probably: but this isn't free, it makes all patterns more costly, and the framework heavier in production.

It seems to me like this is a neat feature for customizing what you want or not want to be canonicalized depending on various needs of your downstream implementation.

I would argue that the canonical form of the IR isn't something meant to be customized with this kind of fine grain control.
Also, relying on some sort of string backdoor does not seem like a great API on the principle as well.

[rscottmanley]

I would argue that the canonical form of the IR isn't something meant to be customized with this kind of fine grain control.
Also, relying on some sort of string backdoor does not seem like a great API on the principle as well.

I agree with everything you're saying from a compiler design perspective.. but at the same time, I don't think this particular pattern is very good candidate for canonicalization and unfortunately feel the need to disable it somehow. For example, loop optimizations that may apply to a fir::DoLoopOp nest, a scf::ForOp nest and a scf::ParallelOp nest.. but have to be engineered differently as only the latter canonicalizes this way. I also would argue that it's more difficult to pick out the fact it's a nest in verbose IR. I assume this pattern is useful for working on scf::ParallelOp but it's less useful for working on "LoopLike" ops.

[rscottmanley]

I am closing this PR. Thank you for your input everyone!

[joker-eph]

I don't think this particular pattern is very good candidate for canonicalization a

That I could agree with: parallelism is often expressed in a hierarchical way for good reasons, the collapsing done here is losing this information and does not seem like a good canonicalization to me.

[Hardcode84]

I don't think this particular pattern is very good candidate for canonicalization a

That I could agree with: parallelism is often expressed in a hierarchical way for good reasons, the collapsing done here is losing this information and does not seem like a good canonicalization to me.

scf.parallel already support arbitrarily nested loops as part of the op itself and loop order is preserved by this pattern (scf.parallel doesn't assign any particular semantics to the loop order, though). So I don't think any info is actually lost within upstream scf.parallel semantics and you can split back into individual loops as needed.