chore: import upstream snapshot with attribution

2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
@@ -0,0 +1,161 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * ************************************************************************
+ * Modified from pytorch_scatter 
+ * Copyright (c) 2020 Matthias Fey <matthias.fey@tu-dortmund.de>
+ * See https://github.com/rusty1s/pytorch_scatter/blob/master/LICENSE for details
+ * ************************************************************************
+ */
+
+#include "TensorInfo.cuh"
+#include "common/dimsHelpers.h"
+#include "reducer.cuh"
+#include "scatterElementsPluginKernel.h"
+#include <thrust/device_vector.h>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+
+#define THREADS 256
+#define BLOCKS(N) (N + THREADS - 1) / THREADS
+
+using detail::TensorInfo;
+using detail::getTensorInfo;
+using nvinfer1::pluginInternal::volume;
+
+template <typename TScalar, ReductionType tReduce>
+__global__ void scatterElements_kernel(const TScalar* updatesData, const TensorInfo<int64_t, int32_t> indexInfo,
+    TScalar* outData, int32_t nE, int32_t nK, int32_t nN, int32_t nbElements)
+{
+
+    int32_t thread_idx = blockIdx.x * blockDim.x + threadIdx.x;
+
+    int32_t b = thread_idx / (nE * nK);
+    int32_t k = thread_idx % nK;
+
+    if (thread_idx < nbElements)
+    {
+        int32_t offset = detail::IndexToOffset<int64_t, int32_t, -1>::get(thread_idx, indexInfo);
+        int64_t idx = indexInfo.data[offset];
+
+        Reducer<TScalar, tReduce>::atomic_write(outData + b * nN * nK + idx * nK + k, updatesData[thread_idx]);
+    }
+}
+
+bool hasBfloat16AtomicAdd()
+{
+  int deviceId;
+  cudaGetDevice(&deviceId);
+  cudaDeviceProp deviceProp;
+  cudaGetDeviceProperties(&deviceProp, deviceId);
+  return deviceProp.major >= 8;
+}
+
+inline uint32_t getElementSize(nvinfer1::DataType t)
+{
+    switch (t)
+    {
+    case nvinfer1::DataType::kINT64: return 8;
+    case nvinfer1::DataType::kINT32:
+    case nvinfer1::DataType::kFLOAT: return 4;
+    case nvinfer1::DataType::kBF16:
+    case nvinfer1::DataType::kHALF: return 2;
+    case nvinfer1::DataType::kBOOL:
+    case nvinfer1::DataType::kUINT8:
+    case nvinfer1::DataType::kINT8:
+    case nvinfer1::DataType::kFP8: return 1;
+    case nvinfer1::DataType::kINT4:
+    case nvinfer1::DataType::kFP4:
+    case nvinfer1::DataType::kE8M0:
+        PLUGIN_FAIL("Unsupported data type");
+    }
+    return 0;
+}
+
+template <typename TScalar>
+void dispatchScatterElementsKernel(void* outDataPtr, void const* dataDataPtr, void const* updatesDataPtr,
+    void const* indicesDataPtr, PluginTensorDesc const& outDesc, PluginTensorDesc const& dataDesc,
+    PluginTensorDesc const& updatesDesc, PluginTensorDesc const& indicesDesc, int64_t axis, ReductionType reduction,
+    cudaStream_t stream)
+{
+    auto updatesNumEl = volume(updatesDesc.dims);
+    auto nB = 1;
+    for (auto i = 0; i < axis; i++)
+    {
+        nB *= updatesDesc.dims.d[i];
+    }
+    auto nE = updatesDesc.dims.d[axis];
+    auto nK = updatesNumEl / (nB * nE);
+    auto nN = outDesc.dims.d[axis];
+
+    auto indexInfo = getTensorInfo<int64_t, int32_t>(indicesDataPtr, indicesDesc);
+
+    auto updatesData = (TScalar*) updatesDataPtr;
+    auto outData = (TScalar*) outDataPtr;
+
+    AT_DISPATCH_REDUCTION_TYPES(reduction, [&] {
+        scatterElements_kernel<TScalar, REDUCE>
+            <<<BLOCKS(updatesNumEl), THREADS, 0, stream>>>(updatesData, indexInfo, outData, nE, nK, nN, updatesNumEl);
+    });
+}
+
+#define DISPATCH_RUN_KERNEL(TYPE)                                                                                      \
+    dispatchScatterElementsKernel<TYPE>(outDataPtr, dataDataPtr, updatesDataPtr, indicesDataPtr, outDesc, dataDesc,    \
+        updatesDesc, indicesDesc, axis, reduction, stream)
+
+void runScatterElementsKernel(void* outDataPtr, void const* dataDataPtr, void const* updatesDataPtr,
+    void const* indicesDataPtr, PluginTensorDesc const& outDesc, PluginTensorDesc const& dataDesc,
+    PluginTensorDesc const& updatesDesc, PluginTensorDesc const& indicesDesc, int64_t axis, ReductionType reduction,
+    cudaStream_t stream)
+
+{
+    auto updatesNumEl = volume(updatesDesc.dims);
+    auto outNumEl = volume(outDesc.dims);
+
+    // copy dataDataPtr data to outDataPtr area first
+    cudaMemcpyAsync(outDataPtr, dataDataPtr, getElementSize(outDesc.type) * outNumEl, cudaMemcpyDeviceToDevice, stream);
+
+    if (updatesNumEl == 0)
+    {
+        return;
+    }
+
+    switch (outDesc.type)
+    {
+    case nvinfer1::DataType::kFLOAT: DISPATCH_RUN_KERNEL(float); break;
+    case nvinfer1::DataType::kHALF: DISPATCH_RUN_KERNEL(__half); break;
+    case nvinfer1::DataType::kINT32: DISPATCH_RUN_KERNEL(int32_t); break;
+    case nvinfer1::DataType::kINT64: DISPATCH_RUN_KERNEL(int64_t); break;
+    case nvinfer1::DataType::kBF16:  DISPATCH_RUN_KERNEL(__nv_bfloat16); break;
+    case nvinfer1::DataType::kBOOL:
+    case nvinfer1::DataType::kUINT8:
+    case nvinfer1::DataType::kINT8:
+    case nvinfer1::DataType::kINT4:
+    case nvinfer1::DataType::kFP8:
+    case nvinfer1::DataType::kFP4:
+    case nvinfer1::DataType::kE8M0:
+      std::ostringstream stream;
+      stream << "Unsupported data type:" << (int)outDesc.type << std::endl;
+      PLUGIN_FAIL(stream.str().c_str());
+      break;
+    }
+}
+
+} // namespace plugin
+} // namespace nvinfer1