[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–60

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

114

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

114

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

114

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–60

Done.

114

Got it. Working on this.

[srishti-pm]

mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp

114

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

114

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

114

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

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

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/include/mlir/Analysis/DataFlow/LivenessAnalysis.h

 ///   (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 has an op with memory effects OR
+///   (1.c) is a non-forwarded call operand.
 ///
 /// 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`.
 struct Liveness : public AbstractSparseLattice {

 public:
   using SparseBackwardDataFlowAnalysis::SparseBackwardDataFlowAnalysis;
 
   void visitOperation(Operation *op, ArrayRef<Liveness *> operands,
                       ArrayRef<const Liveness *> results) override;
 
   void visitBranchOperand(OpOperand &operand) override;
 
+  void visitCallOperand(OpOperand &operand) override;
+
   void setToExitState(Liveness *lattice) override;
 };
 
 //===----------------------------------------------------------------------===//
 // RunLivenessAnalysis
 //===----------------------------------------------------------------------===//
 
 /// Runs liveness analysis on the IR defined by `op`.

mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h

explicit AbstractSparseBackwardDataFlowAnalysis(
       DataFlowSolver &solver, SymbolTableCollection &symbolTable);
 
   /// The operation transfer function. Given the result lattices, this
   /// function is expected to set the operand lattices.
   virtual void visitOperationImpl(
       Operation *op, ArrayRef<AbstractSparseLattice *> operandLattices,
       ArrayRef<const AbstractSparseLattice *> resultLattices) = 0;
 
-  // Visit operands on branch instructions that are not forwarded
+  // 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;
+
   /// Set the given lattice element(s) at control flow exit point(s).
   virtual void setToExitState(AbstractSparseLattice *lattice) = 0;
 
   /// Set the given lattice element(s) at control flow exit point(s).
   void setAllToExitStates(ArrayRef<AbstractSparseLattice *> lattices);
 
   /// Get the lattice element for a value.
   virtual AbstractSparseLattice *getLatticeElement(Value value) = 0;

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 has an op with memory effects OR
+///   (1.c) is a non-forwarded call operand.
 ///
 /// 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`.
 
void LivenessAnalysis::visitBranchOperand(OpOperand &operand) {
   // 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`.
   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);
     }
   } else if (isa<BranchOpInterface>(op)) {
     // When the op is a `BranchOpInterface`, like a `cf.cond_br` or a
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
   // `op` is a `RegionBranchTerminatorOpInterface` or return-like because its
   // non-forwarded operand depends on not only its memory effects/results but
   // also on those of its parent's.
   if (!isa<RegionBranchTerminatorOpInterface>(op) &&
       !op->hasTrait<OpTrait::ReturnLike>())
     return;
   Operation *parentOp = op->getParentOp();
   SmallVector<const Liveness *, 4> parentResultsLiveness;
   for (const Value parentResult : parentOp->getResults())
     parentResultsLiveness.push_back(getLatticeElement(parentResult));
   visitOperation(parentOp, operandLiveness, parentResultsLiveness);
 }
 
+void LivenessAnalysis::visitCallOperand(OpOperand &operand) {
+  // We know (at the moment) and assume (for the future) that `operand` is a
+  // non-forwarded call operand of an op implementing `CallOpInterface`.
+  assert(isa<CallOpInterface>(operand.getOwner()) &&
+         "expected the op to implement `CallOpInterface`");
+
+  // The lattices of the non-forwarded call operands don't get updated like the
+  // forwarded call operands or the non-call operands. Thus they need to be
+  // handled separately. This is where we handle them.
+
+  // This marks values of type (1.c) liveness as "live". A non-forwarded
+  // call operand is live.
+  Liveness *operandLiveness = getLatticeElement(operand.get());
+  propagateIfChanged(operandLiveness, operandLiveness->markLive());
+}
+
 void LivenessAnalysis::setToExitState(Liveness *lattice) {
   // This marks values of type (2) liveness as "live".
   lattice->markLive();
 }
 
 //===----------------------------------------------------------------------===//
 // RunLivenessAnalysis
 //===----------------------------------------------------------------------===//

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);
+        visitCallOperand(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

^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.c) "is a non-forwarded call operand"
+// 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_type_1.c(%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: [callarg2]
+// 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/Analysis/DataFlow/TestSparseBackwardDataFlowAnalysis.cpp

 public:
   using SparseBackwardDataFlowAnalysis::SparseBackwardDataFlowAnalysis;
 
   void visitOperation(Operation *op, ArrayRef<WrittenTo *> operands,
                       ArrayRef<const WrittenTo *> results) override;
 
   void visitBranchOperand(OpOperand &operand) override;
 
+  void visitCallOperand(OpOperand &operand) override;
+
   void setToExitState(WrittenTo *lattice) override { lattice->writes.clear(); }
 };
 
 void WrittenToAnalysis::visitOperation(Operation *op,
                                        ArrayRef<WrittenTo *> operands,
                                        ArrayRef<const WrittenTo *> results) {
   if (auto store = dyn_cast<memref::StoreOp>(op)) {
     SetVector<StringAttr> newWrites;
void WrittenToAnalysis::visitBranchOperand(OpOperand &operand) {
   WrittenTo *lattice = getLatticeElement(operand.get());
   SetVector<StringAttr> newWrites;
   newWrites.insert(
       StringAttr::get(operand.getOwner()->getContext(),
                       "brancharg" + Twine(operand.getOperandNumber())));
   propagateIfChanged(lattice, lattice->addWrites(newWrites));
 }
 
+void WrittenToAnalysis::visitCallOperand(OpOperand &operand) {
+  // Mark call operands as "callarg%d", with %d the operand number.
+  WrittenTo *lattice = getLatticeElement(operand.get());
+  SetVector<StringAttr> newWrites;
+  newWrites.insert(
+      StringAttr::get(operand.getOwner()->getContext(),
+                      "callarg" + Twine(operand.getOperandNumber())));
+  propagateIfChanged(lattice, lattice->addWrites(newWrites));
+}
+
 } // end anonymous namespace
 
 namespace {
 struct TestWrittenToPass
     : public PassWrapper<TestWrittenToPass, OperationPass<>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestWrittenToPass)
 
   StringRef getArgument() const override { return "test-written-to"; }

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);