diff --git a/backends/aoti/common_shims_slim.cpp b/backends/aoti/common_shims_slim.cpp index d9e255bf220..e92c9d319ae 100644 --- a/backends/aoti/common_shims_slim.cpp +++ b/backends/aoti/common_shims_slim.cpp @@ -36,6 +36,19 @@ AOTITorchError aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes) { return Error::Ok; } +AOTITorchError +aoti_torch_get_size(Tensor* tensor, int64_t d, int64_t* ret_size) { + if (tensor == nullptr || ret_size == nullptr) { + return Error::InvalidArgument; + } + const auto ndim = static_cast(tensor->dim()); + if (d < -ndim || d >= ndim) { + return Error::InvalidArgument; + } + *ret_size = tensor->size(d); + return Error::Ok; +} + AOTITorchError aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides) { if (tensor == nullptr || ret_strides == nullptr) { return Error::InvalidArgument; diff --git a/backends/aoti/common_shims_slim.h b/backends/aoti/common_shims_slim.h index 0d60f578605..dea91ff8e67 100644 --- a/backends/aoti/common_shims_slim.h +++ b/backends/aoti/common_shims_slim.h @@ -39,6 +39,11 @@ aoti_torch_get_data_ptr(Tensor* tensor, void** ret_data_ptr); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_sizes(Tensor* tensor, int64_t** ret_sizes); +// Generated code reads the size of a result whose shape depends on the data +// through this, one dim at a time. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_get_size(Tensor* tensor, int64_t d, int64_t* ret_size); + AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_strides(Tensor* tensor, int64_t** ret_strides); diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index 5f2879a7b32..fc89ab581bf 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -255,8 +255,8 @@ set(_aoti_cuda_shim_sources runtime/shims/memory.cpp # Only build CUDA-only operator shims when the CUDA toolchain is available. if(NOT EXECUTORCH_BUILD_ROCM AND CMAKE_CUDA_COMPILER) - list(APPEND _aoti_cuda_shim_sources runtime/shims/int4mm.cu - runtime/shims/sort.cu runtime/shims/rand.cu + list(APPEND _aoti_cuda_shim_sources runtime/shims/index.cu + runtime/shims/int4mm.cu runtime/shims/sort.cu runtime/shims/rand.cu ) endif() diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 0bd1599cd5a..1f55a6099a9 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -839,6 +839,8 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: "at::_ops::sort_stable::call": None, "aoti_torch_cuda_sort_stable": None, "aoti_torch_cuda_randint_low_out": None, + "aoti_torch_cuda_index_Tensor": None, + "aoti_torch_index_put_out": None, } @classmethod diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index 9f48cd08607..3294d1ec262 100644 Binary files a/backends/cuda/runtime/aoti_cuda_shims.lib and b/backends/cuda/runtime/aoti_cuda_shims.lib differ diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib.md b/backends/cuda/runtime/aoti_cuda_shims.lib.md index af2fb65ca13..2a561c057c8 100644 --- a/backends/cuda/runtime/aoti_cuda_shims.lib.md +++ b/backends/cuda/runtime/aoti_cuda_shims.lib.md @@ -10,7 +10,7 @@ pin moves and the generated wrapper starts calling something new: ``` nm --defined-only aoti_cuda_shims.lib \ - | sed -n 's/.* T _\?\(aoti_torch_[a-z_0-9]*\)$/\1/p' \ + | sed -n 's/.* T _\?\(aoti_torch_[A-Za-z_0-9]*\)$/\1/p' \ | sort -u > exports.txt # add the new names to exports.txt, then { echo 'LIBRARY aoti_cuda_shims.dll'; echo 'EXPORTS'; sed 's/^/ /' exports.txt; } \ diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 712ed5ff545..ce429a9a65a 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -87,28 +87,6 @@ constexpr char kUseSharedCudaStream[] = "use_shared_cuda_stream"; constexpr char kEnableCudaGraphForMethod[] = "enable_cuda_graph_for_method"; constexpr int kCudaGraphWarmupSteps = 3; constexpr char kWeightSharingAcrossMethods[] = "weight_sharing_across_methods"; - -// A CUDA graph replay copies an output's raw bytes into its ETensor, which is -// right only when every element sits at the same offset in both. Size-1 dims -// never move an element, so their strides may differ. -bool output_is_raw_copyable(const SlimTensor* output, const Tensor* et_output) { - if (!output->is_contiguous() || - output->dim() != static_cast(et_output->dim())) { - return false; - } - // The ETensor still has its previous shape here, so its strides can be stale - // for dynamic shapes; the default dim order is enough on its own. - if (executorch::runtime::is_contiguous_dim_order( - et_output->dim_order().data(), et_output->dim())) { - return true; - } - for (size_t d = 0; d < output->dim(); d++) { - if (output->size(d) != 1 && output->stride(d) != et_output->strides()[d]) { - return false; - } - } - return true; -} } // anonymous namespace bool CudaGraphState::start_warmup() { @@ -1090,6 +1068,21 @@ class ET_EXPERIMENTAL CudaBackend final handle->method_name.c_str()); return Error::NotSupported; } + // Replay copies raw bytes and never resizes, so the ETensor takes the + // captured shape here and keeps it for every replay. + const Error resize_err = + executorch::ET_RUNTIME_NAMESPACE::resize_tensor( + args[i + n_inputs]->toTensor(), slim_outputs[i]->sizes()); + if (resize_err != Error::Ok) { + std::fill(slim_outputs.begin(), slim_outputs.end(), nullptr); + ET_LOG( + Error, + "CUDA graph: cannot resize output %zu of '%s' to its captured " + "shape; this call fails and later calls run without a graph", + i, + handle->method_name.c_str()); + return resize_err; + } } gerr = cudaGraphInstantiate( diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu new file mode 100644 index 00000000000..97335a75ea0 --- /dev/null +++ b/backends/cuda/runtime/shims/index.cu @@ -0,0 +1,1062 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace executorch::backends::cuda { + +namespace c10_slim = executorch::backends::aoti::slim::c10; +using c10_slim::ScalarType; +using executorch::runtime::Error; + +namespace { + +constexpr int kMaxDims = 16; +constexpr int kThreads = 256; + +// One indexed dimension of the tensor being indexed. Its index is an int64 or +// int32 tensor, or the row-major positions of a mask's true elements, from +// which this dimension's coordinate is (position / mask_divisor) % mask_size. +struct DimIndex { + const void* data; + int64_t strides[kMaxDims]; // along the iteration dims, 0 where broadcast + int64_t size; // of the indexed dimension + int64_t stride; // of the indexed dimension + int64_t mask_divisor; // 0 for an index tensor + int64_t mask_size; + bool is_int32; +}; + +// The iteration space of an indexing call is the shape of the indexed result. +// Each point maps to an element of the indexed tensor and to one of the other +// tensor: the output of a gather, the values of a put. The dims the indices +// produce are [broadcast_start, broadcast_start + broadcast_ndim). +struct IndexPlan { + int ndim; + int64_t sizes[kMaxDims]; + int64_t indexed_strides[kMaxDims]; // 0 along the dims the indices produce + int64_t other_strides[kMaxDims]; + int broadcast_start; + int broadcast_ndim; + int64_t broadcast_numel; + int64_t after_numel; // of the dims after the ones the indices produce + bool from_mask[kMaxDims]; // size is a mask's count, unknown at export + int num_indexed; + DimIndex indexed[kMaxDims]; +}; + +// Adds the offsets of `linear` within dims [first, end). +__device__ __forceinline__ void add_offsets( + const IndexPlan& plan, + int first, + int end, + int64_t linear, + int64_t& indexed_offset, + int64_t& other_offset) { + for (int d = end - 1; d >= first; --d) { + const int64_t coordinate = linear % plan.sizes[d]; + linear /= plan.sizes[d]; + indexed_offset += coordinate * plan.indexed_strides[d]; + other_offset += coordinate * plan.other_strides[d]; + } +} + +// Adds the offsets of a point of the dims the indices produce, including the +// element the indices pick. +__device__ __forceinline__ void add_index_offsets( + const IndexPlan& plan, + int64_t point, + int64_t& indexed_offset, + int64_t& other_offset) { + int64_t index_offsets[kMaxDims]; + for (int i = 0; i < plan.num_indexed; ++i) { + index_offsets[i] = 0; + } + for (int d = plan.broadcast_start + plan.broadcast_ndim - 1; + d >= plan.broadcast_start; + --d) { + const int64_t coordinate = point % plan.sizes[d]; + point /= plan.sizes[d]; + other_offset += coordinate * plan.other_strides[d]; + for (int i = 0; i < plan.num_indexed; ++i) { + index_offsets[i] += coordinate * plan.indexed[i].strides[d]; + } + } + for (int i = 0; i < plan.num_indexed; ++i) { + const DimIndex& index = plan.indexed[i]; + int64_t position = index.is_int32 + ? static_cast(index.data)[index_offsets[i]] + : static_cast(index.data)[index_offsets[i]]; + if (index.mask_divisor != 0) { + position = position / index.mask_divisor % index.mask_size; + } + if (position < 0) { + position += index.size; + } + // Checked on the device, as ATen and Inductor's generated code do, since + // a check on the host would have to wait for the stream. + CUDA_KERNEL_ASSERT_MSG( + position >= 0 && position < index.size, "index out of range"); + indexed_offset += position * index.stride; + } +} + +// Adds the offsets of a point of the other dims, the slice each index picks. +__device__ __forceinline__ void add_slice_offsets( + const IndexPlan& plan, + int64_t point, + int64_t& indexed_offset, + int64_t& other_offset) { + const int end = plan.broadcast_start + plan.broadcast_ndim; + add_offsets( + plan, + end, + plan.ndim, + point % plan.after_numel, + indexed_offset, + other_offset); + add_offsets( + plan, + 0, + plan.broadcast_start, + point / plan.after_numel, + indexed_offset, + other_offset); +} + +__device__ __forceinline__ void locate( + const IndexPlan& plan, + int64_t linear, + int64_t& indexed_offset, + int64_t& other_offset) { + indexed_offset = 0; + other_offset = 0; + const int64_t rest = linear / plan.after_numel; + add_index_offsets( + plan, rest % plan.broadcast_numel, indexed_offset, other_offset); + add_slice_offsets( + plan, + rest / plan.broadcast_numel * plan.after_numel + + linear % plan.after_numel, + indexed_offset, + other_offset); +} + +template +__global__ void gather_kernel( + const IndexPlan plan, + int64_t numel, + const T* __restrict__ self, + T* __restrict__ out) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < numel; + linear += static_cast(gridDim.x) * blockDim.x) { + int64_t self_offset; + int64_t out_offset; + locate(plan, linear, self_offset, out_offset); + out[out_offset] = self[self_offset]; + } +} + +// Adds in AccT, ATen's accumulation type, and rounds once. +template +__device__ __forceinline__ void store(T* out, AccT sum) { + *out = kAccumulate ? static_cast(static_cast(*out) + sum) + : static_cast(sum); +} + +// For indices that cannot repeat an element, which is any mix of masks and +// full slices. +template +__global__ void put_kernel( + const IndexPlan plan, + int64_t numel, + T* __restrict__ out, + const T* __restrict__ values) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < numel; + linear += static_cast(gridDim.x) * blockDim.x) { + int64_t out_offset; + int64_t values_offset; + locate(plan, linear, out_offset, values_offset); + store( + out + out_offset, static_cast(values[values_offset])); + } +} + +__global__ void index_offsets_kernel( + const IndexPlan plan, + int64_t points, + int64_t* __restrict__ out_offsets, + int64_t* __restrict__ values_offsets) { + for (int64_t point = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + point < points; + point += static_cast(gridDim.x) * blockDim.x) { + int64_t out_offset = 0; + int64_t values_offset = 0; + add_index_offsets(plan, point, out_offset, values_offset); + out_offsets[point] = out_offset; + values_offsets[point] = values_offset; + } +} + +// After a stable sort of the index points by the element they pick, the +// points that pick one element are a run in index order. Each thread takes +// one element of the slice of one run, adds the run's values in that order +// or keeps the last, and writes once, so the result does not depend on +// scheduling. +template +__global__ void put_sorted_kernel( + const IndexPlan plan, + int64_t points, + int64_t slice_numel, + const int64_t* __restrict__ out_offsets, + const int64_t* __restrict__ values_offsets, + T* __restrict__ out, + const T* __restrict__ values) { + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < points * slice_numel; + linear += static_cast(gridDim.x) * blockDim.x) { + const int64_t i = linear / slice_numel; + const int64_t offset = out_offsets[i]; + if (i > 0 && out_offsets[i - 1] == offset) { + continue; + } + int64_t out_offset = offset; + int64_t values_slice_offset = 0; + add_slice_offsets( + plan, linear % slice_numel, out_offset, values_slice_offset); + // Starts from the element itself, so the additions associate as they do + // one after another. + AccT result = kAccumulate ? static_cast(out[out_offset]) : AccT{}; + for (int64_t j = i; j < points && out_offsets[j] == offset; ++j) { + const auto value = + static_cast(values[values_offsets[j] + values_slice_offset]); + result = kAccumulate ? static_cast(result + value) : value; + } + out[out_offset] = static_cast(result); + } +} + +// A mask's element at a row-major position is set. +struct MaskIsSet { + const uint8_t* data; + int ndim; + int64_t sizes[kMaxDims]; + int64_t strides[kMaxDims]; + + __device__ bool operator()(int64_t position) const { + int64_t offset = 0; + for (int d = ndim - 1; d >= 0; --d) { + offset += position % sizes[d] * strides[d]; + position /= sizes[d]; + } + return data[offset] != 0; + } +}; + +// Stream-ordered scratch, freed on the same stream when it goes out of scope. +class Scratch { + public: + Scratch(int32_t device, cudaStream_t stream) + : device_(device), stream_(stream) {} + + Scratch(const Scratch&) = delete; + Scratch& operator=(const Scratch&) = delete; + + ~Scratch() { + for (void* block : blocks_) { + (void)CudaAllocator::deallocate_stream_ordered(block, device_, stream_); + } + } + + void* allocate(size_t nbytes) { + auto block = CudaAllocator::allocate_stream_ordered( + std::max(nbytes, 1), device_, stream_); + if (!block.ok()) { + return nullptr; + } + blocks_.push_back(block.get()); + return block.get(); + } + + private: + int32_t device_; + cudaStream_t stream_; + std::vector blocks_; +}; + +// Writes the row-major positions of a mask's true elements to scratch and +// reads their count back, which waits for the stream. +Error mask_positions( + const Tensor& mask, + Scratch& scratch, + cudaStream_t stream, + const int64_t** positions, + int64_t* count) { + const auto numel = static_cast(mask.numel()); + *positions = nullptr; + *count = 0; + if (numel == 0) { + return Error::Ok; + } + // A graph capture cannot record the host read of the count below, so refuse + // before allocating anything. + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamIsCapturing(stream, &capture)); + ET_CHECK_OR_RETURN_ERROR( + capture == cudaStreamCaptureStatusNone, + NotSupported, + "index: a mask index cannot run inside a CUDA graph capture"); + // CUB counts items in an int on CUDA 12. + ET_CHECK_OR_RETURN_ERROR( + numel <= INT_MAX, + NotSupported, + "index: masks with more than %d elements are not supported", + INT_MAX); + MaskIsSet is_set{static_cast(mask.data_ptr())}; + is_set.ndim = static_cast(mask.dim()); + for (int d = 0; d < is_set.ndim; ++d) { + is_set.sizes[d] = mask.size(d); + is_set.strides[d] = mask.stride(d); + } + auto* selected = + static_cast(scratch.allocate(numel * sizeof(int64_t))); + auto* selected_count = + static_cast(scratch.allocate(sizeof(int64_t))); + ET_CHECK_OR_RETURN_ERROR( + selected != nullptr && selected_count != nullptr, + MemoryAllocationFailed, + "index: allocating the mask positions failed"); + thrust::counting_iterator first(0); + size_t temp_bytes = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceSelect::If( + nullptr, + temp_bytes, + first, + selected, + selected_count, + numel, + is_set, + stream)); + void* temp = scratch.allocate(temp_bytes); + ET_CHECK_OR_RETURN_ERROR( + temp != nullptr, + MemoryAllocationFailed, + "index: allocating the mask scan scratch failed"); + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceSelect::If( + temp, + temp_bytes, + first, + selected, + selected_count, + numel, + is_set, + stream)); + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync( + count, selected_count, sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamSynchronize(stream)); + *positions = selected; + return Error::Ok; +} + +struct Advanced { + int64_t dim; // of the indexed tensor + bool from_mask; + DimIndex index; + std::vector sizes; + std::vector strides; +}; + +// Lays out an indexing call the way ATen's index and index_put do: masks +// become one index per dimension they cover, all indices broadcast together, +// and the dimensions they produce replace the indexed ones if those are +// adjacent and go first otherwise. Sets *may_repeat when two points can land +// on one element, which only an index tensor allows. +Error plan_indexing( + const Tensor& indexed, + Tensor* const* indices, + int64_t num_indices, + Scratch& scratch, + cudaStream_t stream, + IndexPlan* plan, + bool* may_repeat) { + const auto ndim = static_cast(indexed.dim()); + ET_CHECK_OR_RETURN_ERROR( + ndim <= kMaxDims, + NotSupported, + "index: tensors with more than %d dims are not supported", + kMaxDims); + ET_CHECK_OR_RETURN_ERROR( + num_indices <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + *may_repeat = false; + std::vector advanced; + int64_t dim = 0; + for (int64_t i = 0; i < num_indices; ++i) { + const Tensor* index = indices[i]; + if (index == nullptr) { + ++dim; + continue; + } + ET_CHECK_OR_RETURN_ERROR( + index->is_cuda(), + NotSupported, + "index: index %" PRId64 " is not on the CUDA device", + i); + const auto dtype = index->dtype(); + if (dtype == ScalarType::Bool || dtype == ScalarType::Byte) { + const auto mask_ndim = static_cast(index->dim()); + // A 0-dim mask covers no dims, so it adds no index, as in ATen. + if (mask_ndim == 0) { + continue; + } + ET_CHECK_OR_RETURN_ERROR( + dim + mask_ndim <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + for (int64_t j = 0; j < mask_ndim; ++j) { + ET_CHECK_OR_RETURN_ERROR( + index->size(j) == indexed.size(dim + j), + InvalidArgument, + "index: mask %" PRId64 + " does not match the indexed shape at dim %" PRId64, + i, + dim + j); + } + const int64_t* positions = nullptr; + int64_t count = 0; + ET_CHECK_OK_OR_RETURN_ERROR( + mask_positions(*index, scratch, stream, &positions, &count)); + // Row-major position p has coordinate p / divisor % size along dim j. + std::vector divisors(mask_ndim, 1); + for (int64_t j = mask_ndim - 1; j > 0; --j) { + divisors[j - 1] = divisors[j] * std::max(index->size(j), 1); + } + for (int64_t j = 0; j < mask_ndim; ++j) { + Advanced entry{dim + j, true}; + entry.index.data = positions; + entry.index.mask_divisor = divisors[j]; + entry.index.mask_size = index->size(j); + entry.sizes = {count}; + entry.strides = {1}; + advanced.push_back(entry); + } + dim += mask_ndim; + } else if (dtype == ScalarType::Long || dtype == ScalarType::Int) { + ET_CHECK_OR_RETURN_ERROR( + dim < ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + ET_CHECK_OR_RETURN_ERROR( + static_cast(index->dim()) <= kMaxDims, + NotSupported, + "index: index tensors with more than %d dims are not supported", + kMaxDims); + Advanced entry{dim, false}; + entry.index.data = index->data_ptr(); + entry.index.is_int32 = dtype == ScalarType::Int; + entry.sizes.assign(index->sizes().begin(), index->sizes().end()); + entry.strides.assign(index->strides().begin(), index->strides().end()); + advanced.push_back(entry); + *may_repeat = true; + ++dim; + } else { + ET_LOG( + Error, + "index: index %" PRId64 + " has dtype %d; indices must be int64, int32, bool or uint8", + i, + static_cast(dtype)); + return Error::InvalidArgument; + } + } + ET_CHECK_OR_RETURN_ERROR( + dim <= ndim, + InvalidArgument, + "index: too many indices for a tensor of %" PRId64 " dims", + ndim); + // ATen's advanced indexing rejects a call with only None and 0-dim masks; + // its plan would cover the whole tensor. + ET_CHECK_OR_RETURN_ERROR( + !advanced.empty(), + InvalidArgument, + "index: at least one index must be a tensor that is not a 0-dim mask"); + + size_t broadcast_ndim = 0; + for (const auto& entry : advanced) { + broadcast_ndim = std::max(broadcast_ndim, entry.sizes.size()); + } + std::vector broadcast(broadcast_ndim, 1); + std::vector masked(broadcast_ndim, false); + std::vector fixed(broadcast_ndim, false); + for (const auto& entry : advanced) { + const size_t lead = broadcast_ndim - entry.sizes.size(); + for (size_t d = 0; d < entry.sizes.size(); ++d) { + if (entry.from_mask) { + masked[lead + d] = true; + } else if (entry.sizes[d] != 1) { + fixed[lead + d] = true; + } + int64_t& size = broadcast[lead + d]; + ET_CHECK_OR_RETURN_ERROR( + entry.sizes[d] == 1 || size == 1 || entry.sizes[d] == size, + InvalidArgument, + "index: the indices cannot be broadcast together"); + if (size == 1) { + size = entry.sizes[d]; + } + } + } + + const bool adjacent = advanced.back().dim - advanced.front().dim + 1 == + static_cast(advanced.size()); + const int64_t first = advanced.front().dim; + std::vector is_indexed(ndim, false); + for (const auto& entry : advanced) { + is_indexed[entry.dim] = true; + } + std::vector plain_dims; + for (int64_t d = 0; d < ndim; ++d) { + if (!is_indexed[d]) { + plain_dims.push_back(d); + } + } + // Dims of the indexed tensor in iteration order, -1 for each broadcast dim. + std::vector order; + for (int64_t d : plain_dims) { + if (adjacent && d < first) { + order.push_back(d); + } + } + const auto broadcast_start = static_cast(order.size()); + order.insert(order.end(), broadcast_ndim, -1); + for (int64_t d : plain_dims) { + if (!adjacent || d > first) { + order.push_back(d); + } + } + ET_CHECK_OR_RETURN_ERROR( + order.size() <= static_cast(kMaxDims), + NotSupported, + "index: results with more than %d dims are not supported", + kMaxDims); + + *plan = IndexPlan{}; + plan->ndim = static_cast(order.size()); + plan->broadcast_start = static_cast(broadcast_start); + plan->broadcast_ndim = static_cast(broadcast_ndim); + plan->broadcast_numel = 1; + for (int64_t size : broadcast) { + plan->broadcast_numel *= size; + } + // As in ATen, an index into a dim of size 0 is out of range unless there are + // no index points. The result is empty, so no kernel would assert it. + for (const auto& entry : advanced) { + ET_CHECK_OR_RETURN_ERROR( + plan->broadcast_numel == 0 || indexed.size(entry.dim) > 0, + InvalidArgument, + "index: index out of range for dim %" PRId64 " of size 0", + entry.dim); + } + plan->after_numel = 1; + for (size_t k = static_cast(broadcast_start) + broadcast_ndim; + k < order.size(); + ++k) { + plan->after_numel *= indexed.size(order[k]); + } + for (size_t k = 0; k < order.size(); ++k) { + if (order[k] >= 0) { + plan->sizes[k] = indexed.size(order[k]); + plan->indexed_strides[k] = indexed.stride(order[k]); + } else { + plan->sizes[k] = broadcast[k - broadcast_start]; + plan->from_mask[k] = + masked[k - broadcast_start] && !fixed[k - broadcast_start]; + } + } + plan->num_indexed = static_cast(advanced.size()); + for (size_t i = 0; i < advanced.size(); ++i) { + const auto& entry = advanced[i]; + DimIndex& index = plan->indexed[i]; + index = entry.index; + index.size = indexed.size(entry.dim); + index.stride = indexed.stride(entry.dim); + const size_t lead = broadcast_ndim - entry.sizes.size(); + for (size_t d = 0; d < entry.sizes.size(); ++d) { + if (entry.sizes[d] != 1) { + index.strides[broadcast_start + lead + d] = entry.strides[d]; + } + } + } + return Error::Ok; +} + +int64_t plan_numel(const IndexPlan& plan) { + int64_t numel = 1; + for (int d = 0; d < plan.ndim; ++d) { + numel *= plan.sizes[d]; + } + return numel; +} + +unsigned int grid_size(int64_t numel) { + return static_cast( + std::min((numel + kThreads - 1) / kThreads, 65535)); +} + +// The meta function decided the layout with each mask's count unknown, so a +// count is never taken to be 1, and the result is only known to have fewer +// than 2 elements when another size is 0. +bool is_one(const IndexPlan& plan, int d) { + return !plan.from_mask[d] && plan.sizes[d] == 1; +} + +bool is_contiguous(const IndexPlan& plan, const int64_t* strides) { + bool any_unknown = false; + int64_t numel = 1; + for (int d = 0; d < plan.ndim; ++d) { + if (!plan.from_mask[d] && plan.sizes[d] == 0) { + return true; + } + any_unknown |= plan.from_mask[d]; + numel *= plan.sizes[d]; + } + if (!any_unknown && numel < 2) { + return true; + } + int64_t expected = 1; + for (int d = plan.ndim - 1; d >= 0; --d) { + if (is_one(plan, d)) { + continue; + } + if (strides[d] != expected) { + return false; + } + expected *= plan.sizes[d]; + } + return true; +} + +bool is_channels_last_contiguous( + const IndexPlan& plan, + const int64_t* strides) { + if (plan.ndim != 4) { + return false; + } + int64_t expected = 1; + for (int d : {1, 3, 2, 0}) { + if (!is_one(plan, d)) { + if (strides[d] != expected) { + return false; + } + expected *= plan.sizes[d]; + } + } + return true; +} + +// The output strides index.Tensor's meta function gives, which Inductor bakes +// into the code that reads the result: those of an elementwise output for an +// input with `strides` (compute_elementwise_output_logical_to_physical_perm in +// torch/_prims_common). +std::vector output_strides(const IndexPlan& plan, bool self_is_empty) { + const int64_t* sizes = plan.sizes; + const int64_t* strides = plan.indexed_strides; + const int ndim = plan.ndim; + std::vector perm(ndim); + std::iota(perm.begin(), perm.end(), 0); + const bool contiguous = is_contiguous(plan, strides); + const bool channels_last = is_channels_last_contiguous(plan, strides); + if (self_is_empty || ndim < 2 || (contiguous && !channels_last)) { + // Row-major. + } else if (channels_last && !contiguous) { + perm = {0, 2, 3, 1}; + } else { + auto ge = [](int64_t a, int64_t b) { + return b == 0 || (a != 0 && (a >= b || a % b == 0)); + }; + auto should_swap = [&](int a, int b) { + const int64_t stride_a = strides[a]; + const int64_t stride_b = strides[b]; + if (stride_a == 0 || stride_b == 0) { + return 0; + } + if (stride_a == stride_b) { + return ge(sizes[b], sizes[a]) ? 0 : 1; + } + if (ge(stride_b, stride_a)) { + return -1; + } + return ge(stride_a, stride_b) ? 1 : 0; + }; + std::reverse(perm.begin(), perm.end()); + for (int i = 1; i < ndim; ++i) { + int dim1 = i; + for (int dim0 = i - 1; dim0 >= 0; --dim0) { + const int comparison = should_swap(perm[dim0], perm[dim1]); + if (comparison > 0) { + std::swap(perm[dim0], perm[dim1]); + dim1 = dim0; + } else if (comparison < 0) { + break; + } + } + } + std::reverse(perm.begin(), perm.end()); + } + std::vector result(ndim); + int64_t expected = 1; + for (int k = ndim - 1; k >= 0; --k) { + result[perm[k]] = expected; + expected *= std::max(sizes[perm[k]], 1); + } + return result; +} + +template +Error gather( + const IndexPlan& plan, + int64_t numel, + const void* self, + void* out, + cudaStream_t stream) { + gather_kernel<<>>( + plan, numel, static_cast(self), static_cast(out)); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + +template +Error put( + const IndexPlan& plan, + int64_t numel, + bool may_repeat, + void* out_data, + const void* values_data, + Scratch& scratch, + cudaStream_t stream) { + auto* out = static_cast(out_data); + const auto* values = static_cast(values_data); + if (!may_repeat) { + put_kernel + <<>>(plan, numel, out, values); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; + } + // One key per index point, so the scratch and the sort grow with the + // number of indices, not with the elements they write. + const int64_t points = plan.broadcast_numel; + // CUB counts items in an int on CUDA 12. + ET_CHECK_OR_RETURN_ERROR( + points <= INT_MAX, + NotSupported, + "index_put: more than %d index points are not supported", + INT_MAX); + auto* offsets = + static_cast(scratch.allocate(4 * points * sizeof(int64_t))); + ET_CHECK_OR_RETURN_ERROR( + offsets != nullptr, + MemoryAllocationFailed, + "index_put: allocating the offsets failed"); + int64_t* out_offsets = offsets; + int64_t* values_offsets = offsets + points; + int64_t* sorted_out_offsets = offsets + 2 * points; + int64_t* sorted_values_offsets = offsets + 3 * points; + index_offsets_kernel<<>>( + plan, points, out_offsets, values_offsets); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + const auto count = static_cast(points); + size_t temp_bytes = 0; + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + nullptr, + temp_bytes, + out_offsets, + sorted_out_offsets, + values_offsets, + sorted_values_offsets, + count, + 0, + static_cast(sizeof(int64_t) * 8), + stream)); + void* temp = scratch.allocate(temp_bytes); + ET_CHECK_OR_RETURN_ERROR( + temp != nullptr, + MemoryAllocationFailed, + "index_put: allocating the sort scratch failed"); + ET_CUDA_CHECK_OR_RETURN_ERROR(cub::DeviceRadixSort::SortPairs( + temp, + temp_bytes, + out_offsets, + sorted_out_offsets, + values_offsets, + sorted_values_offsets, + count, + 0, + static_cast(sizeof(int64_t) * 8), + stream)); + put_sorted_kernel + <<>>( + plan, + points, + numel / points, + sorted_out_offsets, + sorted_values_offsets, + out, + values); + ET_CUDA_KERNEL_LAUNCH_CHECK_OR_RETURN_ERROR(); + return Error::Ok; +} + +} // namespace + +#ifdef __cplusplus +extern "C" { +#endif + +AOTITorchError aoti_torch_cuda_index_Tensor( + Tensor* self, + Tensor* const** indices, + int64_t indices_len_, + Tensor** ret0) { + ET_CHECK_OR_RETURN_ERROR( + self != nullptr && ret0 != nullptr && + (indices != nullptr || indices_len_ == 0), + InvalidArgument, + "aoti_torch_cuda_index_Tensor: null argument"); + ET_CHECK_OR_RETURN_ERROR( + self->is_cuda(), + NotSupported, + "aoti_torch_cuda_index_Tensor: self is not on the CUDA device"); + const auto device = static_cast(self->device_index()); + auto stream_result = getCurrentCUDAStream(self->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), + Internal, + "aoti_torch_cuda_index_Tensor: failed to get CUDA stream"); + cudaStream_t stream = stream_result.get(); + + std::vector index_tensors(indices_len_); + for (int64_t i = 0; i < indices_len_; ++i) { + index_tensors[i] = indices[i] == nullptr ? nullptr : *indices[i]; + } + Scratch scratch(device, stream); + IndexPlan plan; + bool may_repeat = false; + ET_CHECK_OK_OR_RETURN_ERROR(plan_indexing( + *self, + index_tensors.data(), + indices_len_, + scratch, + stream, + &plan, + &may_repeat)); + + const std::vector sizes(plan.sizes, plan.sizes + plan.ndim); + const std::vector strides = output_strides(plan, self->numel() == 0); + Tensor* result = nullptr; + ET_CHECK_OK_OR_RETURN_ERROR(aoti_torch_empty_strided( + plan.ndim, + sizes.data(), + strides.data(), + static_cast(self->dtype()), + static_cast(c10_slim::DeviceType::CUDA), + device, + &result)); + // The caller takes ownership of *ret0 only when this call succeeds. + std::unique_ptr owner(result, [](Tensor* tensor) { + (void)aoti_torch_delete_tensor_object(tensor); + }); + std::copy(strides.begin(), strides.end(), plan.other_strides); + + const int64_t numel = plan_numel(plan); + if (numel > 0) { + const void* src = self->data_ptr(); + void* dst = result->data_ptr(); + switch (self->itemsize()) { + case 1: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 2: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 4: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + case 8: + ET_CHECK_OK_OR_RETURN_ERROR( + gather(plan, numel, src, dst, stream)); + break; + default: + ET_LOG( + Error, + "aoti_torch_cuda_index_Tensor: unsupported element size %zu", + self->itemsize()); + return Error::NotSupported; + } + } + *ret0 = owner.release(); + return Error::Ok; +} + +AOTITorchError aoti_torch_index_put_out( + Tensor* out, + Tensor* self, + Tensor* const* indices, + const uint32_t num_indices, + Tensor* values, + bool accumulate) { + ET_CHECK_OR_RETURN_ERROR( + out != nullptr && values != nullptr && + (indices != nullptr || num_indices == 0), + InvalidArgument, + "aoti_torch_index_put_out: null argument"); + ET_CHECK_OR_RETURN_ERROR( + out == self, + NotSupported, + "aoti_torch_index_put_out: out must be self, as AOTInductor calls it"); + ET_CHECK_OR_RETURN_ERROR( + out->is_cuda() && values->is_cuda(), + NotSupported, + "aoti_torch_index_put_out: tensors must be on the CUDA device"); + ET_CHECK_OR_RETURN_ERROR( + values->dtype() == out->dtype(), + InvalidArgument, + "aoti_torch_index_put_out: values and out must have the same dtype"); + const auto device = static_cast(out->device_index()); + auto stream_result = getCurrentCUDAStream(out->device_index()); + ET_CHECK_OR_RETURN_ERROR( + stream_result.ok(), + Internal, + "aoti_torch_index_put_out: failed to get CUDA stream"); + cudaStream_t stream = stream_result.get(); + + Scratch scratch(device, stream); + IndexPlan plan; + bool may_repeat = false; + ET_CHECK_OK_OR_RETURN_ERROR(plan_indexing( + *out, indices, num_indices, scratch, stream, &plan, &may_repeat)); + + const auto values_ndim = static_cast(values->dim()); + ET_CHECK_OR_RETURN_ERROR( + values_ndim <= plan.ndim, + InvalidArgument, + "aoti_torch_index_put_out: values cannot be broadcast to the indexed shape"); + const int lead = plan.ndim - values_ndim; + for (int d = 0; d < values_ndim; ++d) { + const int64_t size = values->size(d); + ET_CHECK_OR_RETURN_ERROR( + size == 1 || size == plan.sizes[lead + d], + InvalidArgument, + "aoti_torch_index_put_out: values cannot be broadcast to the indexed shape"); + plan.other_strides[lead + d] = size == 1 ? 0 : values->stride(d); + } + + const int64_t numel = plan_numel(plan); + if (numel == 0) { + return Error::Ok; + } + void* dst = out->data_ptr(); + const void* src = values->data_ptr(); + if (!accumulate) { + switch (out->itemsize()) { + case 1: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 2: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 4: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case 8: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + default: + ET_LOG( + Error, + "aoti_torch_index_put_out: unsupported element size %zu", + out->itemsize()); + return Error::NotSupported; + } + } + // The accumulation types are ATen's. + switch (out->dtype()) { + case ScalarType::Float: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Half: + return put<__half, float, true>( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::BFloat16: + return put<__nv_bfloat16, float, true>( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Long: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Int: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Short: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Char: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Byte: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + case ScalarType::Bool: + return put( + plan, numel, may_repeat, dst, src, scratch, stream); + default: + ET_LOG( + Error, + "aoti_torch_index_put_out: unsupported dtype %d", + static_cast(out->dtype())); + return Error::NotSupported; + } +} + +#ifdef __cplusplus +} +#endif + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/shims/index.h b/backends/cuda/runtime/shims/index.h new file mode 100644 index 00000000000..3476eaa8c88 --- /dev/null +++ b/backends/cuda/runtime/shims/index.h @@ -0,0 +1,82 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::cuda { + +using executorch::backends::aoti::AOTITorchError; +using executorch::backends::aoti::Tensor; + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * Advanced indexing, self[indices], with the semantics of aten::index.Tensor. + * + * Inductor calls this fallback when an index is a mask, so the number of true + * elements decides the output size. Reading that count waits for the stream, + * so a CUDA graph cannot capture this call. As in ATen, an index out of range + * triggers a device-side assert, and an index into a dimension of size 0 is + * an error unless the indices select nothing. + * + * @param self Tensor to index (CUDA device). It, each index and the result + * can have up to 16 dims. + * @param indices indices_len_ pointers to index tensors, null for a full + * slice. An index is int64 or int32, or a bool or uint8 mask that covers as + * many dimensions of self as it has. At least one entry must be an index that + * is not a 0-dim mask. + * @param indices_len_ Number of entries in indices + * @param ret0 Output: new tensor, strided the way Inductor expects + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_cuda_index_Tensor( + Tensor* self, + Tensor* const** indices, + int64_t indices_len_, + Tensor** ret0); + +/** + * out[indices] = values, or out[indices] += values when accumulate is set, + * with the semantics of aten::index_put_. + * + * Inductor calls this fallback, with out and self the same tensor, for an + * index that is a mask, for accumulation into a dtype its generated code + * cannot add atomically, and in deterministic mode. Writes that land on the + * same element are applied in index order, and half and bfloat16 values are + * summed in float, so the result is deterministic and small terms are not + * lost. A mask makes this call wait for the stream. Indices are checked as + * for aoti_torch_cuda_index_Tensor. + * + * @param out Tensor to write (CUDA device, up to 16 dims, as for + * aoti_torch_cuda_index_Tensor) + * @param self Must be out; AOTInductor passes the tensor it updates as both + * @param indices num_indices index tensors, null for a full slice, with the + * same dtypes as for aoti_torch_cuda_index_Tensor + * @param num_indices Number of entries in indices + * @param values Values to write, broadcastable to the indexed shape, same + * dtype as out + * @param accumulate If true, add values instead of writing them + * @return AOTITorchError + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_index_put_out( + Tensor* out, + Tensor* self, + Tensor* const* indices, + const uint32_t num_indices, + Tensor* values, + bool accumulate); + +#ifdef __cplusplus +} +#endif + +} // namespace executorch::backends::cuda diff --git a/backends/cuda/runtime/targets.bzl b/backends/cuda/runtime/targets.bzl index 5ff2b15867c..d5635100b62 100644 --- a/backends/cuda/runtime/targets.bzl +++ b/backends/cuda/runtime/targets.bzl @@ -35,6 +35,7 @@ def define_common_targets(is_fbcode = False): name = "runtime_shims", srcs = [ "shims/cuda_guard.cpp", + "shims/index.cu", "shims/int4mm.cu", "shims/memory.cpp", "shims/rand.cu", @@ -43,6 +44,7 @@ def define_common_targets(is_fbcode = False): ], headers = [ "shims/cuda_guard.h", + "shims/index.h", "shims/int4mm.cuh", "shims/int4mm.h", "shims/memory.h", diff --git a/backends/cuda/runtime/test/test_cuda_output_copy.cpp b/backends/cuda/runtime/test/test_cuda_output_copy.cpp index 1952b23a18f..e1b72bd2f17 100644 --- a/backends/cuda/runtime/test/test_cuda_output_copy.cpp +++ b/backends/cuda/runtime/test/test_cuda_output_copy.cpp @@ -12,7 +12,7 @@ namespace slim = ::executorch::backends::aoti::slim; using namespace ::executorch::runtime; -using ::executorch::backends::cuda::copy_slimtensor_to_device_etensor_async; +using namespace ::executorch::backends::cuda; namespace { @@ -115,4 +115,74 @@ TEST_F(CudaOutputCopyTest, RejectsAViewOutsideItsStorage) { Error::InvalidArgument); } +// Whether an output with these sizes and strides can be raw-copied into an +// ETensor that has this shape, dim order and strides from an earlier call. +bool raw_copyable( + std::initializer_list output_sizes, + std::initializer_list output_strides, + std::vector sizes, + std::vector dim_order, + std::vector strides) { + std::vector data(192); + slim::SlimTensor output = slim::from_blob( + data.data(), output_sizes, output_strides, slim::c10::ScalarType::Float); + etensor::TensorImpl impl( + etensor::ScalarType::Float, + static_cast(sizes.size()), + sizes.data(), + data.data(), + dim_order.data(), + strides.data()); + etensor::Tensor etensor(&impl); + return output_is_raw_copyable(&output, &etensor); +} + +TEST(CudaOutputRawCopyTest, AcceptsAnOutputLaidOutLikeItsETensor) { + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, + {48, 1, 12, 3}, + {2, 3, 4, 4}, + {0, 2, 3, 1}, + {48, 1, 12, 3})); + // The ETensor's strides are stale from a smaller shape. + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, {48, 1, 12, 3}, {1, 3, 2, 2}, {0, 2, 3, 1}, {12, 1, 6, 3})); + EXPECT_TRUE(raw_copyable( + {2, 3, 4, 4}, {48, 16, 4, 1}, {1, 3, 2, 2}, {0, 1, 2, 3}, {12, 4, 2, 1})); + // A size-1 dim may have any stride, as with a batch of one. + EXPECT_TRUE(raw_copyable( + {1, 3, 4, 4}, {1, 1, 12, 3}, {1, 3, 4, 4}, {0, 2, 3, 1}, {48, 1, 12, 3})); + // An empty output has nothing to copy, whatever its strides. + EXPECT_TRUE(raw_copyable( + {2, 0, 4, 4}, + {16, 16, 4, 1}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); +} + +TEST(CudaOutputRawCopyTest, RejectsAnOutputLaidOutUnlikeItsETensor) { + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {48, 16, 4, 1}, + {2, 3, 4, 4}, + {0, 2, 3, 1}, + {48, 1, 12, 3})); + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {48, 1, 12, 3}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); + // Every row is followed by a gap, as in a slice of a wider tensor. + EXPECT_FALSE(raw_copyable( + {2, 3, 4, 4}, + {96, 32, 8, 1}, + {2, 3, 4, 4}, + {0, 1, 2, 3}, + {48, 16, 4, 1})); + EXPECT_FALSE(raw_copyable( + {3, 4, 4}, {16, 4, 1}, {1, 3, 4, 4}, {0, 1, 2, 3}, {48, 16, 4, 1})); +} + } // namespace diff --git a/backends/cuda/runtime/utils.h b/backends/cuda/runtime/utils.h index 6ca1c64d493..6dd860e11e1 100644 --- a/backends/cuda/runtime/utils.h +++ b/backends/cuda/runtime/utils.h @@ -299,6 +299,43 @@ inline executorch::runtime::Error _copy_slimtensor_to_etensor_impl( } } // namespace +/** + * Whether copying an output's raw bytes into its ETensor puts every element + * where the ETensor expects it, which is how a CUDA graph replay lands it. + * + * The ETensor may still have its previous shape, so its strides can be stale + * for dynamic shapes. Its expected strides come from its dim order and the + * output's sizes instead, the way resizing it computes them. Size-1 dims never + * move an element, so their strides may differ. + * + * @param output Pointer to the output SlimTensor (must not be null). + * @param etensor Pointer to the destination ETensor (must not be null). + * @return True if the raw copy is correct, including when there is nothing + * to copy. + */ +inline bool output_is_raw_copyable( + const executorch::backends::aoti::slim::SlimTensor* output, + const executorch::runtime::etensor::Tensor* etensor) { + const size_t ndim = output->dim(); + if (ndim != static_cast(etensor->dim())) { + return false; + } + if (output->numel() == 0) { + return true; + } + const auto dim_order = etensor->dim_order(); + int64_t expected_stride = 1; + for (size_t i = ndim; i > 0; --i) { + const int64_t d = dim_order[i - 1]; + const int64_t size = output->size(d); + if (size != 1 && output->stride(d) != expected_stride) { + return false; + } + expected_stride *= size; + } + return true; +} + /** * Copies data from a SlimTensor to an ETensor asynchronously. * diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index d962db09117..96f637a829c 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -6,6 +6,7 @@ import contextlib import os +import tempfile import threading import unittest from typing import Tuple @@ -653,23 +654,63 @@ def test_triton_kernel_mode_off(self): "SDPA kernel export with triton_kernel_mode=OFF failed", ) - def test_missing_scatter_and_index_put_kernels_fail_at_export(self): - """Inductor leaves these int64 ops to ATen kernels the CUDA shim lacks.""" + def test_missing_scatter_kernel_fails_at_export(self): + """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks.""" - class Int64Accumulate(torch.nn.Module): + class Int64ScatterAdd(torch.nn.Module): def forward(self, counts: torch.Tensor, idx: torch.Tensor): - ones = torch.ones_like(idx) - return ( - counts.scatter_add(0, idx, ones), - counts.index_put((idx,), ones, accumulate=True), - ) + return counts.scatter_add(0, idx, torch.ones_like(idx)) inputs = (torch.zeros(8, dtype=torch.long), torch.randint(0, 8, (32,))) with self.assertRaises(RuntimeError) as context: - self._export_to_cuda_with_lower(Int64Accumulate(), inputs) - message = str(context.exception) - self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", message) - self.assertIn("aoti_torch_index_put_out", message) + self._export_to_cuda_with_lower(Int64ScatterAdd(), inputs) + self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", str(context.exception)) + + @unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", + ) + def test_mask_index_and_index_put(self): + """A mask index makes the result size depend on the data, and Inductor + has no atomic add for int64, so these run in the index shims.""" + from executorch.runtime import Runtime + + class MaskIndex(torch.nn.Module): + def forward(self, x, mask, column, idx, counts): + out = x.clone() + out[mask] = x[mask] * 2 + # The mask and column index dims 0 and 2, which are not adjacent. + picked = x[mask, :, column].sum(0) + # A transposed input gives a result that is not row-major. + rows = x.transpose(1, 2)[mask] * 3 + ones = torch.ones_like(idx) + counts = counts.index_put((idx,), ones, accumulate=True) + return out, picked, rows, counts + + model = MaskIndex() + x = torch.randn(4, 5, 6) + column = torch.tensor([2]) + idx = torch.tensor([0, 3, 3, 1, 3]) + # Nonzero counts check that repeated writes add to the old values. + counts = torch.arange(8) * 5 + masks = ( + torch.tensor([True, False, True, True]), + torch.zeros(4, dtype=torch.bool), + torch.ones(4, dtype=torch.bool), + ) + edge = self._export_to_cuda_with_lower( + model, (x, masks[0], column, idx, counts) + ) + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + edge.to_executorch().write_to_file(f) + method = Runtime.get().load_program(path).load_method("forward") + for mask in masks: + inputs = (x, mask, column, idx, counts) + outputs = method.execute(list(inputs)) + for output, expected in zip(outputs, model(*inputs), strict=True): + torch.testing.assert_close(output.cpu(), expected) def test_device_info_propagated_to_cuda_delegate_outputs(self): """ diff --git a/backends/cuda/tests/test_output_stride_rearrange.py b/backends/cuda/tests/test_output_stride_rearrange.py index b1bf91638a7..431bb56a88f 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -282,6 +282,15 @@ def forward(self, x): return self.proj(x)[:, :, :14] +class _Double(nn.Module): + def __init__(self, memory_format): + super().__init__() + self.memory_format = memory_format + + def forward(self, x): + return (x * 2).contiguous(memory_format=self.memory_format) + + class TestStridedOutputWithCudaGraph(unittest.TestCase): """A CUDA graph replay copies raw bytes, so a strided output must not reach it.""" @@ -359,6 +368,20 @@ def test_strided_output_first_seen_at_capture(self): else: self._assert_matches_eager(module, x, out) + def test_output_shape_changes_at_capture(self): + torch.manual_seed(0) + # The capture call has a larger batch than the warmup calls before it. + inputs = [torch.randn(batch, 3, 4, 4) for batch in (1, 1, 1, 2, 2, 2)] + batch = torch.export.Dim("batch", min=1, max=8) + for memory_format in (torch.contiguous_format, torch.channels_last): + module = _Double(memory_format).eval() + outputs = self._run_with_cuda_graph( + module, (inputs[3].cuda(),), inputs, dynamic_shapes=({0: batch},) + ) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(memory_format=memory_format, call=i): + self._assert_matches_eager(module, x, out) + if __name__ == "__main__": unittest.main() diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index 99cba99cfc5..d81d9e9e187 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -41,6 +41,8 @@ "at::_ops::sort_stable::call", "aoti_torch_cuda_sort_stable", "aoti_torch_cuda_randint_low_out", + "aoti_torch_cuda_index_Tensor", + "aoti_torch_index_put_out", } ) diff --git a/runtime/executor/method.cpp b/runtime/executor/method.cpp index 34aca8e3856..c79b02eb52c 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -1528,13 +1528,14 @@ Method::set_output_data_ptr(void* buffer, size_t size, size_t output_idx) { return Error::InvalidArgument; } + // A later call can grow a dynamic output back to its upper bound. ET_CHECK_OR_RETURN_ERROR( - t.nbytes() <= size, + tensor_meta->nbytes() <= size, InvalidArgument, "buffer size: %" ET_PRIsize_t - " is smaller then expected tensor size: %" ET_PRIsize_t, + " is smaller than the largest output size: %" ET_PRIsize_t, size, - t.nbytes()); + tensor_meta->nbytes()); // Set data return internal::set_tensor_data(t, buffer, size); diff --git a/runtime/executor/method.h b/runtime/executor/method.h index 834325f008c..501569285f9 100644 --- a/runtime/executor/method.h +++ b/runtime/executor/method.h @@ -154,8 +154,8 @@ class Method final { * * @param[in] buffer The block of memory to point the specified tensor at. * - * @param[in] size the length of buffer in bytes, must be >= the nbytes of the - * specified tensor. + * @param[in] size the length of buffer in bytes, must be >= the largest + * nbytes the output can have, `MethodMeta::output_tensor_meta(i)->nbytes()`. * * @param[in] output_idx The index of the output to set the data_ptr for. Must * correspond to a tensor, and that tensor must not have had a buffer diff --git a/runtime/executor/test/method_test.cpp b/runtime/executor/test/method_test.cpp index bcbe3e64316..49cf36ad6c7 100644 --- a/runtime/executor/test/method_test.cpp +++ b/runtime/executor/test/method_test.cpp @@ -317,6 +317,33 @@ TEST_F(MethodTest, AliasedIOTest) { } } +TEST_F(MethodTest, SetOutputDataPtrChecksUpperBound) { + // The "cat" model's output is at most (4, 4) and is not memory planned. + ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes); + Result method = programs_["cat"]->load_method("forward", &mmm.get()); + ASSERT_EQ(method.error(), Error::Ok); + + float input[4] = {}; + int32_t sizes[2] = {1, 4}; + uint8_t dim_order[2] = {0, 1}; + int32_t strides[2] = {4, 1}; + executorch::aten::TensorImpl impl( + executorch::aten::ScalarType::Float, 2, sizes, input, dim_order, strides); + ASSERT_EQ( + method->set_input(EValue(executorch::aten::Tensor(&impl)), 0), Error::Ok); + float output[16]; + ASSERT_EQ(method->set_output_data_ptr(output, sizeof(output), 0), Error::Ok); + ASSERT_EQ(method->execute(), Error::Ok); + ASSERT_EQ(method->get_output(0).toTensor().nbytes(), 8 * sizeof(float)); + + // A later call can grow the output back to (4, 4), so a buffer that only + // fits the current (2, 4) output must be refused. + float small_output[8]; + EXPECT_EQ( + method->set_output_data_ptr(small_output, sizeof(small_output), 0), + Error::InvalidArgument); +} + TEST_F(MethodTest, SetInputRejectsOverflowingSizes) { // The "cat" model (ModuleDynamicCatUnallocatedIO) has a 2D input. // set_input validates numel/nbytes overflow before resize_tensor.