From 8a7c7bb426c04284207ccbd37593256b412de217 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Mon, 2 Nov 2026 23:58:57 -0800 Subject: [PATCH] [release/1.6] Fix CUDA indexing, graph execution and output buffers Cherry-pick #24299 with the CUDA graph, input-shape and shim-check follow-ups already carried by #24320. Resolve the Windows import library by retaining the release dtype-element-size export alongside the three indexing/get-size exports. Prepared with OpenAI Codex. --- .ci/scripts/wheel/test_shared_libraries.py | 2 + backends/aoti/common_shims_slim.cpp | 13 + backends/aoti/common_shims_slim.h | 5 + backends/cuda/CMakeLists.txt | 4 +- backends/cuda/cuda_backend.py | 2 + .../cuda/passes/move_cond_predicate_to_cpu.py | 4 + backends/cuda/runtime/aoti_cuda_shims.lib | Bin 45442 -> 47886 bytes backends/cuda/runtime/cuda_backend.cpp | 196 ++- backends/cuda/runtime/cuda_delegate_handle.h | 19 +- backends/cuda/runtime/shims/index.cu | 1064 +++++++++++++++++ backends/cuda/runtime/shims/index.h | 83 ++ backends/cuda/runtime/targets.bzl | 2 + .../cuda/runtime/test/test_cuda_kv_cache.cpp | 7 +- .../runtime/test/test_cuda_output_copy.cpp | 72 +- backends/cuda/runtime/utils.h | 37 + backends/cuda/tests/test_cuda_export.py | 66 +- .../tests/test_output_stride_rearrange.py | 113 +- backends/cuda/tests/test_sort_shim.py | 2 + runtime/executor/method.cpp | 7 +- runtime/executor/method.h | 5 +- runtime/executor/test/method_test.cpp | 27 + 21 files changed, 1636 insertions(+), 94 deletions(-) create mode 100644 backends/cuda/runtime/shims/index.cu create mode 100644 backends/cuda/runtime/shims/index.h diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py index b9ad0149ec1..85afe0b3718 100644 --- a/.ci/scripts/wheel/test_shared_libraries.py +++ b/.ci/scripts/wheel/test_shared_libraries.py @@ -104,6 +104,8 @@ "aoti_torch_cuda_randint_low_out", "aoti_torch_cuda_sort_stable", "aoti_torch_cuda_guard_set_index", + "aoti_torch_cuda_index_Tensor", + "aoti_torch_index_put_out", ) _QUANTIZED_KERNEL_SYMBOLS = ( diff --git a/backends/aoti/common_shims_slim.cpp b/backends/aoti/common_shims_slim.cpp index fa94c55455f..7a5a53ee744 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 29b7fa9b3d7..af275b5514b 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/passes/move_cond_predicate_to_cpu.py b/backends/cuda/passes/move_cond_predicate_to_cpu.py index f876a684019..ba8c36f63b7 100644 --- a/backends/cuda/passes/move_cond_predicate_to_cpu.py +++ b/backends/cuda/passes/move_cond_predicate_to_cpu.py @@ -55,6 +55,10 @@ def forward(self, hidden_states, encoder_hidden_states): In this example if we keep `self.initialized` on GPU, we will need to copy it to CPU for every forward pass. We move the predicate to CPU to avoid device to host copies. This pass is only applicable to models that use torch.cond and its predicate is a constant buffer. + + A method this pass changes keeps a constant in host memory, so the CUDA backend runs it + without a CUDA graph even when one is requested: a graph would keep the branch the capture + call took. """ requires_exported_program = True diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index 0e6a6db6d81971372fbb636f35c9cf6c5ffcff56..aa97264e523c2155b7d27a28f75db138096369e8 100644 GIT binary patch delta 1618 zcmZ|Odr(wW7y$6^EV2>oU2u0rC1gOILXo|f$FgVxLcmC{raVMOL>7~bB2gZasNsVE z#Yp#~0P%sMW`t>eW1~)CuRokJo0@XWX^f_+X_|aYY2(!N_|6_|teNi2Z|}X|p2v5- zd(M_`OIgU>A7-+v zjC6wAN+-WpPJI9%o!)?>5fNlFc>;o0R6Sc!T(zpcwhj~AcAwdNZ)CjLK^1}L&6rbV z^YsVfI~E7dwFJ@(OSK4ECA+R}Y`_DiaBKD5w~E@v2Ed(z6KSuODuUffcefC(4xK=^ ztP9NMdc!bWN%*N_ib(eVx-cPh5@k)uilysug#s_9JMPUQx@ZI=od$-@8D~-R{BOdU z(;~^0(JQbbFof=;rwaJfP3kr29^nI`CV4*&N_P6(aOP8WD7SErACMFLCnY&e{2 zrQUfQ9m&^e?;MV1%(r1pei$|9WeeQ>m_gU`Il0J%8qJ^Y5!h5<(6u;j_gmp)ykRzD zMf^{w6x&#{EDS{bK8>P_xNGs^u_-Ql(b_y%W5*WWE zfks<-iW6lKbhDIG2`<;^T0W1kW0{pk%DCs%3Y}(`bHrQepn+v+0uL`YNb_hl zE0YAaRgEnxXr&#?t2D}8!DE#^sZ;w}Zg-{HP7kc)Vm$bCub`XNEYrO{4{>k)Mw1&e1jem%s zTcCY|(_cg}9gnY%p^MF2!f`D+olD?;ZJWkpf3SnjY~&SvvDryUEgUV`YQqzogJ@ur zOW+%=x+-^K+ay}Gm8;X)?xe3^*Ys-`l`>4>ns{&baU-v`^UcH z&>kyUdwA7h`#5i|p~(RJ#xjv^_azJbeLxirHyqS(hBkpN?8y=M{-BfQ?kf@^5!1A( zbZS4xI*u6Bv6eqd@KH_s*V#jPT1Hpu_x}@>=kq^pGN(0>Y>r z$dAba8@vn)gN-@DqJc7bm>FY~82_3?(Ja##O^i#_=rqKrQH;}f95(rUbA5N;ckk|f zhyT#^+|^y!7fjD^C2suukpR#H5R(tE&8Eb!14#A{l*HEnq(`fiq*wqcMW^h37eI1O zD5=8$Qra11?~4G^e%?pQEDg#BKu(&H%Nt91zbW~*0i;4sDtSUmNgIGv+NeC^1&}Hp zQ=T;dNY!5{wKD)xU9Zwm1|S^|I{3j({HmN<29R3LDs8+->PS|)q5-5JpIGYtOz8~+ zNCW(yGL0*wa$VS{x1wDig)4^~8{=vo&lNbRQ6HnngmF9FGI|6W=aZ<-Y!{K~Y5%t#qoZSJ zaZ8b&Dr0k#$(Z03xR!0>q(=JXK$al)Q?Ck{nm*3U6tw$bNFXPta6Zo&5A$4r{|y&r z3A|UJQBmNiz~({=`8}lqSKJ!?)yEN!$3pofhXqnmDlL1xLfCMtZzr}p@1fmSf)~6R zP5F4Cr8Ew=l$)?B^=JI2EDm>+AHrteX5O_j8?HWNCVfS&z?yOwbyOY|;$h4-#L;(E zT)~k_7p?d430A7&X}LN_pyQZES`8xF5hluN;GC278U;Ezy0gKGnGptxJI)!pp4Vuhoul^~t#tQ9hM=dL z#sw{%N*A=_WJpj?OS-@xUvSY6tz1u2ODq+)X##JyrqYdej_hx z_7IQIr>Bbr7Iiw+O|G7f$M=E;`tYTL0uuuoef~I~*nuNG4qE7T3ye79q?DJrj=}D& zm>m^CW4&D6-#t!s?_IqXIyO)s@S{GBIxBf!_rNyF8RVOpJ+oKIOV?fGs^L#-4~{Hur^Pq90sU`i q^uCu@8Rks*li5PAEM@Z7IIq#~Z*g=rY$D&y6oIxIPMUkWR`(x9%&4mX diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 712ed5ff545..296f6436225 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() { @@ -636,7 +614,18 @@ class ET_EXPERIMENTAL CudaBackend final // Initialize CUDA graph state if enabled for this method. Not on a device // without memory pools, where the method runs without a graph. if (should_use_cuda_graph_for_method(method_name)) { - if (handle->cuda_graph_state.start_warmup()) { + const std::string host_constants = describe_host_constants(handle); + if (!host_constants.empty()) { + // A graph records only the device work of the capture call. Host code, + // such as the choice of a torch.cond branch on a flag in host memory, + // would be fixed at its value from that call. + ET_LOG( + Info, + "CUDA graph requested for method '%s', but %s; running without a " + "CUDA graph", + method_name.c_str(), + host_constants.c_str()); + } else if (handle->cuda_graph_state.start_warmup()) { ET_LOG( Info, "CUDA graph enabled for method '%s' (warmup=%d)", @@ -790,14 +779,26 @@ class ET_EXPERIMENTAL CudaBackend final // Copy new input data (GPU-resident) into static input buffers (D2D) for (size_t i = 0; i < n_inputs; i++) { auto* et_input = &(args[i]->toTensor()); + const auto sizes = et_input->sizes(); + const auto& captured = handle->cuda_graph_state.static_input_sizes[i]; ET_CHECK_OR_RETURN_ERROR( - et_input->nbytes() == - handle->cuda_graph_state.static_input_nbytes[i], + sizes.size() == captured.size(), InvalidArgument, - "CUDA graph replay: input %zu size mismatch (expected %zu, got %zu)", + "CUDA graph replay: input %zu has %zu dims, captured with %zu", i, - handle->cuda_graph_state.static_input_nbytes[i], - et_input->nbytes()); + sizes.size(), + captured.size()); + for (size_t d = 0; d < sizes.size(); d++) { + ET_CHECK_OR_RETURN_ERROR( + sizes[d] == captured[d], + InvalidArgument, + "CUDA graph replay: input %zu has size %lld in dim %zu, " + "captured with %lld", + i, + static_cast(sizes[d]), + d, + static_cast(captured[d])); + } ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpyAsync( handle->cuda_graph_state.static_input_ptrs[i], et_input->const_data_ptr(), @@ -830,13 +831,70 @@ class ET_EXPERIMENTAL CudaBackend final return commit_offgraph_kv_step(handle, kv_step_width); } - // --------------------------------------------------------------- - // Normal path (also used for WARMUP and CAPTURE phases) - // --------------------------------------------------------------- - bool is_capture_step = + const bool is_capture_step = (handle->cuda_graph_state.phase == CudaGraphPhase::Warmup && handle->cuda_graph_state.warmup_remaining == 0); + bool graph_ran = false; + Error error = + run_once(handle, args, n_inputs, n_outputs, is_capture_step, graph_ran); + if (error != Error::Ok && is_capture_step && !graph_ran) { + // A capture records kernels without running them, so until the new graph + // runs, the inputs and any state the program updates are untouched. A + // program that keeps state in host memory, where host code does run + // during a capture, never gets here: init leaves its graph off. + ET_LOG( + Info, + "CUDA graph: capture failed for '%s'; running this call and later " + "ones without a graph", + handle->method_name.c_str()); + // The capture guard turned graphs off, so later calls run eagerly too. + ET_DCHECK(handle->cuda_graph_state.phase == CudaGraphPhase::Disabled); + error = run_once( + handle, + args, + n_inputs, + n_outputs, + /*is_capture_step=*/false, + graph_ran); + } + ET_CHECK_OK_OR_RETURN_ERROR(error); + return commit_offgraph_kv_step(handle, kv_step_width); + } + + private: + // Names a constant of the program in host memory, or says that the + // constants could not be listed, since unknown counts as host memory. + // Empty when every constant is on the device. + static std::string describe_host_constants(cuda::CudaDelegateHandle* handle) { + const char* unlisted = "its constants could not be listed"; + if (handle->extract_constants_map == nullptr) { + return unlisted; + } + std::unordered_map constants; + if (handle->extract_constants_map( + handle->container_handle, + reinterpret_cast(&constants), + /*use_inactive=*/false) != Error::Ok) { + return unlisted; + } + for (const auto& [fqn, tensor] : constants) { + if (reinterpret_cast(tensor)->is_cpu()) { + return "constant '" + fqn + "' is in host memory"; + } + } + return ""; + } + // Runs the program once. On the capture step it captures the CUDA graph + // instead, then launches it once, and sets graph_ran once that launch is + // queued. + Error run_once( + cuda::CudaDelegateHandle* handle, + Span args, + size_t n_inputs, + size_t n_outputs, + bool is_capture_step, + bool& graph_ran) const { // Wrrap them as SlimTensors for AOTI execution without any host<->device // copies. std::vector slim_inputs(n_inputs); @@ -863,25 +921,10 @@ class ET_EXPERIMENTAL CudaBackend final return; } // Free what this attempt pinned and disable graphs for this method, so - // a capture that cannot succeed costs one error rather than one on - // every fourth call for the life of the process. Eager execution is - // correct, just slower. - for (void* ptr : state_->static_input_ptrs) { - (void)cudaFree(ptr); - } - state_->static_input_ptrs.clear(); - state_->static_output_ptrs.clear(); - state_->static_input_nbytes.clear(); - state_->static_output_nbytes.clear(); - // Same order as ~CudaGraphState: the exec depends on the graph. - if (state_->graph_exec != nullptr) { - (void)cudaGraphExecDestroy(state_->graph_exec); - state_->graph_exec = nullptr; - } - if (state_->graph != nullptr) { - (void)cudaGraphDestroy(state_->graph); - state_->graph = nullptr; - } + // a capture that cannot succeed is tried once rather than on every + // later call for the life of the process. Eager execution is correct, + // just slower. + state_->release(); state_->phase = CudaGraphPhase::Disabled; state_->warmup_remaining = 0; (void)cudaGetLastError(); @@ -925,6 +968,8 @@ class ET_EXPERIMENTAL CudaBackend final // through the guard instead of leaking this allocation. handle->cuda_graph_state.static_input_ptrs.push_back(static_ptr); handle->cuda_graph_state.static_input_nbytes.push_back(nbytes); + handle->cuda_graph_state.static_input_sizes.emplace_back( + et_input->sizes().begin(), et_input->sizes().end()); ET_CUDA_CHECK_OR_RETURN_ERROR(cudaMemcpy( static_ptr, @@ -1056,11 +1101,23 @@ class ET_EXPERIMENTAL CudaBackend final } } - ET_CHECK_OR_RETURN_ERROR( - error == Error::Ok, - Internal, - "AOTInductorModelContainerRun failed with error code %d", - error); + if (error != Error::Ok) { + // A failed capture is retried without a graph, which reports a real + // failure at Error. + if (is_capture_step) { + ET_LOG( + Info, + "AOTInductorModelContainerRun failed during CUDA graph capture " + "with error code %d", + error); + } else { + ET_LOG( + Error, + "AOTInductorModelContainerRun failed with error code %d", + error); + } + return Error::Internal; + } if (is_capture_step) { // End capture → instantiate graph @@ -1083,13 +1140,28 @@ class ET_EXPERIMENTAL CudaBackend final slim_outputs[i], &(args[i + n_inputs]->toTensor()))) { std::fill(slim_outputs.begin(), slim_outputs.end(), nullptr); ET_LOG( - Error, + Info, "CUDA graph: output %zu of '%s' is not laid out like its " - "ETensor; this call fails and later calls run without a graph", + "ETensor", i, 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( + Info, + "CUDA graph: cannot resize output %zu of '%s' to its captured " + "shape", + i, + handle->method_name.c_str()); + return resize_err; + } } gerr = cudaGraphInstantiate( @@ -1101,7 +1173,6 @@ class ET_EXPERIMENTAL CudaBackend final Internal, "cudaGraphInstantiate failed: %s", cudaGetErrorString(gerr)); - handle->cuda_graph_state.note_graph_allocations(); // Record static output pointers (stable under graph replay). Releasing // them from slim_outputs here, before the copies below, keeps the cleanup @@ -1127,6 +1198,8 @@ class ET_EXPERIMENTAL CudaBackend final Internal, "cudaGraphLaunch (first replay) failed: %s", cudaGetErrorString(gerr)); + graph_ran = true; + handle->cuda_graph_state.note_graph_allocations(); // Copy capture-step outputs into the planned GPU ET buffers (D2D). for (size_t i = 0; i < n_outputs; i++) { @@ -1143,7 +1216,7 @@ class ET_EXPERIMENTAL CudaBackend final // Last failure point is behind us, so the captured state is now the state // the next call should replay from. capture_guard.disarm(); - return commit_offgraph_kv_step(handle, kv_step_width); + return Error::Ok; } // ----- Normal / WARMUP execution continues here ----- @@ -1200,9 +1273,10 @@ class ET_EXPERIMENTAL CudaBackend final slim_outputs[i] = nullptr; } - return commit_offgraph_kv_step(handle, kv_step_width); + return Error::Ok; } + public: void destroy(DelegateHandle* handle_) const override { if (handle_ == nullptr) { return; diff --git a/backends/cuda/runtime/cuda_delegate_handle.h b/backends/cuda/runtime/cuda_delegate_handle.h index f42f0f2382f..ead44960a46 100644 --- a/backends/cuda/runtime/cuda_delegate_handle.h +++ b/backends/cuda/runtime/cuda_delegate_handle.h @@ -103,8 +103,12 @@ struct CudaWeightStorage { // Phases of the CUDA graph lifecycle for a delegate handle. // // The transition flow is: -// Disabled ──(if CUDA graph is enabled for this method)──▶ Warmup +// Disabled ──(if CUDA graph is enabled for this method, the device +// has memory pools, and every constant is in device +// memory)──▶ Warmup // Warmup ──(after `warmup_remaining` execute() calls)──▶ Replay +// Warmup ──(if the capture fails, or an output is not laid out +// like its ETensor)──▶ Disabled // // - Disabled: CUDA graph is not used for this method. Every execute() runs // eagerly through the normal kernel-launch path. @@ -118,7 +122,9 @@ struct CudaWeightStorage { // - Replay: The captured `graph_exec` is launched on every execute() call. // Inputs are memcpy'd into the static input buffers, the graph is replayed, // and outputs are memcpy'd back from the static output buffers. No tensor -// setup or kernel launches happen on the host hot path. +// setup or kernel launches happen on the host hot path. The graph only +// works for the input shapes it was captured with, so a call whose input +// shapes differ from the capture call fails with InvalidArgument. enum class CudaGraphPhase { Disabled = 0, Warmup = 1, @@ -141,6 +147,9 @@ struct CudaGraphState { std::vector static_output_ptrs; std::vector static_input_nbytes; std::vector static_output_nbytes; + // An input with another shape of the same byte count would replay with the + // captured shape and return a wrong result with no error. + std::vector> static_input_sizes; // Allocations the graph makes and never frees -- the outputs AOTI allocates // while being captured. AutoFreeOnLaunch frees one launch's before the next @@ -163,7 +172,8 @@ struct CudaGraphState { } // Records the graph's outstanding allocations: every memory-allocation node - // without a matching free node. Call once the graph is captured. + // without a matching free node. Call once the graph has launched: only a + // launch allocates them, so freeing them before one fails. void note_graph_allocations() { graph_allocations.clear(); #if !defined(EXECUTORCH_USE_HIP) @@ -227,6 +237,7 @@ struct CudaGraphState { static_output_ptrs.clear(); static_input_nbytes.clear(); static_output_nbytes.clear(); + static_input_sizes.clear(); } // Drops the captured graph so a new one is captured, for when memory the @@ -253,6 +264,7 @@ struct CudaGraphState { static_output_ptrs(std::move(other.static_output_ptrs)), static_input_nbytes(std::move(other.static_input_nbytes)), static_output_nbytes(std::move(other.static_output_nbytes)), + static_input_sizes(std::move(other.static_input_sizes)), graph_allocations(std::move(other.graph_allocations)) { other.graph = nullptr; other.graph_exec = nullptr; @@ -270,6 +282,7 @@ struct CudaGraphState { static_output_ptrs = std::move(other.static_output_ptrs); static_input_nbytes = std::move(other.static_input_nbytes); static_output_nbytes = std::move(other.static_output_nbytes); + static_input_sizes = std::move(other.static_input_sizes); graph_allocations = std::move(other.graph_allocations); other.graph = nullptr; diff --git a/backends/cuda/runtime/shims/index.cu b/backends/cuda/runtime/shims/index.cu new file mode 100644 index 00000000000..4b696cb03b9 --- /dev/null +++ b/backends/cuda/runtime/shims/index.cu @@ -0,0 +1,1064 @@ +/* + * 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. Info, not Error: the backend then retries the + // call without a graph. + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaStreamIsCapturing(stream, &capture)); + if (capture != cudaStreamCaptureStatusNone) { + ET_LOG(Info, "index: a mask index cannot run inside a CUDA graph capture"); + return Error::NotSupported; + } + // 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; + // The sort below takes the count as an int. + 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; + } + } + // Half and bfloat16 add in float, as ATen does. Narrower integers add in + // int64, which wraps to the same result as adding in their own type. + 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..3cdfc64e3f9 --- /dev/null +++ b/backends/cuda/runtime/shims/index.h @@ -0,0 +1,83 @@ +/* + * 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); + +/** + * Writes values into out at indices, with the semantics of aten::index_put_. + * When accumulate is set, every value is added, also when an index repeats. + * + * 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 less is lost to + * rounding, though a term far smaller than the running sum can still vanish. + * 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_kv_cache.cpp b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp index 0c32da2079f..948a53c0c1b 100644 --- a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp @@ -1039,15 +1039,14 @@ TEST_F(CudaKVCacheTest, RetiringAGraphFreesWhatItsLastLaunchAllocated) { graph.graph, cudaGraphInstantiateFlagAutoFreeOnLaunch), cudaSuccess); + ASSERT_EQ(cudaGraphLaunch(graph.graph_exec, stream), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + EXPECT_GE(graph_mem_in_use(), before + kBytes); graph.note_graph_allocations(); // Only the allocation the graph leaves outstanding is its to free. ASSERT_EQ(graph.graph_allocations.size(), 1); EXPECT_EQ(graph.graph_allocations[0], output); - ASSERT_EQ(cudaGraphLaunch(graph.graph_exec, stream), cudaSuccess); - ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); - EXPECT_GE(graph_mem_in_use(), before + kBytes); - graph.release(); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); ASSERT_EQ(cudaDeviceGraphMemTrim(device), cudaSuccess); 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 2f68b4e7341..8e187df4584 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 @@ -690,23 +691,64 @@ 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) + # Negative and int32: the shims wrap the index and read it as int32. + column = torch.tensor([-4], dtype=torch.int32) + 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..2cbde6cc6f8 100644 --- a/backends/cuda/tests/test_output_stride_rearrange.py +++ b/backends/cuda/tests/test_output_stride_rearrange.py @@ -282,6 +282,39 @@ def forward(self, x): return self.proj(x)[:, :, :14] +class _MaskedRows(nn.Module): + def forward(self, x, mask): + return x[mask] * 2 + + +class _MaskedAdd(nn.Module): + def forward(self, x, mask): + return x.index_put((mask,), torch.ones(4), accumulate=True) + + +class _HostFlag(nn.Module): + """A torch.cond on a flag the program keeps in host memory and updates on + every call.""" + + def __init__(self): + super().__init__() + self.register_buffer("flag", torch.tensor(False), persistent=False) + + def forward(self, x): + picked = torch.cond(self.flag, lambda t: t + t, lambda t: t.neg(), (x,)) + self.flag.fill_(x.shape[1] > 1) + return (picked,) + + +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.""" @@ -295,14 +328,16 @@ def setUpClass(cls): def _run_with_cuda_graph(self, module, example, inputs, dynamic_shapes=None): from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner - from executorch.exir import to_edge_transform_and_lower + from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower from executorch.runtime import Runtime with torch.no_grad(): ep = torch.export.export(module, example, dynamic_shapes=dynamic_shapes) spec = CudaBackend.generate_method_name_compile_spec("forward") program = to_edge_transform_and_lower( - ep, partitioner=[CudaPartitioner([spec])] + ep, + partitioner=[CudaPartitioner([spec])], + compile_config=EdgeCompileConfig(_check_ir_validity=False), ).to_executorch() runtime = Runtime.get() option = "enable_cuda_graph_for_method" @@ -320,8 +355,9 @@ def _run_with_cuda_graph(self, module, example, inputs, dynamic_shapes=None): # The option is process wide. runtime.backend_registry.set_option("CudaBackend", {option: ""}) for x in inputs: + args = x if isinstance(x, tuple) else (x,) try: - results.append(method.execute([x])[0]) + results.append(method.execute(list(args))[0]) except RuntimeError as e: results.append(e) return results @@ -354,9 +390,74 @@ def test_strided_output_first_seen_at_capture(self): ) for i, (x, out) in enumerate(zip(inputs, outputs)): with self.subTest(call=i): - if i == 3: - self.assertIsInstance(out, RuntimeError) - else: + self._assert_matches_eager(module, x, out) + + def test_input_shape_changes_after_capture(self): + torch.manual_seed(0) + # The last calls swap two dims, which keeps the byte count, so only the + # shape tells them apart from the captured input. + shapes = [(2, 3, 4, 4)] * 4 + [(2, 4, 3, 4)] * 2 + inputs = [torch.randn(shape) for shape in shapes] + dims = { + 1: torch.export.Dim("c", min=2, max=8), + 2: torch.export.Dim("h", min=2, 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[0].cuda(),), inputs, dynamic_shapes=(dims,) + ) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(memory_format=memory_format, call=i): + if i < 4: + self._assert_matches_eager(module, x, out) + else: + # InvalidArgument, from the replay's input shape check. + self.assertIn("error 0x12", str(out)) + + def test_host_state_runs_without_a_graph(self): + # A graph would keep the branch of the capture call, and a failed + # capture would already have updated the flag before running the call + # again, so this method must not use a graph. + torch.manual_seed(0) + module = _HostFlag().cuda().eval() + eager = _HostFlag().cuda().eval() + eager.load_state_dict(module.state_dict()) + inputs = [torch.randn(1, rows, 64) for rows in (1, 1, 1, 50, 50, 50)] + rows = torch.export.Dim("rows", min=1, max=64) + outputs = self._run_with_cuda_graph( + module, (inputs[3].cuda(),), inputs, dynamic_shapes=({1: rows},) + ) + for i, (x, out) in enumerate(zip(inputs, outputs)): + with self.subTest(call=i): + with torch.no_grad(): + expected = eager(x.cuda())[0].cpu() + torch.testing.assert_close(out, expected) + + def test_mask_runs_without_a_graph(self): + # A mask index reads its count back to the host, which a capture + # cannot record, so the capture call has to run without a graph. + torch.manual_seed(0) + inputs = [(torch.randn(8, 4), torch.rand(8) > 0.5) for _ in range(6)] + for module in (_MaskedRows().eval(), _MaskedAdd().eval()): + outputs = self._run_with_cuda_graph(module, inputs[0], inputs) + for i, (args, out) in enumerate(zip(inputs, outputs)): + with self.subTest(module=type(module).__name__, call=i): + self.assertIsInstance(out, torch.Tensor) + torch.testing.assert_close(out, module(*args)) + + 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) 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 9b8b9e34008..90618c7cf93 100644 --- a/runtime/executor/method.cpp +++ b/runtime/executor/method.cpp @@ -1494,13 +1494,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..9432e2cc891 100644 --- a/runtime/executor/method.h +++ b/runtime/executor/method.h @@ -154,8 +154,9 @@ 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(output_idx)->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 0310c5d648b..ef97f8db657 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.