wehub-resource-sync
175 lines
7.6 KiB
Plaintext
175 lines
7.6 KiB
Plaintext
/*
|
|
* 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
|