ci: add runtime benchmarking to CI (and ViT)

Author: avik-pal

benchmark/Project.toml

@@ -3,7 +3,6 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Boltz = "4544d5e4-abc5-4dea-817f-29e4c205d9c8"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
-LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 

benchmark/benchmarks.jl

@@ -5,6 +5,7 @@ using Boltz, Lux, Random
 
 const SUITE = BenchmarkGroup()
 
+SUITE["runtime"] = BenchmarkGroup()
 SUITE["comptime"] = BenchmarkGroup()
 
 SUITE["comptime"]["basics"] = BenchmarkGroup()
@@ -20,18 +21,53 @@ end setup = begin
     a = Reactant.ConcreteRArray(ones(2, 10))
 end
 
+SUITE["runtime"]["lux neural networks"] = BenchmarkGroup()
 SUITE["comptime"]["lux neural networks"] = BenchmarkGroup()
 
-for depth in [11, 13, 16, 19]
-    SUITE["comptime"]["lux neural networks"]["vgg$depth"] = @benchmarkable Reactant.compile(
-        vgg, (x, ps_concrete, st_concrete)
-    ) setup = begin
-        vgg = Vision.VGG($depth; pretrained=false, batchnorm=false)
+for depth in [11, 13, 16, 19], batchnorm in [false,]#  true] <-- not working yet
+    SUITE["comptime"]["lux neural networks"]["vgg$(depth) bn=$(batchnorm)"] = @benchmarkable begin
+        @compile vgg(x, ps_concrete, st_concrete)
+    end setup = begin
+        vgg = Vision.VGG($depth; pretrained=false, batchnorm=$(batchnorm))
         ps, st = Lux.setup(Random.default_rng(), vgg)
         ps_concrete = Reactant.to_rarray(ps)
         st_concrete = Reactant.to_rarray(Lux.testmode(st))
         x = Reactant.to_rarray(rand(Float32, 224, 224, 3, 16))
     end
+
+    SUITE["runtime"]["lux neural networks"]["vgg$(depth) bn=$(batchnorm) (compiled)"] = @benchmarkable begin
+        vgg_compiled(x, ps_concrete, st_concrete)
+    end setup = begin
+        vgg = Vision.VGG($depth; pretrained=false, batchnorm=$(batchnorm))
+        ps, st = Lux.setup(Random.default_rng(), vgg)
+        ps_concrete = Reactant.to_rarray(ps)
+        st_concrete = Reactant.to_rarray(Lux.testmode(st))
+        x = Reactant.to_rarray(rand(Float32, 224, 224, 3, 16))
+        vgg_compiled = @compile vgg(x, ps_concrete, st_concrete)
+    end
+end
+
+for version in (:tiny, :base)
+    SUITE["comptime"]["lux neural networks"]["ViT $(version)"] = @benchmarkable begin
+        @compile vit(x, ps_concrete, st_concrete)
+    end setup = begin
+        vit = Vision.ViT($(Meta.quot(version)))
+        ps, st = Lux.setup(Random.default_rng(), vit)
+        ps_concrete = Reactant.to_rarray(ps)
+        st_concrete = Reactant.to_rarray(Lux.testmode(st))
+        x = Reactant.to_rarray(rand(Float32, 256, 256, 3, 16))
+    end
+
+    SUITE["runtime"]["lux neural networks"]["ViT $(version) (compiled)"] = @benchmarkable begin
+        vit_compiled(x, ps_concrete, st_concrete)
+    end setup = begin
+        vit = Vision.ViT($(Meta.quot(version)))
+        ps, st = Lux.setup(Random.default_rng(), vit)
+        ps_concrete = Reactant.to_rarray(ps)
+        st_concrete = Reactant.to_rarray(Lux.testmode(st))
+        x = Reactant.to_rarray(rand(Float32, 256, 256, 3, 16))
+        vit_compiled = @compile vit(x, ps_concrete, st_concrete)
+    end
 end
 
 function sumcos(x)