1 files changed, 16 insertions, 15 deletions

diff --git a/measure_startup_oh.c b/measure_startup_oh.c
index 56f9ea7..7c7ba14 100644
--- a/measure_startup_oh.c
+++ b/measure_startup_oh.c
@@ -12,31 +12,32 @@ extern char **environ;
 
 #include "testbench.h"
 
-#define measure_launch_oh "./measure_launch_oh"
+static char* child_argv_default[3] = {"./measure_launch_oh", "1", NULL};
 
 int main(int argc, char **argv) {
-        int *__unused, i;
         struct timespec start;
 
         if (argc == 2 && (!strcmp(argv[1], "--help") || !strcmp(argv[1], "-h"))) {
-                fprintf(stderr, "Usage: %s [environment var to apply to child (optional)]\n", argv[0]);
+                fprintf(stderr, "Usage: %s [-e environment var to apply to child (optional) | prog prog_arg...]\n", argv[0]);
                 return 1;
         }
 
-        static char* child_argv[3];
-        child_argv[0] = measure_launch_oh;
-        child_argv[1] = "1";
-        child_argv[2] = NULL;
-
-        // Use this to apply LD_PRELOAD to the child only
-        if (argv[1])
-                putenv(argv[1]);
+        char **child_argv = child_argv_default;
+        if (argv[1]) {
+                // Change env
+                // (Use this to apply LD_PRELOAD to the child only)
+                if (!strcmp(argv[1], "-e"))
+                        putenv(argv[2]);
+                // Special exec (e.g., ./nvtaskset --gpc-list 0 ./measure_launch_oh 1)
+                else
+                        child_argv = &argv[1];
+        }
 
         clock_gettime(CLOCK_MONOTONIC_RAW, &start);
-        fprintf(stderr, "(%d) Starting '"measure_launch_oh"' at CLOCK_MONOTONIC_RAW = %ld ns\n",
-                        getpid(), time2ns(start));
-        if (execve(measure_launch_oh, child_argv, environ)) {
-                perror("Unable to start '"measure_launch_oh"'");
+        fprintf(stderr, "(%d) Starting '%s' at CLOCK_MONOTONIC_RAW = %ld ns\n",
+                        getpid(), child_argv[0], time2ns(start));
+        if (execve(child_argv[0], child_argv, environ)) {
+                perror("Unable to start subprogram");
                 return 1;
         }