| 210 | |
| 211 | template<DeviceType device_type, typename T> |
| 212 | void WriteSlice(user_op::KernelComputeContext* ctx, const user_op::Tensor* src, |
| 213 | user_op::Tensor* dst, const SliceContext& slice_ctx, |
| 214 | const bool from_large_to_small) { |
| 215 | const user_op::Tensor* large = from_large_to_small ? src : dst; |
| 216 | const user_op::Tensor* small = from_large_to_small ? dst : src; |
| 217 | // Check physical tensor's shape |
| 218 | for (const auto& split_info : slice_ctx.GetSplitInfo()) { |
| 219 | if (split_info.split_axis != SPLIT_AXIS_FOR_NON_SPLIT) { |
| 220 | CHECK_EQ(large->shape_view().At(split_info.split_axis), split_info.upper - split_info.lower) |
| 221 | << "split_info shape mismatch physical tensor shape"; |
| 222 | } |
| 223 | } |
| 224 | |
| 225 | const std::vector<int64_t> start_attr = ctx->Attr<std::vector<int64_t>>("start"); |
| 226 | const std::vector<int64_t> stop_attr = ctx->Attr<std::vector<int64_t>>("stop"); |
| 227 | const std::vector<int64_t> step_attr = ctx->Attr<std::vector<int64_t>>("step"); |
| 228 | const int64_t ndim = start_attr.size(); |
| 229 | std::vector<int64_t> positive_start_vec(ndim); |
| 230 | std::vector<int64_t> positive_stop_vec(ndim); |
| 231 | |
| 232 | // regulate axis number |
| 233 | std::vector<int64_t> logical_dims(ndim); |
| 234 | { |
| 235 | for (int i = 0; i < ndim; i++) { |
| 236 | if (!slice_ctx.IsAxisPushed(i)) { |
| 237 | // axis is not split, logical shape is same as physical shape |
| 238 | logical_dims[i] = large->shape_view().At(i); |
| 239 | } |
| 240 | } |
| 241 | for (const auto& split_info : slice_ctx.GetSplitInfo()) { |
| 242 | if (split_info.split_axis != SPLIT_AXIS_FOR_NON_SPLIT) { |
| 243 | logical_dims[split_info.split_axis] = split_info.logical_length; |
| 244 | } |
| 245 | } |
| 246 | } |
| 247 | for (int i = 0; i < ndim; i++) { |
| 248 | positive_start_vec[i] = RegulateSliceStart(start_attr[i], logical_dims[i]); |
| 249 | positive_stop_vec[i] = RegulateSliceStop(stop_attr[i], logical_dims[i]); |
| 250 | } |
| 251 | |
| 252 | SliceParams large_slice_param; |
| 253 | std::copy(large->stride().begin(), large->stride().end(), large_slice_param.stride); |
| 254 | SliceParams small_slice_param; |
| 255 | std::copy(small->stride().begin(), small->stride().end(), small_slice_param.stride); |
| 256 | ConstructSliceParamsLarge(slice_ctx, positive_start_vec, positive_stop_vec, step_attr, |
| 257 | large->shape_view(), &large_slice_param); |
| 258 | ConstructSliceParamsSmall(slice_ctx, positive_start_vec, positive_stop_vec, step_attr, |
| 259 | small->shape_view(), &small_slice_param); |
| 260 | CHECK_EQ(large_slice_param.elem_cnt(), small_slice_param.elem_cnt()); |
| 261 | if (large_slice_param.ndim == 0 && small_slice_param.ndim == 0) { |
| 262 | // Copy data directly for scalar tensor |
| 263 | AutoMemcpy(ctx->stream(), dst->mut_dptr<T>(), src->dptr<T>(), sizeof(T), src->mem_case(), |
| 264 | dst->mem_case()); |
| 265 | return; |
| 266 | } |
| 267 | if (from_large_to_small) { |
| 268 | if (small_slice_param.elem_cnt() == small->shape_view().elem_cnt()) { |
| 269 | SliceKernelUtil<device_type, T>::Forward(ctx->stream(), large_slice_param, src->dptr<T>(), |
nothing calls this directly
no test coverage detected