diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 99576e938f..a5f3290e37 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -4,6 +4,14 @@

Improvements 🛠

+* Allow to apply instrumentation to each pass within NamedSequenceOp. + [(#1978)](https://github.com/PennyLaneAI/catalyst/pull/1978) + For example: + ``` + quantum-opt ... --apply-transform-sequence --mlir-print-ir-after-all + ``` + should also dump the IR result for each subpass within NamedSequeceOp + * Added `detensorizefunctionboundary` pass to remove scalar tensors across function boundaries and enabled `symbol-dce` pass to remove dead functions, reducing the number of instructions for compilation. [(#1904)](https://github.com/PennyLaneAI/catalyst/pull/1904) diff --git a/mlir/lib/Catalyst/Transforms/ApplyTransformSequencePass.cpp b/mlir/lib/Catalyst/Transforms/ApplyTransformSequencePass.cpp index cfd1900e28..e3dd1ce472 100644 --- a/mlir/lib/Catalyst/Transforms/ApplyTransformSequencePass.cpp +++ b/mlir/lib/Catalyst/Transforms/ApplyTransformSequencePass.cpp @@ -19,20 +19,129 @@ #include -#include "llvm/Support/Debug.h" - +#include "mlir/Dialect/Transform/IR/TransformOps.h" #include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h" +#include "mlir/Dialect/Transform/Transforms/TransformInterpreterUtils.h" #include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/AnalysisManager.h" #include "mlir/Pass/Pass.h" - -#include "Catalyst/IR/CatalystDialect.h" -#include "QEC/IR/QECDialect.h" +#include "mlir/Pass/PassInstrumentation.h" +#include "mlir/Pass/PassManager.h" using namespace llvm; using namespace mlir; -using namespace catalyst; namespace catalyst { + +/// Generate a meaningful name for a transform operation for pass instrumentation +static std::string getTransformOpName(transform::TransformOpInterface transformOp) +{ + std::string baseName = transformOp->getName().getStringRef().str(); + + if (auto applyPassOp = dyn_cast(transformOp.getOperation())) { + if (auto passName = applyPassOp.getPassName(); !passName.empty()) { + return "transform_" + passName.str(); + } + } + + // convert "." to "_" + std::replace(baseName.begin(), baseName.end(), '.', '_'); + return baseName; +} + +/// A fake pass wrapper that represents a single transform operation. Allowing it to be tracked by +/// pass instrumentation. +class TransformOpSubPass : public OperationPass<> { + private: + transform::TransformOpInterface transformOp; + std::string opNameStr; + + public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TransformOpSubPass) + + TransformOpSubPass(transform::TransformOpInterface op) + : OperationPass(TypeID::get()), transformOp(op), + opNameStr(getTransformOpName(op)) + { + } + + void runOnOperation() override + { + llvm_unreachable("TransformOpSubPass should not be executed"); + } + + StringRef getName() const override { return opNameStr; } + StringRef getArgument() const override { return opNameStr; } + StringRef getDescription() const override { return "Transform dialect operation"; } + + std::unique_ptr clonePass() const override + { + return std::make_unique(transformOp); + } + + transform::TransformOpInterface getTransformOp() const { return transformOp; } +}; + +/// Apply transforms with individual subpass tracking by executing each transform operation +/// individually with instrumentation hooks. This implements a custom sequence +/// execution that mirrors the logic in NamedSequenceOp::apply but with instrumentation. +/// Reference to transform::NamedSequenceOp::apply in +/// https://github.com/llvm/llvm-project/blob/334e9bf2dd01fbbfe785624c0de477b725cde6f2/mlir/lib/ +/// Dialect/Transform/IR/TransformOps.cpp#L2378 +LogicalResult applyTransformsWithSubpassTracking(Operation *payload, + transform::NamedSequenceOp namedSequence, + PassInstrumentor *passInstrumentor) +{ + // TODO: We currently only expect to have a single block in the sequence. It may change in the + // future. + assert(namedSequence.getBody().hasOneBlock() && + "Expected exactly one transform op in the sequence block"); + + Block &sequenceBlock = namedSequence.getBody().front(); + if (sequenceBlock.without_terminator().empty()) { + return success(); + } + + transform::TransformState state = + transform::detail::makeTransformStateForTesting(namedSequence->getParentRegion(), payload); + + // Map the entry block argument to the list of operations. + // Note: this is the same implementation as PossibleTopLevelTransformOp but + // without attaching the interface / trait since that is tailored to a + // dangling top-level op that does not get "called". + auto scope = state.make_region_scope(namedSequence.getBody()); + if (failed(transform::detail::mapPossibleTopLevelTransformOpBlockArguments( + state, namedSequence, namedSequence.getBody()))) { + return failure(); + } + + for (Operation &transformOp : sequenceBlock.without_terminator()) { + if (auto transformInterface = dyn_cast(transformOp)) { + auto subPass = std::make_unique(transformInterface); + + // hook before pass + passInstrumentor->runBeforePass(subPass.get(), payload); + + DiagnosedSilenceableFailure result = state.applyTransform(transformInterface); + + if (result.isDefiniteFailure()) { + // hook after pass failed + passInstrumentor->runAfterPassFailed(subPass.get(), payload); + return failure(); + } + + if (result.isSilenceableFailure()) { + (void)result.silence(); + } + + // hook after pass + passInstrumentor->runAfterPass(subPass.get(), payload); + } + } + + return success(); +} + #define GEN_PASS_DEF_APPLYTRANSFORMSEQUENCEPASS #include "Catalyst/Transforms/Passes.h.inc" @@ -79,12 +188,23 @@ struct ApplyTransformSequencePass } }); - // Perform the transform - if (failed(mlir::transform::applyTransforms( - payload, cast(transformer_main_sequence), {}, - mlir::transform::TransformOptions(), false))) { - return signalPassFailure(); - }; + if (PassInstrumentor *passInstrumentor = getAnalysisManager().getPassInstrumentor()) { + // Manually execute the transform sequence with individual subpass tracking + if (auto namedSequence = + dyn_cast(transformer_main_sequence)) { + if (failed(applyTransformsWithSubpassTracking(payload, namedSequence, + passInstrumentor))) { + return signalPassFailure(); + } + } + } + else { + if (failed(transform::applyTransforms( + payload, cast(transformer_main_sequence), {}, + transform::TransformOptions(), false))) { + return signalPassFailure(); + } + } transformer.erase(); } diff --git a/mlir/test/Catalyst/ApplyTransformSequenceInstrumentationTest.mlir b/mlir/test/Catalyst/ApplyTransformSequenceInstrumentationTest.mlir new file mode 100644 index 0000000000..c6fdedf05a --- /dev/null +++ b/mlir/test/Catalyst/ApplyTransformSequenceInstrumentationTest.mlir @@ -0,0 +1,42 @@ +// Copyright 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. + +// RUN: quantum-opt %s --apply-transform-sequence --mlir-print-ir-after-all --split-input-file --verify-diagnostics 2>&1 | FileCheck %s --check-prefix=CHECK-INSTRUMENTATION + +// Instrumentation should show individual transform operations as subpasses +// CHECK-INSTRUMENTATION: transform_cse +// CHECK-INSTRUMENTATION: transform_remove-chained-self-inverse +// CHECK-INSTRUMENTATION: ApplyTransformSequencePass + +module @workflow { + + module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.op<"builtin.module">) { + %0 = transform.apply_registered_pass "cse" to %arg0 : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> + %1 = transform.apply_registered_pass "remove-chained-self-inverse" to %0 : (!transform.op<"builtin.module">) -> !transform.op<"builtin.module"> + transform.yield + } + } + + func.func private @f(%arg0: tensor) -> !quantum.bit { + %c_0 = stablehlo.constant dense<0> : tensor + %extracted = tensor.extract %c_0[] : tensor + %0 = quantum.alloc( 1) : !quantum.reg + %1 = quantum.extract %0[%extracted] : !quantum.reg -> !quantum.bit + %out_qubits = quantum.custom "Hadamard"() %1 : !quantum.bit + %out_qubits_1 = quantum.custom "Hadamard"() %out_qubits : !quantum.bit + return %out_qubits_1 : !quantum.bit + } + +}