fix: temp patch for syrk lowering

Author: avik-pal

src/Ops.jl

@@ -4215,20 +4215,22 @@ end
     alpha::Union{TracedRNumber{T},T},
     beta::Union{TracedRNumber{T},T};
     uplo::Char,
+    output_uplo::Char,
     transpose_a::Char,
     location=mlir_stacktrace("syrk", @__FILE__, @__LINE__),
 ) where {T,N}
     ctx = MLIR.IR.current_context()
-    uplo_attr = MLIR.API.enzymexlaLapackUploAttrGet(ctx, LAPACK_UPLO_MAP[uplo])
 
     res = MLIR.IR.result(
         enzymexla.blas_syrk(
             A.mlir_data,
             C.mlir_data,
             constant(alpha; location).mlir_data,
             constant(beta; location).mlir_data;
-            uplo=uplo_attr,
-            output_uplo=uplo_attr,
+            uplo=MLIR.API.enzymexlaLapackUploAttrGet(ctx, LAPACK_UPLO_MAP[uplo]),
+            output_uplo=MLIR.API.enzymexlaLapackUploAttrGet(
+                ctx, LAPACK_UPLO_MAP[output_uplo]
+            ),
             transpose=MLIR.API.enzymexlaLapackTransposeAttrGet(
                 ctx, LAPACK_TRANSPOSE_MAP[transpose_a]
             ),

src/stdlibs/BLAS.jl

@@ -523,6 +523,7 @@ function BLAS.syrk!(
         alpha,
         beta;
         uplo=uplo,
+        output_uplo='F', # TODO: this is a temporary patch. Why is the default flipped?
         transpose_a=trans,
     )
     copyto!(C, res)

test/integration/blas.jl

@@ -1,8 +1,6 @@
 using LinearAlgebra, Reactant, Test
 using LinearAlgebra: BLAS
 
-const RunningOnTPU = contains(string(Reactant.devices()[1]), "TPU")
-
 @testset "Level 1" begin
     @testset "asum" begin
         x = Reactant.TestUtils.construct_test_array(Float32, 32)
@@ -284,8 +282,7 @@ end
         @jit BLAS.syrk!('U', 'N', 2.0f0, A_ra, 3.0f0, C_ra)
         C_target = copy(C)
         BLAS.syrk!('U', 'N', 2.0f0, A, 3.0f0, C_target)
-        @test UpperTriangular(C_ra) ≈ UpperTriangular(C_target) atol = 1e-3 rtol = 1e-3 broken =
-            !RunningOnTPU
+        @test UpperTriangular(C_ra) ≈ UpperTriangular(C_target) atol = 1e-3 rtol = 1e-3
 
         # test 'L' and 'T'
         A2 = Reactant.TestUtils.construct_test_array(Float32, 16, 16)
@@ -296,8 +293,7 @@ end
         @jit BLAS.syrk!('L', 'T', 2.0f0, A2_ra, 3.0f0, C2_ra)
         C2_target = copy(C2)
         BLAS.syrk!('L', 'T', 2.0f0, A2, 3.0f0, C2_target)
-        @test LowerTriangular(C2_ra) ≈ LowerTriangular(C2_target) atol = 1e-3 rtol = 1e-3 broken =
-            !RunningOnTPU
+        @test LowerTriangular(C2_ra) ≈ LowerTriangular(C2_target) atol = 1e-3 rtol = 1e-3
     end
 
     if isdefined(BLAS, :gemmt!)
@@ -415,8 +411,7 @@ end
             BLAS.trsm('L', 'U', 'N', 'N', 2.0f0, Ainv, B) atol = 1e-3 rtol = 1e-3
 
         @test UpperTriangular(@jit(BLAS.syrk('U', 'N', 2.0f0, A_ra))) ≈
-            UpperTriangular(BLAS.syrk('U', 'N', 2.0f0, A)) atol = 1e-3 rtol = 1e-3 broken =
-            !RunningOnTPU
+            UpperTriangular(BLAS.syrk('U', 'N', 2.0f0, A)) atol = 1e-3 rtol = 1e-3
         @test UpperTriangular(@jit(BLAS.syr2k('U', 'N', 2.0f0, A_ra, B_ra))) ≈
             UpperTriangular(BLAS.syr2k('U', 'N', 2.0f0, A, B)) atol = 1e-3 rtol = 1e-3