LibtorchUtils.C

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://mooseframework.inl.gov
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#ifdef MOOSE_LIBTORCH_ENABLED

#include "LibtorchUtils.h"

#include <algorithm>
#include <functional>
#include <numeric>

namespace LibtorchUtils
{

namespace
{

template <typename DataType>
torch::TensorOptions
tensorOptions()
{
  if constexpr (std::is_same<DataType, double>::value)
    return torch::TensorOptions().dtype(at::kDouble);
  else if constexpr (std::is_same<DataType, float>::value)
    return torch::TensorOptions().dtype(at::kFloat);
  else
    static_assert(Moose::always_false<DataType>,
                  "Tensor conversion is not implemented for the given data type!");
}

template <typename DataType>
void
checkTensorShape(const std::vector<DataType> & vector, c10::IntArrayRef sizes)
{
  const auto expected_numel =
      std::accumulate(sizes.begin(), sizes.end(), int64_t{1}, std::multiplies<int64_t>());

  if (expected_numel != cast_int<int64_t>(vector.size()))
    mooseError("The requested tensor shape is incompatible with the vector size.");
}

} // namespace

template <typename DataType>
torch::Tensor
vectorToTensorCopy(const std::vector<DataType> & vector, c10::IntArrayRef sizes)
{
  checkTensorShape(vector, sizes);

  const auto options = tensorOptions<DataType>();
  auto tensor = torch::empty(sizes, options);

  if (!vector.empty())
    std::copy(vector.begin(), vector.end(), tensor.template data_ptr<DataType>());

  return tensor;
}

// Explicitly instantiate for DataType=Real
template torch::Tensor vectorToTensorCopy<Real>(const std::vector<Real> & vector,
                                                c10::IntArrayRef sizes);

template <typename DataType>
void
vectorToTensor(const std::vector<DataType> & vector, torch::Tensor & tensor, const bool detach)
{
  tensor = vectorToTensorCopy(vector, {long(vector.size()), 1});

  if (detach)
    tensor = tensor.detach();
}

// Explicitly instantiate for DataType=Real
template void
vectorToTensor<Real>(const std::vector<Real> & vector, torch::Tensor & tensor, const bool detach);

template <typename DataType>
torch::Tensor
vectorToTensorView(std::vector<DataType> & vector, c10::IntArrayRef sizes)
{
  checkTensorShape(vector, sizes);

  const auto options = tensorOptions<DataType>();
  if (vector.empty())
    return torch::empty(sizes, options);

  return torch::from_blob(vector.data(), sizes, options);
}

// Explicitly instantiate for DataType=Real
template torch::Tensor vectorToTensorView<Real>(std::vector<Real> & vector, c10::IntArrayRef sizes);

void
moveToLibtorchDevice(torch::Tensor & tensor, const torch::DeviceType device_type)
{
  tensor = tensor.to(device_type);
}

torch::Tensor
toCPUContiguous(const torch::Tensor & tensor)
{
  // CPU accessors can handle strides, but data_ptr()-based reads require dense logical order.
  return tensor.detach().to(tensor.options().device(at::kCPU)).contiguous();
}

template <typename DataType>
void
tensorToVector(torch::Tensor & tensor, std::vector<DataType> & vector)
{
  try
  {
    tensor.data_ptr<DataType>();
  }
  catch (const c10::Error & e)
  {
    mooseError(
        "Cannot cast tensor values to", MooseUtils::prettyCppType<DataType>(), "!\n", e.msg());
  }

  const auto & sizes = tensor.sizes();

  long int max_size = 0;
  for (const auto & dim_size : sizes)
    // We do this comparison because XCode complains if we use std::max
    max_size = dim_size > max_size ? dim_size : max_size;

  mooseAssert(max_size == tensor.numel(), "The given tensor should be one-dimensional!");
  vector = {tensor.data_ptr<DataType>(), tensor.data_ptr<DataType>() + tensor.numel()};
}

// Explicitly instantiate for DataType=Real
template void tensorToVector<Real>(torch::Tensor & tensor, std::vector<Real> & vector);

} // LibtorchUtils namespace

#endif