From 1334ffe46e57b3c414db1747fc636b7cedad1f63 Mon Sep 17 00:00:00 2001
From: Balint Joo <bjoo@nvidia.com>
Date: Thu, 7 Nov 2024 22:35:01 +0000
Subject: [PATCH] Exposed Params to set Matrix Acceleration for Coarse Ops

---
 .../ferm/invert/quda_solvers/quda_mg_utils.h  | 10 +++-
 .../quda_solvers/quda_multigrid_params.cc     | 49 +++++++++++++------
 .../quda_solvers/quda_multigrid_params.h      |  5 ++
 3 files changed, 48 insertions(+), 16 deletions(-)

diff --git a/lib/actions/ferm/invert/quda_solvers/quda_mg_utils.h b/lib/actions/ferm/invert/quda_solvers/quda_mg_utils.h
index 0eb3339c9..e72e10c6a 100644
--- a/lib/actions/ferm/invert/quda_solvers/quda_mg_utils.h
+++ b/lib/actions/ferm/invert/quda_solvers/quda_mg_utils.h
@@ -11,6 +11,7 @@
 #include "chromabase.h"
 
 #include <quda.h>
+#include <quda_arch.h>
 #include "meas/inline/io/named_objmap.h"
 #include "actions/ferm/invert/quda_solvers/syssolver_quda_multigrid_clover_params.h"
 
@@ -225,7 +226,14 @@ namespace Chroma {
 
 				} 
 
-				for (int i=0; i<mg_param.n_level; i++) {
+#ifdef QUDA_MMA_AVAILABLE 
+			  mg_param.dslash_use_mma[0] = QUDA_BOOLEAN_FALSE; // this value is ignored for now, set it to FALSE
+				for (int i=0; i < mg_param.n_level-1; ++i) {
+					mg_param.setup_use_mma[i] = QUDA_BOOLEAN_TRUE;
+					mg_param.dslash_use_mma[i+1] = ip.matrix_accelerate_coarse[i] ? QUDA_BOOLEAN_TRUE : QUDA_BOOLEAN_FALSE; 
+				} 
+#endif
+				for (int i=0; i < mg_param.n_level; i++) {
 					for (int j=0; j< Nd; j++) {
 						if( i < mg_param.n_level-1 ) {
 							mg_param.geo_block_size[i][j] = ip.blocking[i][j];
diff --git a/lib/actions/ferm/invert/quda_solvers/quda_multigrid_params.cc b/lib/actions/ferm/invert/quda_solvers/quda_multigrid_params.cc
index d876cb2c6..5e4e4e94b 100644
--- a/lib/actions/ferm/invert/quda_solvers/quda_multigrid_params.cc
+++ b/lib/actions/ferm/invert/quda_solvers/quda_multigrid_params.cc
@@ -97,20 +97,6 @@ namespace Chroma {
     readArray(paramtop, "SmootherSchwarzCycle", smootherSchwarzCycle, 1);
 
     read(paramtop, "NullVectors", nvec);
-    read(paramtop, "Pre-SmootherApplications", nu_pre);
-    read(paramtop, "Post-SmootherApplications", nu_post);
-    if (nvec.size() != mg_levels-1 ) {
- 
-      QDPIO::cout<<"Warning. There are "<< blocking.size() 
-		 << " blockings but only " << nvec.size() << " sets of NullVectors" << std::endl;
-      QDP_abort(1);
-    }
-    if (nu_pre.size() != mg_levels-1 ) {
- 
-      QDPIO::cout<<"Error. There are "<< (mg_levels-1)  
-		 << " blockings but only " << nu_pre.size() << " sets pre-smoothing iterations" << std::endl;
-      QDP_abort(1);
-    }
 
 		{
 			int paramcount = paramtop.count("NullVectorsBatchSize");
@@ -133,7 +119,39 @@ namespace Chroma {
 				nvec_batch.resize(mg_levels-1);
 				for( int i=0; i < mg_levels-1; i++) nvec_batch[i] = 1;
 			}	
-		}	
+		}
+	
+		{
+			int paramcount = paramtop.count("MatrixAccelerateCoarse");
+			if( paramcount == 1 ) {
+				read(paramtop, "MatrixAccelerateCoarse", matrix_accelerate_coarse);
+				if( matrix_accelerate_coarse.size() != mg_levels - 1 ) {
+					QDPIO::cout << "MatrixAccelerateCoarse, if given, needs to be of size mg_levels -1 "
+											<< "since it applies only to coarse levels\n";
+					QDP_abort(1);
+				}
+			}
+			else {
+				matrix_accelerate_coarse.resize(mg_levels-1);
+				for(int i=0; i < mg_levels-1; i++) matrix_accelerate_coarse[i] = false;
+			}
+		}
+
+
+    read(paramtop, "Pre-SmootherApplications", nu_pre);
+    read(paramtop, "Post-SmootherApplications", nu_post);
+    if (nvec.size() != mg_levels-1 ) {
+ 
+      QDPIO::cout<<"Warning. There are "<< blocking.size() 
+		 << " blockings but only " << nvec.size() << " sets of NullVectors" << std::endl;
+      QDP_abort(1);
+    }
+    if (nu_pre.size() != mg_levels-1 ) {
+ 
+      QDPIO::cout<<"Error. There are "<< (mg_levels-1)  
+		 << " blockings but only " << nu_pre.size() << " sets pre-smoothing iterations" << std::endl;
+      QDP_abort(1);
+    }
 
 
 
@@ -284,6 +302,7 @@ namespace Chroma {
     write(xml, "SchwarzType", p.schwarzType);
     write(xml, "NullVectors", p.nvec);
 		write(xml, "NullVectorsBatchSize", p.nvec_batch);
+		write(xml, "MatrixAccelerateCoarse", p.matrix_accelerate_coarse);
     write(xml, "MultiGridLevels", p.mg_levels);
     write(xml, "GenerateNullSpace", p.generate_nullspace);
     write(xml, "GenerateAllLevels", p.generate_all_levels);
diff --git a/lib/actions/ferm/invert/quda_solvers/quda_multigrid_params.h b/lib/actions/ferm/invert/quda_solvers/quda_multigrid_params.h
index 0614d05a0..66e9b6ec5 100644
--- a/lib/actions/ferm/invert/quda_solvers/quda_multigrid_params.h
+++ b/lib/actions/ferm/invert/quda_solvers/quda_multigrid_params.h
@@ -36,6 +36,7 @@ namespace Chroma
     multi1d<int> nu_post;
     multi1d< multi1d<int> > blocking;
 		multi1d<int> nvec_batch;
+		multi1d<bool> matrix_accelerate_coarse;
     int outer_gcr_nkrylov;
     int precond_gcr_nkrylov;
     std::string cycle_type;
@@ -67,6 +68,10 @@ namespace Chroma
       blocking.resize(mg_levels-1);
       nvec.resize(mg_levels-1);
 			nvec_batch.resize(mg_levels-1);
+			matrix_accelerate_coarse.resize(mg_levels-1);
+		  for(int i=0; i < mg_levels-1; i++) {
+			  matrix_accelerate_coarse[i] = false;
+			}
       nu_pre.resize(mg_levels-1);
       nu_post.resize(mg_levels-1);
       maxIterSubspaceCreate.resize(mg_levels-1);