Apply instrumentation to each pass within NamedSequenceOp

doc/releases/changelog-dev.md

@@ -4,6 +4,14 @@
 
 <h3>Improvements 🛠</h3>
 
+* 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)
 

mlir/lib/Catalyst/Transforms/ApplyTransformSequencePass.cpp

@@ -19,20 +19,130 @@
 
 #include <cassert>
 
-#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(mlir::transform::TransformOpInterface transformOp)
+{
+    std::string baseName = transformOp->getName().getStringRef().str();
+
+    if (auto applyPassOp =
+            dyn_cast<mlir::transform::ApplyRegisteredPassOp>(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:
+    mlir::transform::TransformOpInterface transformOp;
+    std::string opNameStr;
+
+  public:
+    MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TransformOpSubPass)
+
+    TransformOpSubPass(mlir::transform::TransformOpInterface op)
+        : OperationPass(TypeID::get<TransformOpSubPass>()), 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<Pass> clonePass() const override
+    {
+        return std::make_unique<TransformOpSubPass>(transformOp);
+    }
+
+    mlir::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.
+/// ref: transform::NamedSequenceOp::apply in
+/// /path/to/mlir/llvm-project/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+LogicalResult applyTransformsWithSubpassTracking(Operation *payload,
+                                                 mlir::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();
+    }
+
+    mlir::transform::TransformState state = mlir::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(mlir::transform::detail::mapPossibleTopLevelTransformOpBlockArguments(
+            state, namedSequence, namedSequence.getBody()))) {
+        return failure();
+    }
+
+    for (Operation &transformOp : sequenceBlock.without_terminator()) {
+        if (auto transformInterface =
+                dyn_cast<mlir::transform::TransformOpInterface>(transformOp)) {
+            auto subPass = std::make_unique<TransformOpSubPass>(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 +189,23 @@ struct ApplyTransformSequencePass
             }
         });
 
-        // Perform the transform
-        if (failed(mlir::transform::applyTransforms(
-                payload, cast<mlir::transform::TransformOpInterface>(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<mlir::transform::NamedSequenceOp>(transformer_main_sequence)) {
+                if (failed(applyTransformsWithSubpassTracking(payload, namedSequence,
+                                                              passInstrumentor))) {
+                    return signalPassFailure();
+                }
+            }
+        }
+        else {
+            if (failed(mlir::transform::applyTransforms(
+                    payload, cast<mlir::transform::TransformOpInterface>(transformer_main_sequence),
+                    {}, mlir::transform::TransformOptions(), false))) {
+                return signalPassFailure();
+            }
+        }
 
         transformer.erase();
     }

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<f64>) -> !quantum.bit {
+    %c_0 = stablehlo.constant dense<0> : tensor<i64>
+    %extracted = tensor.extract %c_0[] : tensor<i64>
+    %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
+  }
+
+}