[MLIR][analysis] Fix call op handling in sparse backward dataflow

[srishti-pm]

Summary

Currently, data in AbstractSparseBackwardDataFlowAnalysis is
considered to flow one-to-one, in order, from the operands of an op
implementing CallOpInterface to the arguments of the function it is
calling.

This understanding of the data flow is inaccurate. The operands of such
an op that forward to the function arguments are obtained using a
method provided by CallOpInterface called getArgOperands().

This commit fixes this bug by using getArgOperands() instead of
getOperands() to get the mapping from operands to function arguments
because not all operands necessarily forward to the function arguments
and even if they do, they don't necessarily have to be in the order in
which they appear in the op. The operands that don't get forwarded are
handled by the newly introduced visitCallOperand() function, which
works analogous to the visitBranchOperand() function.

This fix is also propagated to liveness analysis that earlier relied on
this incorrect implementation of the sparse backward dataflow analysis
framework and corrects some incorrect assumptions made in it.

Extra cleanup: Improved a comment and removed an unnecessary code line.

Signed-off-by: Srishti Srivastava <[e-mail]>

[matthiaskramm]

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

59

I'm having a bit of difficulty parsing this sentence. Could you rephrase/clarify?

111

Is this still correct if the called block only produces non-live values?

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

441

I know the intention is the same as in unforwarded branch operands, but a call is a not a branch.

Add (or rename to) a new function, "visitNonForwardedOperand", "visitInternalOperand", or some such thing?

[srishti-pm]

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

441

Okay. Shall I just call it visitCallOperand()?

[srishti-pm]

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

111

I think so, yes.

So, this is what I thought:-

A function call affects the stack. So, even when a function doesn't return any live results or have any live arguments, it still affects the memory. And thus, a non-forwarded call op operand should always be live. What do you think?

[matthiaskramm]

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

111

I'd argue that stack/memory are architecture specific, and we're at a level of compilation where we're not reasoning about machine-level semantics yet.

After all, we also do inlining, and that also affects the stack (moreover, the stack level of non-dead code!), but we consider that OK.

With that in mind, I think the correct thing to do here is to derive the liveness of each non-forwarded operand by doing the disjunction (meet) of the liveness of all forwarded operands.

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

441

Not sure. visitCallOperand() makes it sound like this would be called for all operands of a call?

[srishti-pm]

Address comments.

[srishti-pm]

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

441

I thought it would be analogous to the visitBranchOperand() function which is only called for the non-forwarded operands of a branch instruction. Or should I modify the naming of both these functions to visitNonForwardedBranchOperand() and visitNonForwardedCallOperand(), respectively?

// Visit operands on branch instructions that are not forwarded.
virtual void visitBranchOperand(OpOperand &operand) = 0;

// Visit operands on call instructions that are not forwarded.
virtual void visitCallOperand(OpOperand &operand) = 0;

[srishti-pm]

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

59

Done.

111

Got it. Working on this.

[srishti-pm]

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

111

With that in mind, I think the correct thing to do here is to derive the liveness of each non-forwarded operand by doing the disjunction (meet) of the liveness of all forwarded operands.

Actually, shouldn't a non-forwarded call op operand be live when a result of the call op is live or the function it calls have an op with memory effects? Its liveness will be independent of the liveness of the forwarded operands, right?

[matthiaskramm]

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

111

You're right, of course. The results determine the liveness.
And good point about memory effects. Unfortunately, that'll have to be recursive: If the function calls any other function which has memory effects, that also means it's live.

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

441

Ok, that's fair. No objection to your naming proposal!

[srishti-pm]

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

111

Marked all non-forwarded call operands "live" conservatively based on our internal discussion. It won't allow very aggressive optimizations but at least it will be safe for optimization (marking something "live" is safer than marking it "non-live").

[srishti-pm]

Rebased to main.

[srishti-pm]

All comments given so far in this patch have been addressed.

Awaiting more review comments and/or approval.

[jcai19]

mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h

46 ↗

(On Diff #548712)

It's not quite clear to me what is 1.b

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

58

I don't fully understand this condition. It reads as if the value is an operand and a block?

[jcai19]

mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h

46 ↗

(On Diff #548712)

It's not quite clear to me what is 1.b

Please ignore.

[srishti-pm]

Address Jian's comments.

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

58

Fixed the language to make it clearer.

[srishti-pm]

All comments are addressed here and it is approved by two reviewers. Will wait one more day before landing this in case other reviewers would like to suggest changes too :)

[srishti-pm]

Rebased to main.

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

 //===- LivenessAnalysis.cpp - Liveness analysis ---------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include <mlir/Analysis/DataFlow/LivenessAnalysis.h>
 
 #include <mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h>
 #include <mlir/Analysis/DataFlow/DeadCodeAnalysis.h>
 #include <mlir/Analysis/DataFlow/SparseAnalysis.h>
 #include <mlir/Analysis/DataFlowFramework.h>
 #include <mlir/IR/Operation.h>
 #include <mlir/IR/Value.h>
+#include <mlir/Interfaces/CallInterfaces.h>
 #include <mlir/Interfaces/SideEffectInterfaces.h>
 #include <mlir/Support/LLVM.h>
 
 using namespace mlir;
 using namespace mlir::dataflow;
 
 //===----------------------------------------------------------------------===//
 // Liveness

 ///   (1) has memory effects OR
 ///   (2) is returned by a public function OR
 ///   (3) is used to compute a value of type (1) or (2).
 /// It is also to be noted that a value could be of multiple types (1/2/3) at
 /// the same time.
 ///
 /// A value "has memory effects" iff it:
 ///   (1.a) is an operand of an op with memory effects OR
 ///   (1.b) is a non-forwarded branch operand and a block where its op could
-///   take the control has an op with memory effects.
+///   take the control is a callable or has an op with memory effects.
 ///
 /// A value `A` is said to be "used to compute" value `B` iff `B` cannot be
 /// computed in the absence of `A`. Thus, in this implementation, we say that
 /// value `A` is used to compute value `B` iff:
 ///   (3.a) `B` is a result of an op with operand `A` OR
 ///   (3.b) `A` is used to compute some value `C` and `C` is used to compute
 ///   `B`.
 
for (const Liveness *r : results) {
     }
     addDependency(const_cast<Liveness *>(r), op);
   }
 }
 
 void LivenessAnalysis::visitBranchOperand(OpOperand &operand) {
   // We know (at the moment) and assume (for the future) that `operand` is a
   // non-forwarded branch operand of a `RegionBranchOpInterface`,
-  // `BranchOpInterface`, `RegionBranchTerminatorOpInterface` or return-like op.
+  // `BranchOpInterface`, `RegionBranchTerminatorOpInterface`, return-like op,
+  // or `CallOpInterface`.
   Operation *op = operand.getOwner();
   assert((isa<RegionBranchOpInterface>(op) || isa<BranchOpInterface>(op) ||
           isa<RegionBranchTerminatorOpInterface>(op) ||
-          op->hasTrait<OpTrait::ReturnLike>()) &&
+          op->hasTrait<OpTrait::ReturnLike>() || isa<CallOpInterface>(op)) &&
          "expected the op to be `RegionBranchOpInterface`, "
-         "`BranchOpInterface`, `RegionBranchTerminatorOpInterface`, or "
-         "return-like");
+         "`BranchOpInterface`, `RegionBranchTerminatorOpInterface`, "
+         "return-like, or `CallOpInterface`");
 
   // The lattices of the non-forwarded branch operands don't get updated like
   // the forwarded branch operands or the non-branch operands. Thus they need
   // to be handled separately. This is where we handle them.
 
-  // This marks values of type (1.b) liveness as "live". A non-forwarded
-  // branch operand will be live if a block where its op could take the control
-  // has an op with memory effects.
-  // Populating such blocks in `blocks`.
+  // Marking values of type (1.b) liveness as "live"...
+
+  // Since the block to which a `CallOpInterface` with a non-forwarded branch
+  // operand takes control is callable, we simply mark these operands as live.
+  if (isa<CallOpInterface>(op)) {
+    Liveness *operandLiveness = getLatticeElement(operand.get());
+    propagateIfChanged(operandLiveness, operandLiveness->markLive());
+    return;
+  }
+
+  // A non-forwarded branch operand will be live if a block where its op could
+  // take the control has an op with memory effects. Populating such blocks in
+  // `blocks`.
   SmallVector<Block *, 4> blocks;
   if (isa<RegionBranchOpInterface>(op)) {
     // When the op is a `RegionBranchOpInterface`, like an `scf.for` or an
     // `scf.index_switch` op, its branch operand controls the flow into this
     // op's regions.
     for (Region &region : op->getRegions()) {
       for (Block &block : region)
         blocks.push_back(&block);
for (Operation &nestedOp : *block) {
         break;
       }
     }
   }
 
   // Now that we have checked for memory-effecting ops in the blocks of concern,
   // we will simply visit the op with this non-forwarded operand to potentially
   // mark it "live" due to type (1.a/3) liveness.
-  if (operand.getOperandNumber() > 0)
-    return;
   SmallVector<Liveness *, 4> operandLiveness;
   operandLiveness.push_back(getLatticeElement(operand.get()));
   SmallVector<const Liveness *, 4> resultsLiveness;
   for (const Value result : op->getResults())
     resultsLiveness.push_back(getLatticeElement(result));
   visitOperation(op, operandLiveness, resultsLiveness);
 
   // We also visit the parent op with the parent's results and this operand if

mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp

if (auto branch = dyn_cast<BranchOpInterface>(op)) {
     // of the branch operation itself (for example the boolean for if/else).
     for (int index : unaccounted.set_bits()) {
       OpOperand &operand = op->getOpOperand(index);
       visitBranchOperand(operand);
     }
     return;
   }
 
-  // For function calls, connect the arguments of the entry blocks
-  // to the operands of the call op.
+  // For function calls, connect the arguments of the entry blocks to the
+  // operands of the call op that are forwarded to these arguments.
   if (auto call = dyn_cast<CallOpInterface>(op)) {
     Operation *callableOp = call.resolveCallable(&symbolTable);
     if (auto callable = dyn_cast_or_null<CallableOpInterface>(callableOp)) {
+      // Not all operands of a call op forward to arguments. Such operands are
+      // stored in `unaccounted`.
+      BitVector unaccounted(op->getNumOperands(), true);
+
+      OperandRange argOperands = call.getArgOperands();
+      MutableArrayRef<OpOperand> argOpOperands =
+          operandsToOpOperands(argOperands);
       Region *region = callable.getCallableRegion();
       if (region && !region->empty()) {
         Block &block = region->front();
-        for (auto [blockArg, operand] :
-             llvm::zip(block.getArguments(), operandLattices)) {
-          meet(operand, *getLatticeElementFor(op, blockArg));
+        for (auto [blockArg, argOpOperand] :
+             llvm::zip(block.getArguments(), argOpOperands)) {
+          meet(getLatticeElement(argOpOperand.get()),
+               *getLatticeElementFor(op, blockArg));
+          unaccounted.reset(argOpOperand.getOperandNumber());
         }
       }
+      // Handle the operands of the call op that aren't forwarded to any
+      // arguments.
+      for (int index : unaccounted.set_bits()) {
+        OpOperand &opOperand = op->getOpOperand(index);
+        visitBranchOperand(opOperand);
+      }
       return;
     }
   }
 
   // When the region of an op implementing `RegionBranchOpInterface` has a
   // terminator implementing `RegionBranchTerminatorOpInterface` or a
   // return-like terminator, the region's successors' arguments flow back into
   // the "successor operands" of this terminator.

mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir

 // RUN: mlir-opt -split-input-file -test-liveness-analysis %s 2>&1 | FileCheck %s
 
 // Positive test: Type (1.a) "is an operand of an op with memory effects"
 // zero is live because it is stored in memory.
 // CHECK-LABEL: test_tag: zero:
 // CHECK-NEXT:  result #0: live
 func.func @test_1_type_1.a(%arg0: memref<i32>) {
   %c0_i32 = arith.constant {tag = "zero"} 0 : i32
   memref.store %c0_i32, %arg0[] : memref<i32>
   return
 }
 
 // -----
 
 // Positive test: Type (1.b) "is a non-forwarded branch operand and a block
-// where its op could take the control has an op with memory effects"
+// where its op could take the control is a callable or has an op with memory
+// effects"
 // %arg2 is live because it can make the control go into a block with a memory
 // effecting op.
 // CHECK-LABEL: test_tag: br:
 // CHECK-NEXT:  operand #0: live
 // CHECK-NEXT:  operand #1: live
 // CHECK-NEXT:  operand #2: live
 func.func @test_2_RegionBranchOpInterface_type_1.b(%arg0: memref<i32>, %arg1: memref<i32>, %arg2: i1) {
   %c0_i32 = arith.constant 0 : i32
   cf.cond_br %arg2, ^bb1(%c0_i32 : i32), ^bb2(%c0_i32 : i32) {tag = "br"}
 ^bb1(%0 : i32):
   memref.store %0, %arg0[] : memref<i32>
   cf.br ^bb3
 ^bb2(%1 : i32):
   memref.store %1, %arg1[] : memref<i32>
   cf.br ^bb3
 ^bb3:
   return
 }
 
 // -----
 
 // Positive test: Type (1.b) "is a non-forwarded branch operand and a block
-// where its op could take the control has an op with memory effects"
+// where its op could take the control is a callable or has an op with memory
+// effects"
 // %arg0 is live because it can make the control go into a block with a memory
 // effecting op.
 // CHECK-LABEL: test_tag: flag:
 // CHECK-NEXT:  operand #0: live
 func.func @test_3_BranchOpInterface_type_1.b(%arg0: i32, %arg1: memref<i32>, %arg2: memref<i32>) {
   %c0_i32 = arith.constant 0 : i32
   cf.switch %arg0 : i32, [
     default: ^bb1,
     42: ^bb2
   ] {tag = "flag"}
 ^bb1:
   memref.store %c0_i32, %arg1[] : memref<i32>
   cf.br ^bb3
 ^bb2:
   memref.store %c0_i32, %arg2[] : memref<i32>
   cf.br ^bb3
 ^bb3:
   return
 }
 
 // -----
 
+func.func private @private(%arg0 : i32, %arg1 : i32) {
+  func.return
+}
+
+// Positive test: Type (1.b) "is a non-forwarded branch operand and a block
+// where its op could take the control is a callable or has an op with memory
+// effects"
+// CHECK-LABEL: test_tag: call
+// CHECK-LABEL:  operand #0: not live
+// CHECK-LABEL:  operand #1: not live
+// CHECK-LABEL:  operand #2: live
+func.func @test_4_CallOpInterface_type_1.b(%arg0: i32, %arg1: i32, %device: i32, %m0: memref<i32>) {
+  test.call_on_device @private(%arg0, %arg1), %device {tag = "call"} : (i32, i32, i32) -> ()
+  return
+}
+
+// -----
+
 // Positive test: Type (2) "is returned by a public function"
 // zero is live because it is returned by a public function.
 // CHECK-LABEL: test_tag: zero:
 // CHECK-NEXT:  result #0: live
-func.func @test_4_type_2() -> (f32){
+func.func @test_5_type_2() -> (f32){
   %0 = arith.constant {tag = "zero"} 0.0 : f32
   return %0 : f32
 }
 
 // -----
 
 // Positive test: Type (3) "is used to compute a value of type (1) or (2)"
 // %arg1 is live because the scf.while has a live result and %arg1 is a

 // to non-live values %0#1 and %arg7.
 // CHECK-LABEL: test_tag: condition:
 // CHECK-NEXT:  operand #0: live
 // CHECK-NEXT:  operand #1: live
 // CHECK-NEXT:  operand #2: not live
 // CHECK-NEXT:  operand #3: live
 // CHECK-LABEL: test_tag: add:
 // CHECK-NEXT:  operand #0: live
-func.func @test_5_RegionBranchTerminatorOpInterface_type_3(%arg0: memref<i32>, %arg1: i1) -> (i32) {
+func.func @test_6_RegionBranchTerminatorOpInterface_type_3(%arg0: memref<i32>, %arg1: i1) -> (i32) {
   %c0_i32 = arith.constant 0 : i32
   %c1_i32 = arith.constant 1 : i32
   %c2_i32 = arith.constant 2 : i32
   %0:3 = scf.while (%arg2 = %c0_i32, %arg3 = %c1_i32, %arg4 = %c2_i32, %arg5 = %c2_i32) : (i32, i32, i32, i32) -> (i32, i32, i32) {
     memref.store %arg4, %arg0[] : memref<i32>
     scf.condition(%arg1) {tag = "condition"} %arg2, %arg3, %arg5 : i32, i32, i32
   } do {
   ^bb0(%arg6: i32, %arg7: i32, %arg8: i32):

 // CHECK-LABEL: test_tag: ten:
 // CHECK-NEXT:  result #0: live
 // CHECK-LABEL: test_tag: one:
 // CHECK-NEXT:  result #0: live
 // CHECK-LABEL: test_tag: x:
 // CHECK-NEXT:  result #0: live
 // CHECK-LABEL: test_tag: y:
 // CHECK-NEXT:  result #0: not live
-func.func @test_6_type_3(%arg0: memref<i32>) {
+func.func @test_7_type_3(%arg0: memref<i32>) {
   %c0 = arith.constant {tag = "zero"} 0 : index
   %c10 = arith.constant {tag = "ten"} 10 : index
   %c1 = arith.constant {tag = "one"} 1 : index
   %x = arith.constant {tag = "x"} 0 : i32
   %y = arith.constant {tag = "y"} 1 : i32
   %0:2 = scf.for %arg1 = %c0 to %c10 step %c1 iter_args(%arg2 = %x, %arg3 = %y) -> (i32, i32) {
     %1 = arith.addi %x, %x : i32
     %2 = func.call @private0(%1) : (i32) -> i32

 // CHECK-NEXT:  result #0: live
 // CHECK-LABEL: test_tag: in_private2:
 // CHECK-NEXT:  operand #0: live
 // CHECK-NEXT:  result #0: live
 // CHECK-LABEL: test_tag: final:
 // CHECK-NEXT:  operand #0: live
 // CHECK-NEXT:  operand #1: live
 // CHECK-NEXT:  result #0: live
-func.func @test_7_type_3(%arg: i32) -> (i32) {
+func.func @test_8_type_3(%arg: i32) -> (i32) {
   %0 = func.call @private1(%arg) : (i32) -> i32
   %final = arith.muli %0, %arg {tag = "final"} : i32
   return %final : i32
 }
 
 // -----
 
 // Negative test: None of the types (1), (2), or (3)
 // zero is not live because it has no effect outside the program: it doesn't
 // affect the memory or the program output.
 // CHECK-LABEL: test_tag: zero:
 // CHECK-NEXT:  result #0: not live
 // CHECK-LABEL: test_tag: one:
 // CHECK-NEXT:  result #0: live
-func.func @test_8_negative() -> (f32){
+func.func @test_9_negative() -> (f32){
   %0 = arith.constant {tag = "zero"} 0.0 : f32
   %1 = arith.constant {tag = "one"} 1.0 : f32
   return %1 : f32
 }
 
 // -----
 
 // Negative test: None of the types (1), (2), or (3)
 // %1 is not live because it has no effect outside the program: it doesn't
 // affect the memory or the program output. Even though it is returned by the
 // function `@private_1`, it is never used by the caller.
 // Note that this test clearly shows how this liveness analysis utility differs
 // from the existing liveness utility present at
 // llvm-project/mlir/include/mlir/Analysis/Liveness.h. The latter marks %1 as
 // live as it exists the block of function `@private_1`, simply because it is
 // computed inside and returned by the block, irrespective of whether or not it
 // is used by the caller.
 // CHECK-LABEL: test_tag: one:
 // CHECK:  result #0: not live
 func.func private @private_1() -> (i32, i32) {
   %0 = arith.constant 0 : i32
   %1 = arith.addi %0, %0 {tag = "one"} : i32
   return %0, %1 : i32, i32
 }
-func.func @test_9_negative() -> (i32) {
+func.func @test_10_negative() -> (i32) {
   %0:2 = func.call @private_1() : () -> (i32, i32)
   return %0#0 : i32
 }

mlir/test/Analysis/DataFlow/test-written-to.mlir

 // CHECK-NEXT:  llvm.call @decl(%arg0) : (i64) -> ()
 // CHECK-NEXT:  llvm.return
 
 llvm.func @decl(i64)
 
 llvm.func @func(%lb : i64) -> () {
   llvm.call @decl(%lb) : (i64) -> ()
   llvm.return
-} 
+}
+
+// -----
+
+func.func private @callee(%arg0 : i32, %arg1 : i32) -> i32 {
+  func.return %arg0 : i32
+}
+
+// CHECK-LABEL: test_tag: a
+// CHECK-LABEL:  operand #0: [b]
+// CHECK-LABEL:  operand #1: []
+// CHECK-LABEL:  operand #2: [brancharg2]
+// CHECK-LABEL:  result #0: [b]
+func.func @test_call_on_device(%arg0: i32, %arg1: i32, %device: i32, %m0: memref<i32>) {
+  %0 = test.call_on_device @callee(%arg0, %arg1), %device {tag = "a"} : (i32, i32, i32) -> (i32)
+  memref.store %0, %m0[] {tag_name = "b"} : memref<i32>
+  return
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

 Operation::operand_range TestCallAndStoreOp::getArgOperands() {
   return getCalleeOperands();
 }
 
 MutableOperandRange TestCallAndStoreOp::getArgOperandsMutable() {
   return getCalleeOperandsMutable();
 }
 
+CallInterfaceCallable TestCallOnDeviceOp::getCallableForCallee() {
+  return getCallee();
+}
+
+void TestCallOnDeviceOp::setCalleeFromCallable(CallInterfaceCallable callee) {
+  setCalleeAttr(callee.get<SymbolRefAttr>());
+}
+
+Operation::operand_range TestCallOnDeviceOp::getArgOperands() {
+  return getForwardedOperands();
+}
+
+MutableOperandRange TestCallOnDeviceOp::getArgOperandsMutable() {
+  return getForwardedOperandsMutable();
+}
+
 void TestStoreWithARegion::getSuccessorRegions(
     std::optional<unsigned> index, ArrayRef<Attribute> operands,
     SmallVectorImpl<RegionSuccessor> &regions) {
   if (!index) {
     regions.emplace_back(&getBody(), getBody().front().getArguments());
   } else {
     regions.emplace_back();
   }

mlir/test/lib/Dialect/Test/TestOps.td

def TestCallAndStoreOp : TEST_Op<"call_and_store",
   let results = (outs
     Variadic<AnyType>:$results
   );
   let assemblyFormat =
     "$callee `(` $callee_operands `)` `,` $address attr-dict "
     "`:` functional-type(operands, results)";
 }
 
+def TestCallOnDeviceOp : TEST_Op<"call_on_device",
+    [DeclareOpInterfaceMethods<CallOpInterface>]> {
+  let arguments = (ins
+    SymbolRefAttr:$callee,
+    Variadic<AnyType>:$forwarded_operands,
+    AnyType:$non_forwarded_device_operand
+  );
+  let results = (outs
+    Variadic<AnyType>:$results
+  );
+  let assemblyFormat =
+    "$callee `(` $forwarded_operands `)` `,` $non_forwarded_device_operand "
+    "attr-dict `:` functional-type(operands, results)";
+}
+
 def TestStoreWithARegion : TEST_Op<"store_with_a_region",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface>,
      SingleBlock]> {
   let arguments = (ins
     Arg<AnyMemRef, "", [MemWrite]>:$address,
     BoolAttr:$store_before_region
   );
   let regions = (region AnyRegion:$body);