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             : #pragma once
      10             : 
      11             : #include "TensorOperatorBase.h"
      12             : 
      13             : /**
      14             :  * Compute group with internal dependency resolution
      15             :  */
      16             : class ComputeGroup : public TensorOperatorBase
      17             : {
      18             : public:
      19             :   static InputParameters validParams();
      20             : 
      21             :   ComputeGroup(const InputParameters & parameters);
      22             : 
      23             :   virtual void init() override;
      24             : 
      25             :   virtual void computeBuffer() override;
      26             : 
      27             :   virtual void updateDependencies() override;
      28             : 
      29           8 :   std::size_t getComputeCount() const { return _compute_count; }
      30             : 
      31             : protected:
      32             :   /// nested tensor computes
      33             :   std::vector<std::shared_ptr<TensorOperatorBase>> _computes;
      34             : 
      35             :   /// for diagnostic purposes we can make sure that every requested buffer is defined
      36             :   typedef std::vector<std::tuple<const torch::Tensor *, std::string, std::string>>
      37             :       CheckedTensorList;
      38             :   std::vector<CheckedTensorList> _checked_tensors;
      39             : 
      40             :   bool _visited;
      41             : 
      42             :   std::size_t _compute_count;
      43             : };