fix: copy result attribtues when extracting single ops from secret.generic

When canonicalizing, secret.generic ops are hoisted when they are fully plaintext. But this doesn't copy the attributes of the op onto the new hoisted generic's result attrs, and then those aren't copied to the old generic's new input.

This results in a segfault when running in debug mode, due to a segfault in printing secret.generic arg attributes (that expect the arg attr dict to have arg attrs for each i in the arg list).

PiperOrigin-RevId: 889323354

Author: asraa

lib/Dialect/Secret/IR/BUILD

@@ -102,6 +102,7 @@ cc_library(
         ":types_inc_gen",
         "@heir//lib/Dialect:HEIRInterfaces",
         "@heir//lib/Kernel",
+        "@heir//lib/Utils:AttributeUtils",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ControlFlowInterfaces",
         "@llvm-project//mlir:IR",

lib/Dialect/Secret/IR/SecretOps.cpp

@@ -9,6 +9,7 @@
 
 #include "lib/Dialect/Secret/IR/SecretPatterns.h"
 #include "lib/Dialect/Secret/IR/SecretTypes.h"
+#include "lib/Utils/AttributeUtils.h"
 #include "llvm/include/llvm/ADT/STLExtras.h"          // from @llvm-project
 #include "llvm/include/llvm/ADT/Sequence.h"           // from @llvm-project
 #include "llvm/include/llvm/ADT/SmallVector.h"        // from @llvm-project
@@ -401,13 +402,15 @@ GenericOp GenericOp::removeYieldedValues(ValueRange yieldedValuesToRemove,
            "Cannot remove a value that is not yielded");
   }
 
+  SmallVector<Attribute> newResultAttrs;
   SmallVector<int, 4> indicesToErase;
   for (unsigned int i = 0; i < getYieldOp()->getNumOperands(); ++i) {
     if (std::find(yieldedValuesToRemove.begin(), yieldedValuesToRemove.end(),
                   getYieldOp()->getOperand(i)) != yieldedValuesToRemove.end()) {
       indicesToErase.push_back(i);
     } else {
       remainingResults.push_back(getResult(i));
+      newResultAttrs.push_back(getResultAttrDict(i));
     }
   }
 
@@ -416,6 +419,11 @@ GenericOp GenericOp::removeYieldedValues(ValueRange yieldedValuesToRemove,
   for (int i : llvm::reverse(indicesToErase)) {
     getYieldOp().getValuesMutable().erase(i);
   }
+  // Update the result attr dictionary to remove the deleted results
+  if (this->getAllResultAttrsAttr()) {
+    this->setResultAttrsAttr(
+        ArrayAttr::get(this->getContext(), newResultAttrs));
+  }
 
   auto newResultTypes = llvm::to_vector<4>(
       llvm::map_range(yieldOp.getValues().getTypes(), [](Type t) -> Type {
@@ -435,10 +443,13 @@ GenericOp GenericOp::removeYieldedValues(ArrayRef<int> yieldedIndicesToRemove,
            "Cannot remove an index that is out of range");
   }
 
+  SmallVector<Attribute> newResultAttrs;
+  newResultAttrs.reserve(getNumResults() - yieldedIndicesToRemove.size());
   for (size_t i = 0; i < getYieldOp()->getNumOperands(); ++i) {
     if (std::find(yieldedIndicesToRemove.begin(), yieldedIndicesToRemove.end(),
                   i) == yieldedIndicesToRemove.end()) {
       remainingResults.push_back(getResult(i));
+      newResultAttrs.push_back(getResultAttrDict(i));
     }
   }
 
@@ -447,6 +458,11 @@ GenericOp GenericOp::removeYieldedValues(ArrayRef<int> yieldedIndicesToRemove,
   for (int i : llvm::reverse(yieldedIndicesToRemove)) {
     getYieldOp().getValuesMutable().erase(i);
   }
+  // Update the result attr dictionary to remove the deleted results
+  if (this->getAllResultAttrsAttr()) {
+    this->setResultAttrsAttr(
+        ArrayAttr::get(this->getContext(), newResultAttrs));
+  }
 
   auto newResultTypes = llvm::to_vector<4>(
       llvm::map_range(yieldOp.getValues().getTypes(), [](Type t) -> Type {
@@ -526,25 +542,61 @@ GenericOp GenericOp::extractOpBeforeGeneric(Operation* opToExtract,
         auto* newOp = b.clone(*opToExtract, mp);
         YieldOp::create(b, loc, newOp->getResults());
       });
-  LLVM_DEBUG({
-    llvm::dbgs() << "After adding new single-op generic:\n";
-    newGeneric->getParentOp()->dump();
-  });
-
+  // Set the result attrs of the new generic to be the attrs of the opToExtract.
+  for (auto attr : opToExtract->getAttrs()) {
+    auto attrName = attr.getName();
+    auto attrValue = attr.getValue();
+    if (!attrName.getValue().contains(".") || !attrValue) continue;
+    if (opToExtract->getNumResults() == 1) {
+      setAttributeAssociatedWith(newGeneric.getResult(0), attrName, attrValue);
+    } else if (auto arrayValueAttr = dyn_cast<ArrayAttr>(attrValue);
+               arrayValueAttr.size() == opToExtract->getNumResults()) {
+      // If the opToExtract has multiple result, each result attr has an array
+      // attr elements corresponding to each result.
+      for (auto [i, value] : llvm::enumerate(arrayValueAttr)) {
+        setAttributeAssociatedWith(newGeneric.getResult(i), attrName, value);
+      }
+    }
+  }
   // Once the op is split off into a new generic op, we need to add new
   // operands to the old generic op, add new corresponding block arguments, and
   // replace all uses of the opToExtract's results with the created block
+  // arguments. We also need to copy arg attrs corresponding to the new block
   // arguments.
+  auto oldGenericArgAttrs = this->getAllOperandAttrsAttr();
+  SmallVector<Attribute> newGenericArgAttrs =
+      oldGenericArgAttrs ? llvm::to_vector(oldGenericArgAttrs)
+                         : SmallVector<Attribute>(this->getNumOperands(),
+                                                  ::mlir::DictionaryAttr::get(
+                                                      this->getContext(), {}));
+  for (OpResult newOperand : newGeneric.getResults()) {
+    auto attrDict = newGeneric.getResultAttrDict(newOperand.getResultNumber());
+    newGenericArgAttrs.push_back(
+        attrDict ? attrDict
+                 : ::mlir::DictionaryAttr::get(this->getContext(), {}));
+  }
   SmallVector<Value> oldGenericNewBlockArgs;
   rewriter.modifyOpInPlace(*this, [&]() {
     getInputsMutable().append(newGeneric.getResults());
     for (auto ty : opToExtract->getResultTypes()) {
       BlockArgument arg = getBody()->addArgument(ty, opToExtract->getLoc());
       oldGenericNewBlockArgs.push_back(arg);
     }
+    if (!llvm::all_of(newGenericArgAttrs, [](Attribute attr) {
+          if (attr == nullptr) return true;
+          auto dictAttr = dyn_cast<DictionaryAttr>(attr);
+          return !dictAttr || dictAttr.empty();
+        })) {
+      this->setOperandAttrsAttr(
+          ArrayAttr::get(this->getContext(), newGenericArgAttrs));
+    }
   });
   rewriter.replaceOp(opToExtract, oldGenericNewBlockArgs);
 
+  LLVM_DEBUG({
+    llvm::dbgs() << "After adding new single-op generic:\n";
+    newGeneric->getParentOp()->dump();
+  });
   return newGeneric;
 }
 

lib/Dialect/Secret/IR/SecretPatterns.cpp

@@ -178,6 +178,17 @@ LogicalResult RemoveUnusedGenericArgs::matchAndRewrite(
         body->eraseArgument(i);
         op.getOperation()->eraseOperand(i);
       });
+      auto attrs = op.getAllOperandAttrsAttr();
+      if (attrs) {
+        SmallVector<Attribute> attrList;
+        for (auto [j, attr] : llvm::enumerate(attrs)) {
+          if (j != i) {
+            attrList.push_back(attr);
+          }
+        }
+        op.setOperandAttrsAttr(ArrayAttr::get(op.getContext(), attrList));
+      }
+
       // Ensure the next iteration uses the right arg number
       --i;
     } else if (llvm::any_of(arg.getUsers(), [&](Operation* user) {
@@ -254,6 +265,16 @@ LogicalResult RemoveNonSecretGenericArgs::matchAndRewrite(
         body->eraseArgument(i);
         op.getOperation()->eraseOperand(i);
       });
+      auto attrs = op.getAllOperandAttrsAttr();
+      if (attrs) {
+        SmallVector<Attribute> attrList;
+        for (auto [j, attr] : llvm::enumerate(attrs)) {
+          if (j != i) {
+            attrList.push_back(attr);
+          }
+        }
+        op.setOperandAttrsAttr(ArrayAttr::get(op.getContext(), attrList));
+      }
       i--;
     }
   }

tests/Dialect/Secret/Transforms/canonicalize/generic_attrs.mlir

@@ -0,0 +1,27 @@
+// RUN: heir-opt --canonicalize %s | FileCheck %s
+
+// This is a regression test specifically testing that the result attrs of the assign_layout that is hoisted out is preserved.
+
+#kernel = #secret.kernel<name = "MatvecDiagonal", force = false>
+#layout = #tensor_ext.layout<"{ [i0, i1, i2, i3] -> [ct, slot] : i0 = 0 and ct = 0 and (-25i1 - 5i2 - i3 + slot) mod 128 = 0 and 0 <= i1 <= 3 and 0 <= i2 <= 4 and 0 <= i3 <= 1023 - 25i1 - 5i2 and i3 <= 4 and 0 <= slot <= 1023 and 1024*floor((-128 + 25i1 + 5i2 + i3)/1024) <= -1024 + 25i1 + 5i2 + i3 }">
+#layout1 = #tensor_ext.layout<"{ [i0, i1, i2, i3] -> [ct, slot] : i0 = 0 and i1 = 0 and ct = 0 and (-10i2 - i3 + slot) mod 128 = 0 and 0 <= i2 <= 9 and 0 <= i3 <= 9 and 0 <= slot <= 1023 }">
+#layout2 = #tensor_ext.layout<"{ [i0, i1, i2, i3] -> [ct, slot] : exists (e0, e1, e2, e3, e4, e6: i1 = 0 and 128e6 = -10i2 - i3 + ct + slot - 20e0 - 2e1 and 0 <= i0 <= 3 and 0 <= i2 <= 1 and 0 <= i3 <= 1 and 0 <= ct <= 15 and 0 <= slot <= 1023 and 0 <= e0 <= 4 and 0 <= e1 <= 35 - 5i2 - 10e0 and e1 <= 4 and 0 <= e2 <= 1 and 16*floor((slot)/16) >= slot - 2e3 and slot - 4e2 - 2e3 + 2e4 <= 16*floor((slot)/16) <= 1 + slot - 4e2 - 2e3 + 2e4 and 16*floor((slot)/16) <= 3 + slot - 2e3 and -1 + 25i0 + 4slot + 5e0 + e1 - 2e2 - 8e3 <= 64*floor((slot)/16) <= 25i0 + 4slot + 5e0 + e1 - 2e2 - 8e3 and 50i0 + 7slot + 10e0 + 2e1 - 4e2 - 16e3 + 4e4 <= 112*floor((slot)/16) <= 1 + 50i0 + 7slot + 10e0 + 2e1 - 4e2 - 16e3 + 4e4) }">
+#layout3 = #tensor_ext.layout<"{ [i0, i1, i2, i3] -> [ct, slot] : exists (e1, e2, e3: i0 = 0 and ct = 0 and 0 <= i1 <= 3 and 0 <= i2 <= 4 and 0 <= i3 <= 4 and 0 <= slot <= 99 and 128*floor((-1 - 25i1 - 5i2 - i3)/128) <= -29 - 25i1 - 5i2 - i3 and 0 <= e1 <= 4 and slot <= 2e2 <= 3 + slot and 124 + 25i1 + 5i2 + i3 + 25slot + 128*floor((-1 - 25i1 - 5i2 - i3)/128) - 5e1 <= 50e2 <= 128 + 25i1 + 5i2 + i3 + 25slot + 128*floor((-1 - 25i1 - 5i2 - i3)/128) - 5e1 and -1 - slot + 4e1 + 2e2 <= 2e3 <= -slot + 4e1 + 2e2 and -257 - 50i1 - 10i2 - 2i3 - 49slot - 256*floor((-1 - 25i1 - 5i2 - i3)/128) + 10e1 + 100e2 <= 10e3 <= -256 - 50i1 - 10i2 - 2i3 - 49slot - 256*floor((-1 - 25i1 - 5i2 - i3)/128) + 10e1 + 100e2) }">
+module {
+  // CHECK: @conv2d_nchw
+  // CHECK: tensor_ext.assign_layout
+  // CHECK-SAME: tensor_ext.layout
+  // CHECK: secret.generic
+  // CHECK: return
+  func.func @conv2d_nchw(%arg0: !secret.secret<tensor<4x1x2x2xf32>> {tensor_ext.layout = #layout2}) -> (!secret.secret<tensor<4x1x2x2xf32>> {tensor_ext.layout = #layout2}) {
+    %cst = arith.constant dense<0.000000e+00> : tensor<1x4x5x5xf32>
+    %cst_0 = arith.constant dense<2.500000e-01> : tensor<4x1x2x2xf32>
+    %0 = secret.generic(%arg0: !secret.secret<tensor<4x1x2x2xf32>> {tensor_ext.layout = #layout2}) {
+    ^body(%input0: tensor<4x1x2x2xf32>):
+      %1 = tensor_ext.assign_layout %cst_0 {layout = #layout2, tensor_ext.layout = #layout2} : tensor<4x1x2x2xf32>
+      %2 = arith.addf %input0, %1 {tensor_ext.layout = #layout2} : tensor<4x1x2x2xf32>
+      secret.yield %2 : tensor<4x1x2x2xf32>
+    } -> (!secret.secret<tensor<4x1x2x2xf32>> {tensor_ext.layout = #layout2})
+    return %0 : !secret.secret<tensor<4x1x2x2xf32>>
+  }
+}