wehub-resource-sync
152 lines
4.4 KiB
Plaintext
152 lines
4.4 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_TENSOR_INFO_H
|
|
#define TRT_SCATTER_ELEMENTS_TENSOR_INFO_H
|
|
|
|
#include "common/plugin.h"
|
|
|
|
namespace nvinfer1
|
|
{
|
|
namespace plugin
|
|
{
|
|
namespace detail
|
|
{
|
|
|
|
static constexpr int32_t kMAX_TENSORINFO_DIMS = 25;
|
|
|
|
// CUDA kernel argument that defines tensor layout
|
|
template <typename TScalar, typename TIndex>
|
|
struct TensorInfo
|
|
{
|
|
TensorInfo();
|
|
TensorInfo(const TScalar* p, int32_t dim, TIndex sz[kMAX_TENSORINFO_DIMS], TIndex st[kMAX_TENSORINFO_DIMS]);
|
|
|
|
// Contiguous tensors of more than one dimension are collapsed down
|
|
// to one tensor
|
|
__host__ __device__ inline bool isContiguous() const
|
|
{
|
|
return (dims == 1 && strides[0] == 1);
|
|
}
|
|
|
|
const TScalar* data;
|
|
TIndex sizes[kMAX_TENSORINFO_DIMS];
|
|
TIndex strides[kMAX_TENSORINFO_DIMS];
|
|
int32_t dims;
|
|
};
|
|
|
|
// Creates TensorInfo object from PluginTensorDesc and data address
|
|
template <typename TScalar, typename TIndex>
|
|
TensorInfo<TScalar, TIndex> getTensorInfo(const void* d, PluginTensorDesc const& t)
|
|
{
|
|
TIndex sz[kMAX_TENSORINFO_DIMS];
|
|
TIndex st[kMAX_TENSORINFO_DIMS];
|
|
|
|
int32_t dims = t.dims.nbDims;
|
|
for (int32_t i = 0; i < dims; ++i)
|
|
{
|
|
sz[i] = t.dims.d[i];
|
|
}
|
|
for (int32_t i = dims; i < kMAX_TENSORINFO_DIMS; ++i)
|
|
{
|
|
sz[i] = static_cast<TIndex>(0);
|
|
}
|
|
// calculate strides
|
|
st[dims - 1] = 1;
|
|
for (int32_t i = dims - 2; i >= 0; --i)
|
|
{
|
|
st[i] = st[i + 1] * sz[i + 1];
|
|
}
|
|
return TensorInfo<TScalar, TIndex>(reinterpret_cast<const TScalar*>(d), dims, sz, st);
|
|
}
|
|
|
|
template <typename TScalar, typename TIndex>
|
|
TensorInfo<TScalar, TIndex>::TensorInfo()
|
|
{
|
|
data = nullptr;
|
|
dims = 0;
|
|
}
|
|
|
|
template <typename TScalar, typename TIndex>
|
|
TensorInfo<TScalar, TIndex>::TensorInfo(
|
|
const TScalar* p, int32_t dim, TIndex sz[kMAX_TENSORINFO_DIMS], TIndex st[kMAX_TENSORINFO_DIMS])
|
|
{
|
|
data = p;
|
|
dims = dim;
|
|
for (int32_t i = 0; i < dim; ++i)
|
|
{
|
|
sizes[i] = sz[i];
|
|
strides[i] = st[i];
|
|
}
|
|
}
|
|
|
|
// Translate a linear index for the apply to a T* offset;
|
|
// specialized on `Dims` to reduce nvcc compilation time
|
|
template <typename TScalar, typename TIndex, int tDims>
|
|
struct IndexToOffset
|
|
{
|
|
static __host__ __device__ TIndex get(TIndex linearId, const TensorInfo<TScalar, TIndex>& info)
|
|
{
|
|
|
|
TIndex offset = 0;
|
|
|
|
// Uses static dims
|
|
for (int32_t i = tDims - 1; i > 0; --i)
|
|
{
|
|
TIndex curDimIndex = linearId % info.sizes[i];
|
|
TIndex curDimOffset = curDimIndex * info.strides[i];
|
|
offset += curDimOffset;
|
|
linearId /= info.sizes[i];
|
|
}
|
|
|
|
return offset + linearId * info.strides[0];
|
|
}
|
|
};
|
|
|
|
// Uses dynamic (runtime) instead of static (compiletime) dims
|
|
template <typename TScalar, typename TIndex>
|
|
struct IndexToOffset<TScalar, TIndex, -1>
|
|
{
|
|
static inline __host__ __device__ TIndex get(TIndex linearId, const TensorInfo<TScalar, TIndex>& info)
|
|
{
|
|
|
|
TIndex offset = 0;
|
|
|
|
for (int32_t i = info.dims - 1; i > 0; --i)
|
|
{
|
|
TIndex curDimIndex = linearId % info.sizes[i];
|
|
TIndex curDimOffset = curDimIndex * info.strides[i];
|
|
offset += curDimOffset;
|
|
linearId /= info.sizes[i];
|
|
}
|
|
|
|
return offset + linearId * info.strides[0];
|
|
}
|
|
};
|
|
|
|
} // namespace detail
|
|
} // namespace plugin
|
|
} // namespace nvinfer1
|
|
|
|
#endif // TRT_SCATTER_ELEMENTS_TENSOR_INFO_H
|