feat: new jll version + new compiler passes (#1791)

* feat: new jll version + new compiler passes

* fix: use newer libtpu builds

* fix: try always downloading libtpu on ci

* feat: version check for libtpu

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <wmoses@mit.edu>", "Valentin Churavy <vchuravy@mit.edu>", "Sergio Sánchez Ramírez <sergio.sanchez.ramirez@bsc.es>", "Paul Berg <paul@plutojl.org>", "Avik Pal <avikpal@mit.edu>", "Mosè Giordano <mose@gnu.org>"]
-version = "0.2.171"
+version = "0.2.172"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -105,7 +105,7 @@ PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.16"
-Reactant_jll = "0.0.253"
+Reactant_jll = "0.0.254"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/Compiler.jl

@@ -910,6 +910,7 @@ function optimization_passes(
         "remove_no_ops_from_while_loop",
         "while_is_copy_simplify",
         "split_variadic_scatter_op",
+        "dynamic_slice_simplify",
     ]
 
     if !compile_options.disable_auto_batching_passes
@@ -1124,9 +1125,11 @@ function optimization_passes(
         if AGGRESSIVE_PROPAGATION[]
             push!(transform_passes_list, "reshape_slice(0)")
             push!(transform_passes_list, "reshape_elementwise(0)")
+            push!(transform_passes_list, "reshape_dynamic_slice(0)")
         else
             push!(transform_passes_list, "reshape_slice(1)")
             push!(transform_passes_list, "reshape_elementwise(1)")
+            push!(transform_passes_list, "reshape_dynamic_slice(1)")
         end
     elseif compile_options.reshape_propagate === :down
         append!(

src/accelerators/TPU.jl

@@ -10,6 +10,9 @@ using unzip_jll: unzip
 const libtpu_dir = Ref{Union{Nothing,String}}(nothing)
 const RUNNING_IN_CLOUD_TPU_VM = Ref(false)
 
+const LIBTPU_VERSION = "0.0.28.dev20251027"
+const LIBTPU_SO = "libtpu-$(replace(string(LIBTPU_VERSION), '.' => '_')).so"
+
 function __init__()
     @static if !Sys.isapple()
         if !Reactant.precompiling() && has_tpu()
@@ -32,18 +35,18 @@ end
 
 get_libtpu_dir() = libtpu_dir[]
 
-get_libtpu_path() = joinpath(get_libtpu_dir(), "libtpu.so")
+get_libtpu_path() = joinpath(get_libtpu_dir(), LIBTPU_SO)
 
 function download_libtpu_if_needed(path=nothing)
     path === nothing && (path = get_libtpu_dir())
     @assert path !== nothing "libtpu_dir is not set!"
 
-    libtpu_path = joinpath(path, "libtpu.so")
+    libtpu_path = joinpath(path, LIBTPU_SO)
     if !isfile(libtpu_path)
         zip_file_path = joinpath(path, "tpu.zip")
         tmp_dir = joinpath(path, "tmp")
         Downloads.download(
-            "https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20250811+nightly-py3-none-manylinux_2_31_x86_64.whl",
+            "https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu/libtpu-0.0.28.dev20251027+nightly-cp314-cp314t-manylinux_2_31_x86_64.whl",
             zip_file_path,
         )
         run(`$(unzip()) -qq $(zip_file_path) -d $(tmp_dir)`)

test/ops.jl

@@ -791,12 +791,12 @@ end
 
 @testset "acos" begin
     x = Reactant.to_rarray(Float32[-1.0, 0.0, 1.0])
-    @test acos.(Array(x)) ≈ @jit(Ops.acos(x)) broken = RunningOnTPU
+    @test acos.(Array(x)) ≈ @jit(Ops.acos(x))
 end
 
 @testset "acosh" begin
     x = Reactant.to_rarray(Float32[1.0, 10.0])
-    @test acosh.(Array(x)) ≈ @jit(Ops.acosh(x)) broken = RunningOnTPU
+    @test acosh.(Array(x)) ≈ @jit(Ops.acosh(x))
 end
 
 @testset "asin" begin