diff --git a/.github/workflows/linux_cuda_no_cudnn.yml b/.github/workflows/linux_cuda_no_cudnn.yml index abe92589aa705..5adcf016ea47d 100644 --- a/.github/workflows/linux_cuda_no_cudnn.yml +++ b/.github/workflows/linux_cuda_no_cudnn.yml @@ -93,7 +93,7 @@ jobs: ldd /build/Release/Release/libonnxruntime_providers_cuda.so | tee /tmp/ldd.txt ! grep -i cudnn /tmp/ldd.txt' - - name: Run no-cuDNN CUDA EP smoke test + - name: Run no-cuDNN CUDA EP op smoke test run: | docker run --rm --gpus all \ -v "${{ runner.temp }}/Release:/build/Release" \ @@ -102,6 +102,25 @@ jobs: PATH=/opt/python/cp312-cp312/bin:$PATH LD_LIBRARY_PATH=/usr/local/cuda-13.0/lib64:${LD_LIBRARY_PATH:-} export PATH LD_LIBRARY_PATH + + # The build image contains cuDNN. Remove its runtime libraries from this disposable + # test container so the smoke test proves that the provider works when cuDNN is + # physically unavailable, rather than merely relying on enable_cudnn=0. + for root in /usr /opt /lib /lib64; do + if [ -e "$root" ]; then + find "$root" -name "libcudnn*.so*" -print -delete 2>/dev/null || true + fi + done + ldconfig + if ldconfig -p | grep -qi cudnn; then + echo "cuDNN remains available in the dynamic linker cache" >&2 + exit 1 + fi + if find /usr /opt /lib /lib64 -name "libcudnn*.so*" -print -quit 2>/dev/null | grep -q .; then + echo "cuDNN runtime libraries remain in the no-cuDNN test container" >&2 + exit 1 + fi + WHEEL_PATH=$(find /build/Release/Release/dist -type f -name "onnxruntime_gpu-*.whl" | head -n 1) if [ -z "$WHEEL_PATH" ]; then echo "No built onnxruntime GPU wheel found under /build/Release/Release/dist" >&2 @@ -111,21 +130,114 @@ jobs: python -m pip install --no-cache-dir --force-reinstall --no-deps numpy onnx "$WHEEL_PATH" python - <<"PY" import numpy as np - import onnx import onnxruntime as ort from onnx import TensorProto, helper - x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]) - y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3]) - node = helper.make_node("Add", ["x", "x"], ["y"]) - graph = helper.make_graph([node], "cuda_no_cudnn_smoke", [x], [y]) - model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) - model.ir_version = 10 - providers = [("CUDAExecutionProvider", {"enable_cudnn": "0"})] - sess = ort.InferenceSession(model.SerializeToString(), providers=providers) - data = np.arange(6, dtype=np.float32).reshape(2, 3) - result = sess.run(None, {"x": data})[0] - np.testing.assert_allclose(result, data + data) - print("CUDA no-cuDNN smoke test passed") + + + def run_op(name, model, feeds, expected, rtol=1e-4, atol=1e-4): + model.ir_version = 10 + options = ort.SessionOptions() + options.add_session_config_entry("session.disable_cpu_ep_fallback", "1") + sess = ort.InferenceSession(model.SerializeToString(), sess_options=options, providers=providers) + got = sess.run(None, feeds)[0] + np.testing.assert_allclose(got, expected, rtol=rtol, atol=atol) + print("[no-cuDNN] " + name + " passed") + + + def make_model(nodes, inputs, outputs, opset, initializers=None): + graph = helper.make_graph(nodes, "no_cudnn_smoke", inputs, outputs, initializer=initializers or []) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) + + + def f32_in(shape): + return helper.make_tensor_value_info("x", TensorProto.FLOAT, shape) + + + data = np.random.rand(2, 3).astype(np.float32) + + # Elementwise baseline (does not depend on cuDNN). + run_op( + "Add", + make_model( + [helper.make_node("Add", ["x", "x"], ["y"])], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + 21, + ), + {"x": data}, + data + data, + ) + + # LogSoftmax over the last axis (softmax kernel, no cuDNN). + shifted = data - data.max(axis=1, keepdims=True) + log_softmax = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True)) + run_op( + "LogSoftmax", + make_model( + [helper.make_node("LogSoftmax", ["x"], ["y"], axis=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])], + 13, + ), + {"x": data}, + log_softmax, + ) + + # ReduceMean over the last axis (matrix-reduction path, no cuDNN). + run_op( + "ReduceMean", + make_model( + [helper.make_node("ReduceMean", ["x"], ["y"], axes=[1], keepdims=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 1])], + 13, + ), + {"x": data}, + data.mean(axis=1, keepdims=True), + ) + + # ReduceSum over a middle axis, which requires the general no-cuDNN path. + reduce_sum_data = np.random.rand(2, 3, 4).astype(np.float32) + run_op( + "ReduceSum", + make_model( + [helper.make_node("ReduceSum", ["x", "axes"], ["y"], keepdims=1)], + [f32_in([2, 3, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 1, 4])], + 13, + initializers=[helper.make_tensor("axes", TensorProto.INT64, [1], [1])], + ), + {"x": reduce_sum_data}, + reduce_sum_data.sum(axis=1, keepdims=True), + ) + + # ArgMax over the last axis (custom cuDNN-free kernel). + run_op( + "ArgMax", + make_model( + [helper.make_node("ArgMax", ["x"], ["y"], axis=1, keepdims=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.INT64, [2, 1])], + 13, + ), + {"x": data}, + np.argmax(data, axis=1).reshape(2, 1).astype(np.int64), + ) + + # ArgMin over the last axis (custom cuDNN-free kernel). + run_op( + "ArgMin", + make_model( + [helper.make_node("ArgMin", ["x"], ["y"], axis=1, keepdims=1)], + [f32_in([2, 3])], + [helper.make_tensor_value_info("y", TensorProto.INT64, [2, 1])], + 13, + ), + {"x": data}, + np.argmin(data, axis=1).reshape(2, 1).astype(np.int64), + ) + + print("CUDA no-cuDNN op smoke tests passed") PY' diff --git a/docs/cuda_plugin_ep/QUICK_START.md b/docs/cuda_plugin_ep/QUICK_START.md index ab7b4308f13e5..d6593d8e22d18 100644 --- a/docs/cuda_plugin_ep/QUICK_START.md +++ b/docs/cuda_plugin_ep/QUICK_START.md @@ -27,7 +27,7 @@ For local Linux CUDA 13 validation, use the no-cuDNN helper script. It keeps `CU bash .env/cuda_130_plugin_no_cudnn.sh --build --test_plugin ``` -The test mode sets `ORT_TEST_CUDA_PLUGIN_EP=1` and `ORT_TEST_CUDA_PLUGIN_NO_CUDNN=1`, which passes `enable_cudnn=0` to plugin sessions and skips plugin tests for operators that still require cuDNN, such as Conv, ConvTranspose, BatchNormalization, InstanceNormalization, LRN, ArgMax, reductions, Einsum, and cuDNN-backed pooling paths. +The test mode sets `ORT_TEST_CUDA_PLUGIN_EP=1` and `ORT_TEST_CUDA_PLUGIN_NO_CUDNN=1`, which passes `enable_cudnn=0` to plugin sessions and skips plugin tests for operators that still require cuDNN, such as Conv, ConvTranspose, BatchNormalization, InstanceNormalization, LRN, Einsum, and cuDNN-backed pooling paths. ## Minimum ONNX Runtime Version diff --git a/onnxruntime/core/providers/cuda/cuda_kernel.h b/onnxruntime/core/providers/cuda/cuda_kernel.h index 85d629764da48..6152b5f888eb2 100644 --- a/onnxruntime/core/providers/cuda/cuda_kernel.h +++ b/onnxruntime/core/providers/cuda/cuda_kernel.h @@ -133,6 +133,10 @@ class CudaKernel : public OpKernel { return RequireCudnnHandle(GetCudnnHandle(static_cast(ctx->GetComputeStream()))); } + inline cudnnHandle_t TryGetCudnnHandle(OpKernelContext* ctx) const { + return GetCudnnHandle(static_cast(ctx->GetComputeStream())); + } + static inline cudnnHandle_t GetCudnnHandle(onnxruntime::CudaStream* stream) { return stream ? stream->cudnn_handle_ : nullptr; } diff --git a/onnxruntime/core/providers/cuda/math/softmax.h b/onnxruntime/core/providers/cuda/math/softmax.h index c0c0818042c15..09b6eeec834b5 100644 --- a/onnxruntime/core/providers/cuda/math/softmax.h +++ b/onnxruntime/core/providers/cuda/math/softmax.h @@ -47,7 +47,7 @@ class Softmax final : public CudaKernel { } } - log_softmax_ = info.GetKernelDef().OpName() == "LogSoftmax"; + log_softmax_ = node.OpType() == "LogSoftmax"; } Status ComputeInternal(OpKernelContext* context) const override; diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index b1cc748be754a..7c6cdd0e4dc80 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h @@ -1064,8 +1064,30 @@ class CudaKernel : public OpKernel { std::string("cuDNN is unavailable or disabled for CUDA Plugin Execution Provider: ") + onnxruntime::cuda::CudnnLibrary::Get().Error())); } - if (handle != nullptr && stream != nullptr) { - CUDNN_CALL_THROW(cudnnSetStream(handle, stream)); + // Bind the shared handle to the current compute stream. cudaStream_t 0/nullptr is the default + // stream, which is still a valid stream to bind, so do this unconditionally to avoid leaving + // the handle bound to a stale stream from a previous call. + CUDNN_CALL_THROW(cudnnSetStream(handle, stream)); + return handle; + } + + cudnnHandle_t TryGetCudnnHandle(OpKernelContext* ctx) const { + auto stream = Stream(ctx); + auto handle = GetCudnnHandle(stream); + if (handle != nullptr) { + return handle; + } + + handle = DefaultCudnnHandle(); + if (handle != nullptr) { + // Bind the shared handle to the current compute stream. cudaStream_t 0/nullptr is the default + // stream, which is still a valid stream to bind, so do this unconditionally to avoid leaving + // the handle bound to a stale stream from a previous call. + // Keep this accessor non-throwing: if the stream cannot be bound, treat it as "no cuDNN handle" + // so callers can fall back to a cuDNN-free path instead of failing. + if (!CUDNN_CALL(cudnnSetStream(handle, stream)).IsOK()) { + return nullptr; + } } return handle; } diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu index ed97507f87641..afc2d9b061a55 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu @@ -4,6 +4,8 @@ #include "core/providers/cuda/reduction/reduction_functions.h" #include +#include +#include #include #include @@ -513,6 +515,236 @@ INSTANTIATE_REDUCE_MATRIX_COLUMNS(double); INSTANTIATE_REDUCE_MATRIX_COLUMNS(BFloat16); #undef INSTANTIATE_REDUCE_MATRIX_COLUMNS +namespace detail { +constexpr int kMaxReduceRank = 16; + +struct ReduceSumNdMetadata { + int output_segment_count{}; + int reduction_segment_count{}; + int64_t output_segment_sizes[kMaxReduceRank]{}; + int64_t output_segment_strides[kMaxReduceRank]{}; + int64_t reduction_segment_sizes[kMaxReduceRank]{}; + int64_t reduction_segment_strides[kMaxReduceRank]{}; + int64_t output_count{}; + int64_t reduction_count{}; +}; + +template +struct SumState { + T sum{}; + + __device__ __forceinline__ void Add(T value) { sum += value; } + __device__ __forceinline__ T Result() const { return sum; } +}; + +template <> +struct SumState { + double sum{}; + double correction{}; + + // Neumaier compensation is used instead of Kahan so that independently accumulated + // thread partials can be merged without losing a small value between large values. + __device__ __forceinline__ void Add(double value) { + const double next = __dadd_rn(sum, value); + const double error = fabs(sum) >= fabs(value) + ? __dadd_rn(__dsub_rn(sum, next), value) + : __dadd_rn(__dsub_rn(value, next), sum); + correction = __dadd_rn(correction, error); + sum = next; + } + + __device__ __forceinline__ double Result() const { return __dadd_rn(sum, correction); } +}; + +template +struct MergeSumState { + __device__ __forceinline__ SumState operator()(SumState lhs, const SumState& rhs) const { + lhs.Add(rhs.sum); + if constexpr (std::is_same_v) { + lhs.Add(rhs.correction); + } + return lhs; + } +}; + +template +__device__ __forceinline__ T CastReduceSumResult(TAccum value) { + if constexpr (std::is_integral_v) { + const double value_as_double = static_cast(value); + const double max_value = static_cast(std::numeric_limits::max()); + const double min_value = static_cast(std::numeric_limits::min()); + if (value_as_double >= max_value) return std::numeric_limits::max(); + if (value_as_double <= min_value) return std::numeric_limits::min(); + } + return static_cast(value); +} + +template +__global__ void reduce_sum_nd_kernel(const T* input, T* output, ReduceSumNdMetadata metadata) { + using TAccum = std::conditional_t, double, AccumulationType_t>; + using BlockReduce = cub::BlockReduce, BlockSize>; + __shared__ typename BlockReduce::TempStorage reduce_storage; + __shared__ int64_t input_base; + + // One cooperative block reduces each output. Grid-striding keeps the launch bounded for + // large outputs while retaining enough blocks to saturate the device. + for (int64_t output_index = blockIdx.x; + output_index < metadata.output_count; + output_index += gridDim.x) { + if (threadIdx.x == 0) { + int64_t remaining = output_index; + int64_t base = 0; + for (int segment = metadata.output_segment_count - 1; segment >= 0; --segment) { + const int64_t coordinate = segment == 0 ? remaining : remaining % metadata.output_segment_sizes[segment]; + if (segment != 0) remaining /= metadata.output_segment_sizes[segment]; + base += coordinate * metadata.output_segment_strides[segment]; + } + input_base = base; + } + __syncthreads(); + + SumState thread_sum{}; + for (int64_t reduction_index = threadIdx.x; reduction_index < metadata.reduction_count; + reduction_index += BlockSize) { + int64_t remaining = reduction_index; + int64_t input_index = input_base; + // Adjacent reduced dimensions are collapsed on the host, so this loop performs + // divisions only at reduced/non-reduced boundaries rather than once per rank. + for (int segment = metadata.reduction_segment_count - 1; segment >= 0; --segment) { + const int64_t coordinate = segment == 0 ? remaining : remaining % metadata.reduction_segment_sizes[segment]; + if (segment != 0) remaining /= metadata.reduction_segment_sizes[segment]; + input_index += coordinate * metadata.reduction_segment_strides[segment]; + } + thread_sum.Add(static_cast(input[input_index])); + } + + const SumState block_sum = BlockReduce(reduce_storage).Reduce(thread_sum, MergeSumState{}); + if (threadIdx.x == 0) { + output[output_index] = CastReduceSumResult(block_sum.Result()); + } + __syncthreads(); // reduce_storage is reused by the next grid-stride iteration. + } +} +} // namespace detail + +template +Status reduce_sum_nd(cudaStream_t stream, const T* input, T* output, + gsl::span dims, gsl::span axes) { + ORT_RETURN_IF_NOT(dims.size() <= detail::kMaxReduceRank, + "The general CUDA ReduceSum kernel supports ranks up to ", detail::kMaxReduceRank, "."); + + detail::ReduceSumNdMetadata metadata; + const int rank = gsl::narrow_cast(dims.size()); + std::array strides{}; + SafeInt stride = 1; + for (int axis = rank - 1; axis >= 0; --axis) { + ORT_RETURN_IF_NOT(dims[axis] > 0, "ReduceSum dimensions must be positive."); + strides[axis] = static_cast(stride); + stride *= dims[axis]; + } + + std::array reduced{}; + if (axes.empty()) { + for (int axis = 0; axis < rank; ++axis) reduced[axis] = true; + } else { + for (int64_t axis : axes) { + if (axis < 0) axis += rank; + ORT_RETURN_IF_NOT(axis >= 0 && axis < rank, "ReduceSum axis is out of range."); + ORT_RETURN_IF_NOT(!reduced[axis], "ReduceSum axes must not contain duplicates."); + reduced[axis] = true; + } + } + + SafeInt output_count = 1; + SafeInt reduction_count = 1; + for (int axis = 0; axis < rank;) { + const bool is_reduced = reduced[axis]; + SafeInt segment_size = 1; + int last_axis = axis; + do { + segment_size *= dims[axis]; + last_axis = axis++; + } while (axis < rank && reduced[axis] == is_reduced); + + if (is_reduced) { + const int segment = metadata.reduction_segment_count++; + metadata.reduction_segment_sizes[segment] = static_cast(segment_size); + metadata.reduction_segment_strides[segment] = strides[last_axis]; + reduction_count *= static_cast(segment_size); + } else { + const int segment = metadata.output_segment_count++; + metadata.output_segment_sizes[segment] = static_cast(segment_size); + metadata.output_segment_strides[segment] = strides[last_axis]; + output_count *= static_cast(segment_size); + } + } + metadata.output_count = static_cast(output_count); + metadata.reduction_count = static_cast(reduction_count); + + constexpr int block_size = 256; + constexpr int max_blocks = 65535; + const int grid_size = static_cast(std::min(max_blocks, metadata.output_count)); + detail::reduce_sum_nd_kernel<<>>(input, output, metadata); + return CUDA_CALL(cudaGetLastError()); +} + +#define INSTANTIATE_REDUCE_SUM_ND(T) \ + template Status reduce_sum_nd(cudaStream_t stream, const T* input, T* output, \ + gsl::span dims, gsl::span axes) +INSTANTIATE_REDUCE_SUM_ND(half); +INSTANTIATE_REDUCE_SUM_ND(float); +INSTANTIATE_REDUCE_SUM_ND(double); +INSTANTIATE_REDUCE_SUM_ND(BFloat16); +INSTANTIATE_REDUCE_SUM_ND(int32_t); +INSTANTIATE_REDUCE_SUM_ND(int64_t); +#undef INSTANTIATE_REDUCE_SUM_ND + +namespace detail { +template +__global__ void arg_min_max_last_axis_kernel(const TIn* input, int64_t* output, int m, int n) { + const int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= m) return; + + const int64_t row_offset = static_cast(row) * n; + TIn best_value = input[row_offset]; + int64_t best_index = 0; + for (int i = 1; i < n; ++i) { + const TIn value = input[row_offset + i]; + if constexpr (IsArgMax) { + if (value > best_value) { + best_value = value; + best_index = i; + } + } else { + if (value < best_value) { + best_value = value; + best_index = i; + } + } + } + + output[row] = best_index; +} +} // namespace detail + +template +Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* output, int m, int n) { + // The kernel reads input[row_offset] unconditionally, so a non-empty reduction axis is required. + if (m == 0 || n <= 0) return Status::OK(); + constexpr int block_size = 256; + const int grid_size = (m + block_size - 1) / block_size; + detail::arg_min_max_last_axis_kernel<<>>(input, output, m, n); + return CUDA_CALL(cudaGetLastError()); +} + +#define INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(T) \ + template Status arg_min_max_last_axis(cudaStream_t stream, const T* input, int64_t* output, int m, int n); \ + template Status arg_min_max_last_axis(cudaStream_t stream, const T* input, int64_t* output, int m, int n) +INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(half); +INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(float); +INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(double); +#undef INSTANTIATE_ARG_MIN_MAX_LAST_AXIS + } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h index 30b0afe4fe799..2e73a73a353f7 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.h +++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.h @@ -103,6 +103,19 @@ Status reduce_matrix_rows(cudaStream_t stream, const TIn* input, TOut* output, i template Status reduce_matrix_columns(cudaStream_t stream, const TIn* input, TOut* output, int m, int n, void* buffer, size_t buffer_size); +/** + * Computes ReduceSum for an arbitrary set of axes. This is the general fallback + * for axis layouts that cannot be flattened into one of the optimized matrix + * reductions above. + */ +template +Status reduce_sum_nd(cudaStream_t stream, const T* input, T* output, + gsl::span dims, gsl::span axes); + +/** Computes ArgMax/ArgMin indices over the last dimension in a row-major matrix. */ +template +Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* output, int m, int n); + /** Apply unary elementwise division. */ template void UnaryDiv(cudaStream_t stream, const T* input, T* output, T denominator, size_t count); diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index a8019cda5c411..ed6195c0bbb8b 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -299,10 +299,6 @@ Status PrepareForReduce(const Tensor* X, const int64_t rank = gsl::narrow(input_shape.NumDimensions()); prepare_reduce_metadata.input_count = input_shape.Size(); - if (rank > 8) { - return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "cuDNN only supports up to 8-D tensors in reduction"); - } - const auto input_dims = input_shape.GetDims(); std::vector reduced(rank, false); if (axes.size() > 0) { @@ -481,6 +477,41 @@ Status ReduceComputeCore(const AllocatorPtr& gpu_allocator, const CudaKernel* ke } } + if constexpr (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES) { + if (axes.size() == 1) { + const int64_t rank = input_shape.NumDimensions(); + const int64_t axis = HandleNegativeAxis(axes[0], rank); + if (axis == rank - 1) { + const int64_t m = input_shape.SizeToDimension(axis); + const int64_t n = input_shape[axis]; + if (n > 0 && m <= std::numeric_limits::max() && n <= std::numeric_limits::max()) { + if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MAX) { + return arg_min_max_last_axis(stream, reinterpret_cast(input.Data()), + output.MutableData(), gsl::narrow_cast(m), + gsl::narrow_cast(n)); + } + if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MIN) { + return arg_min_max_last_axis(stream, reinterpret_cast(input.Data()), + output.MutableData(), gsl::narrow_cast(m), + gsl::narrow_cast(n)); + } + } + } + } + } + + // Preserve the optimized matrix reductions above and the existing cuDNN path when available. + // Without cuDNN, use a general CUDA kernel for plain ReduceSum axis layouts that cannot be + // represented as a contiguous matrix reduction. + if constexpr (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_NO_INDICES) { + if ((cudnn_handle == nullptr || input_shape.NumDimensions() > 8) && + cudnn_reduce_op == CUDNN_REDUCE_TENSOR_ADD && + !calculate_log && !calculate_sqt && !log_sum_exp) { + return reduce_sum_nd(stream, reinterpret_cast(input.Data()), + reinterpret_cast(output.MutableData()), input_shape.GetDims(), axes); + } + } + // This reduction keep adding values to this buffer. If a non-zero value, say 1000, is here, the sum will start with 1000. // Therefore zeroing out the memory is required CUDA_RETURN_IF_ERROR(cudaMemsetAsync(output.MutableDataRaw(), 0, output.SizeInBytes(), stream)); @@ -785,7 +816,7 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe const bool fast_reduction = fast_reduction_ && !ctx->GetUseDeterministicCompute(); return ReduceComputeCore(AllocatorPtr{}, this, *X, prepare_reduce_metadata, *Y, cudnn_reduce_op, axes, calculate_log_, calculate_sqt_, log_sum_exp_, fast_reduction, - Stream(ctx), GetComputeStream(ctx), GetCudnnHandle(ctx)); + Stream(ctx), GetComputeStream(ctx), TryGetCudnnHandle(ctx)); } #define SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(T) \ @@ -797,9 +828,8 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe const Tensor* X = ctx->Input(0); \ TensorShapeVector axes; \ size_t num_inputs = ctx->InputCount(); \ - if (num_inputs == 2) { \ - const Tensor* axes_tensor = ctx->Input(1); \ - ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); \ + const Tensor* axes_tensor = num_inputs == 2 ? ctx->Input(1) : nullptr; \ + if (axes_tensor != nullptr) { \ ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor."); \ auto nDims = static_cast(axes_tensor->Shape()[0]); \ const auto* data = axes_tensor->Data(); \ @@ -858,6 +888,14 @@ Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnRe return Status::OK(); \ } \ \ + if constexpr (std::is_same_v || std::is_same_v) { \ + if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_ADD && \ + !calculate_log_ && !calculate_sqt_ && !log_sum_exp_) { \ + return reduce_sum_nd(Stream(ctx), reinterpret_cast(X->Data()), \ + reinterpret_cast(Y->MutableData()), X->Shape().GetDims(), axes); \ + } \ + } \ + \ CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes(), Stream(ctx))); \ \ size_t indices_bytes = 0; \ diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc index 95f5ae30016e2..13d6016d09de9 100644 --- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc +++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc @@ -2863,6 +2863,31 @@ TEST(ReductionOpTest, ReduceSum_int64) { test.Run(); } +#if defined(USE_CUDA) +TEST(ReductionOpTest, ReduceSum_int64_omitted_optional_axes) { + OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {3}, {1, 2, 3}); + test.AddOptionalInputEdge(); + test.AddOutput("reduced", {}, {6}); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +TEST(ReductionOpTest, ReduceSum_int64_cancellation) { + OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)0); + const int64_t large = int64_t{1} << 53; + test.AddInput("data", {3}, {large, 1, -large}); + test.AddInput("axes", {1}, {0}); + test.AddOutput("reduced", {}, {1}); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} +#endif + TEST(ReductionOpTest, ReduceSum_default_axes_keepdims) { OpTester test("ReduceSum"); test.AddAttribute("keepdims", (int64_t)1); diff --git a/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc b/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc index 593255b9e9c23..410527f0a6ea6 100644 --- a/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc @@ -3,7 +3,9 @@ #include "gtest/gtest.h" +#include #include +#include #include "core/providers/cuda/shared_inc/cuda_utils.h" #include "core/common/optional.h" @@ -238,6 +240,97 @@ TEST(ReductionFunctionsTest, ReduceColumnsToColumnRepeated) { TestReduceColumnsToColumnRepeated(17, 8192, 100, 2e-4f); } +TEST(ReductionFunctionsTest, ReduceSumNdMiddleAndMultipleAxes) { + const std::vector dims{2, 3, 4, 2}; + const std::vector axes{1, 3}; + std::vector input(48); + std::iota(input.begin(), input.end(), 1.0f); + std::vector expected(8, 0.0f); + for (int64_t d0 = 0; d0 < dims[0]; ++d0) { + for (int64_t d1 = 0; d1 < dims[1]; ++d1) { + for (int64_t d2 = 0; d2 < dims[2]; ++d2) { + for (int64_t d3 = 0; d3 < dims[3]; ++d3) { + expected[d0 * dims[2] + d2] += input[((d0 * dims[1] + d1) * dims[2] + d2) * dims[3] + d3]; + } + } + } + } + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + CheckDeviceValues(expected.size(), d_output.get(), expected.data(), 1e-6f); +} + +TEST(ReductionFunctionsTest, ReduceSumNdIntegerSaturation) { + const std::vector dims{2, 3, 2}; + const std::vector axes{1}; + const int32_t big = 1'100'000'000; + const std::vector input(12, big); + const std::vector expected(4, std::numeric_limits::max()); + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(int32_t), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + std::vector actual(expected.size()); + cudaMemcpy(actual.data(), d_output.get(), actual.size() * sizeof(int32_t), cudaMemcpyDeviceToHost); + EXPECT_EQ(actual, expected); +} + +TEST(ReductionFunctionsTest, ReduceSumNdLargeReductionSmallOutput) { + const std::vector dims{2, 131072, 3}; + const std::vector axes{1}; + std::vector input(TensorShape(dims).Size(), 1.0f); + const std::vector expected(6, 131072.0f); + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + CheckDeviceValues(expected.size(), d_output.get(), expected.data(), 0.0f); +} + +TEST(ReductionFunctionsTest, ReduceSumNdInt64Cancellation) { + const std::vector dims{1, 3, 1}; + const std::vector axes{1}; + const int64_t large = int64_t{1} << 53; + const std::vector input{large, 1, -large}; + const std::vector expected{1}; + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(int64_t), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + std::vector actual(expected.size()); + cudaMemcpy(actual.data(), d_output.get(), actual.size() * sizeof(int64_t), cudaMemcpyDeviceToHost); + EXPECT_EQ(actual, expected); +} + +TEST(ReductionFunctionsTest, ReduceSumNdRank9) { + const std::vector dims{2, 2, 2, 2, 2, 2, 2, 2, 2}; + const std::vector axes{1, 3, 5, 7}; + std::vector input(TensorShape(dims).Size(), 1.0f); + const std::vector expected(32, 16.0f); + + auto d_input = AllocateDeviceMemory(input.size()); + auto d_output = AllocateDeviceMemory(expected.size()); + cudaMemcpy(d_input.get(), input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice); + + ASSERT_STATUS_OK(reduce_sum_nd(0, d_input.get(), d_output.get(), dims, axes)); + ASSERT_TRUE(CUDA_CALL(cudaDeviceSynchronize()).IsOK()); + CheckDeviceValues(expected.size(), d_output.get(), expected.data(), 0.0f); +} + TEST(ReductionFunctionsTest, BufferOffsets) { const int m = 2048; const int n = 1024; diff --git a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py index caff34b1c79d4..c95a03dc50d8e 100644 --- a/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py +++ b/onnxruntime/test/python/transformers/test_cuda_plugin_ep.py @@ -1804,6 +1804,25 @@ def expected(f): result = _run_model_test(target_device, "Softmax", model, feed, expected) self.assertEqual(result, TEST_PASS, "Softmax test failed") + def test_op_log_softmax(self): + target_device = get_cuda_plugin_device() + model = _make_simple_model( + "LogSoftmax", + [("X", TensorProto.FLOAT, [2, 5])], + [("Y", TensorProto.FLOAT, [2, 5])], + attrs={"axis": 1}, + opset=13, + ) + feed = {"X": np.random.rand(2, 5).astype(np.float32)} + + def expected(f): + x = f["X"] + shifted = x - np.max(x, axis=1, keepdims=True) + return shifted - np.log(np.sum(np.exp(shifted), axis=1, keepdims=True)) + + result = _run_model_test(target_device, "LogSoftmax", model, feed, expected) + self.assertEqual(result, TEST_PASS, "LogSoftmax test failed") + def test_op_relu(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -1900,7 +1919,6 @@ def test_op_flatten(self): result = _run_model_test(target_device, "Flatten", model, feed, lambda f: f["X"].reshape(2, 12)) self.assertEqual(result, TEST_PASS, "Flatten test failed") - @requires_cudnn def test_op_argmax(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -1916,6 +1934,21 @@ def test_op_argmax(self): ) self.assertEqual(result, TEST_PASS, "ArgMax test failed") + def test_op_argmin(self): + target_device = get_cuda_plugin_device() + model = _make_simple_model( + "ArgMin", + [("X", TensorProto.FLOAT, [3, 5])], + [("Y", TensorProto.INT64, [3, 1])], + attrs={"axis": 1, "keepdims": 1}, + opset=13, + ) + feed = {"X": np.random.rand(3, 5).astype(np.float32)} + result = _run_model_test( + target_device, "ArgMin", model, feed, lambda f: np.argmin(f["X"], axis=1).reshape(3, 1) + ) + self.assertEqual(result, TEST_PASS, "ArgMin test failed") + def test_op_topk(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -2030,23 +2063,21 @@ def expected(f): result = _run_model_test(target_device, "ConvTranspose", model, feed, expected) self.assertEqual(result, TEST_PASS, "ConvTranspose test failed") - @requires_cudnn def test_op_reduce_mean(self): target_device = get_cuda_plugin_device() model = _make_simple_model( "ReduceMean", [("X", TensorProto.FLOAT, [3, 4, 5])], - [("Y", TensorProto.FLOAT, [3, 1, 5])], - attrs={"axes": [1], "keepdims": 1}, + [("Y", TensorProto.FLOAT, [3, 4, 1])], + attrs={"axes": [2], "keepdims": 1}, opset=13, ) feed = {"X": np.random.rand(3, 4, 5).astype(np.float32)} result = _run_model_test( - target_device, "ReduceMean", model, feed, lambda f: np.mean(f["X"], axis=1, keepdims=True) + target_device, "ReduceMean", model, feed, lambda f: np.mean(f["X"], axis=2, keepdims=True) ) self.assertEqual(result, TEST_PASS, "ReduceMean test failed") - @requires_cudnn def test_op_reduce_sum(self): target_device = get_cuda_plugin_device() model = _make_simple_model( @@ -2064,6 +2095,32 @@ def test_op_reduce_sum(self): ) self.assertEqual(result, TEST_PASS, "ReduceSum test failed") + def _run_reduce_sum_integer_last_axis(self, onnx_dtype, np_dtype): + # Mirrors the qwen attention_mask usage: a rank-2 integer ReduceSum over the last axis. + # Integer ReduceSum takes a specialized path that does not use the float matrix fast path, + # so without cuDNN it must fall back to the general native kernel. Covering it here keeps the + # no-cuDNN CI honest for the exact layout that broke real models. + target_device = get_cuda_plugin_device() + model = _make_simple_model( + "ReduceSum", + [("X", onnx_dtype, [2, 8]), ("axes", TensorProto.INT64, [1])], + [("Y", onnx_dtype, [2, 1])], + attrs={"keepdims": 1}, + opset=13, + ) + axes_init = helper.make_tensor("axes", TensorProto.INT64, [1], [1]) + model.graph.initializer.append(axes_init) + feed = {"X": np.arange(16, dtype=np_dtype).reshape(2, 8)} + return _run_model_test(target_device, "ReduceSum", model, feed, lambda f: np.sum(f["X"], axis=1, keepdims=True)) + + def test_op_reduce_sum_int64_last_axis(self): + result = self._run_reduce_sum_integer_last_axis(TensorProto.INT64, np.int64) + self.assertEqual(result, TEST_PASS, "ReduceSum int64 test failed") + + def test_op_reduce_sum_int32_last_axis(self): + result = self._run_reduce_sum_integer_last_axis(TensorProto.INT32, np.int32) + self.assertEqual(result, TEST_PASS, "ReduceSum int32 test failed") + def test_op_gather_nd(self): target_device = get_cuda_plugin_device() model = _make_simple_model(