Home
last modified time | relevance | path

Searched defs:output_layouts (Results 1 – 25 of 45) sorted by relevance

12

/external/tensorflow/tensorflow/dtensor/mlir/expansions/
Dsplit_spmd_expander.cc103 llvm::DenseMap<int, Layout> output_layouts(split_op.getNumResults()); in ComputeLayoutForward() local
114 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
156 llvm::DenseMap<int, Layout> output_layouts(split_v_op.getNumResults()); in ComputeLayoutForward() local
167 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dreplicated_spmd_expander.cc46 const std::vector<Layout>& output_layouts) { in ReplicatedRelayoutOperandsAndOutputs()
121 llvm::DenseMap<int, Layout> output_layouts(op->getNumResults()); in ComputeLayoutForward() local
132 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dtrivial_spmd_expander.h40 const llvm::DenseMap<int, Layout>& output_layouts) override { in ComputeLayoutBackward()
57 const llvm::DenseMap<int, Layout>& output_layouts) override { in ComputeLayoutBackward()
Dtrivial_spmd_expander.cc84 llvm::DenseMap<int, Layout> output_layouts(op->getNumResults()); in ComputeLayoutForward() local
96 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dnullary_spmd_expander.cc119 llvm::DenseMap<int, Layout> output_layouts; in ComputeLayoutForward() local
139 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dcontrol_flow_spmd_expander.cc58 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
97 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Ddataparallel_spmd_expander.cc151 const std::vector<Layout>& output_layouts, in IntermediateBatchLayout()
208 const std::vector<Layout>& output_layouts) { in RelayoutOperandsAndOutputs()
312 llvm::DenseMap<int, Layout> output_layouts; in ComputeLayoutForward() local
341 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dmeta_spmd_expander.cc173 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
233 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
313 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
502 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
836 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
927 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
1112 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
1241 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dsave_restore_spmd_expander.h56 const llvm::DenseMap<int, Layout>& output_layouts) override { in ComputeLayoutBackward()
Dtensorlist_setitem_spmd_expander.cc69 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dscatter_spmd_expander.cc183 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in TensorScatterOpComputeLayoutBackward()
261 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dtensorlist_getitem_spmd_expander.cc56 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Drange_spmd_expander.cc48 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Ddtensor_op_spmd_expander.cc228 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
383 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
542 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dsparse_to_dense_spmd_expander.cc51 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dsave_restore_spmd_expander.cc836 const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutForward()
870 llvm::DenseMap<int, Layout> output_layouts(restore_v2.getNumResults()); in ComputeLayoutForward() local
884 llvm::DenseMap<int, Layout> output_layouts(restore_v2.getNumResults()); in ComputeLayoutForward() local
903 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dtensorlist_reserve_spmd_expander.cc72 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Didentity_n_spmd_expander.cc83 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dqr_spmd_expander.cc105 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dfill_spmd_expander.cc102 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dtop_k_spmd_expander.cc94 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dcumsum_spmd_expander.cc115 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dslice_spmd_expander.cc559 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
673 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
818 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
958 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
Dexpanddims_spmd_expander.cc103 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
/external/tensorflow/tensorflow/dtensor/mlir/
Dspmd_expander.cc158 const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutForward()
163 mlir::Operation* op, const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()
170 const llvm::DenseMap<int, Layout>& output_layouts) { in ComputeLayoutBackward()

12