perf: FNOs

Author: avik-pal

perf/neuraloperators/main.jl

@@ -36,4 +36,26 @@ function run_deeponet_benchmarks()
     return nothing
 end
 
+function run_fno_benchmarks()
+    @info "Running FNO benchmarks"
+
+    model = FourierNeuralOperator((16, 16), 3, 8, 64)
+    ps, st = xdev(Lux.setup(Random.default_rng(), model))
+    x = xdev(rand(Float32, 64, 64, 1, 256))
+
+    primal_timings = Reactant.with_config(;
+        dot_general_precision=PrecisionConfig.HIGH,
+        convolution_precision=PrecisionConfig.HIGH,
+    ) do
+        benchmark_nn_primal(
+            model, x, ps, st; disable_scatter_gather_bench=true, disable_pad_bench=true
+        )
+    end
+
+    pretty_print_table(permutedims(hcat([[t...] for t in primal_timings]...), (2, 1)))
+
+    return nothing
+end
+
+run_fno_benchmarks()
 run_deeponet_benchmarks()