Implement fastflops counter

bsc-quantic · Jan 18, 2024 · 32e21ac · 32e21ac
1 parent 0766f76
commit 32e21ac
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/src/Counters.jl b/src/Counters.jl
@@ -17,6 +17,24 @@ flops(sexpr::SizedEinExpr) =
 
 flops(expr::EinExpr, size) = flops(SizedEinExpr(expr, size))
 
+function fastflops(sexpr::SizedEinExpr)
+    if nargs(sexpr) == 0 || nargs(sexpr) == 1 && isempty(suminds(sexpr))
+        return 0
+    end
+
+    mapreduce(
+        log ∘ Base.Fix1(getindex, sexpr.size),
+        +,
+        Iterators.flatten((head(sexpr), suminds(sexpr))),
+        init = zero(Float64),
+    ) |>
+    exp |>
+    round |>
+    BigInt
+end
+
+fastflops(expr::EinExpr, size) = fastflops(SizedEinExpr(expr, size))
+
 """
     removedsize(path::EinExpr)