Implement staging the main harness loop of PolyBenchLauncher into Python.

Author: Andrija Kolic

sdk/mx.sdk/mx_sdk_benchmark.py

@@ -492,7 +492,7 @@ def __init__(
             else:
                 self._native_mode: bool = suite.is_native_mode(bmSuiteArgs)
                 self._post_processors: List[DataPointsPostProcessor] = suite._get_post_processors(
-                    self.benchmark, self._native_mode
+                    self.benchmark, self._native_mode, self.bmSuiteArgs
                 )
 
         @property
@@ -526,6 +526,11 @@ def run(self, benchmarks, bmSuiteArgs) -> DataPoints:
                 self._image_cache.add(self._current_image)
             return datapoints
 
+    def use_stage_aware_benchmark_mixin_intercept_run(self):
+        if self.jvm(self.execution_context.bmSuiteArgs) == "cpython":
+            return True
+        return False
+
     def _resolve_current_benchmark(self, benchmarks) -> ResolvedPolybenchBenchmark:
         if benchmarks is None or len(benchmarks) != 1:
             mx.abort(f"Must specify one benchmark at a time (given: {benchmarks})")
@@ -575,6 +580,9 @@ def _set_image_context(self, resolved_benchmark: ResolvedPolybenchBenchmark, bm_
 
     def _base_image_name(self) -> Optional[str]:
         """Overrides the image name used to build/run the image."""
+        if self.jvm(self.execution_context.bmSuiteArgs) == "cpython":
+            benchmark_sanitized = self.execution_context.benchmark.replace("/", "-").replace(".", "-")
+            return f"{benchmark_sanitized}-staged-benchmark"
         assert self._current_image, "Image should have been set already"
         return self._current_image.full_executable_name()
 
@@ -659,6 +667,9 @@ def _infer_host_vm_config(self, bm_suite_args, dims):
                 return existing_config
             return dims["host-vm-config"] + "-" + edition
         else:
+            non_graal_vms = ["cpython"]
+            if self.jvm(bm_suite_args) in non_graal_vms:
+                return self.jvmConfig(bm_suite_args)
             # assume config used when building a GraalVM distribution
             return "graal-enterprise-libgraal-pgo" if edition == "ee" else "graal-core-libgraal"
 
@@ -847,11 +858,11 @@ def process_datapoints(self, datapoints: DataPoints) -> DataPoints:
             fork_info = self._suite.execution_context.fork_info
             # When running non-native benchmarks there is no concept of stages, there is only a single bench suite run.
             # So we store a pretend "run" stage to indicate that all datapoints (for this fork) have already been produced.
-            current_stage = (
-                Stage.from_string("run")
-                if not self._suite.execution_context.native_mode
-                else self._suite.stages_info.current_stage
-            )
+            current_stage = Stage.from_string("run")
+            try:
+                current_stage = self._suite.stages_info.current_stage
+            except AttributeError:
+                pass
             # Preserve this fork-stage's datapoints for eventual post-processing.
             self._fork_stage_datapoints[(fork_info.current_fork_index, current_stage)] = datapoints
             if (
@@ -959,20 +970,40 @@ def update_fn(dp):
             )
             super().__init__(suite, selector_fn, key_fn, field, update_fn, lower_percentile, upper_percentile)
 
+    class GraalSpecificFieldsRemoverPostProcessor(DataPointsPostProcessor):
+        """
+        Removes all platform Graal specific fields from all the datapoints.
+        Used for cleaning up the bench results of a benchmark that runs on
+        a different platform (e.g. CPython).
+        The removed fields include:
+            * The "guest-vm" and "guest-vm-config" fields.
+            * All the "platform.*" fields.
+        """
+
+        def process_datapoints(self, datapoints: DataPoints) -> DataPoints:
+            return [{k: v for k, v in dp.items() if self._should_be_kept(k)} for dp in datapoints]
+
+        def _should_be_kept(self, key) -> bool:
+            return key not in ["guest-vm", "guest-vm-config"] and not key.startswith("platform.")
+
     def post_processors(self) -> List[DataPointsPostProcessor]:
         return self.execution_context.post_processors
 
-    def _get_post_processors(self, benchmark: str, native_mode: bool):
+    def _get_post_processors(self, benchmark: str, native_mode: bool, bm_suite_args: List[str]):
+        post_processors = []
+        if self.jvm(bm_suite_args) == "cpython":
+            post_processors += [PolybenchBenchmarkSuite.GraalSpecificFieldsRemoverPostProcessor()]
         if native_mode:
-            return [
+            post_processors += [
                 PolybenchBenchmarkSuite.NativeModeBenchmarkRenamingPostProcessor(self),
                 PolybenchBenchmarkSuite.NativeModeBuildSummaryPostProcessor(self, benchmark),
                 PolybenchBenchmarkSuite.NativeModeBenchmarkSummaryPostProcessor(self, benchmark),
             ]
         else:
-            return [
+            post_processors += [
                 PolybenchBenchmarkSuite.ServerModeBenchmarkSummaryPostProcessor(self, benchmark),
             ]
+        return post_processors
 
     @staticmethod
     def _get_metric_name(bench_output) -> Optional[str]:

truffle/mx.truffle/mx_polybench/model.py

@@ -24,13 +24,29 @@
  */
 package org.graalvm.polybench;
 
+import org.graalvm.polybench.ast.Decl.Compound;
+import org.graalvm.polybench.ast.Decl.Subroutine;
+import org.graalvm.polybench.ast.Decl.Variable;
+import org.graalvm.polybench.ast.Expr.Atom.Int;
+import org.graalvm.polybench.ast.Expr.Atom.Null;
+import org.graalvm.polybench.ast.Expr.BinaryOp;
+import org.graalvm.polybench.ast.Expr.ConstructorCall;
+import org.graalvm.polybench.ast.Expr.ListLengthCall;
+import org.graalvm.polybench.ast.Expr.Reference.Ident;
+import org.graalvm.polybench.ast.Operator;
+import org.graalvm.polybench.ast.Stat;
+import org.graalvm.polybench.ast.Stat.Assign;
+import org.graalvm.polybench.ast.Stat.Block;
+import org.graalvm.polybench.ast.Stat.Return;
+import org.graalvm.polybench.ast.Tree.Program.Builder;
+
 import java.util.Optional;
 
 /**
  * Summarizes the results of a polybench benchmark with an average value computed from iteration
  * datapoints.
  */
-class AverageSummary implements Summary {
+class AverageSummary implements Summary, StageableClass {
     @Override
     public Optional<Double> postprocess(double[] results) {
         double sum = 0;
@@ -44,4 +60,43 @@ public Optional<Double> postprocess(double[] results) {
     public String toString() {
         return "AverageSummary{}";
     }
+
+    @Override
+    public boolean stageClass(Builder programBuilder) {
+        Stat thenBranch;
+        Stat elseBranch;
+        try (Compound.Builder classBuilder = new Compound.Builder(programBuilder, "AverageSummary", null)) {
+            try (Subroutine.Builder methodBuilder = new Subroutine.Builder(classBuilder, "postprocess", new Variable[]{Variable.of("results")})) {
+                methodBuilder.append(Variable.of("sum", new Int(0)));
+                try (Stat.Foreach.Builder foreachBuilder = new Stat.Foreach.Builder(methodBuilder, Variable.of("result"), new Ident("results"))) {
+                    foreachBuilder.append(new Assign(new Ident("sum"),
+                                    new BinaryOp(new Ident("sum"), new Ident("result"), Operator.PLUS)));
+                }
+                methodBuilder.append(Variable.of("resultsLength", new ListLengthCall(new Ident("results"))));
+                // if (resultsLength > 0)
+                try (Block.Builder thenBranchBuilder = new Block.Builder(null)) {
+                    thenBranchBuilder.append(new Return(new BinaryOp(new Ident("sum"), new Ident("resultsLength"), Operator.DIV)));
+                    thenBranch = thenBranchBuilder.build();
+                }
+                try (Block.Builder elseBranchBuilder = new Block.Builder(null)) {
+                    elseBranchBuilder.append(new Return(new Null()));
+                    elseBranch = elseBranchBuilder.build();
+                }
+                methodBuilder.append(new Stat.If(new BinaryOp(new Ident("resultsLength"), new Int(0), Operator.GREATER_THAN), thenBranch, elseBranch));
+            }
+            return true;
+        } catch (UnsupportedOperationException e) {
+            return false;
+        }
+    }
+
+    @Override
+    public boolean stageSingletonInitialization(Builder programBuilder, String singleton) {
+        try {
+            programBuilder.append(Variable.of(singleton, ConstructorCall.of(new Ident("AverageSummary"))));
+            return true;
+        } catch (UnsupportedOperationException e) {
+            return false;
+        }
+    }
 }

truffle/src/org.graalvm.polybench/src/org/graalvm/polybench/AverageSummary.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020, 2021, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2020, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -46,12 +46,16 @@ public class Config {
     Metric metric;
     boolean evalSourceOnlyDefault;
     Summary summary;
+    Language stagingLanguage;
+    String stagingFilePath;
+    boolean logStagedProgram;
+    String stagedProgramLauncher;
 
     final List<String> unrecognizedArguments = new ArrayList<>();
 
     private static final int UNINITIALIZED_ITERATIONS = -1;
-    private static final int DEFAULT_WARMUP = 20;
-    private static final int DEFAULT_ITERATIONS = 30;
+    public static final int DEFAULT_WARMUP = 20;
+    public static final int DEFAULT_ITERATIONS = 30;
 
     /**
      * Multi-context runs related configuration.

truffle/src/org.graalvm.polybench/src/org/graalvm/polybench/Config.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.  Oracle designates this
+ * particular file as subject to the "Classpath" exception as provided
+ * by Oracle in the LICENSE file that accompanied this code.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package org.graalvm.polybench;
+
+enum Language {
+    PYTHON
+}

truffle/src/org.graalvm.polybench/src/org/graalvm/polybench/Language.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.  Oracle designates this
+ * particular file as subject to the "Classpath" exception as provided
+ * by Oracle in the LICENSE file that accompanied this code.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package org.graalvm.polybench;
+
+import org.graalvm.polybench.ast.Tree.Program;
+
+import java.io.IOException;
+
+/**
+ * Generates a program in the target language from an AST.
+ */
+public interface LanguageGenerator {
+    byte[] generate(Program program) throws IOException;
+}

truffle/src/org.graalvm.polybench/src/org/graalvm/polybench/LanguageGenerator.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.  Oracle designates this
+ * particular file as subject to the "Classpath" exception as provided
+ * by Oracle in the LICENSE file that accompanied this code.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package org.graalvm.polybench;
+
+class LanguageGeneratorFactory {
+    public static LanguageGenerator getLanguageGenerator(Language language) {
+        return switch (language) {
+            case PYTHON -> new PythonGenerator();
+        };
+    }
+}

truffle/src/org.graalvm.polybench/src/org/graalvm/polybench/LanguageGeneratorFactory.java

@@ -132,5 +132,4 @@ public Optional<Double> reportAfterAll() {
 
     public void reset() {
     }
-
 }