Skip to content

Commit 22824b6

Browse files
paul0403mehrdad2m
andauthored
Update llvm, 2025 Q3 (#1916)
**Context:** Update llvm, mhlo and enzyme, 2025 Q3. The latest pair of good versions, indicated by mhlo, is tensorflow/mlir-hlo@1dd2e71 ``` mhlo=1dd2e71331014ae0373f6bf900ce6be393357190 llvm=f8cb7987c64dcffb72414a40560055cb717dbf74 ``` For Enzyme, we go to the latest release https://github.com/EnzymeAD/Enzyme/releases/tag/v0.0.186 ``` enzyme=v0.0.186 ``` with commit `8c1a596158f6194f10e8ffd56a1660a61c54337e` **Description of the Change:** Miscellaneous: 1. `GreedyRewriteConfig.stuff = blah` -> `GreedyRewriteConfig.setStuff(blah)` llvm/llvm-project#137122 2. llvm gep op `inbounds` attribute is subsumed under a gep sign wrap enum flag llvm/llvm-project#137272 3. `arith::Constant[Int, Float]Op` builders now have the same argument order as other ops (output type first, then arguments) llvm/llvm-project#144636 (note that Enzyme also noticed this EnzymeAD/Enzyme#2379 😆 ) 4. The `lookupOrCreateFn` functions now take in a builder instead of instantiating a new one llvm/llvm-project#136421 5. `getStridedElementPtr` now takes in `rewriter` as the first argument (instead of the last), like all the other utils llvm/llvm-project#138984 6. The following functions now return a `LogicalResult`, and will be caught by warnings as errors as `-Wunused-result`: - `func::FuncOp.[insert, erase]Argument(s)` llvm/llvm-project#137130 - `getBackwardSlice()` llvm/llvm-project#140961 Things related to `transform.apply_registered_pass` op: 1. It now takes in a `dynamic_options` llvm/llvm-project#142683. We don't need to use this as all our pass options are static. 2. The options it takes in are now dictionaries instead of strings llvm/llvm-project#143159 Bufferization: 1. `bufferization.to_memref` op is renamed to `bufferization.to_buffer` llvm/llvm-project#137180 3. `bufferization.to_tensor` op's builder now needs the result type to be explicit llvm/llvm-project#142986. This is also needed by a patched mhlo pass. 4. The `getBuffer()` methods take in a new arg for `BufferizationState` llvm/llvm-project#141019, llvm/llvm-project#141466 5. `UnknownTypeConverterFn` in bufferization options now takes in just a type instead of a full value llvm/llvm-project#144658 **Related GitHub Issues:** [sc-95176] [sc-95664] --------- Co-authored-by: Mehrdad Malek <[email protected]>
1 parent 7d46225 commit 22824b6

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

51 files changed

+645
-288
lines changed

.dep-versions

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22
# To update JAX version alongside compatible dependency tags, run the following script:
33
# python3 .github/workflows/set_dep_versions.py {JAX_version}
44
jax=0.6.2
5-
mhlo=617a9361d186199480c080c9e8c474a5e30c22d1
6-
llvm=179d30f8c3fddd3c85056fd2b8e877a4a8513158
7-
enzyme=v0.0.180
5+
mhlo=1dd2e71331014ae0373f6bf900ce6be393357190
6+
llvm=f8cb7987c64dcffb72414a40560055cb717dbf74
7+
enzyme=v0.0.186
88

99
# Always remove custom PL/LQ versions before release.
1010

.github/workflows/build-wheel-linux-arm64.yaml

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,12 @@ jobs:
108108
ref: ${{ needs.constants.outputs.llvm_version }}
109109
path: ${{ github.workspace }}/mlir/llvm-project
110110

111+
- name: Patch LLVM Source
112+
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
113+
run: |
114+
cd $GITHUB_WORKSPACE/mlir/llvm-project
115+
git apply $GITHUB_WORKSPACE/mlir/patches/llvm-bufferization-segfault.patch
116+
111117
- name: Clone MHLO Submodule
112118
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
113119
uses: actions/checkout@v4
@@ -122,6 +128,7 @@ jobs:
122128
cd $GITHUB_WORKSPACE/mlir/mlir-hlo
123129
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-remove-shardy.patch
124130
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-add-back-necessary-passes.patch
131+
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-rename-sort.patch
125132
126133
- name: Clone Enzyme Submodule
127134
if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
@@ -134,9 +141,8 @@ jobs:
134141
- name: Patch Enzyme Source
135142
if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
136143
run: |
137-
export TARGET_FILE=$GITHUB_WORKSPACE/mlir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
138-
export PATCH_FILE=$GITHUB_WORKSPACE/mlir/patches/enzyme-nvvm-fabs-intrinsics.patch
139-
patch -p1 $TARGET_FILE $PATCH_FILE
144+
cd $GITHUB_WORKSPACE/mlir/Enzyme
145+
git apply $GITHUB_WORKSPACE/mlir/patches/enzyme-nvvm-fabs-intrinsics.patch
140146
141147
# Cache external project builds
142148
- name: Restore LLVM Build

.github/workflows/build-wheel-linux-x86_64.yaml

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,12 @@ jobs:
127127
ref: ${{ needs.constants.outputs.llvm_version }}
128128
path: ${{ github.workspace }}/mlir/llvm-project
129129

130+
- name: Patch LLVM Source
131+
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
132+
run: |
133+
cd $GITHUB_WORKSPACE/mlir/llvm-project
134+
git apply $GITHUB_WORKSPACE/mlir/patches/llvm-bufferization-segfault.patch
135+
130136
- name: Clone MHLO Submodule
131137
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
132138
uses: actions/checkout@v4
@@ -141,6 +147,7 @@ jobs:
141147
cd $GITHUB_WORKSPACE/mlir/mlir-hlo
142148
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-remove-shardy.patch
143149
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-add-back-necessary-passes.patch
150+
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-rename-sort.patch
144151
145152
- name: Clone Enzyme Submodule
146153
if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
@@ -153,9 +160,8 @@ jobs:
153160
- name: Patch Enzyme Source
154161
if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
155162
run: |
156-
export TARGET_FILE=$GITHUB_WORKSPACE/mlir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
157-
export PATCH_FILE=$GITHUB_WORKSPACE/mlir/patches/enzyme-nvvm-fabs-intrinsics.patch
158-
patch -p1 $TARGET_FILE $PATCH_FILE
163+
cd $GITHUB_WORKSPACE/mlir/Enzyme
164+
git apply $GITHUB_WORKSPACE/mlir/patches/enzyme-nvvm-fabs-intrinsics.patch
159165
160166
# Cache external project builds
161167
- name: Restore LLVM Build

.github/workflows/build-wheel-macos-arm64.yaml

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,12 @@ jobs:
113113
ref: ${{ needs.constants.outputs.llvm_version }}
114114
path: ${{ github.workspace }}/mlir/llvm-project
115115

116+
- name: Patch LLVM Source
117+
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
118+
run: |
119+
cd $GITHUB_WORKSPACE/mlir/llvm-project
120+
git apply $GITHUB_WORKSPACE/mlir/patches/llvm-bufferization-segfault.patch
121+
116122
- name: Clone MHLO Submodule
117123
if: steps.cache-mhlo-source.outputs.cache-hit != 'true'
118124
uses: actions/checkout@v4
@@ -127,6 +133,7 @@ jobs:
127133
cd $GITHUB_WORKSPACE/mlir/mlir-hlo
128134
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-remove-shardy.patch
129135
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-add-back-necessary-passes.patch
136+
git apply $GITHUB_WORKSPACE/mlir/patches/mhlo-rename-sort.patch
130137
131138
- name: Clone Enzyme Submodule
132139
if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
@@ -139,9 +146,8 @@ jobs:
139146
- name: Patch Enzyme Source
140147
if: steps.cache-enzyme-source.outputs.cache-hit != 'true'
141148
run: |
142-
export TARGET_FILE=$GITHUB_WORKSPACE/mlir/Enzyme/enzyme/Enzyme/TypeAnalysis/TypeAnalysis.cpp
143-
export PATCH_FILE=$GITHUB_WORKSPACE/mlir/patches/enzyme-nvvm-fabs-intrinsics.patch
144-
patch -p1 $TARGET_FILE $PATCH_FILE
149+
cd $GITHUB_WORKSPACE/mlir/Enzyme
150+
git apply $GITHUB_WORKSPACE/mlir/patches/enzyme-nvvm-fabs-intrinsics.patch
145151
146152
# Cache external project builds
147153
- name: Restore LLVM Build

Makefile

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -282,6 +282,9 @@ clean-plugin:
282282
clean-llvm:
283283
$(MAKE) -C mlir clean-llvm
284284

285+
reset-llvm:
286+
$(MAKE) -C mlir reset-llvm
287+
285288
clean-mhlo:
286289
$(MAKE) -C mlir clean-mhlo
287290

doc/releases/changelog-dev.md

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,16 @@
1212
* The JAX version used by Catalyst is updated to 0.6.2.
1313
[(#1897)](https://github.com/PennyLaneAI/catalyst/pull/1897)
1414

15+
* The version of LLVM, mlir-hlo, and Enzyme used by Catalyst has been updated.
16+
[(#1916)](https://github.com/PennyLaneAI/catalyst/pull/1916)
17+
18+
The LLVM version has been updated to
19+
[commit f8cb798](https://github.com/llvm/llvm-project/tree/f8cb7987c64dcffb72414a40560055cb717dbf74).
20+
The mlir-hlo version has been updated to
21+
[commit 1dd2e71](https://github.com/tensorflow/mlir-hlo/tree/1dd2e71331014ae0373f6bf900ce6be393357190).
22+
The Enzyme version has been updated to
23+
[v0.0.186](https://github.com/EnzymeAD/Enzyme/releases/tag/v0.0.186).
24+
1525
<h3>Deprecations 👋</h3>
1626

1727
<h3>Bug fixes 🐛</h3>

frontend/catalyst/jax_extras/lowering.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from __future__ import annotations
1717

1818
import logging
19+
import textwrap
1920

2021
import jax
2122
from jax._src.dispatch import jaxpr_replicas
@@ -38,6 +39,7 @@
3839

3940
import catalyst
4041
from catalyst.logging import debug_logger
42+
from catalyst.utils.exceptions import CompileError
4143
from catalyst.utils.patching import Patcher
4244

4345
# pylint: disable=protected-access
@@ -165,3 +167,57 @@ def custom_lower_jaxpr_to_module(
165167
worklist += [*op.body.operations]
166168

167169
return ctx.module, ctx.context
170+
171+
172+
def get_mlir_attribute_from_pyval(value):
173+
"""
174+
Given a value of any type, construct an mlir attribute of corresponding type.
175+
176+
We set up the context and location outside because recursive calls to this function
177+
will segfault if multiple `Context()`s are instantiated.
178+
"""
179+
180+
attr = None
181+
match value:
182+
case bool():
183+
attr = ir.BoolAttr.get(value)
184+
185+
case int():
186+
if -9223372036854775808 <= value < 0: # 2**63
187+
attr = ir.IntegerAttr.get(ir.IntegerType.get_signed(64), value)
188+
elif 0 <= value < 18446744073709551616: # = 2**64
189+
attr = ir.IntegerAttr.get(ir.IntegerType.get_signless(64), value)
190+
else:
191+
raise CompileError(
192+
textwrap.dedent(
193+
"""
194+
Large interger attributes currently not supported in MLIR,
195+
see https://github.com/llvm/llvm-project/issues/128072
196+
"""
197+
)
198+
)
199+
200+
case float():
201+
attr = ir.FloatAttr.get(ir.F64Type.get(), value)
202+
203+
case str():
204+
attr = ir.StringAttr.get(value)
205+
206+
case list() | tuple():
207+
element_attrs = [get_mlir_attribute_from_pyval(elem) for elem in value]
208+
attr = ir.ArrayAttr.get(element_attrs)
209+
210+
case dict():
211+
named_attrs = {}
212+
for k, v in value.items():
213+
if not isinstance(k, str):
214+
raise CompileError(
215+
f"Dictionary keys for MLIR DictionaryAttr must be strings, got: {type(k)}"
216+
)
217+
named_attrs[k] = get_mlir_attribute_from_pyval(v)
218+
attr = ir.DictAttr.get(named_attrs)
219+
220+
case _:
221+
raise CompileError(f"Cannot convert Python type {type(value)} to an MLIR attribute.")
222+
223+
return attr

frontend/catalyst/jax_primitives_utils.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -280,7 +280,11 @@ def transform_named_sequence_lowering(jax_ctx: mlir.LoweringRuleContext, pipelin
280280
for _pass in pipeline:
281281
options = _pass.get_options()
282282
apply_registered_pass_op = ApplyRegisteredPassOp(
283-
result=transform_mod_type, target=target, pass_name=_pass.name, options=options
283+
result=transform_mod_type,
284+
target=target,
285+
pass_name=_pass.name,
286+
options=options,
287+
dynamic_options={},
284288
)
285289
target = apply_registered_pass_op.result
286290
transform_yield_op = YieldOp(operands_=[]) # pylint: disable=unused-variable

frontend/catalyst/passes/pass_api.py

Lines changed: 11 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import pennylane as qml
2121

22+
from catalyst.jax_extras.lowering import get_mlir_attribute_from_pyval
2223
from catalyst.tracing.contexts import EvaluationContext
2324

2425
PipelineDict: TypeAlias = dict[str, dict[str, str]]
@@ -286,23 +287,18 @@ def __init__(self, name: str, *options: list[str], **valued_options: dict[str, s
286287

287288
def get_options(self):
288289
"""
289-
Stringify options according to what mlir-opt expects.
290-
291-
ApplyRegisteredPassOp expects options to be a single StringAttr
292-
which follows the same format as the one used with mlir-opt.
293-
294-
https://mlir.llvm.org/docs/Dialects/Transform/#transformapply_registered_pass-transformapplyregisteredpassop
295-
296-
Options passed to a pass are specified via the syntax {option1=value1 option2=value2 ...},
297-
i.e., use space-separated key=value pairs for each option.
290+
Build a dictionary mapping option names to MLIR attributes.
291+
ApplyRegisteredPassOp expects options to be a dictionary from strings to attributes.
292+
See https://github.com/llvm/llvm-project/pull/143159
293+
"""
294+
options_dict = {}
295+
for option in self.options:
296+
options_dict[str(option)] = get_mlir_attribute_from_pyval(True)
298297

299-
https://mlir.llvm.org/docs/Tutorials/MlirOpt/#running-a-pass-with-options
298+
for option, value in self.valued_options.items():
299+
options_dict[str(option)] = get_mlir_attribute_from_pyval(value)
300300

301-
Experimentally we found that single-options also work without values.
302-
"""
303-
retval = " ".join(f"{str(option)}" for option in self.options)
304-
retval2 = " ".join(f"{str(key)}={str(value)}" for key, value in self.valued_options.items())
305-
return " ".join([retval, retval2]).strip()
301+
return options_dict
306302

307303
def __repr__(self):
308304
return (

frontend/test/lit/test_mlir_plugin.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ def test_pass_options():
106106
"""Is the option in the generated MLIR?"""
107107

108108
@qjit(target="mlir")
109-
# CHECK: options = "an-option maxValue=1"
109+
# CHECK: options = {"an-option" = true, "maxValue" = 1 : i64}
110110
@catalyst.passes.apply_pass("some-pass", "an-option", maxValue=1)
111111
@qml.qnode(qml.device("null.qubit", wires=1))
112112
def example():

0 commit comments

Comments
 (0)