perf: add optimize option to reactant runs

LuxDL · Nov 3, 2024 · 4a02032 · 4a02032
1 parent cd80c8d
commit 4a02032
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/perf/resnet/lux_reactant.jl b/perf/resnet/lux_reactant.jl
@@ -18,6 +18,11 @@ function parse_commandline()
             help = "Model size"
             arg_type = Int
             default = 50
+
+        "--optimize"
+            help = "Optimization Level for Reactant"
+            arg_type = String
+            default = "all"
     end
     #! format: on
 
@@ -43,7 +48,9 @@ function main()
 
         x = rand(Float32, 224, 224, 3, b) |> dev
 
-        model_compiled = Reactant.compile(model, (x, ps_ra, st_ra); sync=true)
+        model_compiled = Reactant.compile(
+            model, (x, ps_ra, st_ra); sync=true, optimize=Symbol(parsed_args["optimize"])
+        )
 
         timings[b] = @belapsed $(model_compiled)($(x), $(ps_ra), $(st_ra))