From 7014eb323600deafefed8e307ebdd099c551dba6 Mon Sep 17 00:00:00 2001
From: Gregory Becker <becker33@llnl.gov>
Date: Thu, 23 Feb 2023 15:40:34 -0800
Subject: [PATCH] matrices: broadcast key combinatorially applies to all nodes
 in matrix

---
 lib/spack/spack/schema/env.py     |  4 ++++
 lib/spack/spack/spec_list.py      | 26 ++++++++++++++++++++++----
 lib/spack/spack/test/spec_list.py |  9 +++++++++
 3 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/lib/spack/spack/schema/env.py b/lib/spack/spack/schema/env.py
index 88c1ef1c85b..b96958559ba 100644
--- a/lib/spack/spack/schema/env.py
+++ b/lib/spack/spack/schema/env.py
@@ -30,6 +30,10 @@
                         "type": "array",
                         "items": {"type": "array", "items": {"type": "string"}},
                     },
+                    "broadcast": {
+                        "type": "array",
+                        "items": {"type": "array", "items": {"type": "string"}},
+                    },
                     "exclude": {"type": "array", "items": {"type": "string"}},
                 },
             },
diff --git a/lib/spack/spack/spec_list.py b/lib/spack/spack/spec_list.py
index 9efa8cbb50f..c5f1de0667c 100644
--- a/lib/spack/spack/spec_list.py
+++ b/lib/spack/spack/spec_list.py
@@ -186,9 +186,13 @@ def _expand_matrix_constraints(matrix_config):
                 new_row.append([r])
         expanded_rows.append(new_row)
 
+    # TODO someday: allow matrices inside `broadcast`
+    broadcast_rows = matrix_config.get("broadcast", [])
     excludes = matrix_config.get("exclude", [])  # only compute once
     sigil = matrix_config.get("sigil", "")
 
+    broadcast_constraints = list(itertools.product(*broadcast_rows))
+
     results = []
     for combo in itertools.product(*expanded_rows):
         # Construct a combined spec to test against excludes
@@ -212,15 +216,29 @@ def _expand_matrix_constraints(matrix_config):
         if any(test_spec.satisfies(x) for x in excludes):
             continue
 
-        if sigil:
-            flat_combo[0] = Spec(sigil + str(flat_combo[0]))
+        # If no broadcast, this is [(,)].
+        # It will run once, as required, and apply no constraints
+        for broadcast_combo in broadcast_constraints:
+            final_combo = [_apply_broadcast(spec.copy(), broadcast_combo) for spec in flat_combo]
 
-        # Add to list of constraints
-        results.append(flat_combo)
+            if sigil:
+                final_combo[0] = Spec(sigil + str(final_combo[0]))
+
+            # Add to list of constraints
+            results.append(final_combo)
 
     return results
 
 
+def _apply_broadcast(spec, constraints):
+    if constraints:
+        for node in spec.traverse():
+            if node.name:
+                for constraint in constraints:
+                    node.constrain(constraint)
+    return spec
+
+
 def _sigilify(item, sigil):
     if isinstance(item, dict):
         if sigil:
diff --git a/lib/spack/spack/test/spec_list.py b/lib/spack/spack/test/spec_list.py
index 7df87542b0d..ff62f6bd6de 100644
--- a/lib/spack/spack/test/spec_list.py
+++ b/lib/spack/spack/test/spec_list.py
@@ -214,3 +214,12 @@ def test_spec_list_constraints_with_structure(
         speclist = SpecList("specs", [matrix])
         assert len(speclist.specs) == 1
         assert libdwarf_spec in speclist.specs[0]
+
+    def test_spec_list_broadcast(self, mock_packages):
+        matrix = {"matrix": [["mpileaks"], ["^callpath"]], "broadcast": [["%gcc", "%clang"]]}
+        speclist = SpecList("specs", [matrix])
+
+        assert len(speclist) == 2
+        for spec in speclist:
+            for node in spec.traverse():
+                assert node.compiler.name == spec.compiler.name