Returns the required transient allocation size (in bytes) for a given array, or 0 if it's not a transient array.
| 160 | // Returns the required transient allocation size (in bytes) for a given array, |
| 161 | // or 0 if it's not a transient array. |
| 162 | std::size_t TransientArraySize(const Model& model, const string& array_name, |
| 163 | std::size_t transient_data_alignment) { |
| 164 | if (!IsAllocatableTransientArray(model, array_name)) { |
| 165 | return 0; |
| 166 | } |
| 167 | const auto& array = &model.GetArray(array_name); |
| 168 | CHECK(array->has_shape()) |
| 169 | << "Array '" << array_name << "' doesn't have a shape"; |
| 170 | if (array->data_type == ArrayDataType::kNone) { |
| 171 | // Catch a typical issue at the moment with RNN states |
| 172 | for (const auto& rnn_state : model.flags.rnn_states()) { |
| 173 | if (rnn_state.state_array() == array_name) { |
| 174 | LOG(FATAL) |
| 175 | << "A RNN state array, " << array_name << ", still does not " |
| 176 | << "have a known data type after all graph transformations have " |
| 177 | << "run."; |
| 178 | } |
| 179 | } |
| 180 | LOG(FATAL) << "An array, " << array_name << ", still does not " |
| 181 | << "have a known data type after all graph transformations have " |
| 182 | << "run."; |
| 183 | } |
| 184 | const std::size_t elem_size = ElementSize(array->data_type); |
| 185 | const std::size_t raw_size = |
| 186 | elem_size * RequiredBufferSizeForShape(array->shape()); |
| 187 | const std::size_t rounded_size = |
| 188 | RoundUpToNextMultipleOf(raw_size, transient_data_alignment); |
| 189 | return rounded_size; |
| 190 | } |
| 191 | |
| 192 | // Allocates an array: call this for every array just before the first |
| 193 | // op where it is used. |
no test coverage detected