Fix support for scf.for in helpers to set/find attrs

This change updates the attribute set/get helpers so that they
can operate on RegionBranchOpInterface (provided the same op
also implements OperandAndResultAttrInterface). This is only
hard for region block arguments, and it required a rather unclear
method to map from a region's block argument to the corresponding
OpOperand of the parent op, in order to figure out which
of the parent op's operands are forwarded to the block argument
via control flow.

PiperOrigin-RevId: 885723146

Author: j2kun

lib/Utils/AttributeUtils.cpp

@@ -11,12 +11,23 @@
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"     // from @llvm-project
 #include "mlir/include/mlir/IR/Value.h"                 // from @llvm-project
 #include "mlir/include/mlir/IR/Visitors.h"              // from @llvm-project
+#include "mlir/include/mlir/Interfaces/ControlFlowInterfaces.h"  // from @llvm-project
 #include "mlir/include/mlir/Interfaces/FunctionInterfaces.h"  // from @llvm-project
 #include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project
 
 namespace mlir {
 namespace heir {
 
+// https://discourse.llvm.org/t/regionbranchopinterface-agnostically-connect-a-region-block-argument-to-parent-opoperand/90199
+static RegionBranchInverseSuccessorMapping invertRegionBranchSuccessorMapping(
+    const RegionBranchSuccessorMapping& operandToInputs) {
+  RegionBranchInverseSuccessorMapping inputToOperands;
+  for (const auto& [operand, inputs] : operandToInputs) {
+    for (Value input : inputs) inputToOperands[input].push_back(operand);
+  }
+  return inputToOperands;
+}
+
 int getOperandNumber(Operation* op, Value value) {
   for (OpOperand& operand : op->getOpOperands()) {
     if (operand.get() == value) return operand.getOperandNumber();
@@ -49,6 +60,23 @@ Attribute findAttributeForBlockArgument(BlockArgument blockArg,
         if (operandNumber == -1) return nullptr;
         return argAttrInterface.getOperandAttr(operandNumber, attrName);
       })
+      .Case<RegionBranchOpInterface>(
+          [&](RegionBranchOpInterface op) -> Attribute {
+            auto argAttrInterface =
+                dyn_cast<OperandAndResultAttrInterface>(op.getOperation());
+            if (!argAttrInterface) return nullptr;
+
+            RegionBranchSuccessorMapping mapping;
+            op.getSuccessorOperandInputMapping(mapping,
+                                               RegionBranchPoint::parent());
+            auto inverseMapping = invertRegionBranchSuccessorMapping(mapping);
+            if (inverseMapping.count(blockArg)) {
+              return argAttrInterface.getOperandAttr(
+                  inverseMapping.lookup(blockArg)[0]->getOperandNumber(),
+                  attrName);
+            }
+            return nullptr;
+          })
       .Case<OperandAndResultAttrInterface>([&](auto op) -> Attribute {
         return op.getOperandAttr(blockArg.getArgNumber(), attrName);
       })
@@ -141,6 +169,21 @@ void setAttributeForBlockArgument(BlockArgument blockArg, StringRef attrName,
         if (operandNumber == -1) return;
         return argAttrInterface.setOperandAttr(operandNumber, attrName, attr);
       })
+      .Case<RegionBranchOpInterface>([&](RegionBranchOpInterface op) {
+        auto argAttrInterface =
+            dyn_cast<OperandAndResultAttrInterface>(op.getOperation());
+        if (!argAttrInterface) return;
+
+        RegionBranchSuccessorMapping mapping;
+        op.getSuccessorOperandInputMapping(mapping,
+                                           RegionBranchPoint::parent());
+        auto inverseMapping = invertRegionBranchSuccessorMapping(mapping);
+        if (inverseMapping.count(blockArg)) {
+          return argAttrInterface.setOperandAttr(
+              inverseMapping.lookup(blockArg)[0]->getOperandNumber(), attrName,
+              attr);
+        }
+      })
       .Case<OperandAndResultAttrInterface>([&](auto op) {
         return op.setOperandAttr(blockArg.getArgNumber(), attrName, attr);
       });

lib/Utils/BUILD

@@ -61,6 +61,7 @@ cc_library(
         "@heir//lib/Dialect:HEIRInterfaces",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:ControlFlowInterfaces",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:FunctionInterfaces",
         "@llvm-project//mlir:IR",