ALOHA: merge multi-coupling routines, and make FD gauge routines self-contained

.github/workflows/IOtest.yml

@@ -64,6 +64,17 @@ jobs:
             ./tests/test_manager.py testIO_test_pptt_fks_loonly testIO_test_pptt_fksreal testIO_test_pptt_fksrealew testIO_test_ppw_fksall testIO_test_tdecay_fksreal testIO_test_wprod_fksew -t0
 
 
+  # FD gauge output reference files (fortran DHELAS + madmatrix HelAmps)
+  IOtest_fd_gauge:
+    runs-on: ubuntu-24.04
+    steps:
+      - uses: actions/checkout@v5
+      - name: test FD gauge output IO tests
+        run: |
+            cd $GITHUB_WORKSPACE
+            ./tests/test_manager.py testIO_FDgauge_standalone_fortran testIO_FDgauge_madmatrix -pA -t0
+
+
   # C++ exporter IOTests (from unittest_20)
   IOtest_cpp_write:
     runs-on: ubuntu-24.04

.github/workflows/acceptancetest_mg7.yml

@@ -58,6 +58,23 @@ jobs:
             cd $GITHUB_WORKSPACE
             ./tests/test_manager.py test_mg7_ufo_aloha -pA -t0 -l INFO
 
+  # FD gauge: the madmatrix and the fortran standalone must compute the same
+  # value, and the scalar and vectorised madmatrix backends must agree with each
+  # other (the FD wavefunctions are written once per backend in helas_fd.h)
+  acceptancetest_mg7_fd_gauge:
+    needs: build_madspace
+    runs-on: ubuntu-24.04
+    if: github.event_name == 'push' || github.event.pull_request.head.repo.fork == true
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/checkout_mg5
+      - uses: ./.github/actions/install_madspace
+      - uses: ./.github/actions/restore-pip-cache
+      - name: test FD gauge madmatrix vs fortran
+        run: |
+            cd $GITHUB_WORKSPACE
+            ./tests/test_manager.py test_standalone_mg7_fd_vs_fortran test_standalone_mg7_fd_simd_lanes -pA -t0 -l INFO
+
   acceptancetest_mg7_output_directory:
     needs: build_madspace
     runs-on: ubuntu-24.04

aloha/create_aloha.py

@@ -59,7 +59,11 @@
 class AbstractRoutine(object):
     """ store the result of the computation of Helicity Routine
     this is use for storing and passing to writer """
-
+
+    # builders of the merged multiple coupling routines, shared by all the
+    # outgoing of a given set of Lorentz structures (see get_combined_routine)
+    combined_builder = {}
+
     def __init__(self, expr, outgoing, spins, name, infostr, model, denom=None):
         """ store the information """
 
@@ -84,10 +88,134 @@ def add_symmetry(self, outgoing):
 
     def add_combine(self, lor_list):
         """add a combine rule """
-        
+
         if lor_list not in self.combined:
             self.combined.append(lor_list)
 
+    def get_combined_routine(self, lor_names):
+        """Return the AbstractRoutine associated to the merged structure
+              Coup(1) * <structure of self> + Coup(i+1) * <structure of lor_names[i]>
+        i.e. the routine that a writer can output as a single subroutine (one
+        coupling argument per structure) instead of a wrapper calling each
+        single structure routine in turn.
+
+        The expression is built exactly like the one of a single coupling
+        routine, so the momenta, the propagator denominator and the temporary
+        variables are shared by all the structures.
+        Return None if such a merge is not possible (unknown Lorentz structure,
+        structures acting on different spins, loop routine, ...). In that case
+        the caller has to fall back on the wrapper form.
+        """
+
+        if not hasattr(self, 'combined_routine'):
+            self.combined_routine = {}
+
+        key = tuple(lor_names)
+        if key not in self.combined_routine:
+            try:
+                self.combined_routine[key] = self.compute_combined_routine(lor_names)
+            except Exception as error:
+                logger.debug('can not merge the routines %s (%s): %s',
+                             self.name, ','.join(lor_names), error)
+                self.combined_routine[key] = None
+
+        routine = self.combined_routine[key]
+        if routine is not None:
+            # the same merged routine is written once per tag set (the MP pass
+            # adds the 'MP' tag on the fly) -> keep the tag in sync.
+            routine.tag = list(self.tag)
+        return routine
+
+    def compute_combined_routine(self, lor_names):
+        """Compute the merged routine of get_combined_routine (no caching)"""
+
+        if self.model is None:
+            return None
+        if any(t.startswith('L') for t in self.tag):
+            # loop routines have their own (already explicit) combine mechanism
+            return None
+
+        l_lorentz = [getattr(self.model.lorentz, name)
+                     for name in (self.name,) + tuple(lor_names)]
+        if any(lor.spins != l_lorentz[0].spins for lor in l_lorentz[1:]):
+            return None
+
+        conjg = tuple(int(t[1:]) for t in self.tag if t.startswith('C'))
+        # the model is part of the key: two models can define different
+        # structures under the same Lorentz name
+        key = (self.model, tuple(lor.name for lor in l_lorentz), conjg)
+        if key in self.combined_builder:
+            builder = self.combined_builder[key]
+        else:
+            builder = CombineRoutineBuilder(l_lorentz, self.model)
+            if conjg:
+                builder = builder.define_conjugate_builder(conjg)
+            # the kernel is computed once and re-used for each outgoing
+            self.combined_builder[key] = builder
+
+        routine = builder.compute_routine(self.outgoing, list(self.tag))
+        # the merge does not change which particles are identical
+        routine.symmetries = list(self.symmetries)
+        routine.tag = list(self.tag)
+        return routine
+
+    def get_combined_routines(self, lor_names):
+        """Return the routines to sum -- one per Lorentz structure, this one
+        first -- when the structures do not act on the same spins, so that
+        get_combined_routine can not build a single expression for them.
+
+        This is the FD gauge situation: a massive vector and its Goldstone are
+        the very same wavefunction (components 1-4 and 5 of one object), so the
+        structures of a vertex mix V and S on a given leg. They still write
+        into one output and read one wavefunction per leg, so a writer can
+        assemble them into a single routine.
+        Return None when even that is not possible.
+        """
+
+        if not hasattr(self, 'combined_routines'):
+            self.combined_routines = {}
+
+        key = tuple(lor_names)
+        if key not in self.combined_routines:
+            try:
+                self.combined_routines[key] = self.compute_combined_routines(lor_names)
+            except Exception as error:
+                logger.debug('can not assemble the routines %s (%s): %s',
+                             self.name, ','.join(lor_names), error)
+                self.combined_routines[key] = None
+
+        routines = self.combined_routines[key]
+        if routines is not None:
+            for routine in routines:
+                routine.tag = list(self.tag)
+        return routines
+
+    def compute_combined_routines(self, lor_names):
+        """Compute the routines of get_combined_routines (no caching)"""
+
+        if self.model is None:
+            return None
+        if any(t.startswith('L') for t in self.tag):
+            return None
+
+        conjg = tuple(int(t[1:]) for t in self.tag if t.startswith('C'))
+        routines = [self]
+        for name in lor_names:
+            lorentz = getattr(self.model.lorentz, name)
+            if len(lorentz.spins) != len(self.spins):
+                # not the same legs: nothing to assemble them into
+                return None
+            key = (self.model, name, conjg)
+            if key in self.combined_builder:
+                builder = self.combined_builder[key]
+            else:
+                builder = AbstractRoutineBuilder(lorentz, self.model)
+                if conjg:
+                    builder = builder.define_conjugate_builder(conjg)
+                self.combined_builder[key] = builder
+            routines.append(builder.compute_routine(self.outgoing, list(self.tag)))
+        return routines
+
     def write(self, output_dir, language='Fortran', mode='self', combine=True, options=None, **opt):
         """ write the content of the object """
         # Set loop_mode based on the tag of this specific routine
@@ -699,7 +827,15 @@ def __init__(self, l_lorentz, model=None):
         self.outgoing = None
         self.lorentz_expr = []
         for i, lor in enumerate(l_lorentz):
-            self.lorentz_expr.append( 'Coup(%s) * (%s)' % (i+1, lor.structure))
+            structure = lor.structure
+            # AbstractRoutineBuilder.__init__ only expanded the formfactors of
+            # l_lorentz[0], and that expansion is overwritten here -> redo it
+            # for each structure entering the combination.
+            if getattr(lor, 'formfactors', None):
+                for formf in lor.formfactors:
+                    pat = re.compile(r'\b%s\b' % formf.name)
+                    structure = pat.sub('(%s)' % formf.value, structure)
+            self.lorentz_expr.append( 'Coup(%s) * (%s)' % (i+1, structure))
         self.lorentz_expr = ' + '.join(self.lorentz_expr)
         self.routine_kernel = None
         self.contracted = {}

aloha/template_files/madmatrix/helas_fd.h

@@ -493,14 +493,13 @@
 
     // FD gauge
      const cxtype_sv cI = cxmake( 0 + fptype_sv{ 0 },  1 + fptype_sv{ 0 }  );
-#ifdef MGONGPU_CPPSIMD
-    fptype_sv nA[5];
-    fptype_sv nB[5];
-#endif
     fptype_sv n[5];
     fptype_sv nk;
-    const fptype_sv zero{0.};
-    const fptype_sv one{1.};
+    // NB: broadcast to every SIMD lane. 'fptype_sv one{1.}' sets the first
+    // lane only (the others are zero filled), which left n, and then nk, at
+    // zero in all lanes but the first one
+    const fptype_sv zero = 0. + fptype_sv{ 0 };
+    const fptype_sv one = 1. + fptype_sv{ 0 };
 
     if( vmass != 0. )
     {
@@ -569,7 +568,7 @@
         w[1] = cxmake( -vmass/nk * n[1], zero );
         w[2] = cxmake( -vmass/nk * n[2], zero );
         w[3] = cxmake( -vmass/nk * n[3], zero );
-        w[4] = static_cast<fptype>(nsv)*cI;
+        w[4] = -static_cast<fptype>(nsv)*cI; // as in fortran vxxxxx (vc%W(5) = -nsv*ci) and in the SIMD branch below
       }
 
 #else
@@ -604,37 +603,33 @@
       w[2] = cxternary( mask, vcA_4, cxternary( maskB, vcB1_4, vcB2_4 ) );
       w[3] = cxternary( mask, vcA_5, vcB_5 );
 
-      //FD gauge
-      //branch A
-      nA[0] = fpternary( pvec0 >= zero , one , -one);
-      nA[1] = -pvec1/pp;
-      nA[2] = -pvec2/pp;
-      nA[3] = -pvec3/pp;
-      nA[4] = zero;
-
-      //branch B
-      nB[0] = nA[0];
-      nB[1] = zero;
-      nB[2] = zero;
-      nB[3] = -nA[0];
-
-      const fptype_sv b_A = fpternary(pp > zero, one , zero);
-      const fptype_sv b_B = fpternary(pp <= zero , one , zero);
-
-      n[0] = nA[0]*b_A + nB[0]*b_B;
-      n[1] = nA[1]*b_A + nB[1]*b_B;
-      n[2] = nA[2]*b_A + nB[2]*b_B;
-      n[3] = nA[3]*b_A + nB[3]*b_B;
-      n[4] = nA[4];
+      //FD gauge: same two branches as the scalar code above, selected lane by
+      //lane. The division uses ppDENOM (=1 where pp==0) so that the lanes that
+      //do not take it are not poisoned: a select discards the other value, but
+      //a nan surviving an arithmetic combination (nan*0 is nan) would not be.
+      const bool_v maskFD = ( pp > zero );
+      n[0] = fpternary( pvec0 >= zero , one , -one );
+      n[1] = fpternary( maskFD, -pvec1 / ppDENOM, zero );
+      n[2] = fpternary( maskFD, -pvec2 / ppDENOM, zero );
+      n[3] = fpternary( maskFD, -pvec3 / ppDENOM, -n[0] );
+      n[4] = zero;
 
       nk = n[0]*pvec0 - n[1]*pvec1 - n[2]*pvec2 - n[3]*pvec3;
 
-      const bool_v mask3 = { (abs(nhel) == 1 ? 1 : 0) }; // first element replicated
-      w[0] = cxternary( mask3, w[0], cxmake( -vmass/nk * n[0], zero));
-      w[1] = cxternary( mask3, w[1], cxmake( -vmass/nk * n[1], zero));
-      w[2] = cxternary( mask3, w[2], cxmake( -vmass/nk * n[2], zero));
-      w[3] = cxternary( mask3, w[3], cxmake( -vmass/nk * n[3], zero));
-      w[4] = cxternary( mask3, cxzero_sv(), -static_cast<fptype>(nsv)*cI);
+      // nhel is a scalar: no need for a per-lane mask here (and a bool_v built
+      // from a single value would only set the first lane)
+      if ( abs(nhel) == 1 )
+      {
+        w[4] = cxzero_sv();
+      }
+      else
+      {
+        w[0] = cxmake( -vmass/nk * n[0], zero );
+        w[1] = cxmake( -vmass/nk * n[1], zero );
+        w[2] = cxmake( -vmass/nk * n[2], zero );
+        w[3] = cxmake( -vmass/nk * n[3], zero );
+        w[4] = -static_cast<fptype>(nsv)*cI;
+      }
 #endif
     }
     else
@@ -1025,7 +1020,7 @@
       n[0] = fpternary( q[0].real() >= 0.f , one , -one );
       n[1] = zero;
       n[2] = zero;
-      n[3] = fpternary( q[0].real() >= 0.f , -one , one); //possible error in Fortran
+      n[3] = fpternary( q[0].real() >= 0.f , -one , one); // -sign(q0), as in fortran and python define_gauge_dir
       n[4] = zero;
     }
 #else
@@ -1034,7 +1029,7 @@
     n[0] = fpternary( q[0].real() >= 0.f , one , -one);
     n[1] = fpternary( qsign , -q[1].real() / qabs , zero );
     n[2] = fpternary( qsign , -q[2].real() / qabs , zero );
-    n[3] = fpternary( qsign , -q[3].real() / qabs , fpternary( q[0].real() >= 0.f , one , -one));
+    n[3] = fpternary( qsign , -q[3].real() / qabs , fpternary( q[0].real() >= 0.f , -one , one)); // same gauge as the branch above
     n[4] = zero;
 #endif
  }