cnpy_utils.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); 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 _LBANN_CNPY_UTILS_HPP_
#define _LBANN_CNPY_UTILS_HPP_

#include "cnpy.h"
#include "lbann/utils/exception.hpp"
#include <string>
#include <vector>

namespace lbann {
namespace cnpy_utils {

size_t compute_cnpy_array_offset(const cnpy::NpyArray& na,
                                 std::vector<size_t> indices);

template <typename T>
inline size_t ptr_offset(const cnpy::NpyArray& na, std::vector<size_t> indices)
{
  if ((sizeof(T) != na.word_size) && (sizeof(T) != 1u)) {
    throw lbann_exception(
      std::string("cnpy_utils::ptr_offset() :") +
      "The data type is not consistent with the word size of the array.");
  }
  return (compute_cnpy_array_offset(na, indices) *
          ((sizeof(T) == 1u) ? na.word_size : 1u));
}

template <typename T>
inline T& data(const cnpy::NpyArray& na, const std::vector<size_t> indices)
{
  return *(reinterpret_cast<T*>(&(*na.data_holder)[0]) +
           ptr_offset<T>(na, indices));
}

template <typename T>
inline T* data_ptr(const cnpy::NpyArray& na, const std::vector<size_t> indices)
{
  return (reinterpret_cast<T*>(&(*na.data_holder)[0]) +
          ptr_offset<T>(na, indices));
}

template <>
inline void* data_ptr<void>(const cnpy::NpyArray& na,
                            const std::vector<size_t> indices)
{
  return data_ptr<uint8_t>(na, indices);
}

void shrink_to_fit(cnpy::NpyArray& na, size_t sz);

std::string show_shape(const cnpy::NpyArray& na);

} // end of namespace cnpy_utils
} // end of namespace lbann

#endif // _LBANN_CNPY_UTILS_HPP_