Reject zero divisors in CUDA Mod kernels

onnxruntime/core/providers/cuda/cu_inc/common.cuh

@@ -501,6 +501,9 @@ __device__ __inline__ half _Gelu(half a) {
 
 template <typename T>
 __device__ __inline__ T _Mod(T a, T b) {
+  if (b == T(0)) {
+    return T(0);
+  }
   T r = a % b;
   T zero = T(0);
   if ((r > zero && b < zero) || (r < zero && b > zero)) {
@@ -511,6 +514,9 @@ __device__ __inline__ T _Mod(T a, T b) {
 
 template <typename T>
 __device__ __inline__ T _Fmod(T a, T b) {
+  if (b == T(0)) {
+    return T(0);
+  }
   return a % b;
 }
 

onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc

@@ -482,6 +482,35 @@ Status Mod::ComputeInternal(OpKernelContext* context) const {
                   element_type == on::TensorProto_DataType_INT64 || element_type == on::TensorProto_DataType_UINT32 ||
                   element_type == on::TensorProto_DataType_UINT64,
               "Non-fmod can support integer types only.");
+  if (element_type == on::TensorProto_DataType_INT32 || element_type == on::TensorProto_DataType_INT64 ||
+      element_type == on::TensorProto_DataType_UINT32 || element_type == on::TensorProto_DataType_UINT64) {
+    const size_t divisor_count = prepare.rhs_tensor->Shape().Size();
+    if (divisor_count > 0) {
+      auto has_zero_buffer = GetScratchBuffer<int>(1, GetComputeStream(context));
+      CUDA_RETURN_IF_ERROR(cudaMemsetAsync(has_zero_buffer.get(), 0, sizeof(int), Stream(context)));
+      switch (element_type) {
+        case on::TensorProto_DataType_INT32:
+          CheckZeroDivisor(Stream(context), prepare.rhs_tensor->Data<int32_t>(), divisor_count, has_zero_buffer.get());
+          break;
+        case on::TensorProto_DataType_INT64:
+          CheckZeroDivisor(Stream(context), prepare.rhs_tensor->Data<int64_t>(), divisor_count, has_zero_buffer.get());
+          break;
+        case on::TensorProto_DataType_UINT32:
+          CheckZeroDivisor(Stream(context), prepare.rhs_tensor->Data<uint32_t>(), divisor_count, has_zero_buffer.get());
+          break;
+        case on::TensorProto_DataType_UINT64:
+          CheckZeroDivisor(Stream(context), prepare.rhs_tensor->Data<uint64_t>(), divisor_count, has_zero_buffer.get());
+          break;
+      }
+
+      auto has_zero = AllocateBufferOnCPUPinned<int>(1);
+      CUDA_RETURN_IF_ERROR(
+          cudaMemcpyAsync(has_zero.get(), has_zero_buffer.get(), sizeof(int), cudaMemcpyDeviceToHost, Stream(context)));
+      CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(Stream(context)));
+      ORT_RETURN_IF(*has_zero != 0, "Integer modulo by zero");
+    }
+  }
+
 #define CASE_MOD_ELEMENT_TYPE(name, onnx_type, data_type)                                                           \
   case onnx_type: {                                                                                                 \
     Impl_##name<typename ToCudaType<data_type>::MappedType>(                                                        \

onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.cu

@@ -10,6 +10,26 @@
 namespace onnxruntime {
 namespace cuda {
 
+template <typename T>
+__global__ void CheckZeroDivisorKernel(const T* divisor_data, size_t count, int* has_zero) {
+  const size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (idx < count && divisor_data[idx] == T{0}) {
+    atomicExch(has_zero, 1);
+  }
+}
+
+template <typename T>
+void CheckZeroDivisor(cudaStream_t stream, const T* divisor_data, size_t count, int* has_zero) {
+  constexpr int threads_per_block = GridDim::maxThreadsPerBlock;
+  const int blocks_per_grid = static_cast<int>(CeilDiv(count, threads_per_block));
+  CheckZeroDivisorKernel<<<blocks_per_grid, threads_per_block, 0, stream>>>(divisor_data, count, has_zero);
+}
+
+template void CheckZeroDivisor<int32_t>(cudaStream_t, const int32_t*, size_t, int*);
+template void CheckZeroDivisor<int64_t>(cudaStream_t, const int64_t*, size_t, int*);
+template void CheckZeroDivisor<uint32_t>(cudaStream_t, const uint32_t*, size_t, int*);
+template void CheckZeroDivisor<uint64_t>(cudaStream_t, const uint64_t*, size_t, int*);
+
 #define BINARY_ELEMENTWISE_IMPL(name)                      \
   BINARY_ELEMENTWISE_IMPL_DECLARATION(name) {              \
     BinaryElementWiseImpl(stream,                          \

onnxruntime/core/providers/cuda/math/binary_elementwise_ops_impl.h

@@ -8,6 +8,9 @@
 namespace onnxruntime {
 namespace cuda {
 
+template <typename T>
+void CheckZeroDivisor(cudaStream_t stream, const T* divisor_data, size_t count, int* has_zero);
+
 // These macros simplifies coding. To add a new op with following steps:
 // 1. Add a new entry in BINARY_OPS() list
 // 2. (optional) Define templated single element operator in binary_elementwise_ops_impl.cu

onnxruntime/test/providers/cpu/math/element_wise_ops_test.cc

@@ -4831,6 +4831,20 @@ TEST(ModOpTest, Mod_int32_by_zero_constant_initializer) {
            {}, nullptr, &execution_providers);
 }
 
+#ifdef USE_CUDA
+TEST(ModOpTest, Mod_int64_by_zero_CUDA) {
+  OpTester test("Mod", ModOp_ver);
+  test.AddInput<int64_t>("X", {3}, {-3, 4, 7});
+  test.AddInput<int64_t>("Y", {3}, {0, 2, 3});
+  test.AddOutput<int64_t>("Z", {3}, {0, 0, 0});
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCudaExecutionProvider());
+  test.Run(OpTester::ExpectResult::kExpectFailure,
+           "Integer modulo by zero",
+           {}, nullptr, &execution_providers);
+}
+#endif
+
 TEST(BitShiftOpTest, SimpleLeft) {
   OpTester test("BitShift", 11);
   test.AddAttribute("direction", "LEFT");