| 227 | } |
| 228 | |
| 229 | af_err af_assign_gen(af_array* out, const af_array lhs, const dim_t ndims, |
| 230 | const af_index_t* indexs, const af_array rhs_) { |
| 231 | try { |
| 232 | ARG_ASSERT(2, (ndims > 0 && ndims <= AF_MAX_DIMS)); |
| 233 | ARG_ASSERT(3, (indexs != NULL)); |
| 234 | |
| 235 | int track = 0; |
| 236 | vector<af_seq> seqs(AF_MAX_DIMS, af_span); |
| 237 | for (dim_t i = 0; i < ndims; i++) { |
| 238 | if (indexs[i].isSeq) { |
| 239 | track++; |
| 240 | seqs[i] = indexs[i].idx.seq; |
| 241 | } |
| 242 | } |
| 243 | |
| 244 | af_array rhs = rhs_; |
| 245 | if (track == static_cast<int>(ndims)) { |
| 246 | // all indexs are sequences, redirecting to af_assign |
| 247 | return af_assign_seq(out, lhs, ndims, seqs.data(), rhs); |
| 248 | } |
| 249 | |
| 250 | ARG_ASSERT(1, (lhs != 0)); |
| 251 | ARG_ASSERT(4, (rhs != 0)); |
| 252 | |
| 253 | const ArrayInfo& lInfo = getInfo(lhs); |
| 254 | const ArrayInfo& rInfo = getInfo(rhs); |
| 255 | const dim4& lhsDims = lInfo.dims(); |
| 256 | const dim4& rhsDims = rInfo.dims(); |
| 257 | af_dtype lhsType = lInfo.getType(); |
| 258 | af_dtype rhsType = rInfo.getType(); |
| 259 | |
| 260 | if (rhsDims.ndims() == 0) { return af_retain_array(out, lhs); } |
| 261 | |
| 262 | if (lhsDims.ndims() == 0) { |
| 263 | return af_create_handle(out, 0, nullptr, lhsType); |
| 264 | } |
| 265 | |
| 266 | if (ndims == 1 && ndims != static_cast<dim_t>(lInfo.ndims())) { |
| 267 | af_array tmp_in = 0; |
| 268 | af_array tmp_out = 0; |
| 269 | AF_CHECK(af_flat(&tmp_in, lhs)); |
| 270 | AF_CHECK(af_assign_gen(&tmp_out, tmp_in, ndims, indexs, rhs_)); |
| 271 | AF_CHECK( |
| 272 | af_moddims(out, tmp_out, lInfo.ndims(), lInfo.dims().get())); |
| 273 | AF_CHECK(af_release_array(tmp_in)); |
| 274 | // This can run into a double free issue if tmp_in == tmp_out |
| 275 | // The condition ensures release only if both are different |
| 276 | // Issue found on Tegra X1 |
| 277 | if (tmp_in != tmp_out) { AF_CHECK(af_release_array(tmp_out)); } |
| 278 | return AF_SUCCESS; |
| 279 | } |
| 280 | |
| 281 | ARG_ASSERT(1, (lhsType == rhsType)); |
| 282 | ARG_ASSERT(1, (lhsDims.ndims() >= rhsDims.ndims())); |
| 283 | |
| 284 | af_array output = 0; |
| 285 | if (*out != lhs) { |
| 286 | int count = 0; |
no test coverage detected