MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / SetPiece

Method SetPiece

tensorflow/compiler/xla/literal.cc:121–146  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

119 : Literal(shape, /*allocate_arrays=*/true) {}
120
121void Literal::SetPiece(const Shape& shape, Piece* piece, bool allocate_arrays) {
122 if (shape.IsTuple()) {
123 for (int i = 0; i < ShapeUtil::TupleElementCount(shape); ++i) {
124 const Shape& subshape = shape.tuple_shapes(i);
125
126 auto child_piece = Piece();
127 child_piece.set_subshape(&subshape);
128
129 SetPiece(subshape, &child_piece, allocate_arrays);
130
131 piece->emplace_back(std::move(child_piece));
132 }
133 } else if (shape.IsArray()) {
134 if (allocate_arrays) {
135 // Literals can be used as DMA targets, which can require alignment. We
136 // force a 16-byte minimum alignment.
137 constexpr int kMinimumAlignment = 16;
138 piece->set_buffer(static_cast<char*>(tensorflow::port::AlignedMalloc(
139 piece->size_bytes(), kMinimumAlignment)));
140 }
141 } else {
142 // If the shape is neither an array nor tuple, then it must be
143 // zero-sized. Otherwise, some memory needs to be allocated for it.
144 CHECK_EQ(piece->size_bytes(), 0);
145 }
146}
147
148Literal::Literal(const Shape& shape, bool allocate_arrays)
149 : MutableLiteralBase() {

Callers

nothing calls this directly

Calls 9

set_subshapeMethod · 0.80
PieceClass · 0.70
AlignedMallocFunction · 0.50
IsTupleMethod · 0.45
tuple_shapesMethod · 0.45
emplace_backMethod · 0.45
IsArrayMethod · 0.45
set_bufferMethod · 0.45
size_bytesMethod · 0.45

Tested by

no test coverage detected