Searched refs:buffer_tree (Results 1 – 12 of 12) sorted by relevance
336 se::Stream* stream, const ShapeTree<MaybeOwningDeviceMemory>& buffer_tree) { in WriteRootTupleIndexTable() argument337 TF_RET_CHECK(buffer_tree.shape().IsTuple()); in WriteRootTupleIndexTable()338 if (ShapeUtil::TupleElementCount(buffer_tree.shape()) == 0) { in WriteRootTupleIndexTable()342 buffer_tree.element({}).AsDeviceMemoryBase(); in WriteRootTupleIndexTable()343 TF_RET_CHECK(GetByteSizeRequirement(buffer_tree.shape()) == in WriteRootTupleIndexTable()347 for (int64_t i = 0; i < ShapeUtil::TupleElementCount(buffer_tree.shape()); in WriteRootTupleIndexTable()349 elements.push_back(buffer_tree.element({i}).AsDeviceMemoryBase()); in WriteRootTupleIndexTable()351 return WriteSingleTupleIndexTable(stream, elements, buffer_tree.shape(), in WriteRootTupleIndexTable()
138 ShapeTree<MaybeOwningDeviceMemory> buffer_tree( in ExecutionInputsFromScopedShapedBuffers() local146 *buffer_tree.mutable_element(index) = se::OwningDeviceMemory{ in ExecutionInputsFromScopedShapedBuffers()150 *buffer_tree.mutable_element(index) = execution_input_buffer; in ExecutionInputsFromScopedShapedBuffers()153 execution_inputs.emplace_back(std::move(buffer_tree)); in ExecutionInputsFromScopedShapedBuffers()
228 const ShapeTree<MaybeOwningDeviceMemory>& buffer_tree);
67 ShapeTree<se::ScopedDeviceMemory<uint8_t>> buffer_tree(literal_shape); in TransferLiteralToInfeed() local68 for (auto& leaf : buffer_tree.leaves()) { in TransferLiteralToInfeed()86 EnqueueDestination(std::move(buffer_tree)); in TransferLiteralToInfeed()
114103 struct buffer_tree { struct114129 struct buffer_tree trees[0]; argument
104308 struct buffer_tree { struct104334 struct buffer_tree trees[0]; argument
128981 struct rb_root buffer_tree; member
57244 name: "buffer_tree"
57309 name: "buffer_tree"
270692 name: "buffer_tree"