Home
last modified time | relevance | path

Searched defs:input_layouts (Results 1 – 25 of 44) sorted by relevance

12

/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/expansions/
H A Dmeta_spmd_expander.cc166 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
222 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
290 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
318 llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands()); in ComputeLayoutBackward() local
454 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
518 llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands()); in ComputeLayoutBackward() local
803 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
906 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
933 llvm::DenseMap<int, Layout> input_layouts(transpose->getNumOperands()); in ComputeLayoutBackward() local
1084 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
[all …]
H A Dslice_spmd_expander.cc544 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
563 llvm::DenseMap<int, Layout> input_layouts(slice_op.getNumOperands()); in ComputeLayoutBackward() local
645 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
684 llvm::DenseMap<int, Layout> input_layouts(strided_slice_op.getNumOperands()); in ComputeLayoutBackward() local
778 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
829 llvm::DenseMap<int, Layout> input_layouts(strided_slice_op.getNumOperands()); in ComputeLayoutBackward() local
929 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
969 llvm::DenseMap<int, Layout> input_layouts( in ComputeLayoutBackward() local
H A Dsplit_spmd_expander.cc101 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
118 llvm::DenseMap<int, Layout> input_layouts(split_op.getNumOperands()); in ComputeLayoutBackward() local
154 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
171 llvm::DenseMap<int, Layout> input_layouts(split_v_op.getNumOperands()); in ComputeLayoutBackward() local
H A Deinsum_spmd_expander.cc44 std::vector<Layout> input_layouts(op->getNumOperands()); in ExpandOp() local
250 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
350 llvm::DenseMap<int, Layout> input_layouts(num_inputs); in ComputeLayoutBackward() local
388 const std::vector<Layout>& input_layouts, mlir::Operation* op, in MaybeRelayoutInputs()
H A Dscatter_spmd_expander.cc155 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in TensorScatterOpComputeLayoutForward()
197 llvm::DenseMap<int, Layout> input_layouts(scatter_op.getNumOperands()); in TensorScatterOpComputeLayoutBackward() local
246 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
H A Drange_spmd_expander.cc40 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
51 llvm::DenseMap<int, Layout> input_layouts; in ComputeLayoutBackward() local
H A Ddtensor_op_spmd_expander.cc222 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
377 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
530 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
H A Dtrivial_spmd_expander.h34 const llvm::DenseMap<int, Layout>& input_layouts) override { in ComputeLayoutForward()
51 const llvm::DenseMap<int, Layout>& input_layouts) override { in ComputeLayoutForward()
H A Dgather_spmd_expander.cc169 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
243 llvm::DenseMap<int, Layout> input_layouts(gather_op.getNumOperands()); in ComputeLayoutBackward() local
469 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
H A Dtrivial_spmd_expander.cc81 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
99 llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands()); in ComputeLayoutBackward() local
H A Dcontrol_flow_spmd_expander.cc50 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
90 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
H A Dreplicated_spmd_expander.cc118 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
135 llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands()); in ComputeLayoutBackward() local
H A Delementwise_spmd_expander.cc124 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
150 llvm::DenseMap<int, Layout> input_layouts; in ComputeLayoutBackward() local
H A Dbias_add_spmd_expander.cc117 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
153 llvm::DenseMap<int, Layout> input_layouts; in ComputeLayoutBackward() local
H A Dresource_spmd_expander.cc220 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
248 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts, in ComputeLayoutBackward()
H A Drandom_op_spmd_expander.cc356 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
373 llvm::DenseMap<int, Layout> input_layouts; in ComputeLayoutBackward() local
H A Dconv_spmd_expander.cc417 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
477 llvm::DenseMap<int, Layout> input_layouts(op->getNumOperands()); in ComputeLayoutBackward() local
H A Ddataparallel_spmd_expander.cc309 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
344 llvm::DenseMap<int, Layout> input_layouts; in ComputeLayoutBackward() local
H A Dsoftmax_spmd_expander.cc354 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
623 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
H A Dsave_restore_spmd_expander.cc697 std::vector<Layout> input_layouts; in ExpandDTensorRestoreV2Op() local
835 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts, in ComputeLayoutForward()
H A Dsave_restore_spmd_expander.h48 const llvm::DenseMap<int, Layout>& input_layouts) override { in ComputeLayoutForward()
H A Dtensorlist_setitem_spmd_expander.cc55 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dspmd_expander.cc151 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts) { in ComputeLayoutForward()
157 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts, in ComputeLayoutForward()
169 mlir::Operation* op, const llvm::DenseMap<int, Layout>& input_layouts, in ComputeLayoutBackward()
/aosp_15_r20/external/pytorch/torch/distributed/tensor/parallel/
H A Dstyle.py89 def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh): argument
191 def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh): argument
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/cc/
H A Ddtensor_ops.cc92 std::vector<std::string> input_layouts; in __anon3c2861b20202() local

12