Retain attributes of original scf::ForOp when folding

[nimiwio]

Summary

ForOpIterArgsFolder and ForOpTensorCastFolder construct a new ForOp, but don't retain any pre-existing attributes.

[mehdi_amini]

It is a bit tricky, because we can't really know what these attributes are here for. For example they could refer to individual operands of the loop, but we just shuffled them (removed some).

[nimiwio]

In D111837#3066067, @mehdi_amini wrote:

It is a bit tricky, because we can't really know what these attributes are here for. For example they could refer to individual operands of the loop, but we just shuffled them (removed some).

Good point, I had not considered attribute references to operands of the loop.. For context: my situation is that an attribute was added to a ForOp and is being silently dropped during canonicalization (it's triggering ForOpIterArgsFolder).

Would it be possible to try and detect and handle this situation, e.g. keep track of what we modify/remove, and then update/omit any attributes when we copy them over from the original ForOp? I'm not very familiar with all the ways in which attributes may refer to loop operands, so not sure if something like that would be feasible. I'm not really even sure of a way to hack around this for my situation, besides just not running canonicalization, but that seems wrong..

[mehdi_amini]

Yes it is messy... Unfortunately our situation with attributes isn't great: we really don't have a good story for this in general :(

[bondhugula]

mlir/lib/Dialect/SCF/SCF.cpp

644

I think we need support on derived op types for a getIntrinsicAttributes() and getExternalAttributes(). In this case, since the op's iter args, results, and iter operands are being trimmed and the op doesn't have any intrinsic attributes that correspond to the cross-iter value flow/init and yield, the compiler developer of this pattern knows that it's correct to propagate all external attributes as is. An inherent assumption is that the transforms and rewrites on the ops have the guarantee that any user-defined and external attributes by design don't have an impact on the core semantics of the op and it should be correct for them to propagate them as is. Had those attributes been core attributes, they'd have to be appropriately handled. A corollary of this is that those attaching user-defined attributes shouldn't expect any special handling and treatment on propagation of those.

So what we need is a getExternalAttributes() and just propagate those. In this case, scf.for has zero intrinsic attributes and so this should be fine. But please a comment/warning on a future update once proper support is available.

[bondhugula]

This revision itself looks good to me modulo the code comment to be added.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

644

I don't think I'm in agreement on the safety here: the safe things generic transformations should do is discard unknown attributes.

[bondhugula]

mlir/lib/Dialect/SCF/SCF.cpp

644

And I would invert it and argue the other way: the *safer* or equally safe thing to do is to *preserve* external attributes. Besides safety or if all else is the same on safety, propagating prevents loss of information and provides the desired functionality to pass pipeline developers that are using these external attributes to propagate short-lived information.

I hope you mean the same thing by "unknown" as I mean by "external" -- these are attributes that are not part of an op's definition and one can't expect any loss is safety in propagating those. If one argues that propagating them is unsafe, one could argue that dropping them is more unsafe.

[mehdi_amini]

mlir/lib/Dialect/SCF/SCF.cpp

644

I don't think I adhere to your "inversion" right now, but let's keep it for Discourse.

[nimiwio]

Thank you both for the discussion.. I'm not quite sure how to proceed on this though. As things stand right now, it seems like it is not safe for a pass to add custom attributes to a ForOp. Is this true for any Ops belonging to a Dialect not "owned" by a pass? The notion of "external" vs "intrinsic" attributes seems to make sense to me but if there's not agreement on that, is there some other way that we can propagate custom information about an Op besides through attributes? I'm not necessarily attached to this change, but I'm just not sure of a workable alternative right now.

[bondhugula]

Discourse discussion here: https://llvm.discourse.group/t/on-querying-an-operations-intrinsic-core-vs-external-user-defined-attributes/4498/6

mlir/lib/Dialect/SCF/SCF.cpp

644

Discussion here: https://llvm.discourse.group/t/on-querying-an-operations-intrinsic-core-vs-external-user-defined-attributes/4498/6

mlir/lib/Dialect/SCF/SCF.cpp

LogicalResult matchAndRewrite(scf::ForOp forOp,
     }
 
     if (!canonicalize)
       return failure();
 
     scf::ForOp newForOp = rewriter.create<scf::ForOp>(
         forOp.getLoc(), forOp.lowerBound(), forOp.upperBound(), forOp.step(),
         newIterArgs);
+    newForOp->setAttrs(forOp->getAttrs());
     Block &newBlock = newForOp.region().front();
 
     // Replace the null placeholders with newly constructed values.
     newBlockTransferArgs[0] = newBlock.getArgument(0); // iv
     for (unsigned idx = 0, collapsedIdx = 0, e = newResultValues.size();
          idx != e; ++idx) {
       Value &blockTransferArg = newBlockTransferArgs[1 + idx];
       Value &newResultVal = newResultValues[idx];
for (OpOperand &opOperand : forOp.getIterOpOperands()) {
     }
     newIterOperands.push_back(opOperand.get());
   }
 
   // 2. Create the new forOp shell.
   scf::ForOp newForOp = rewriter.create<scf::ForOp>(
       forOp.getLoc(), forOp.lowerBound(), forOp.upperBound(), forOp.step(),
       newIterOperands);
+  newForOp->setAttrs(forOp->getAttrs());
   Block &newBlock = newForOp.region().front();
   SmallVector<Value, 4> newBlockTransferArgs(newBlock.getArguments().begin(),
                                              newBlock.getArguments().end());
 
   // 3. Inject an incoming cast op at the beginning of the block for the bbArg
   // corresponding to the `replacement` value.
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(&newBlock, newBlock.begin());

mlir/test/Dialect/SCF/canonicalize.mlir

 // CHECK:           [[V0:%.*]] = select {{.*}}, [[C0]], [[C2]]
 // CHECK:           [[V1:%.*]] = select {{.*}}, [[C1]], [[C3]]
 // CHECK:           return [[V0]], [[V1]] : index
 
 // -----
 
 func private @make_i32() -> i32
 
-func @for_yields_2(%lb : index, %ub : index, %step : index) -> i32 {
+func @for_yields_1(%lb : index, %ub : index, %step : index) -> i32 {
   %a = call @make_i32() : () -> (i32)
   %b = scf.for %i = %lb to %ub step %step iter_args(%0 = %a) -> i32 {
     scf.yield %0 : i32
   }
   return %b : i32
 }
 
-// CHECK-LABEL:   func @for_yields_2
+// CHECK-LABEL:   func @for_yields_1
 //  CHECK-NEXT:     %[[R:.*]] = call @make_i32() : () -> i32
 //  CHECK-NEXT:     return %[[R]] : i32
 
+func @for_yields_1_attr(%lb : index, %ub : index, %step : index) -> i32 {
+  %a = call @make_i32() : () -> (i32)
+  %b = scf.for %i = %lb to %ub step %step iter_args(%0 = %a) -> i32 {
+    %c = call @make_i32() : () -> (i32)
+    scf.yield %0 : i32
+  } {someAttr = "someVal"}
+  return %b : i32
+}
+
+// CHECK-LABEL:   func @for_yields_1_attr
+//  CHECK-NEXT:     %[[a:.*]] = call @make_i32() : () -> i32
+//  CHECK-NEXT:     scf.for {{.*}} {
+//  CHECK-NEXT:       %[[c:.*]] = call @make_i32() : () -> i32
+//  CHECK-NEXT:     } {someAttr = "someVal"}
+//  CHECK-NEXT:     return %[[a]] : i32
+
 func @for_yields_3(%lb : index, %ub : index, %step : index) -> (i32, i32, i32) {
   %a = call @make_i32() : () -> (i32)
   %b = call @make_i32() : () -> (i32)
   %r:3 = scf.for %i = %lb to %ub step %step iter_args(%0 = %a, %1 = %a, %2 = %b) -> (i32, i32, i32) {
     %c = call @make_i32() : () -> (i32)
     scf.yield %0, %c, %2 : i32, i32, i32
   }
   return %r#0, %r#1, %r#2 : i32, i32, i32
func @matmul_on_tensors(%t0: tensor<32x1024xf32>, %t1: tensor<1024x1024xf32>) -> tensor<1024x1024xf32> {
   %c32 = arith.constant 32 : index
   %c1024 = arith.constant 1024 : index
 //   CHECK-NOT: tensor.cast
 //       CHECK: %[[FOR_RES:.*]] = scf.for {{.*}} iter_args(%[[ITER_T0:.*]] = %[[T0]]) -> (tensor<32x1024xf32>) {
 //       CHECK:   %[[CAST:.*]] = tensor.cast %[[ITER_T0]] : tensor<32x1024xf32> to tensor<?x?xf32>
 //       CHECK:   %[[DONE:.*]] = call @do(%[[CAST]]) : (tensor<?x?xf32>) -> tensor<?x?xf32>
 //       CHECK:   %[[UNCAST:.*]] = tensor.cast %[[DONE]] : tensor<?x?xf32> to tensor<32x1024xf32>
 //       CHECK:   scf.yield %[[UNCAST]] : tensor<32x1024xf32>
+//       CHECK: {someAttr = "someVal"}
   %0 = tensor.cast %t0 : tensor<32x1024xf32> to tensor<?x?xf32>
   %1 = scf.for %i = %c0 to %c1024 step %c32 iter_args(%iter_t0 = %0) -> (tensor<?x?xf32>) {
     %2 = call @do(%iter_t0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
     scf.yield %2 : tensor<?x?xf32>
-  }
+  } {someAttr = "someVal"}
 //   CHECK-NOT: tensor.cast
 //       CHECK: %[[RES:.*]] = tensor.insert_slice %[[FOR_RES]] into %[[T1]][0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32>
 //       CHECK: return %[[RES]] : tensor<1024x1024xf32>
   %2 = tensor.cast %1 : tensor<?x?xf32> to tensor<32x1024xf32>
   %res = tensor.insert_slice %2 into %t1[0, 0] [32, 1024] [1, 1] : tensor<32x1024xf32> into tensor<1024x1024xf32>
   return %res : tensor<1024x1024xf32>
 }