Scatter lowering

.github/workflows/check-catalyst.yaml

@@ -104,6 +104,14 @@ jobs:
     steps:
     - name: Checkout Catalyst repo
       uses: actions/checkout@v3
+
+    - name: Cache MHLO Source
+      id: cache-mhlo-source
+      uses: actions/cache@v3
+      with:
+        path: mlir/mlir-hlo
+        key: ${{ runner.os }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-default-source
+        enableCrossOsArchive: True
 
     - name: Cache MHLO Build
       id: cache-mhlo
@@ -131,14 +139,6 @@ jobs:
         key: ${{ runner.os }}-llvm-${{ needs.constants.outputs.llvm_version }}-default-build-opt
         fail-on-cache-miss: True
 
-    - name: Cache MHLO Source
-      id: cache-mhlo-source
-      uses: actions/cache@v3
-      with:
-        path: mlir/mlir-hlo
-        key: ${{ runner.os }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-default-source
-        enableCrossOsArchive: True
-
     - name: Clone MHLO Submodule
       if: |
         steps.cache-mhlo.outputs.cache-hit != 'true' ||
@@ -220,7 +220,6 @@ jobs:
         make enzyme
 
   quantum:
-    name: Quantum Dialects Build
     needs: [constants, mhlo, llvm]
     runs-on: ubuntu-latest
 
@@ -242,6 +241,14 @@ jobs:
         enableCrossOsArchive: True
         fail-on-cache-miss: True
 
+    - name: Get Cached LLVM Build
+      id: cache-llvm-build
+      uses: actions/cache@v3
+      with:
+        path: llvm-build
+        key: ${{ runner.os }}-llvm-${{ needs.constants.outputs.llvm_version }}-default-build-opt
+        fail-on-cache-miss: True
+
     - name: Get Cached MHLO Source
       id: cache-mhlo-source
       uses: actions/cache@v3
@@ -250,13 +257,13 @@ jobs:
         key: ${{ runner.os }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-default-source
         enableCrossOsArchive: True
         fail-on-cache-miss: True
-
-    - name: Get Cached LLVM Build
-      id: cache-llvm-build
+    
+    - name: Get Cached MHLO Build
+      id: cache-mhlo
       uses: actions/cache@v3
       with:
-        path: llvm-build
-        key: ${{ runner.os }}-llvm-${{ needs.constants.outputs.llvm_version }}-default-build-opt
+        path: mhlo-build
+        key: ${{ runner.os }}-mhlo-${{ needs.constants.outputs.mhlo_version }}-default-build
         fail-on-cache-miss: True
 
     - name: Get Cached MHLO Build

doc/changelog.md

@@ -2,6 +2,24 @@
 
 <h3>New features</h3>
 
+* Add lowering to tensor dialect for MHLO scatter. It unlocks indexing and updating jax arrays.
+  [(#273)](https://github.com/PennyLaneAI/catalyst/pull/273)
+
+  ```python
+
+  @qjit
+  def add_multiply(l: jax.core.ShapedArray((3,), dtype=float), idx: int):
+      res = l.at[idx].multiply(3)
+      res2 = l.at[idx].add(2)
+      return res + res2
+
+  res = add_multiply(jnp.array([0, 1, 2]), 2)
+  ```
+
+  ```pycon
+  >>> res
+  [0, 2, 10]
+
 * Catalyst users can now use Python for loop statements in their programs without having to
   explicitly use the functional `catalyst.for_loop` form!
   [#258](https://github.com/PennyLaneAI/catalyst/pull/258)
@@ -56,7 +74,8 @@ This release contains contributions from (in alphabetical order):
 Ali Asadi,
 Erick Ochoa Lopez,
 Jacob Mai Peng,
-Sergei Mironov.
+Sergei Mironov,
+Romain Moyard.
 
 # Release 0.3.0
 

frontend/catalyst/compiler.py

@@ -112,6 +112,7 @@ def get_lib_path(project, env_var):
     (
         "QuantumCompilationPass",
         [
+            "scatter-lowering",
             "lower-gradients",
             "adjoint-lowering",
         ],

frontend/test/pytest/test_scatter.py

@@ -0,0 +1,46 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test integration for the lowering of MHLO scatter."""
+
+import jax
+import numpy as np
+from jax import numpy as jnp
+
+from catalyst import qjit
+
+
+def test_add_multiply():
+    """Test to index a jax array and have operations on it."""
+
+    @qjit
+    def add_multiply(l: jax.core.ShapedArray((3,), dtype=float), idx: int):
+        res = l.at[idx].multiply(3)
+        res2 = l.at[idx].add(2)
+        return res + res2
+
+    res = add_multiply(jnp.array([0, 1, 2]), 2)
+    assert np.allclose(res, [0, 2, 10])
+
+
+def test_multiple_index():
+    """Test to index a jax array and have operations on it."""
+
+    @qjit
+    def multiple_index_multiply(l: jax.core.ShapedArray((3,), dtype=float)):
+        res = l.at[1:3].multiply(3)
+        return res
+
+    res = multiple_index_multiply(jnp.array([0, 1, 2]))
+    assert np.allclose(res, [0, 3, 6])

mlir/CMakeLists.txt

@@ -47,6 +47,7 @@ set(ALL_MHLO_PASSES
 
 list(APPEND CMAKE_MODULE_PATH "${MLIR_CMAKE_DIR}")
 list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")
+list(APPEND CMAKE_MODULE_PATH "${MHLO_CMAKE_DIR}")
 include(TableGen)
 include(AddLLVM)
 include(AddMLIR)
@@ -57,8 +58,14 @@ if(QUANTUM_ENABLE_BINDINGS_PYTHON)
   mlir_configure_python_dev_packages()
 endif()
 
+list(GET MHLO_INCLUDE_DIRS 1 MLIRHLO_DIR)
+list(GET MHLO_INCLUDE_DIRS 2 MLIRHLO_BUILD_DIR)
+
 include_directories(${LLVM_INCLUDE_DIRS}
                     ${MLIR_INCLUDE_DIRS}
+                    ${MHLO_INCLUDE_DIRS}
+                    ${MLIRHLO_DIR}/stablehlo
+                    ${MLIRHLO_BUILD_DIR}/stablehlo
                     ${PROJECT_SOURCE_DIR}/include
                     ${PROJECT_BINARY_DIR}/include)
 link_directories(${LLVM_BUILD_LIBRARY_DIR})

mlir/include/Quantum/Transforms/Passes.h

@@ -25,5 +25,6 @@ std::unique_ptr<mlir::Pass> createQuantumConversionPass();
 std::unique_ptr<mlir::Pass> createEmitCatalystPyInterfacePass();
 std::unique_ptr<mlir::Pass> createCopyGlobalMemRefPass();
 std::unique_ptr<mlir::Pass> createAdjointLoweringPass();
+std::unique_ptr<mlir::Pass> createScatterLoweringPass();
 
 } // namespace catalyst

mlir/include/Quantum/Transforms/Passes.td

@@ -68,4 +68,15 @@ def AdjointLoweringPass : Pass<"adjoint-lowering"> {
     let constructor = "catalyst::createAdjointLoweringPass()";
 }
 
+def ScatterLoweringPass : Pass<"scatter-lowering"> {
+    let summary = "Lower scatter op from Stable HLO to loops.";
+
+    let dependentDialects = [
+        "index::IndexDialect",
+        "mhlo::MhloDialect"
+    ];
+
+    let constructor = "catalyst::createScatterLoweringPass()";
+}
+
 #endif // QUANTUM_PASSES

mlir/include/Quantum/Transforms/Patterns.h

@@ -24,6 +24,7 @@ void populateBufferizationLegality(mlir::TypeConverter &, mlir::ConversionTarget
 void populateBufferizationPatterns(mlir::TypeConverter &, mlir::RewritePatternSet &);
 void populateQIRConversionPatterns(mlir::TypeConverter &, mlir::RewritePatternSet &);
 void populateAdjointPatterns(mlir::RewritePatternSet &);
+void populateScatterPatterns(mlir::RewritePatternSet &);
 
 } // namespace quantum
 } // namespace catalyst

mlir/lib/Catalyst/Transforms/RegisterAllPasses.cpp

@@ -22,6 +22,7 @@ void catalyst::registerAllCatalystPasses()
     mlir::registerPass(catalyst::createGradientBufferizationPass);
     mlir::registerPass(catalyst::createGradientLoweringPass);
     mlir::registerPass(catalyst::createGradientConversionPass);
+    mlir::registerPass(catalyst::createScatterLoweringPass);
     mlir::registerPass(catalyst::createAdjointLoweringPass);
     mlir::registerPass(catalyst::createQuantumBufferizationPass);
     mlir::registerPass(catalyst::createQuantumConversionPass);

mlir/lib/Quantum/Transforms/CMakeLists.txt

@@ -9,6 +9,8 @@ file(GLOB SRC
     cp_global_buffers.cpp
     adjoint_lowering.cpp
     AdjointPatterns.cpp
+    scatter_lowering.cpp
+    ScatterPatterns.cpp
 )
 
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)