| 142 | } |
| 143 | |
| 144 | int32_t FlattenConcat::enqueue( |
| 145 | int32_t batchSize, void const* const* inputs, void* const* outputs, void*, cudaStream_t stream) noexcept |
| 146 | { |
| 147 | try |
| 148 | { |
| 149 | PLUGIN_ASSERT(mConcatAxisID != 0); |
| 150 | // mCHW is the first input tensor |
| 151 | auto numConcats = static_cast<int32_t>(pluginInternal::volume(mCHW, /*start*/ 0, /*stop*/ mConcatAxisID - 1)); |
| 152 | |
| 153 | // Num concats will be proportional to number of samples in a batch |
| 154 | if (!mIgnoreBatch) |
| 155 | { |
| 156 | numConcats *= batchSize; |
| 157 | } |
| 158 | |
| 159 | auto* output = static_cast<float*>(outputs[0]); |
| 160 | int32_t offset = 0; |
| 161 | for (int32_t i = 0; i < mNumInputs; ++i) |
| 162 | { |
| 163 | auto const* input = static_cast<float const*>(inputs[i]); |
| 164 | for (int32_t n = 0; n < numConcats; ++n) |
| 165 | { |
| 166 | auto status = cublasScopy(mCublas, mInputConcatAxis[i], input + n * mInputConcatAxis[i], 1, |
| 167 | output + (n * mOutputConcatAxis + offset), 1); |
| 168 | |
| 169 | if (status != CUBLAS_STATUS_SUCCESS) |
| 170 | { |
| 171 | return STATUS_FAILURE; |
| 172 | } |
| 173 | } |
| 174 | offset += mInputConcatAxis[i]; |
| 175 | } |
| 176 | |
| 177 | return STATUS_SUCCESS; |
| 178 | } |
| 179 | catch (std::exception const& e) |
| 180 | { |
| 181 | caughtError(e); |
| 182 | } |
| 183 | return -1; |
| 184 | } |
| 185 | |
| 186 | size_t FlattenConcat::getSerializationSize() const noexcept |
| 187 | { |
nothing calls this directly
no test coverage detected