Program Listing for File TensorLayout.h

↰ Return to documentation for file (nvcv_types/include/nvcv/TensorLayout.h)

/*
 * SPDX-FileCopyrightText: Copyright (c) 2022-2023 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.
 */

#ifndef NVCV_TENSORLAYOUT_H
#define NVCV_TENSORLAYOUT_H

#include "Export.h"
#include "Status.h"
#include "detail/CompilerUtils.h"

#include <stdint.h>
#include <string.h>

#ifdef __cplusplus
extern "C"
{
#endif

#define NVCV_TENSOR_MAX_RANK (15)

typedef struct NVCVTensorLayoutRec
{
    // Not to be used directly.
    char    data[NVCV_TENSOR_MAX_RANK + 1]; // +1 for '\0'
    int32_t rank;
} NVCVTensorLayout;

typedef enum
{
    NVCV_TLABEL_BATCH   = 'N',
    NVCV_TLABEL_CHANNEL = 'C',
    NVCV_TLABEL_FRAME   = 'F',
    NVCV_TLABEL_DEPTH   = 'D',
    NVCV_TLABEL_HEIGHT  = 'H',
    NVCV_TLABEL_WIDTH   = 'W'
} NVCVTensorLabel;

#ifdef __cplusplus
#    define NVCV_TENSOR_LAYOUT_MAKE(layout) \
        NVCVTensorLayout                    \
        {                                   \
            layout, sizeof(layout) - 1      \
        }
#else
#    define NVCV_TENSOR_LAYOUT_MAKE(layout)            \
        {                                              \
            .data = layout, .rank = sizeof(layout) - 1 \
        }
#endif

NVCV_CONSTEXPR static const NVCVTensorLayout NVCV_TENSOR_NONE = NVCV_TENSOR_LAYOUT_MAKE("");

#define NVCV_DETAIL_DEF_TLAYOUT(LAYOUT) \
    NVCV_CONSTEXPR static const NVCVTensorLayout NVCV_TENSOR_##LAYOUT = NVCV_TENSOR_LAYOUT_MAKE(#LAYOUT);
#include "TensorLayoutDef.inc"
#undef NVCV_DETAIL_DEF_TLAYOUT

// clang-format off
NVCV_CONSTEXPR static const NVCVTensorLayout NVCV_TENSOR_IMPLICIT[7] =
{
    // Can't use the NVCV_TENSOR_* identifiers directly,
    // clang complains they are not compile-time constants.
    // We must resort to the make macros instead.
    NVCV_TENSOR_LAYOUT_MAKE(""), // none
    NVCV_TENSOR_LAYOUT_MAKE("W"),
    NVCV_TENSOR_LAYOUT_MAKE("HW"),
    NVCV_TENSOR_LAYOUT_MAKE("NHW"),
    NVCV_TENSOR_LAYOUT_MAKE("NCHW"),
    NVCV_TENSOR_LAYOUT_MAKE("NCDHW"),
    NVCV_TENSOR_LAYOUT_MAKE("NCFDHW"),
};
// clang-format on

NVCV_PUBLIC NVCVStatus nvcvTensorLayoutMake(const char *descr, NVCVTensorLayout *layout);

NVCV_PUBLIC NVCVStatus nvcvTensorLayoutMakeRange(const char *beg, const char *end, NVCVTensorLayout *layout);

NVCV_PUBLIC NVCVStatus nvcvTensorLayoutMakeFirst(NVCVTensorLayout in, int32_t n, NVCVTensorLayout *layout);

NVCV_PUBLIC NVCVStatus nvcvTensorLayoutMakeLast(NVCVTensorLayout in, int32_t n, NVCVTensorLayout *layout);

NVCV_PUBLIC NVCVStatus nvcvTensorLayoutMakeSubRange(NVCVTensorLayout in, int32_t beg, int32_t end,
                                                    NVCVTensorLayout *layout);

inline static int32_t nvcvTensorLayoutFindDimIndex(NVCVTensorLayout layout, char dimLabel, int idxStart)
{
    if (idxStart < 0)
    {
        idxStart = layout.rank + idxStart;
    }

    int n = layout.rank - idxStart;
    if (n > 0)
    {
        void *p = memchr(layout.data + idxStart, dimLabel, n);
        if (p != NULL)
        {
            return (int32_t)((char *)p - (char *)layout.data);
        }
    }
    return -1;
}

NVCV_CONSTEXPR inline static char nvcvTensorLayoutGetLabel(NVCVTensorLayout layout, int idx)
{
    // Must be all a single statement for C++11 compatibility
    return idx < 0 ? (0 <= layout.rank + idx && layout.rank + idx < layout.rank ? layout.data[layout.rank + idx] : '\0')
                   : (0 <= idx && idx < layout.rank ? layout.data[idx] : '\0');
}

NVCV_CONSTEXPR inline static int32_t nvcvTensorLayoutGetNumDim(NVCVTensorLayout layout)
{
    return layout.rank;
}

inline static int32_t nvcvTensorLayoutCompare(NVCVTensorLayout a, NVCVTensorLayout b)
{
    if (a.rank == b.rank)
    {
        return memcmp(a.data, b.data, a.rank);
    }
    else
    {
        return a.rank - b.rank;
    }
}

inline static int32_t nvcvTensorLayoutStartsWith(NVCVTensorLayout layout, NVCVTensorLayout test)
{
    if (test.rank <= layout.rank)
    {
        return memcmp(test.data, layout.data, test.rank) == 0;
    }
    else
    {
        return 0;
    }
}

inline static int32_t nvcvTensorLayoutEndsWith(NVCVTensorLayout layout, NVCVTensorLayout test)
{
    if (test.rank <= layout.rank)
    {
        return memcmp(test.data, layout.data + layout.rank - test.rank, test.rank) == 0;
    }
    else
    {
        return 0;
    }
}

NVCV_CONSTEXPR inline static const char *nvcvTensorLayoutGetName(const NVCVTensorLayout *layout)
{
    return layout == NULL ? "" : layout->data;
}

#ifdef __cplusplus
}
#endif

#endif // NVCV_TENSORLAYOUT_H