Migrate catalyst dialect to new one-shot bufferization

doc/releases/changelog-dev.md

@@ -115,6 +115,7 @@
   The new mlir bufferization interface is required by jax 0.4.29 or higher.
   [(#1027)](https://github.com/PennyLaneAI/catalyst/pull/1027)
   [(#1686)](https://github.com/PennyLaneAI/catalyst/pull/1686)
+  [(#1708)](https://github.com/PennyLaneAI/catalyst/pull/1708)
 
 <h3>Documentation 📝</h3>
 

frontend/catalyst/pipelines.py

@@ -227,12 +227,13 @@
         "empty-tensor-to-alloc-tensor",
         "func.func(bufferization-bufferize)",
         "func.func(tensor-bufferize)",
-        "catalyst-bufferize",  # Must be run before -- func.func(linalg-bufferize)
+        #"catalyst-bufferize",  # Must be run before -- func.func(linalg-bufferize)
+        "one-shot-bufferize{dialect-filter=catalyst}",  # Must be run before -- func.func(linalg-bufferize)
         "func.func(linalg-bufferize)",
         "func.func(tensor-bufferize)",
         "one-shot-bufferize{dialect-filter=quantum}",
         "func-bufferize",
-        "func.func(finalizing-bufferize)",
+        #"func.func(finalizing-bufferize)",
         "canonicalize",  # Remove dead memrefToTensorOp's
         "gradient-postprocess",
         # introduced during gradient-bufferize of callbacks

mlir/include/Catalyst/Transforms/BufferizableOpInterfaceImpl.h

@@ -0,0 +1,23 @@
+// Copyright 2024-2025 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.
+
+#pragma once
+
+using namespace mlir;
+
+namespace catalyst {
+
+void registerBufferizableOpInterfaceExternalModels(mlir::DialectRegistry &registry);
+
+} // namespace catalyst

mlir/include/Catalyst/Transforms/Patterns.h

@@ -21,6 +21,7 @@
 
 namespace catalyst {
 
+// TODO: remove buffer pattern
 void populateBufferizationPatterns(mlir::TypeConverter &, mlir::RewritePatternSet &);
 
 void populateScatterPatterns(mlir::RewritePatternSet &);

mlir/include/Quantum/Transforms/Patterns.h

@@ -21,8 +21,6 @@
 namespace catalyst {
 namespace quantum {
 
-void populateBufferizationLegality(mlir::TypeConverter &, mlir::ConversionTarget &);
-void populateBufferizationPatterns(mlir::TypeConverter &, mlir::RewritePatternSet &);
 void populateQIRConversionPatterns(mlir::TypeConverter &, mlir::RewritePatternSet &);
 void populateAdjointPatterns(mlir::RewritePatternSet &);
 void populateSelfInversePatterns(mlir::RewritePatternSet &);

mlir/lib/Catalyst/IR/CatalystDialect.cpp

@@ -14,6 +14,7 @@
 
 #include "Catalyst/IR/CatalystDialect.h"
 #include "Catalyst/IR/CatalystOps.h"
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h" // needed for generated type parser
 #include "mlir/Interfaces/FunctionImplementation.h"
@@ -40,6 +41,8 @@ void CatalystDialect::initialize()
 #define GET_OP_LIST
 #include "Catalyst/IR/CatalystOps.cpp.inc"
         >();
+
+    declarePromisedInterfaces<bufferization::BufferizableOpInterface, PrintOp, CustomCallOp, CallbackCallOp, CallbackOp>();
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/Catalyst/Transforms/BufferizableOpInterfaceImpl.cpp

@@ -0,0 +1,285 @@
+// Copyright 2024-2025 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.
+
+#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "Catalyst/IR/CatalystOps.h"
+#include "Catalyst/Transforms/BufferizableOpInterfaceImpl.h"
+
+using namespace mlir;
+using namespace mlir::bufferization;
+using namespace catalyst;
+
+namespace {
+
+/// Bufferization of catalyst.print. Get memref of printOp.val.
+struct PrintOpInterface
+    : public bufferization::BufferizableOpInterface::ExternalModel<PrintOpInterface,
+                                                    PrintOp> {
+    bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                                const bufferization::AnalysisState &state) const {
+        return true;
+    }
+
+    bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                                 const bufferization::AnalysisState &state) const {
+        return false;
+    }
+
+    bufferization::AliasingValueList getAliasingValues(Operation *op,
+                                        OpOperand &opOperand,
+                                        const bufferization::AnalysisState &state) const {
+        return {};
+    }
+
+    LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                            const bufferization::BufferizationOptions &options) const {
+        auto printOp = cast<PrintOp>(op);
+        if (printOp.getVal()) {
+            FailureOr<Value> source = getBuffer(rewriter, printOp.getVal(), options);
+            if (failed(source)) {
+                return failure();
+            }
+            bufferization::replaceOpWithNewBufferizedOp<PrintOp>(rewriter, op, *source,
+                            printOp.getConstValAttr(), printOp.getPrintDescriptorAttr());
+        }
+        return success();
+    }
+};
+
+/// Bufferization of catalyst.custom_call. Mainly get buffers for arguments.
+struct CustomCallOpInterface
+    : public bufferization::BufferizableOpInterface::ExternalModel<CustomCallOpInterface,
+                                                    CustomCallOp> {
+    bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                                const bufferization::AnalysisState &state) const {
+        return true;
+    }
+
+    bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                                 const bufferization::AnalysisState &state) const {
+        return true;
+    }
+
+    bufferization::AliasingValueList getAliasingValues(Operation *op,
+                                        OpOperand &opOperand,
+                                        const bufferization::AnalysisState &state) const {
+        return {};
+    }
+
+    LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                            const bufferization::BufferizationOptions &options) const {
+        auto customCallOp = cast<CustomCallOp>(op);
+
+        // Add bufferized arguments
+        SmallVector<Value> bufferArgs;
+        ValueRange operands = customCallOp.getOperands();
+        for (Value operand : operands) {
+            FailureOr<Value> opBuffer = getBuffer(rewriter, operand, options);
+            if (failed(opBuffer)) {
+                return failure();
+            }
+            bufferArgs.push_back(*opBuffer);
+        }
+
+        // Add bufferized return values to the arguments
+        ValueRange results = customCallOp.getResults();
+        for (Value result : results) {
+            Type resultType = result.getType();
+            RankedTensorType tensorType = dyn_cast<RankedTensorType>(resultType);
+            if (!tensorType) {
+                return failure();
+            }
+            auto options = bufferization::BufferizationOptions();
+            FailureOr<Value> tensorAlloc = bufferization::allocateTensorForShapedValue(
+                rewriter, op->getLoc(), result, options, false);
+            MemRefType memrefType =
+                MemRefType::get(tensorType.getShape(), tensorType.getElementType());
+            auto newBuffer =
+                rewriter.create<bufferization::ToMemrefOp>(op->getLoc(), memrefType, *tensorAlloc);
+            bufferArgs.push_back(newBuffer);
+        }
+
+         // Add the initial number of arguments
+        int32_t numArguments = static_cast<int32_t>(customCallOp.getNumOperands());
+        DenseI32ArrayAttr numArgumentsDenseAttr = rewriter.getDenseI32ArrayAttr({numArguments});
+
+        // Create an updated custom call operation
+        rewriter.create<CustomCallOp>(op->getLoc(), TypeRange{}, bufferArgs,
+                                      customCallOp.getCallTargetName(), numArgumentsDenseAttr);
+        size_t startIndex = bufferArgs.size() - customCallOp.getNumResults();
+        SmallVector<Value> bufferResults(bufferArgs.begin() + startIndex, bufferArgs.end());
+        bufferization::replaceOpWithBufferizedValues(rewriter, op, bufferResults);
+
+        return success();
+    }
+};
+
+struct CallbackOpInterface
+    : public bufferization::BufferizableOpInterface::ExternalModel<CallbackOpInterface,
+                                                                   CallbackOp> {
+    bool hasTensorSemantics(Operation *op) const
+    {
+        auto isaTensor = llvm::IsaPred<TensorType>;
+
+        // A function has tensor semantics if it has tensor arguments/results.
+        auto callbackOp = cast<CallbackOp>(op);
+        bool hasTensorArg = any_of(callbackOp.getArgumentTypes(), isaTensor);
+        bool hasTensorResult = any_of(callbackOp.getResultTypes(), isaTensor);
+        if (hasTensorArg || hasTensorResult)
+            return true;
+
+        return false;
+    }
+
+    bufferization::AliasingOpOperandList
+    getAliasingOpOperands(Operation *op, Value value,
+                          const bufferization::AnalysisState &state) const
+    {
+        return {};
+    }
+
+    LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                            const bufferization::BufferizationOptions &options) const
+    {
+        auto callbackOp = cast<CallbackOp>(op);
+
+        auto argTys = callbackOp.getArgumentTypes();
+        auto retTys = callbackOp.getResultTypes();
+        SmallVector<Type> emptyRets;
+        SmallVector<Type> args(argTys.begin(), argTys.end());
+        args.insert(args.end(), retTys.begin(), retTys.end());
+        SmallVector<Type> bufferArgs;
+        for (Type ty : args) {
+            auto tensorType = dyn_cast<RankedTensorType>(ty);
+            if (!tensorType) {
+                bufferArgs.push_back(ty);
+            }
+            else {
+                bufferArgs.push_back(
+                    MemRefType::get(tensorType.getShape(), tensorType.getElementType()));
+            }
+        }
+        auto callbackTy = rewriter.getFunctionType(bufferArgs, emptyRets);
+        rewriter.modifyOpInPlace(op, [&] { callbackOp.setFunctionType(callbackTy); });
+
+        return success();
+    }
+};
+
+void convertTypes(SmallVector<Type> inTypes, SmallVector<Type> &convertedResults){
+    // See https://github.com/llvm/llvm-project/pull/114155/files
+    for (Type inType : inTypes) {
+        if (isa<TensorType>(inType)){
+            convertedResults.push_back(bufferization::getMemRefTypeWithStaticIdentityLayout(cast<TensorType>(inType)));
+        } else {
+            convertedResults.push_back(inType);
+        }
+    }
+}
+
+struct CallbackCallOpInterface
+    : public bufferization::BufferizableOpInterface::ExternalModel<CallbackCallOpInterface,
+                                                                   CallbackCallOp> {
+    bool bufferizesToAllocation(Operation *op, Value value) const { return true; }
+
+    bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
+                                const bufferization::AnalysisState &state) const
+    {
+        return true;
+    }
+
+    bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
+                                 const bufferization::AnalysisState &state) const
+    {
+        // We can safely say false because CallbackCallOp's memrefs
+        // will be put in a JAX array and JAX arrays are immutable.
+        //
+        //    Unlike NumPy arrays, JAX arrays are always immutable.
+        //
+        // https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html
+        return false;
+    }
+
+    bufferization::AliasingValueList
+    getAliasingValues(Operation *op, OpOperand &opOperand,
+                      const bufferization::AnalysisState &state) const
+    {
+        return {};
+    }
+
+    LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
+                            const bufferization::BufferizationOptions &options) const
+    {
+        auto callOp = cast<CallbackCallOp>(op);
+
+        SmallVector<Type> convertedResults;
+        convertTypes(SmallVector<Type>(callOp.getResultTypes()), convertedResults);
+        if (callOp->getNumResults() != convertedResults.size()) {
+            return failure();
+        }
+
+        SmallVector<Value> newInputs;
+        auto operands = callOp.getOperands();
+        for (Value operand : operands) {
+            FailureOr<Value> opBuffer = getBuffer(rewriter, operand, options);
+            if (failed(opBuffer)) {
+                return failure();
+            }
+            newInputs.push_back(*opBuffer);
+        }
+
+        auto results = callOp.getResults();
+        auto loc = callOp->getLoc();
+        SmallVector<Value> outmemrefs;
+        for (auto result : results) {
+            FailureOr<Value> tensorAlloc =
+                bufferization::allocateTensorForShapedValue(rewriter, loc, result, options, false);
+            if (failed(tensorAlloc)) {
+                return failure();
+            }
+
+            auto tensor = *tensorAlloc;
+            RankedTensorType tensorTy = cast<RankedTensorType>(tensor.getType());
+            auto shape = tensorTy.getShape();
+            auto elementTy = tensorTy.getElementType();
+            auto memrefType = MemRefType::get(shape, elementTy);
+            auto toMemrefOp = rewriter.create<bufferization::ToMemrefOp>(loc, memrefType, tensor);
+            auto memref = toMemrefOp.getResult();
+            outmemrefs.push_back(memref);
+            newInputs.push_back(memref);
+        }
+
+        SmallVector<Type> emptyRets;
+        rewriter.create<CallbackCallOp>(loc, emptyRets, callOp.getCallee(), newInputs);
+        bufferization::replaceOpWithBufferizedValues(rewriter, op, outmemrefs);
+        return success();
+    }
+};
+
+} // namespace
+
+void catalyst::registerBufferizableOpInterfaceExternalModels(
+    DialectRegistry &registry) {
+    registry.addExtension(+[](MLIRContext *ctx, CatalystDialect *dialect) {
+        CustomCallOp::attachInterface<CustomCallOpInterface>(*ctx);
+        PrintOp::attachInterface<PrintOpInterface>(*ctx);
+        CallbackOp::attachInterface<CallbackOpInterface>(*ctx);
+        CallbackCallOp::attachInterface<CallbackCallOpInterface>(*ctx);
+    });
+}

mlir/lib/Catalyst/Transforms/CMakeLists.txt

@@ -4,8 +4,9 @@ file(GLOB SRC
     ApplyTransformSequencePass.cpp
     ArrayListToMemRefPass.cpp
     AsyncUtils.cpp
-    BufferizationPatterns.cpp
-    catalyst_bufferize.cpp
+    BufferizableOpInterfaceImpl.cpp
+    BufferizationPatterns.cpp  // remove
+    catalyst_bufferize.cpp  // remove
     catalyst_to_llvm.cpp
     DetectQNodes.cpp
     DetensorizeSCFPass.cpp

mlir/lib/Driver/CompilerDriver.cpp

@@ -56,6 +56,7 @@
 #include "llvm/Transforms/IPO/GlobalDCE.h"
 
 #include "Catalyst/IR/CatalystDialect.h"
+#include "Catalyst/Transforms/BufferizableOpInterfaceImpl.h"
 #include "Catalyst/Transforms/Passes.h"
 #include "Driver/CatalystLLVMTarget.h"
 #include "Driver/CompilerDriver.h"
@@ -964,6 +965,7 @@ int QuantumDriverMainFromCL(int argc, char **argv)
     registerLLVMTranslations(registry);
 
     // Register bufferization interfaces
+    catalyst::registerBufferizableOpInterfaceExternalModels(registry);
     catalyst::quantum::registerBufferizableOpInterfaceExternalModels(registry);
 
     // Register and parse command line options.