Line data    Source code

       1             : //* This file is part of the MOOSE framework
       2             : //* https://mooseframework.inl.gov
       3             : //*
       4             : //* All rights reserved, see COPYRIGHT for full restrictions
       5             : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
       6             : //*
       7             : //* Licensed under LGPL 2.1, please see LICENSE for details
       8             : //* https://www.gnu.org/licenses/lgpl-2.1.html
       9             : 
      10             : #ifdef MOOSE_LIBTORCH_ENABLED
      11             : 
      12             : #include "LibtorchUtils.h"
      13             : 
      14             : #include <algorithm>
      15             : #include <functional>
      16             : #include <numeric>
      17             : 
      18             : namespace LibtorchUtils
      19             : {
      20             : 
      21             : namespace
      22             : {
      23             : 
      24             : template <typename DataType>
      25             : torch::TensorOptions
      26          53 : tensorOptions()
      27             : {
      28             :   if constexpr (std::is_same<DataType, double>::value)
      29          53 :     return torch::TensorOptions().dtype(at::kDouble);
      30             :   else if constexpr (std::is_same<DataType, float>::value)
      31             :     return torch::TensorOptions().dtype(at::kFloat);
      32             :   else
      33             :     static_assert(Moose::always_false<DataType>,
      34             :                   "Tensor conversion is not implemented for the given data type!");
      35             : }
      36             : 
      37             : template <typename DataType>
      38             : void
      39          54 : checkTensorShape(const std::vector<DataType> & vector, c10::IntArrayRef sizes)
      40             : {
      41             :   const auto expected_numel =
      42          54 :       std::accumulate(sizes.begin(), sizes.end(), int64_t{1}, std::multiplies<int64_t>());
      43             : 
      44          54 :   if (expected_numel != cast_int<int64_t>(vector.size()))
      45           1 :     mooseError("The requested tensor shape is incompatible with the vector size.");
      46          53 : }
      47             : 
      48             : } // namespace
      49             : 
      50             : template <typename DataType>
      51             : torch::Tensor
      52          51 : vectorToTensorCopy(const std::vector<DataType> & vector, c10::IntArrayRef sizes)
      53             : {
      54          51 :   checkTensorShape(vector, sizes);
      55             : 
      56          51 :   const auto options = tensorOptions<DataType>();
      57          51 :   auto tensor = torch::empty(sizes, options);
      58             : 
      59          51 :   if (!vector.empty())
      60          51 :     std::copy(vector.begin(), vector.end(), tensor.template data_ptr<DataType>());
      61             : 
      62         102 :   return tensor;
      63           0 : }
      64             : 
      65             : // Explicitly instantiate for DataType=Real
      66             : template torch::Tensor vectorToTensorCopy<Real>(const std::vector<Real> & vector,
      67             :                                                 c10::IntArrayRef sizes);
      68             : 
      69             : template <typename DataType>
      70             : void
      71          51 : vectorToTensor(const std::vector<DataType> & vector, torch::Tensor & tensor, const bool detach)
      72             : {
      73          51 :   tensor = vectorToTensorCopy(vector, {long(vector.size()), 1});
      74             : 
      75          51 :   if (detach)
      76           0 :     tensor = tensor.detach();
      77          51 : }
      78             : 
      79             : // Explicitly instantiate for DataType=Real
      80             : template void
      81             : vectorToTensor<Real>(const std::vector<Real> & vector, torch::Tensor & tensor, const bool detach);
      82             : 
      83             : template <typename DataType>
      84             : torch::Tensor
      85           3 : vectorToTensorView(std::vector<DataType> & vector, c10::IntArrayRef sizes)
      86             : {
      87           3 :   checkTensorShape(vector, sizes);
      88             : 
      89           2 :   const auto options = tensorOptions<DataType>();
      90           2 :   if (vector.empty())
      91           1 :     return torch::empty(sizes, options);
      92             : 
      93           1 :   return torch::from_blob(vector.data(), sizes, options);
      94             : }
      95             : 
      96             : // Explicitly instantiate for DataType=Real
      97             : template torch::Tensor vectorToTensorView<Real>(std::vector<Real> & vector, c10::IntArrayRef sizes);
      98             : 
      99             : void
     100           0 : moveToLibtorchDevice(torch::Tensor & tensor, const torch::DeviceType device_type)
     101             : {
     102           0 :   tensor = tensor.to(device_type);
     103           0 : }
     104             : 
     105             : torch::Tensor
     106           2 : toCPUContiguous(const torch::Tensor & tensor)
     107             : {
     108             :   // CPU accessors can handle strides, but data_ptr()-based reads require dense logical order.
     109           2 :   return tensor.detach().to(tensor.options().device(at::kCPU)).contiguous();
     110             : }
     111             : 
     112             : template <typename DataType>
     113             : void
     114           1 : tensorToVector(torch::Tensor & tensor, std::vector<DataType> & vector)
     115             : {
     116             :   try
     117             :   {
     118           1 :     tensor.data_ptr<DataType>();
     119             :   }
     120           0 :   catch (const c10::Error & e)
     121             :   {
     122           0 :     mooseError(
     123           0 :         "Cannot cast tensor values to", MooseUtils::prettyCppType<DataType>(), "!\n", e.msg());
     124             :   }
     125             : 
     126           1 :   const auto & sizes = tensor.sizes();
     127             : 
     128           1 :   long int max_size = 0;
     129           3 :   for (const auto & dim_size : sizes)
     130             :     // We do this comparison because XCode complains if we use std::max
     131           2 :     max_size = dim_size > max_size ? dim_size : max_size;
     132             : 
     133             :   mooseAssert(max_size == tensor.numel(), "The given tensor should be one-dimensional!");
     134           2 :   vector = {tensor.data_ptr<DataType>(), tensor.data_ptr<DataType>() + tensor.numel()};
     135           1 : }
     136             : 
     137             : // Explicitly instantiate for DataType=Real
     138             : template void tensorToVector<Real>(torch::Tensor & tensor, std::vector<Real> & vector);
     139             : 
     140             : } // LibtorchUtils namespace
     141             : 
     142             : #endif