| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/ |
| H A D | replicated_spmd_expander.cc | 47 TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ReplicatedRelayoutOperandsAndOutputs()
119 TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
133 TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | broadcast_to_spmd_expander.cc | 85 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ExpandOp()
130 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
168 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | gather_spmd_expander.cc | 171 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
241 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
391 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ExpandOp()
471 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
513 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | meta_spmd_expander.cc | 314 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
458 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
503 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
931 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
1030 ExtractDeviceMeshEnclosingCluster(one_hot_op)); in ExpandOp()
1091 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
1116 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
1234 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
|
| H A D | range_spmd_expander.cc | 41 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
49 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | trivial_spmd_expander.cc | 82 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
97 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | fill_spmd_expander.cc | 94 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
103 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | cumsum_spmd_expander.cc | 102 TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
116 TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | dataparallel_spmd_expander.cc | 209 TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in RelayoutOperandsAndOutputs()
310 TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
342 TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | segmentation_spmd_expander.cc | 36 TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
59 TF_ASSIGN_OR_RETURN(const auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | expanddims_spmd_expander.cc | 76 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
108 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | split_spmd_expander.cc | 116 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
169 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | save_restore_spmd_expander.cc | 390 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ExpandSaveV2Op()
500 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ExpandMergeV2Op()
544 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ExpandRestoreV2OpHelper()
850 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
|
| H A D | scatter_spmd_expander.cc | 156 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in TensorScatterOpComputeLayoutForward()
184 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in TensorScatterOpComputeLayoutBackward()
|
| H A D | slice_spmd_expander.cc | 561 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
676 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
821 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
961 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | random_op_spmd_expander.cc | 357 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
372 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | tensorlist_setitem_spmd_expander.cc | 56 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
|
| H A D | sparse_to_dense_spmd_expander.cc | 45 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
|
| H A D | save_restore_spmd_expander.h | 49 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
|
| H A D | tensorlist_reserve_spmd_expander.cc | 64 TF_ASSIGN_OR_RETURN(Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
|
| H A D | dtensor_op_spmd_expander.cc | 347 ExtractDeviceMeshEnclosingCluster(recv_op)); in ExpandOp()
404 ExtractDeviceMeshEnclosingCluster(send_op)); in ExpandOp()
|
| H A D | conv_spmd_expander.cc | 418 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
476 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | einsum_spmd_expander.cc | 253 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutForward()
308 TF_ASSIGN_OR_RETURN(auto mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| H A D | control_flow_spmd_expander.cc | 101 TF_ASSIGN_OR_RETURN(const Mesh mesh, ExtractDeviceMeshEnclosingCluster(op)); in ComputeLayoutBackward()
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/ |
| H A D | layout_parsing.h | 80 StatusOr<Mesh> ExtractDeviceMeshEnclosingCluster(mlir::Operation* op);
|