Skip to content

Commit

Permalink
fix primitive_type for complex (#193)
Browse files Browse the repository at this point in the history 
* fix primitive_type for complex

* add simple complex runtime test
  • Loading branch information
@Pangoraw
Pangoraw authored Oct 23, 2024
1 parent c507e37 commit 35a51ad
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 2 deletions.
4 changes: 2 additions & 2 deletions src/XLA.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -233,8 +233,8 @@ end

@inline primitive_type(::Type{Float64}) = 12

@inline primitive_type(::Type{Complex{Float32}}) = 24
@inline primitive_type(::Type{Complex{Float64}}) = 25
@inline primitive_type(::Type{Complex{Float32}}) = 15
@inline primitive_type(::Type{Complex{Float64}}) = 18

function ArrayFromHostBuffer(client::Client, array::Array{T,N}, device) where {T,N}
sizear = Int64[s for s in reverse(size(array))]
Expand Down
7 changes: 7 additions & 0 deletions test/basic.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -438,3 +438,10 @@ end
@test size(f(y)) == size(x)
@test eltype(f(y)) == eltype(x)
end

@testset "Complex runtime: $CT" for CT in (ComplexF32, ComplexF64)
a = Reactant.to_rarray(ones(CT, 2))
b = Reactant.to_rarray(ones(CT, 2))
c = Reactant.compile(+, (a, b))(a, b)
@test c == ones(CT, 2) + ones(CT, 2)
end

1 comment on commit 35a51ad

@github-actions
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reactant.jl Benchmarks

Benchmark suite Current: 35a51ad Previous: c507e37 Ratio
ViT base (256 x 256 x 3 x 32)/forward/CUDA/Reactant (optimize = :after_enzyme) 1233718266 ns 1236544173 ns 1.00
ViT base (256 x 256 x 3 x 32)/forward/CUDA/Reactant 1272290171 ns 1242881142 ns 1.02
ViT base (256 x 256 x 3 x 32)/forward/CUDA/Reactant (optimize = :before_enzyme) 1414092991 ns 1235658910 ns 1.14
ViT base (256 x 256 x 3 x 32)/forward/CUDA/Reactant (optimize = :only_enzyme) 2507794592 ns 2350423218 ns 1.07
ViT base (256 x 256 x 3 x 32)/forward/CUDA/Lux 218244448 ns 217352004 ns 1.00
ViT base (256 x 256 x 3 x 32)/forward/CPU/Reactant (optimize = :after_enzyme) 7057019589 ns 6794443727 ns 1.04
ViT base (256 x 256 x 3 x 32)/forward/CPU/Reactant 5136395044 ns 5221674868 ns 0.98
ViT base (256 x 256 x 3 x 32)/forward/CPU/Reactant (optimize = :before_enzyme) 5202106378 ns 5093314901 ns 1.02
ViT base (256 x 256 x 3 x 32)/forward/CPU/Reactant (optimize = :only_enzyme) 6826278112 ns 6885416164 ns 0.99
ViT base (256 x 256 x 3 x 32)/forward/CPU/Lux 33666785970 ns 33707808014 ns 1.00
ViT small (256 x 256 x 3 x 4)/forward/CUDA/Reactant (optimize = :after_enzyme) 1308209141 ns 1266808372 ns 1.03
ViT small (256 x 256 x 3 x 4)/forward/CUDA/Reactant 1354474554.5 ns 1273495665.5 ns 1.06
ViT small (256 x 256 x 3 x 4)/forward/CUDA/Reactant (optimize = :before_enzyme) 1312596491 ns 1333702124.5 ns 0.98
ViT small (256 x 256 x 3 x 4)/forward/CUDA/Reactant (optimize = :only_enzyme) 2573078535 ns 2512672037 ns 1.02
ViT small (256 x 256 x 3 x 4)/forward/CUDA/Lux 8770241 ns 8492497 ns 1.03
ViT small (256 x 256 x 3 x 4)/forward/CPU/Reactant (optimize = :after_enzyme) 1571674207.5 ns 1623142223 ns 0.97
ViT small (256 x 256 x 3 x 4)/forward/CPU/Reactant 1540710272.5 ns 1600931077 ns 0.96
ViT small (256 x 256 x 3 x 4)/forward/CPU/Reactant (optimize = :before_enzyme) 1548078395 ns 1598434969 ns 0.97
ViT small (256 x 256 x 3 x 4)/forward/CPU/Reactant (optimize = :only_enzyme) 2744115264 ns 2912977764 ns 0.94
ViT small (256 x 256 x 3 x 4)/forward/CPU/Lux 2213143753 ns 2178291674 ns 1.02
ViT tiny (256 x 256 x 3 x 32)/forward/CUDA/Reactant (optimize = :after_enzyme) 1270044962 ns 1191252514 ns 1.07
ViT tiny (256 x 256 x 3 x 32)/forward/CUDA/Reactant 1234436938 ns 1267623530.5 ns 0.97
ViT tiny (256 x 256 x 3 x 32)/forward/CUDA/Reactant (optimize = :before_enzyme) 1294952462.5 ns 1254309112.5 ns 1.03
ViT tiny (256 x 256 x 3 x 32)/forward/CUDA/Reactant (optimize = :only_enzyme) 2715588750 ns 2522324425 ns 1.08
ViT tiny (256 x 256 x 3 x 32)/forward/CUDA/Lux 22668152 ns 21024085 ns 1.08
ViT tiny (256 x 256 x 3 x 32)/forward/CPU/Reactant (optimize = :after_enzyme) 2134055837 ns 2219640487 ns 0.96
ViT tiny (256 x 256 x 3 x 32)/forward/CPU/Reactant 2138464342 ns 2218982964 ns 0.96
ViT tiny (256 x 256 x 3 x 32)/forward/CPU/Reactant (optimize = :before_enzyme) 2124865845 ns 2252600530 ns 0.94
ViT tiny (256 x 256 x 3 x 32)/forward/CPU/Reactant (optimize = :only_enzyme) 3323398587 ns 3553615972 ns 0.94
ViT tiny (256 x 256 x 3 x 32)/forward/CPU/Lux 6167302164 ns 5451831944 ns 1.13
ViT tiny (256 x 256 x 3 x 4)/forward/CUDA/Reactant (optimize = :after_enzyme) 1255998954 ns 1219179894 ns 1.03
ViT tiny (256 x 256 x 3 x 4)/forward/CUDA/Reactant 1314721136.5 ns 1266307789 ns 1.04
ViT tiny (256 x 256 x 3 x 4)/forward/CUDA/Reactant (optimize = :before_enzyme) 1279651269 ns 1292867362.5 ns 0.99
ViT tiny (256 x 256 x 3 x 4)/forward/CUDA/Reactant (optimize = :only_enzyme) 2554634529 ns 2395083528 ns 1.07
ViT tiny (256 x 256 x 3 x 4)/forward/CUDA/Lux 7369398 ns 7182253.5 ns 1.03
ViT tiny (256 x 256 x 3 x 4)/forward/CPU/Reactant (optimize = :after_enzyme) 1402852118 ns 1478653783 ns 0.95
ViT tiny (256 x 256 x 3 x 4)/forward/CPU/Reactant 1403963317 ns 1473188659 ns 0.95
ViT tiny (256 x 256 x 3 x 4)/forward/CPU/Reactant (optimize = :before_enzyme) 1400054223.5 ns 1458846498 ns 0.96
ViT tiny (256 x 256 x 3 x 4)/forward/CPU/Reactant (optimize = :only_enzyme) 2604456962 ns 2745979768 ns 0.95
ViT tiny (256 x 256 x 3 x 4)/forward/CPU/Lux 1231519358 ns 1076729677.5 ns 1.14
ViT tiny (256 x 256 x 3 x 16)/forward/CUDA/Reactant (optimize = :after_enzyme) 1307469859.5 ns 1307047788 ns 1.00
ViT tiny (256 x 256 x 3 x 16)/forward/CUDA/Reactant 1271107436.5 ns 1277025925 ns 1.00
ViT tiny (256 x 256 x 3 x 16)/forward/CUDA/Reactant (optimize = :before_enzyme) 1331859396 ns 1226039283.5 ns 1.09
ViT tiny (256 x 256 x 3 x 16)/forward/CUDA/Reactant (optimize = :only_enzyme) 2683313614 ns 2531008275 ns 1.06
ViT tiny (256 x 256 x 3 x 16)/forward/CUDA/Lux 12276652 ns 11116198 ns 1.10
ViT tiny (256 x 256 x 3 x 16)/forward/CPU/Reactant (optimize = :after_enzyme) 1688567666 ns 1763174402 ns 0.96
ViT tiny (256 x 256 x 3 x 16)/forward/CPU/Reactant 1688900936 ns 1762263171 ns 0.96
ViT tiny (256 x 256 x 3 x 16)/forward/CPU/Reactant (optimize = :before_enzyme) 1687635701 ns 1755205670 ns 0.96
ViT tiny (256 x 256 x 3 x 16)/forward/CPU/Reactant (optimize = :only_enzyme) 2910962217 ns 3044952264 ns 0.96
ViT tiny (256 x 256 x 3 x 16)/forward/CPU/Lux 2991919710 ns 3217666637.5 ns 0.93
ViT small (256 x 256 x 3 x 16)/forward/CUDA/Reactant (optimize = :after_enzyme) 1290428008 ns 1396484090 ns 0.92
ViT small (256 x 256 x 3 x 16)/forward/CUDA/Reactant 1278577051 ns 1263883093.5 ns 1.01
ViT small (256 x 256 x 3 x 16)/forward/CUDA/Reactant (optimize = :before_enzyme) 1247880063.5 ns 1266345521 ns 0.99
ViT small (256 x 256 x 3 x 16)/forward/CUDA/Reactant (optimize = :only_enzyme) 2595166528 ns 2519876095 ns 1.03
ViT small (256 x 256 x 3 x 16)/forward/CUDA/Lux 27317118 ns 25484298.5 ns 1.07
ViT small (256 x 256 x 3 x 16)/forward/CPU/Reactant (optimize = :after_enzyme) 2166869626 ns 2255146771 ns 0.96
ViT small (256 x 256 x 3 x 16)/forward/CPU/Reactant 2153166296 ns 2293550587 ns 0.94
ViT small (256 x 256 x 3 x 16)/forward/CPU/Reactant (optimize = :before_enzyme) 2169198327 ns 2308771124 ns 0.94
ViT small (256 x 256 x 3 x 16)/forward/CPU/Reactant (optimize = :only_enzyme) 3389215278 ns 3541388172 ns 0.96
ViT small (256 x 256 x 3 x 16)/forward/CPU/Lux 7315062267 ns 8606549877 ns 0.85
ViT small (256 x 256 x 3 x 32)/forward/CUDA/Reactant (optimize = :after_enzyme) 1303537244 ns 1246919786 ns 1.05
ViT small (256 x 256 x 3 x 32)/forward/CUDA/Reactant 1357078784 ns 1246468994 ns 1.09
ViT small (256 x 256 x 3 x 32)/forward/CUDA/Reactant (optimize = :before_enzyme) 1310042021 ns 1242420953 ns 1.05
ViT small (256 x 256 x 3 x 32)/forward/CUDA/Reactant (optimize = :only_enzyme) 2603025150 ns 2496367360 ns 1.04
ViT small (256 x 256 x 3 x 32)/forward/CUDA/Lux 52794954.5 ns 50142997 ns 1.05
ViT small (256 x 256 x 3 x 32)/forward/CPU/Reactant (optimize = :after_enzyme) 3030272839 ns 3110834290 ns 0.97
ViT small (256 x 256 x 3 x 32)/forward/CPU/Reactant 2959738629 ns 3087052147 ns 0.96
ViT small (256 x 256 x 3 x 32)/forward/CPU/Reactant (optimize = :before_enzyme) 2929830073 ns 3073238853 ns 0.95
ViT small (256 x 256 x 3 x 32)/forward/CPU/Reactant (optimize = :only_enzyme) 4299042483 ns 4465557597 ns 0.96
ViT small (256 x 256 x 3 x 32)/forward/CPU/Lux 9053699739 ns 12448453806 ns 0.73
ViT base (256 x 256 x 3 x 16)/forward/CUDA/Reactant (optimize = :after_enzyme) 1303409186 ns 1207653927 ns 1.08
ViT base (256 x 256 x 3 x 16)/forward/CUDA/Reactant 1299005328 ns 1252669878.5 ns 1.04
ViT base (256 x 256 x 3 x 16)/forward/CUDA/Reactant (optimize = :before_enzyme) 1327270461 ns 1284419122 ns 1.03
ViT base (256 x 256 x 3 x 16)/forward/CUDA/Reactant (optimize = :only_enzyme) 2424433409 ns 2435990194 ns 1.00
ViT base (256 x 256 x 3 x 16)/forward/CUDA/Lux 71019572 ns 67827398 ns 1.05
ViT base (256 x 256 x 3 x 16)/forward/CPU/Reactant (optimize = :after_enzyme) 3188826413 ns 3295855546 ns 0.97
ViT base (256 x 256 x 3 x 16)/forward/CPU/Reactant 3252391187 ns 3257027916 ns 1.00
ViT base (256 x 256 x 3 x 16)/forward/CPU/Reactant (optimize = :before_enzyme) 3190391621 ns 3317925170 ns 0.96
ViT base (256 x 256 x 3 x 16)/forward/CPU/Reactant (optimize = :only_enzyme) 4502893678 ns 4627513897 ns 0.97
ViT base (256 x 256 x 3 x 16)/forward/CPU/Lux 13177363170 ns 13704380374 ns 0.96
ViT base (256 x 256 x 3 x 4)/forward/CUDA/Reactant (optimize = :after_enzyme) 1264446460 ns 1250112499 ns 1.01
ViT base (256 x 256 x 3 x 4)/forward/CUDA/Reactant 1275969531 ns 1225627385 ns 1.04
ViT base (256 x 256 x 3 x 4)/forward/CUDA/Reactant (optimize = :before_enzyme) 1236289619.5 ns 1253439196.5 ns 0.99
ViT base (256 x 256 x 3 x 4)/forward/CUDA/Reactant (optimize = :only_enzyme) 2377926514 ns 2466672353 ns 0.96
ViT base (256 x 256 x 3 x 4)/forward/CUDA/Lux 20589563.5 ns 19404003 ns 1.06
ViT base (256 x 256 x 3 x 4)/forward/CPU/Reactant (optimize = :after_enzyme) 1840153244 ns 1913014192 ns 0.96
ViT base (256 x 256 x 3 x 4)/forward/CPU/Reactant 1822459856 ns 1918443489 ns 0.95
ViT base (256 x 256 x 3 x 4)/forward/CPU/Reactant (optimize = :before_enzyme) 1811923850 ns 1919983927 ns 0.94
ViT base (256 x 256 x 3 x 4)/forward/CPU/Reactant (optimize = :only_enzyme) 3057262107 ns 3210069008 ns 0.95
ViT base (256 x 256 x 3 x 4)/forward/CPU/Lux 5829780775 ns 3218575500.5 ns 1.81

This comment was automatically generated by workflow using github-action-benchmark.

Please sign in to comment.