better driver scripts

PRONTOLab · Oct 6, 2024 · a9b3f97 · a9b3f97
1 parent 7a49094
commit a9b3f97
Show file tree

Hide file tree

Showing 4 changed files with 203 additions and 97 deletions.
diff --git a/example/fpopt-logged-driver-generator.py b/example/fpopt-logged-driver-generator.py
@@ -4,11 +4,12 @@
 import random
 import numpy as np
 
-num_samples_per_func = 100
+DEFAULT_NUM_SAMPLES = 10000
 default_regex = "ex\\d+"
 
 np.random.seed(42)
 
+
 def parse_bound(bound):
     if "/" in bound:
         numerator, denominator = map(float, bound.split("/"))
@@ -44,32 +45,59 @@ def parse_c_file(filepath, func_regex):
     return functions
 
 
-def create_driver_function(functions):
-    driver_code = ["int main() {"]
+def create_driver_function(functions, num_samples_per_func):
+    driver_code = [
+        "#include <iostream>",
+        "#include <random>",
+        "",
+        "int main() {",
+        "    std::mt19937 gen(42);",
+        "",
+    ]
     driver_code.append("    initializeLogger();")
-    driver_code.append("    volatile double res;")
 
     for func_name, bounds, params, return_type in functions:
-        print(f"Generating driver code for {func_name}")
-        for _ in range(num_samples_per_func):
-            call_params = []
-            for param in params:
-                param_tokens = param.strip().split()
-                if len(param_tokens) >= 2:
-                    param_name = param_tokens[-1]
-                else:
-                    exit(f"Cannot parse parameter: {param}")
-                try:
-                    min_val = bounds[param_name]["min"]
-                    max_val = bounds[param_name]["max"]
-                except KeyError:
-                    exit(
-                        f"WARNING: Bounds not found for {param_name} in function {func_name}, manually specify the bounds."
-                    )
-                random_value = np.random.uniform(min_val, max_val)
-                call_params.append(str(random_value))
-            driver_code.append(f"    res = __enzyme_autodiff<{return_type}>((void *) {func_name}, {', '.join(call_params)});")
+        for param in params:
+            param_tokens = param.strip().split()
+            if len(param_tokens) >= 2:
+                param_name = param_tokens[-1]
+            else:
+                exit(f"Cannot parse parameter: {param}")
+            try:
+                min_val = bounds[param_name]["min"]
+                max_val = bounds[param_name]["max"]
+            except KeyError:
+                exit(
+                    f"WARNING: Bounds not found for {param_name} in function {func_name}, manually specify the bounds."
+                )
+            dist_name = f"{func_name}_{param_name}_dist"
+            driver_code.append(f"    std::uniform_real_distribution<double> {dist_name}({min_val}, {max_val});")
+    driver_code.append("")
+
+    driver_code.append("    double res = 0.;")
+    driver_code.append("")
 
+    for func_name, bounds, params, return_type in functions:
+        driver_code.append(f"    for (int i = 0; i < {num_samples_per_func}; ++i) {{")
+
+        call_params = []
+        for param in params:
+            param_tokens = param.strip().split()
+            if len(param_tokens) >= 2:
+                param_name = param_tokens[-1]
+            else:
+                exit(f"Cannot parse parameter: {param}")
+            dist_name = f"{func_name}_{param_name}_dist"
+            param_value = f"{dist_name}(gen)"
+            call_params.append(param_value)
+
+        driver_code.append(
+            f"    res += __enzyme_autodiff<{return_type}>((void *) {func_name}, {', '.join(call_params)});"
+        )
+        driver_code.append("    }")
+        driver_code.append("")
+
+    driver_code.append('    std::cout << "Sum: " << res << std::endl;')
     driver_code.append("    printLogger();")
     driver_code.append("    destroyLogger();")
     driver_code.append("    return 0;")
@@ -79,16 +107,17 @@ def create_driver_function(functions):
 
 def main():
     if len(sys.argv) < 2:
-        exit("Usage: script.py <filepath> [function_regex]")
+        exit("Usage: script.py <filepath> [function_regex] [num_samples_per_func (default: 10000)]")
 
     filepath = sys.argv[1]
     func_regex = sys.argv[2] if len(sys.argv) > 2 else default_regex
+    num_samples_per_func = int(sys.argv[3]) if len(sys.argv) > 3 else DEFAULT_NUM_SAMPLES
 
     if len(sys.argv) <= 2:
         print(f"WARNING: No regex provided for target function names. Using default regex: {default_regex}")
 
     functions = parse_c_file(filepath, func_regex)
-    driver_code = create_driver_function(functions)
+    driver_code = create_driver_function(functions, num_samples_per_func)
     new_filepath = os.path.splitext(filepath)[0] + "-logged.cpp"
 
     with open(filepath, "r") as original_file:

diff --git a/example/fpopt-original-driver-generator.py b/example/fpopt-original-driver-generator.py
@@ -4,11 +4,12 @@
 import random
 import numpy as np
 
-num_samples_per_func = 100
+DEFAULT_NUM_SAMPLES = 10000
 default_regex = "ex\\d+"
 
 np.random.seed(42)
 
+
 def parse_bound(bound):
     if "/" in bound:
         numerator, denominator = map(float, bound.split("/"))
@@ -44,48 +45,75 @@ def parse_c_file(filepath, func_regex):
     return functions
 
 
-def create_driver_function(functions):
-    driver_code = ["int main() {"]
-    driver_code.append("    volatile double res;")
+def create_driver_function(functions, num_samples_per_func):
+    driver_code = [
+        "#include <iostream>",
+        "#include <random>",
+        "",
+        "int main() {",
+        "    std::mt19937 gen(42);",
+        "",
+    ]
 
     for func_name, bounds, params, return_type in functions:
-        print(f"Generating driver code for {func_name}")
-        for _ in range(num_samples_per_func):
-            call_params = []
-            for param in params:
-                param_tokens = param.strip().split()
-                if len(param_tokens) >= 2:
-                    param_name = param_tokens[-1]
-                else:
-                    exit(f"Cannot parse parameter: {param}")
-                try:
-                    min_val = bounds[param_name]["min"]
-                    max_val = bounds[param_name]["max"]
-                except KeyError:
-                    exit(
-                        f"WARNING: Bounds not found for {param_name} in function {func_name}, manually specify the bounds."
-                    )
-                random_value = np.random.uniform(min_val, max_val)
-                call_params.append(str(random_value))
-            driver_code.append(f"    res = {func_name}({', '.join(call_params)});")
+        for param in params:
+            param_tokens = param.strip().split()
+            if len(param_tokens) >= 2:
+                param_name = param_tokens[-1]
+            else:
+                exit(f"Cannot parse parameter: {param}")
+            try:
+                min_val = bounds[param_name]["min"]
+                max_val = bounds[param_name]["max"]
+            except KeyError:
+                exit(
+                    f"WARNING: Bounds not found for {param_name} in function {func_name}, manually specify the bounds."
+                )
+            dist_name = f"{func_name}_{param_name}_dist"
+            driver_code.append(f"    std::uniform_real_distribution<double> {dist_name}({min_val}, {max_val});")
+    driver_code.append("")
+
+    driver_code.append("    double res = 0.;")
+    driver_code.append("")
 
+    for func_name, bounds, params, return_type in functions:
+        driver_code.append(f"    for (int i = 0; i < {num_samples_per_func}; ++i) {{")
+
+        call_params = []
+        for param in params:
+            param_tokens = param.strip().split()
+            if len(param_tokens) >= 2:
+                param_name = param_tokens[-1]
+            else:
+                exit(f"Cannot parse parameter: {param}")
+            dist_name = f"{func_name}_{param_name}_dist"
+            param_value = f"{dist_name}(gen)"
+            call_params.append(param_value)
+
+        driver_code.append(f"        res += {func_name}({', '.join(call_params)});")
+        driver_code.append("        std::cout << res << std::endl;")
+        driver_code.append("    }")
+        driver_code.append("")
+
+    driver_code.append('    std::cout << "Sum: " << res << std::endl;')
     driver_code.append("    return 0;")
     driver_code.append("}")
     return "\n".join(driver_code)
 
 
 def main():
     if len(sys.argv) < 2:
-        exit("Usage: script.py <filepath> [function_regex]")
+        exit("Usage: script.py <filepath> [function_regex] [num_samples_per_func (default: 10000)]")
 
     filepath = sys.argv[1]
     func_regex = sys.argv[2] if len(sys.argv) > 2 else default_regex
+    num_samples_per_func = int(sys.argv[3]) if len(sys.argv) > 3 else DEFAULT_NUM_SAMPLES
 
     if len(sys.argv) <= 2:
         print(f"WARNING: No regex provided for target function names. Using default regex: {default_regex}")
 
     functions = parse_c_file(filepath, func_regex)
-    driver_code = create_driver_function(functions)
+    driver_code = create_driver_function(functions, num_samples_per_func)
     new_filepath = os.path.splitext(filepath)[0] + ".cpp"
 
     with open(filepath, "r") as original_file:

diff --git a/experiments/fpopt-logged-driver-generator.py b/experiments/fpopt-logged-driver-generator.py
@@ -46,33 +46,58 @@ def parse_c_file(filepath, func_regex):
 
 
 def create_driver_function(functions, num_samples_per_func):
-    driver_code = ["int main() {"]
+    driver_code = [
+        "#include <iostream>",
+        "#include <random>",
+        "",
+        "int main() {",
+        "    std::mt19937 gen(42);",
+        "",
+    ]
     driver_code.append("    initializeLogger();")
-    driver_code.append("    volatile double res;")
 
     for func_name, bounds, params, return_type in functions:
-        print(f"Generating driver code for {func_name}")
-        for _ in range(num_samples_per_func):
-            call_params = []
-            for param in params:
-                param_tokens = param.strip().split()
-                if len(param_tokens) >= 2:
-                    param_name = param_tokens[-1]
-                else:
-                    exit(f"Cannot parse parameter: {param}")
-                try:
-                    min_val = bounds[param_name]["min"]
-                    max_val = bounds[param_name]["max"]
-                except KeyError:
-                    exit(
-                        f"WARNING: Bounds not found for {param_name} in function {func_name}, manually specify the bounds."
-                    )
-                random_value = np.random.uniform(min_val, max_val)
-                call_params.append(str(random_value))
-            driver_code.append(
-                f"    res = __enzyme_autodiff<{return_type}>((void *) {func_name}, {', '.join(call_params)});"
-            )
+        for param in params:
+            param_tokens = param.strip().split()
+            if len(param_tokens) >= 2:
+                param_name = param_tokens[-1]
+            else:
+                exit(f"Cannot parse parameter: {param}")
+            try:
+                min_val = bounds[param_name]["min"]
+                max_val = bounds[param_name]["max"]
+            except KeyError:
+                exit(
+                    f"WARNING: Bounds not found for {param_name} in function {func_name}, manually specify the bounds."
+                )
+            dist_name = f"{func_name}_{param_name}_dist"
+            driver_code.append(f"    std::uniform_real_distribution<double> {dist_name}({min_val}, {max_val});")
+    driver_code.append("")
+
+    driver_code.append("    double res = 0.;")
+    driver_code.append("")
 
+    for func_name, bounds, params, return_type in functions:
+        driver_code.append(f"    for (int i = 0; i < {num_samples_per_func}; ++i) {{")
+
+        call_params = []
+        for param in params:
+            param_tokens = param.strip().split()
+            if len(param_tokens) >= 2:
+                param_name = param_tokens[-1]
+            else:
+                exit(f"Cannot parse parameter: {param}")
+            dist_name = f"{func_name}_{param_name}_dist"
+            param_value = f"{dist_name}(gen)"
+            call_params.append(param_value)
+
+        driver_code.append(
+            f"    res += __enzyme_autodiff<{return_type}>((void *) {func_name}, {', '.join(call_params)});"
+        )
+        driver_code.append("    }")
+        driver_code.append("")
+
+    driver_code.append('    std::cout << "Sum: " << res << std::endl;')
     driver_code.append("    printLogger();")
     driver_code.append("    destroyLogger();")
     driver_code.append("    return 0;")

diff --git a/experiments/fpopt-original-driver-generator.py b/experiments/fpopt-original-driver-generator.py
@@ -46,32 +46,56 @@ def parse_c_file(filepath, func_regex):
 
 
 def create_driver_function(functions, num_samples_per_func):
-    driver_code = ["#include <iostream>"]
-    driver_code.append("int main() {")
-    driver_code.append("    volatile double res;")
+    driver_code = [
+        "#include <iostream>",
+        "#include <random>",
+        "",
+        "int main() {",
+        "    std::mt19937 gen(42);",
+        "",
+    ]
 
     for func_name, bounds, params, return_type in functions:
-        print(f"Generating driver code for {func_name}")
-        for _ in range(num_samples_per_func):
-            call_params = []
-            for param in params:
-                param_tokens = param.strip().split()
-                if len(param_tokens) >= 2:
-                    param_name = param_tokens[-1]
-                else:
-                    exit(f"Cannot parse parameter: {param}")
-                try:
-                    min_val = bounds[param_name]["min"]
-                    max_val = bounds[param_name]["max"]
-                except KeyError:
-                    exit(
-                        f"WARNING: Bounds not found for {param_name} in function {func_name}, manually specify the bounds."
-                    )
-                random_value = np.random.uniform(min_val, max_val)
-                call_params.append(str(random_value))
-            driver_code.append(f"    res = {func_name}({', '.join(call_params)});")
-
-    driver_code.append("    std::cout << res << std::endl;")
+        for param in params:
+            param_tokens = param.strip().split()
+            if len(param_tokens) >= 2:
+                param_name = param_tokens[-1]
+            else:
+                exit(f"Cannot parse parameter: {param}")
+            try:
+                min_val = bounds[param_name]["min"]
+                max_val = bounds[param_name]["max"]
+            except KeyError:
+                exit(
+                    f"WARNING: Bounds not found for {param_name} in function {func_name}, manually specify the bounds."
+                )
+            dist_name = f"{func_name}_{param_name}_dist"
+            driver_code.append(f"    std::uniform_real_distribution<double> {dist_name}({min_val}, {max_val});")
+    driver_code.append("")
+
+    driver_code.append("    double res = 0.;")
+    driver_code.append("")
+
+    for func_name, bounds, params, return_type in functions:
+        driver_code.append(f"    for (int i = 0; i < {num_samples_per_func}; ++i) {{")
+
+        call_params = []
+        for param in params:
+            param_tokens = param.strip().split()
+            if len(param_tokens) >= 2:
+                param_name = param_tokens[-1]
+            else:
+                exit(f"Cannot parse parameter: {param}")
+            dist_name = f"{func_name}_{param_name}_dist"
+            param_value = f"{dist_name}(gen)"
+            call_params.append(param_value)
+
+        driver_code.append(f"        res += {func_name}({', '.join(call_params)});")
+        driver_code.append("        std::cout << res << std::endl;")
+        driver_code.append("    }")
+        driver_code.append("")
+
+    driver_code.append('    std::cout << "Sum: " << res << std::endl;')
     driver_code.append("    return 0;")
     driver_code.append("}")
     return "\n".join(driver_code)