Skip to content

[mlir] Update docs for Greedy Pattern Rewrite Driver(NFC) #126701

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.

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

Merged
merged 2 commits into from
Feb 17, 2025

Conversation

CoTinker
Copy link
Contributor

The applyOpPatternsAndFold is deprecated, use applyOpPatternsGreedily instead.

The `applyOpPatternsAndFold` is deprecated, use `applyOpPatternsGreedily` instead.
@llvmbot llvmbot added mlir:core MLIR Core Infrastructure mlir labels Feb 11, 2025
@llvmbot
Copy link
Member

llvmbot commented Feb 11, 2025

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

Changes

The applyOpPatternsAndFold is deprecated, use applyOpPatternsGreedily instead.


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

2 Files Affected:

  • (modified) mlir/docs/PatternRewriter.md (+1-1)
  • (modified) mlir/lib/Reducer/ReductionTreePass.cpp (+7-6)
diff --git a/mlir/docs/PatternRewriter.md b/mlir/docs/PatternRewriter.md
index d15e7e5a80678e3..9df46472990101e 100644
--- a/mlir/docs/PatternRewriter.md
+++ b/mlir/docs/PatternRewriter.md
@@ -361,7 +361,7 @@ This driver comes in two fashions:
 *   `applyPatternsGreedily` ("region-based driver") applies patterns to
     all ops in a given region or a given container op (but not the container op
     itself). I.e., the worklist is initialized with all containing ops.
-*   `applyOpPatternsAndFold` ("op-based driver") applies patterns to the
+*   `applyOpPatternsGreedily` ("op-based driver") applies patterns to the
     provided list of operations. I.e., the worklist is initialized with the
     specified list of ops.
 
diff --git a/mlir/lib/Reducer/ReductionTreePass.cpp b/mlir/lib/Reducer/ReductionTreePass.cpp
index 2d2744bfc2732d9..ef32adbab557769 100644
--- a/mlir/lib/Reducer/ReductionTreePass.cpp
+++ b/mlir/lib/Reducer/ReductionTreePass.cpp
@@ -56,13 +56,14 @@ static void applyPatterns(Region &region,
       opsInRange.push_back(&op.value());
   }
 
-  // `applyOpPatternsAndFold` may erase the ops so we can't do the pattern
-  // matching in above iteration. Besides, erase op not-in-range may end up in
-  // invalid module, so `applyOpPatternsAndFold` should come before that
-  // transform.
+  // `applyOpPatternsGreedily` with folding may erase the ops so we can't do the
+  // pattern matching in above iteration. Besides, erase op not-in-range may end
+  // up in invalid module, so `applyOpPatternsGreedily` with folding should come
+  // before that transform.
   for (Operation *op : opsInRange) {
-    // `applyOpPatternsAndFold` returns whether the op is convered. Omit it
-    // because we don't have expectation this reduction will be success or not.
+    // `applyOpPatternsGreedily` with folding returns whether the op is
+    // convered. Omit it because we don't have expectation this reduction will
+    // be success or not.
     GreedyRewriteConfig config;
     config.strictMode = GreedyRewriteStrictness::ExistingOps;
     (void)applyOpPatternsGreedily(op, patterns, config);

@CoTinker
Copy link
Contributor Author

Ping~

@joker-eph joker-eph merged commit 4e41e9a into llvm:main Feb 17, 2025
11 checks passed
@CoTinker CoTinker deleted the patterns branch February 18, 2025 01:31
paul0403 added a commit to PennyLaneAI/catalyst that referenced this pull request Jun 3, 2025
**Context:**
We update the llvm version tagged by jax 0.6.0:
```
mhlo=617a9361d186199480c080c9e8c474a5e30c22d1
llvm=179d30f8c3fddd3c85056fd2b8e877a4a8513158
```

We also update Enzyme to the latest version, which is 0.0.180, at commit
`db0181320d6e425ee963bd496ed0d8dbb615be18`


**Description of the Change:**

Firstly, jax recently moved from the Google github organization to its
own jax-ml organization.
This means the urls, and the retrieval method for the underlying llvm
and mhlo git commit tags, needs to be updated. (Thanks @mehrdad2m !)

Now on to the actual changes. I will list the changes in increasing
complexity.
1. The new enzyme cmake target is `EnzymeStatic-21` (from 20)

2. Enzyme works with a later llvm then our target, so it has some llvm
intrinsics unknown to the one we are targeting. We patch them away. They
do not concern us since they are all intrinsics for nvidia backends.

3. `applyPatternsAndFoldGreedily` is removed. Drop-in replacement is
`applyPatternsGreedily`.
llvm/llvm-project#104649,
llvm/llvm-project#126701

4. ops with `CallOpInterface` must have two new optional attributes
`arg_attrs` and `res_attrs`
llvm/llvm-project#123176

5. `CallInterfaceCallable` objects now must be directly casted to the
callee `SymbolRefAttr`, i.e. `callee.get<SymbolRefAttr>()` ->
`cast<SymbolRefAttr>(callee)`
llvm/llvm-project@35e8989

6. The `lookupOrCreateFn` family of functions now return
`FailureOr<funcop>` instead of just `funcop`, so a `.value()` needs to
be used to retrieve the underlying `funcop`.
llvm/llvm-project@e84f6b6

7. The cpp api for `OneShotBufferizePassOptions` no longer needs
complicated lambdas for the type converter options. They can be set with
the `mlir::bufferization::LayoutMapOption::IdentityLayoutMap` options
directly.

8. The individual `match` and `rewrite` methods in pattern rewrites are
removed. Use the two-in-one `matchAndRewrite` instead.
llvm/llvm-project#129861

9. For rewrite patterns with 1-to-N convertions, a new `macthAndRewrite`
overload with `OneToNOpAdaptor` must be used. For us, this is only the
`catalyst.list*` ops. llvm/llvm-project#116470

10. The lowering of `cf::AssertOp` to llvm was split from the
overall`--covert-cf-to-llvm` pass. We need to manually call this
separate pattern for cf.assert duriing quantum to llvm dialect lowering,
where we also convert cf to llvm.
https://github.com/llvm/llvm-project/pull/120431/files 

11. The new mhlo depends on a
[shardy](https://github.com/openxla/shardy) dialect. Shardy is built
with bazel, not cmake. Building shardy ourselves would be very difficult
(not having bazel in our build ecosystem is a hard constraint, cc @mlxd
), and also not necessary (we just use mhlo for their "standard"
passes). We thus patch out all shardy components.

12. Three necessary passes were removed in mhlo:
`mhlo-legalize-control-flow`, `mhlo-legalize-to-std`,
`hlo-legalize-sort`
tensorflow/mlir-hlo@4a640be#diff-ef0d7e30da19a396ba036405a9ef636f8b1be194618b0a90f4602671fc2ef34d

tensorflow/mlir-hlo@2a5e267#diff-f8c7cb07b43593403e00e0dbf9983f0186b4eb70368cc99af3b924061f1ea46f
- Alongside the removal of `mhlo-legalize-to-std`, the cmake target
`MhloToStandard` was removed too.
  We simply patch them back for now. 
  
**For the above two points, note that there will be an overall migration
to the stablehlo repo, as mhlo is sunseting.
Therefore, spending too much time on this isn't necessary, so we just
patch.**

13. The new pattern applicator (`applyPatternsGreedily`) is more
aggressive in dead code elimination,
and is eliminating dead `Value`s in the adjoint gradient method.
The `nodealloc` function we generate for adjoint gradient lowering used
to only return the qreg, not the expval result. This causes the expval
op to be eliminated since it has no users.
This further causes wrong gradient results, since the entire program,
all ops included (regardless
of dead or not), impacts the gradient through chain rule.
To avoid this, we return the expval result as well.
In doing this, we implicitly assume that differentiated qnodes can only
return expval.
Although this assumption is true and also restricted by frontend,
ideally we should not have it hard coded.
We leave this as a TODO for a future feature.

14. The old `--buffer-deallocation` pass is removed. Intended
replacement is `--buffer-deallocation-pipeline`.
This migration is very complicated. We simply add back the old buffer
deallocation pass in the catalyst dialect as a util for now.
We will revisit this in #1778 . 


mlir lit test updates:
1. `bufferization.to_tensor/memref` updated assembly format
2. gradient adjoint lowering test returns both qreg and expval
3. Some inverse unrealized conversion cast pairs are canceled by the new
pattern rewriter.
4. `llvm.mlir.undef` is deprecated, use `llvm.mlir.poison` instead.
llvm/llvm-project#125629


**Benefits:**
Up to date with upstream versions.


[sc-92017]

---------

Co-authored-by: Tzung-Han Juang <tzunghan.juang@gmail.com>
Co-authored-by: Ritu Thombre <42207923+ritu-thombre99@users.noreply.github.com>
Co-authored-by: Mehrdad Malekmohammadi <mehrdad.malek@xanadu.ai>
Co-authored-by: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com>
Co-authored-by: David Ittah <dime10@users.noreply.github.com>
Co-authored-by: Joey Carter <joseph.carter@xanadu.ai>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
mlir:core MLIR Core Infrastructure mlir
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants