[Mlir-commits] [mlir] [MLIR][Transform] Allow ApplyRegisteredPassOp to take options as a param (PR #142683)

[Rolf Morel]

================
@@ -786,26 +828,124 @@ DiagnosedSilenceableFailure transform::ApplyRegisteredPassOp::applyToOne(
     return emitDefiniteFailure()
            << "unknown pass or pass pipeline: " << getPassName();
 
-  // Create pass manager and run the pass or pass pipeline.
+  // Create pass manager and add the pass or pass pipeline.
   PassManager pm(getContext());
-  if (failed(info->addToPipeline(pm, getOptions(), [&](const Twine &msg) {
+  if (failed(info->addToPipeline(pm, options, [&](const Twine &msg) {
         emitError(msg);
         return failure();
       }))) {
     return emitDefiniteFailure()
            << "failed to add pass or pass pipeline to pipeline: "
            << getPassName();
   }
-  if (failed(pm.run(target))) {
-    auto diag = emitSilenceableError() << "pass pipeline failed";
-    diag.attachNote(target->getLoc()) << "target op";
-    return diag;
+
+  auto targets = SmallVector<Operation *>(state.getPayloadOps(getTarget()));
+  for (Operation *target : targets) {
+    // Make sure that this transform is not applied to itself. Modifying the
+    // transform IR while it is being interpreted is generally dangerous. Even
+    // more so when applying passes because they may perform a wide range of IR
+    // modifications.
+    DiagnosedSilenceableFailure payloadCheck =
+        ensurePayloadIsSeparateFromTransform(*this, target);
+    if (!payloadCheck.succeeded())
+      return payloadCheck;
+
+    // Run the pass or pass pipeline on the current target operation.
+    if (failed(pm.run(target))) {
+      auto diag = emitSilenceableError() << "pass pipeline failed";
+      diag.attachNote(target->getLoc()) << "target op";
+      return diag;
+    }
   }
 
-  results.push_back(target);
+  // The applied pass will have directly modified the payload IR(s).
+  results.set(llvm::cast<OpResult>(getResult()), targets);
   return DiagnosedSilenceableFailure::success();
 }
 
+static ParseResult parseApplyRegisteredPassOptions(
+    OpAsmParser &parser, ArrayAttr &options,
+    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &dynamicOptions) {
+  auto dynamicOptionMarker = UnitAttr::get(parser.getContext());
+  SmallVector<Attribute> optionsArray;
+
+  auto parseOperandOrString = [&]() -> OptionalParseResult {
+    OpAsmParser::UnresolvedOperand operand;
+    OptionalParseResult parsedOperand = parser.parseOptionalOperand(operand);
+    if (parsedOperand.has_value()) {
+      if (failed(parsedOperand.value()))
+        return failure();
+
+      dynamicOptions.push_back(operand);
+      optionsArray.push_back(
+          dynamicOptionMarker); // Placeholder for knowing where to
+                                // inject the dynamic option-as-param.
+      return success();
+    }
+
+    StringAttr stringAttr;
+    OptionalParseResult parsedStringAttr =
+        parser.parseOptionalAttribute(stringAttr);
+    if (parsedStringAttr.has_value()) {
+      if (failed(parsedStringAttr.value()))
+        return failure();
+      optionsArray.push_back(stringAttr);
+      return success();
+    }
+
+    return std::nullopt;
+  };
+
+  OptionalParseResult parsedOptionsElement = parseOperandOrString();
+  while (parsedOptionsElement.has_value()) {
+    if (failed(parsedOptionsElement.value()))
+      return failure();
+    parsedOptionsElement = parseOperandOrString();
+  }
+
+  if (optionsArray.empty()) {
+    return parser.emitError(parser.getCurrentLocation())
+           << "expected at least one option (either a string or a param)";
+  }
+  options = parser.getBuilder().getArrayAttr(optionsArray);
+  return success();
+}
+
+static void printApplyRegisteredPassOptions(OpAsmPrinter &printer,
+                                            Operation *op, ArrayAttr options,
+                                            ValueRange dynamicOptions) {
+  size_t currentDynamicOptionIdx = 0;
+  for (auto [idx, optionAttr] : llvm::enumerate(options)) {
+    if (idx > 0)
+      printer << " "; // Interleave options separator.
+
+    if (isa<UnitAttr>(optionAttr))
+      printer.printOperand(dynamicOptions[currentDynamicOptionIdx++]);
+    else if (auto strAttr = dyn_cast<StringAttr>(optionAttr))
+      printer.printAttribute(strAttr);
+    else
+      assert(false && "each option should be either a StringAttr or UnitAttr");
----------------
rolfmorel wrote:

Done.

https://github.com/llvm/llvm-project/pull/142683