| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | py_executable.cc | 75 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers; in PjRtExecute() local
78 TF_ASSIGN_OR_RETURN(output_buffers, executable_->Execute({args}, options_)); in PjRtExecute()
82 outputs.reserve(output_buffers[0].size()); in PjRtExecute()
83 for (auto& buffer : output_buffers[0]) { in PjRtExecute()
92 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers; in Execute() local
98 TF_ASSIGN_OR_RETURN(output_buffers, in Execute()
103 outputs.reserve(output_buffers[0].size()); in Execute()
104 for (auto& buffer : output_buffers[0]) { in Execute()
114 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers; in ExecuteOnLocalDevices() local
123 TF_ASSIGN_OR_RETURN(output_buffers, in ExecuteOnLocalDevices()
[all …]
|
| D | outfeed_receiver.cc | 412 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers, in SendShutdownOutfeedHeader()
|
| D | outfeed_receiver_test.cc | 50 std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> output_buffers, in CompileAndExecute()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | outfeed_thunk.cc | 42 ShapeTree<std::unique_ptr<OutfeedBuffer>>* output_buffers = in ExecuteOnStream() local
52 const int64 leaf_count = output_buffers->leaf_count(); in ExecuteOnStream()
57 auto output_leaf_it = output_buffers->leaf_begin(); in ExecuteOnStream()
79 ShapeUtil::GetSubshape(output_buffers->shape(), shape_index); in ExecuteOnStream()
|
| D | ir_emission_utils.cc | 801 auto output_buffers = fusion.getOutputBuffers(); in CanEmitFusedDynamicUpdateSliceInPlaceForGpu() local
802 CHECK_EQ(1, output_buffers.size()); in CanEmitFusedDynamicUpdateSliceInPlaceForGpu()
811 auto maybe_rhs = GetAllocationSliceForMlir(output_buffers[0], allocations); in CanEmitFusedDynamicUpdateSliceInPlaceForGpu()
|
| /external/tensorflow/tensorflow/core/tpu/kernels/ |
| D | tpu_execute_op.cc | 439 auto output_buffers = in AllocateOutputTensors() local
442 xla::Shape output_device_shape = output_buffers->buffers.on_device_shape(); in AllocateOutputTensors()
446 stream, &output_buffers->buffers, &output_device_shape)); in AllocateOutputTensors()
475 *buffer = output_buffers->buffers.buffers().element(out_index); in AllocateOutputTensors()
476 *output_buffers->owned_buffers.mutable_element(out_index) = false; in AllocateOutputTensors()
586 return std::move(output_buffers); in AllocateOutputTensors()
765 std::unique_ptr<OutputBuffers> output_buffers, in DoWork()
774 std::make_shared<xla::Literal>(output_buffers->buffers.on_host_shape()); in DoWork()
776 stream, output_buffers->buffers, literal.get(), in DoWork()
|
| /external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
| D | cpu_runtime.cc | 485 std::vector<std::vector<absl::Span<T>>> output_buffers; in DoAllReduce() local
487 output_buffers.reserve(num_participants); in DoAllReduce()
496 output_buffers.emplace_back(); in DoAllReduce()
500 output_buffers.back(); in DoAllReduce()
531 output_buffers[participant_idx][buffer_idx][idx] = out; in DoAllReduce()
664 void** input_buffers, void** output_buffers) { in __xla_cpu_runtime_AllReduce() argument
692 output_buffers[i], xla::ShapeUtil::ByteSizeOf(subshape)); in __xla_cpu_runtime_AllReduce()
|
| D | cpu_runtime.h | 178 void** input_buffers, void** output_buffers);
|
| D | ir_emitter.cc | 1167 llvm::Value* output_buffers = in HandleAllReduceMultipleReplica() local
1195 /*output_buffers=*/b_.CreateBitCast(output_buffers, i8_ptr_type)}, in HandleAllReduceMultipleReplica()
1239 llvm::Value* output_buffers = in HandleAllToAll() local
1256 /*destination_buffers=*/b_.CreateBitCast(output_buffers, i8_ptr_type)}, in HandleAllToAll()
|
| /external/llvm-project/mlir/include/mlir/Dialect/Linalg/IR/ |
| D | LinalgTraits.h | 106 return concreteOp.output_buffers().size() + in getNumOutputs()
|
| D | LinalgStructuredOpsInterface.td | 883 /// Returns all the operands past the inputs, output_buffers and
|
| D | LinalgStructuredOps.td | 499 Variadic<AnyMemRef>:$output_buffers,
|
| /external/llvm-project/mlir/lib/Dialect/Linalg/Transforms/ |
| D | Bufferize.cpp | 198 SmallVector<Value, 2> newOutputBuffers(adaptor.output_buffers().begin(), in matchAndRewrite()
199 adaptor.output_buffers().end()); in matchAndRewrite()
|
| D | DropUnitDims.cpp | 346 insertReshapes(op.output_buffers()); in matchAndRewrite()
|
| /external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/ |
| D | tpu_client.cc | 757 std::vector<std::vector<std::unique_ptr<PyTpuBuffer>>> output_buffers; in ExecuteShardedOnLocalDevices() local
780 TF_ASSIGN_OR_RETURN(output_buffers, ExecuteOnLocalDevices(arg_buffers)); in ExecuteShardedOnLocalDevices()
781 int num_output_buffers = output_buffers[0].size(); in ExecuteShardedOnLocalDevices()
788 std::move(output_buffers[computation][buffer_id])); in ExecuteShardedOnLocalDevices()
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ |
| D | legalize_to_linalg.cc | 186 SmallVector<Value, 4> output_buffers; in matchAndRewrite() local
188 output_buffers.append(args.begin() + num_inputs, args.end()); in matchAndRewrite()
193 output_buffers.push_back( in matchAndRewrite()
198 output_buffers, [](Value v) { return getElementTypeOrSelf(v); })); in matchAndRewrite()
208 loc, op_result_types, inputs, output_buffers, indexing_maps, in matchAndRewrite()
|
| /external/llvm-project/mlir/lib/Dialect/Linalg/IR/ |
| D | LinalgOps.cpp | 502 if (op.inputs().size() + op.output_buffers().size() + in verifyGenericOp()
1614 if (!op.output_buffers().empty()) in printCommonStructuredOpParts()
1615 p << " outs(" << op.output_buffers() << " : " in printCommonStructuredOpParts()
1616 << op.output_buffers().getTypes() << ")"; in printCommonStructuredOpParts()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | layout_assignment.cc | 151 PointsToSet::BufferSet* output_buffers = GetBufferSet(instruction); in OperandBufferForwarded() local
154 return absl::c_any_of(*output_buffers, [&](const LogicalBuffer* b) { in OperandBufferForwarded()
|