test: raising to syrk

Author: avik-pal

src/stdlibs/LinearAlgebra.jl

@@ -46,6 +46,21 @@ function __init__()
             (BLAS.@blasfunc(dgesvj_), :enzymexla_lapack_dgesvj_),
             (BLAS.@blasfunc(cgesvj_), :enzymexla_lapack_cgesvj_),
             (BLAS.@blasfunc(zgesvj_), :enzymexla_lapack_zgesvj_),
+            # syrk
+            (BLAS.@blasfunc(ssyrk_), :enzymexla_blas_ssyrk_),
+            (BLAS.@blasfunc(dsyrk_), :enzymexla_blas_dsyrk_),
+            (BLAS.@blasfunc(csyrk_), :enzymexla_blas_csyrk_),
+            (BLAS.@blasfunc(zsyrk_), :enzymexla_blas_zsyrk_),
+            # trmm
+            (BLAS.@blasfunc(strmm_), :enzymexla_blas_strmm_),
+            (BLAS.@blasfunc(dtrmm_), :enzymexla_blas_dtrmm_),
+            (BLAS.@blasfunc(ctrmm_), :enzymexla_blas_ctrmm_),
+            (BLAS.@blasfunc(ztrmm_), :enzymexla_blas_ztrmm_),
+            # symm
+            (BLAS.@blasfunc(ssymm_), :enzymexla_blas_ssymm_),
+            (BLAS.@blasfunc(dsymm_), :enzymexla_blas_dsymm_),
+            (BLAS.@blasfunc(csymm_), :enzymexla_blas_csymm_),
+            (BLAS.@blasfunc(zsymm_), :enzymexla_blas_zsymm_),
         ]
             sym = Libdl.dlsym(libblastrampoline_handle, cname)
             @ccall MLIR.API.mlir_c.EnzymeJaXMapSymbol(

test/integration/linear_algebra.jl

@@ -723,3 +723,24 @@ end
     @jit LinearAlgebra.normalize!(x_ra)
     @test x_ra ≈ x
 end
+
+raise_to_syrk(x, y) = 3 .* (x * transpose(x)) .+ 5 .* y
+raise_to_syrk2(x, y) = 3 .* (transpose(x) * x) .+ 5 .* y
+
+@testset "syrk optimizations" begin
+    @testset for elty in (Float32, Float64, ComplexF32, ComplexF64)
+        x = Reactant.TestUtils.construct_test_array(elty, 4, 5)
+        y1 = Reactant.TestUtils.construct_test_array(elty, 4, 4)
+        y2 = Reactant.TestUtils.construct_test_array(elty, 5, 5)
+        x_ra = Reactant.to_rarray(x)
+
+        @testset for (fn, y) in ((raise_to_syrk, y1), (raise_to_syrk2, y2))
+            y_ra = Reactant.to_rarray(y)
+
+            hlo = @code_hlo optimize=:before_jit fn(x_ra, y_ra)
+            @test occursin("enzymexla.blas.syrk", repr(hlo))
+
+            @test @jit(fn(x_ra, y_ra)) ≈ fn(x, y)
+        end
+    end
+end