args_to_remap.insert(std::make_pair(op, return_to_operand));
  });

  // Find all callers
  module->walk([&](CallOpInterface op) {
    auto callable = op.getCallableForCallee();
    mlir::SymbolRefAttr sym = callable.dyn_cast<mlir::SymbolRefAttr>();
    if (!sym) return;
    Operation* func = mlir::SymbolTable::lookupNearestSymbolFrom(op, sym);
    if (!args_to_erase.count(func)) return;