Line data    Source code

       1             : /**********************************************************************/
       2             : /*                    DO NOT MODIFY THIS HEADER                       */
       3             : /*             Swift, a Fourier spectral solver for MOOSE             */
       4             : /*                                                                    */
       5             : /*            Copyright 2024 Battelle Energy Alliance, LLC            */
       6             : /*                        ALL RIGHTS RESERVED                         */
       7             : /**********************************************************************/
       8             : 
       9             : #include "MooseError.h"
      10             : #include "PlainTensorBuffer.h"
      11             : #include "SwiftUtils.h"
      12             : #include "DomainAction.h"
      13             : 
      14             : registerMooseObject("SwiftApp", PlainTensorBuffer);
      15             : 
      16             : InputParameters
      17        1214 : PlainTensorBuffer::validParams()
      18             : {
      19             :   InputParameters params = TensorBuffer<torch::Tensor>::validParams();
      20        2428 :   params.addParam<std::vector<int64_t>>("value_dimensions", {}, "Optional value dimensions");
      21        1214 :   return params;
      22           0 : }
      23             : 
      24         874 : PlainTensorBuffer::PlainTensorBuffer(const InputParameters & parameters)
      25         874 :   : TensorBuffer<torch::Tensor>(parameters)
      26             : {
      27         874 : }
      28             : 
      29             : void
      30         860 : PlainTensorBuffer::init()
      31             : {
      32        2580 :   _u = torch::zeros(_domain.getValueShape(getParam<std::vector<int64_t>>("value_dimensions")),
      33         860 :                     MooseTensor::floatTensorOptions());
      34         860 : }
      35             : 
      36             : void
      37       14384 : PlainTensorBuffer::makeCPUCopy()
      38             : {
      39       14384 :   if (!_u.defined())
      40             :     return;
      41             : 
      42       14384 :   if (_cpu_copy_requested)
      43             :   {
      44        1672 :     if (_u.is_cpu())
      45        2992 :       _u_cpu = _u.clone().contiguous();
      46             :     else
      47         352 :       _u_cpu = _u.cpu().contiguous();
      48             :   }
      49             : }