From 32e21ac2ace7670d1a504fd8fbe0622a83450eb2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?=
 <sergio.sanchez.ramirez+git@bsc.es>
Date: Thu, 18 Jan 2024 12:44:40 +0100
Subject: [PATCH] Implement `fastflops` counter

---
 src/Counters.jl | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/src/Counters.jl b/src/Counters.jl
index 6b9427b..d23cf39 100644
--- 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)