Detect intrinsic FP functions (instead of higher-level ones)

triscale-innov · Jun 22, 2019 · b070980 · b070980
1 parent 439a8cb
commit b070980
Showing 1 changed file with 25 additions and 13 deletions.
diff --git a/src/overdub.jl b/src/overdub.jl
@@ -2,25 +2,37 @@ using Cassette
 
 Cassette.@context CounterCtx;
 
-ops = [
-    (:add32, +, Float32, Float32),
-    (:sub32, -, Float32, Float32),
-    (:mul32, *, Float32, Float32),
-    (:div32, /, Float32, Float32),
-    (:add64, +, Float64, Float64),
-    (:sub64, -, Float64, Float64),
-    (:mul64, *, Float64, Float64),
-    (:div64, /, Float64, Float64),
-]
+const ops = (
+    (:add32, Core.Intrinsics.add_float, Float32),
+    (:sub32, Core.Intrinsics.sub_float, Float32),
+    (:mul32, Core.Intrinsics.mul_float, Float32),
+    (:div32, Core.Intrinsics.div_float, Float32),
+    (:add64, Core.Intrinsics.add_float, Float64),
+    (:sub64, Core.Intrinsics.sub_float, Float64),
+    (:mul64, Core.Intrinsics.mul_float, Float64),
+    (:div64, Core.Intrinsics.div_float, Float64),
+)
 
 @eval mutable struct Counter
     $((:($(op[1]) ::Int) for op in ops)...)
     Counter() = new($((0 for _ in 1:length(ops))...))
 end
 
-for (name, op, typ1, typ2) in ops
-    @eval function Cassette.prehook(ctx::CounterCtx, op::typeof($op), a::$typ1, b::$typ2)
-        ctx.metadata.$name +=1
+for typ1 in (Float32, Float64)
+    @eval function Cassette.prehook(ctx::CounterCtx,
+                                    op::Core.IntrinsicFunction,
+                                    ::$typ1,
+                                    ::$typ1)
+        $(Expr(:block,
+               (map(ops) do (name, op, typ2)
+                  typ1 == typ2 || return :nothing
+                  quote
+                    if op == $op
+                       ctx.metadata.$name += 1
+                       return
+                    end
+                  end
+                end)...))
     end
 end