non ml compile

Author: ajnonaka

Exec/GNUmakefile

@@ -1,5 +1,6 @@
 # AMREX_HOME defines the directory in which we will find all the AMReX code.
-AMREX_HOME = ../../amrex
+AMREX_HOME ?= ../../amrex
+SUNDIALS_HOME ?= ../../sundials/instdir
 
 DEBUG        = FALSE
 USE_MPI      = TRUE
@@ -12,11 +13,12 @@ USE_FFT      = TRUE
 USE_ML       = FALSE
 TINY_PROFILE = FALSE
 PROFILE      = FALSE
-
 USE_SUNDIALS = FALSE
-SUNDIALS_HOME ?= ../../sundials/instdir
 
 ifeq ($(USE_ML),TRUE)
+
+  CPPFLAGS += -DAMREX_USE_ML
+
   # Define a macro for the C++ preprocessor
   DEFINES += -DML_ENABLE -D_GLIBCXX_USE_CXX11_ABI=1
 

Source/Demag_ml.cpp

@@ -1,3 +1,5 @@
+#ifdef AMREX_USE_ML
+
 // MagneX_ML_Infer_Dynamic.cpp
 #include "MagneX.H"
 #include <torch/script.h>
@@ -318,4 +320,6 @@ void RunMLDemagOnBox(
 
     // 5) Unpack back to MultiFab
     UnpackTensorToHfieldDynamic(denorm, H_demagfield, mfi, bx, b.expected_spatial);
-}
+}
+
+#endif

Source/MagneX.H

@@ -176,13 +176,7 @@ void WritePlotfile(MultiFab& Ms,
                    const Real& time,
                    const int& plt_step);
 
-
-// void CalculateH_demag_ML(const Array< MultiFab, AMREX_SPACEDIM> &   Mfield,
-//                          torch::jit::script::Module& x_norm_module,
-//                          torch::jit::script::Module& ml_module,
-//                          torch::jit::script::Module& y_norm_module,
-//                         Array< MultiFab, AMREX_SPACEDIM> &   H_demagfield);
-
+#ifdef AMREX_USE_ML
 
 at::Tensor PackMfieldToTensorDynamic(
   const amrex::Array<amrex::MultiFab, AMREX_SPACEDIM>& Mfield,
@@ -216,3 +210,5 @@ void MoveModuleToDevice(torch::jit::script::Module& m,
 
 // Read expected_spatial = [nx, ny, nz] from normalizer module
 amrex::IntVect GetExpectedSpatial(torch::jit::script::Module& x_norm_module);
+
+#endif

Source/MagneX.cpp

@@ -234,6 +234,11 @@ void InitializeMagneXNamespace() {
 
     ml_enable = 0;
     pp.query("ml_enable",ml_enable);
+#ifndef AMREX_USE_ML
+    if (ml_enable == 1) {
+        amrex::Abort("ml_enable=1 requires USE_ML=TRUE");
+    }
+#endif
 
     diag_type = -1;
     pp.query("diag_type",diag_type);

Source/Make.package

@@ -1,6 +1,7 @@
 CEXE_sources += Checkpoint.cpp
 CEXE_sources += ComputeLLGRHS.cpp
 CEXE_sources += Demagnetization.cpp
+CEXE_sources += Demag_ml.cpp
 CEXE_sources += Diagnostics.cpp
 CEXE_sources += EffectiveAnisotropyField.cpp
 CEXE_sources += EffectiveDMIField.cpp
@@ -15,4 +16,3 @@ CEXE_headers += CartesianAlgorithm_K.H
 CEXE_headers += Demagnetization.H
 CEXE_headers += MagneX.H
 CEXE_headers += MagneX_namespace.H
-CEXE_sources += Demag_ml.cpp

Source/main.cpp

@@ -1,6 +1,5 @@
 #include "MagneX.H"
 #include "Demagnetization.H"
-#include <torch/script.h>
 #include <AMReX_MultiFab.H>
 #include <AMReX_VisMF.H>
 #include <AMReX_ParmParse.H>
@@ -11,8 +10,12 @@
 
 #include <cmath>
 
+#ifdef AMREX_USE_ML
+#include <torch/script.h>
 #include <ATen/cuda/CUDAContext.h>  // for at::cuda::setDevice
 #include <c10/cuda/CUDAGuard.h>
+#endif
+
 using namespace amrex;
 using namespace MagneX;
 
@@ -60,10 +63,11 @@ void main_main ()
     Array<MultiFab, AMREX_SPACEDIM> LLG_RHS;
     Array<MultiFab, AMREX_SPACEDIM> LLG_RHS_pre;
     Array<MultiFab, AMREX_SPACEDIM> LLG_RHS_avg;
+#ifdef AMREX_USE_ML
     torch::jit::script::Module ml_module;
     torch::jit::script::Module x_norm_module;
     torch::jit::script::Module y_norm_module;
-
+#endif
 
     // Declare variables for hysteresis
     Real normalized_Mx;
@@ -114,6 +118,7 @@ void main_main ()
     // **********************************
     // // LOAD PYTORCH MODEL
     if (ml_enable == 1) {
+#ifdef AMREX_USE_ML
         BL_PROFILE_VAR("LoadPytorch",LoadPytorch);
 
         // Load pytorch module via torch script
@@ -164,6 +169,7 @@ void main_main ()
                     << expected_spatial[2] << "\n";
 
         Print() << "Model loaded.\n";
+#endif
     }
     else {
         Print() << "ML disabled. Skipping model load.\n";
@@ -449,7 +455,7 @@ void main_main ()
             if (demag_coupling == 1) {
                 // demag_solver.CalculateH_demag(Mfield_old, H_demagfield);
                 if (ml_enable == 1) {
-                    // CalculateH_demag_ML(Mfield_old, x_norm_module, ml_module, y_norm_module, H_demagfield);
+#ifdef AMREX_USE_ML
                     for (amrex::MFIter mfi(Mfield_old[0], amrex::TilingIfNotGPU());
                             mfi.isValid(); ++mfi)
                         {
@@ -460,14 +466,14 @@ void main_main ()
                                 Mfield_old, mfi, bx, expected_spatial, device_id
                             );
 
-                            //
                             at::Tensor norm  = NormalizeInput(M_cuda_f32, x_norm_module);
                             at::Tensor pred  = MLForwardOnly(norm, ml_module);
                             at::Tensor denorm_f64 = DenormalizeOutput(pred, y_norm_module);
 
                             // unpack: tensor -> MultiFab
                             UnpackTensorToHfieldDynamic(denorm_f64, H_demagfield, mfi, bx, expected_spatial);
                     }
+#endif
                 } else {
                     // demag_solver.CalculateH_demag(Mfield_old, H_demagfield);
                     amrex::Gpu::streamSynchronize();
@@ -599,8 +605,7 @@ void main_main ()
                 if (demag_coupling == 1) {
                     // demag_solver.CalculateH_demag(Mfield, H_demagfield);
                     if (ml_enable == 1) {
-                        // CalculateH_demag_ML(Mfield, x_norm_module, ml_module, y_norm_module, H_demagfield);
-
+#ifdef AMREX_USE_ML
                         for (amrex::MFIter mfi(Mfield_old[0], amrex::TilingIfNotGPU());
                             mfi.isValid(); ++mfi)
                         {
@@ -618,6 +623,7 @@ void main_main ()
                             // unpack: tensor -> MultiFab
                             UnpackTensorToHfieldDynamic(denorm_f64, H_demagfield, mfi, bx, expected_spatial);
                         }
+#endif
                     } else {
                         // demag_solver.CalculateH_demag(Mfield, H_demagfield);
                         amrex::Gpu::streamSynchronize();
@@ -810,7 +816,7 @@ void main_main ()
                     if (fast_demag==1) {
                         // demag_solver.CalculateH_demag(ar_state, H_demagfield);
                         if (ml_enable == 1) {
-                            CalculateH_demag_ML(ar_state, x_norm_module, ml_module, y_norm_module, H_demagfield);
+                            amrex::Abort("add ML demag to fast dynamics");
                         } else {
                             demag_solver.CalculateH_demag(ar_state, H_demagfield);
                         }