wehub/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plugin/scatterElementsPlugin/reducer.cuh
+174
View File
@@ -0,0 +1,174 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 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
* ************************************************************************
*/
#ifndef TRT_SCATTER_ELEMENTS_REDUCER_H
#define TRT_SCATTER_ELEMENTS_REDUCER_H
#include <limits>
#include "atomics.cuh"
#include "scatterElementsPluginKernel.h"
namespace nvinfer1
{
namespace plugin
{
#define AT_DISPATCH_REDUCTION_TYPES(reduce, ...) \
[&] { \
switch (reduce) \
{ \
case ReductionType::kSUM: \
{ \
static constexpr ReductionType REDUCE = ReductionType::kSUM; \
return __VA_ARGS__(); \
} \
case ReductionType::kMEAN: \
{ \
static constexpr ReductionType REDUCE = ReductionType::kMEAN; \
return __VA_ARGS__(); \
} \
case ReductionType::kMUL: \
{ \
static constexpr ReductionType REDUCE = ReductionType::kMUL; \
return __VA_ARGS__(); \
} \
case ReductionType::kMIN: \
{ \
static constexpr ReductionType REDUCE = ReductionType::kMIN; \
return __VA_ARGS__(); \
} \
case ReductionType::kMAX: \
{ \
static constexpr ReductionType REDUCE = ReductionType::kMAX; \
return __VA_ARGS__(); \
} \
} \
}()
template <typename TScalar, ReductionType tReduce>
struct Reducer
{
static inline __host__ __device__ TScalar init()
{
if (tReduce == ReductionType::kMUL)
{
return TScalar(1);
}
else if (tReduce == ReductionType::kMIN)
{
return std::numeric_limits<TScalar>::max();
}
else if (tReduce == ReductionType::kMAX)
{
return std::numeric_limits<TScalar>::lowest();
}
else
{
return TScalar(0);
}
}
static inline __host__ __device__ void update(TScalar* val, TScalar newVal)
{
if (tReduce == ReductionType::kSUM || tReduce == ReductionType::kMEAN)
{
*val = *val + newVal;
}
else if (tReduce == ReductionType::kMUL)
{
*val = *val * newVal;
}
else if ((tReduce == ReductionType::kMIN && newVal < *val) || (tReduce == ReductionType::kMAX && newVal > *val))
{
*val = newVal;
}
}
static inline __host__ __device__ void update(TScalar* val, TScalar newVal, int64_t* arg, int64_t newArg)
{
if (tReduce == ReductionType::kSUM || tReduce == ReductionType::kMEAN)
{
*val = *val + newVal;
}
else if (tReduce == ReductionType::kMUL)
{
*val = *val * newVal;
}
else if ((tReduce == ReductionType::kMIN && newVal < *val) || (tReduce == ReductionType::kMAX && newVal > *val))
{
*val = newVal;
*arg = newArg;
}
}
static inline __host__ __device__ void write(
TScalar* address, TScalar val, int64_t* argAddress, int64_t arg, int count)
{
if (tReduce == ReductionType::kSUM || tReduce == ReductionType::kMUL)
{
*address = val;
}
else if (tReduce == ReductionType::kMEAN)
{
*address = val / (TScalar) (count > 0 ? count : 1);
}
else if (tReduce == ReductionType::kMIN || tReduce == ReductionType::kMAX)
{
if (count > 0)
{
*address = val;
*argAddress = arg;
}
else
{
*address = (TScalar) 0;
}
}
}
static inline __device__ void atomic_write(TScalar* address, TScalar val)
{
if (tReduce == ReductionType::kSUM || tReduce == ReductionType::kMEAN)
{
atomAdd(address, val);
}
else if (tReduce == ReductionType::kMUL)
{
atomMul(address, val);
}
else if (tReduce == ReductionType::kMIN)
{
atomMin(address, val);
}
else if (tReduce == ReductionType::kMAX)
{
atomMax(address, val);
}
}
};
} // namespace plugin
} // namespace nvinfer1
#endif // TRT_SCATTER_ELEMENTS_REDUCER_H