| 34 | |
| 35 | template <typename T, typename Context> |
| 36 | void SetKernel(const Context& dev_ctx, |
| 37 | const DenseTensor& x, |
| 38 | const DenseTensor& source, |
| 39 | const std::vector<int64_t>& dims, |
| 40 | const std::vector<int64_t>& stride, |
| 41 | int64_t offset, |
| 42 | DenseTensor* out) { |
| 43 | auto meta = out->meta(); |
| 44 | meta.dims = DDim(dims.data(), static_cast<int>(dims.size())); |
| 45 | meta.strides = DDim(stride.data(), static_cast<int>(stride.size())); |
| 46 | meta.offset = offset; |
| 47 | if (x.numel() == 0 || source.numel() == 0) { |
| 48 | int64_t out_numel = 1; |
| 49 | for (auto d : dims) { |
| 50 | out_numel *= d; |
| 51 | } |
| 52 | if (source.numel() == 0 && x.numel() != 0) { |
| 53 | // Source is empty but x has storage. Reuse x's storage and apply |
| 54 | // the user-specified meta, matching PyTorch behavior. |
| 55 | if (out_numel == 0) { |
| 56 | // Output has 0 elements — no storage needed, just set meta. |
| 57 | out->set_meta(meta); |
| 58 | out->ShareInplaceVersionCounterWith(x); |
| 59 | return; |
| 60 | } |
| 61 | // If the strided view requires more storage than x provides, |
| 62 | // allocate a larger zero-filled buffer and copy x's data into it |
| 63 | // to avoid out-of-bounds reads on elements beyond x's allocation. |
| 64 | int64_t required_size = ComputeRequiredStorageSize(dims, stride, offset); |
| 65 | if (required_size > x.numel()) { |
| 66 | DenseTensor tmp; |
| 67 | std::vector<int64_t> alloc_shape = {required_size}; |
| 68 | Full<T, Context>(dev_ctx, alloc_shape, 0, &tmp); |
| 69 | if (dev_ctx.GetPlace().GetType() == phi::AllocationType::CPU) { |
| 70 | std::memcpy(tmp.data<T>(), x.data<T>(), x.numel() * sizeof(T)); |
| 71 | } else { |
| 72 | memory_utils::Copy(dev_ctx.GetPlace(), |
| 73 | tmp.data<T>(), |
| 74 | dev_ctx.GetPlace(), |
| 75 | x.data<T>(), |
| 76 | x.numel() * sizeof(T), |
| 77 | nullptr); |
| 78 | } |
| 79 | out->clear(); |
| 80 | *out = DenseTensor{tmp.Holder(), meta}; |
| 81 | } else { |
| 82 | out->set_meta(meta); |
| 83 | } |
| 84 | } else if (source.numel() == 0 && x.numel() == 0 && out_numel != 0) { |
| 85 | // Both x and source are 0-size but user wants non-zero shape. |
| 86 | // Allocate zero-filled storage to avoid null pointer access. |
| 87 | int64_t required_size = ComputeRequiredStorageSize(dims, stride, offset); |
| 88 | DenseTensor tmp; |
| 89 | std::vector<int64_t> alloc_shape = {required_size}; |
| 90 | Full<T, Context>(dev_ctx, alloc_shape, 0, &tmp); |
| 91 | out->clear(); |
| 92 | *out = DenseTensor{tmp.Holder(), meta}; |
| 93 | } else if (source.numel() != 0) { |
nothing calls this directly
no test coverage detected