Wrap the Givens rotation methods

Author: kshyatt

lib/cublas/wrappers.jl

@@ -295,6 +295,62 @@ for (fname, fname_64, elty, sty) in ((:cublasSrot_v2, :cublasSrot_v2_64, :Float3
     end
 end
 
+## rotg
+for (fname, elty) in ((:cublasSrotg_v2, :Float32),
+                      (:cublasDrotg_v2, :Float64),
+                      (:cublasCrotg_v2, :ComplexF32),
+                      (:cublasZrotg_v2, :ComplexF64),
+                     )
+    @eval begin
+        function rotg!(a::$elty, b::$elty)
+            c = Ref{real($elty)}(zero(real($elty)))
+            s = Ref{$elty}(zero($elty))
+            ref_a = Ref(a)
+            ref_b = Ref(b)
+            $fname(handle(), ref_a, ref_b, c, s)
+            ref_a[], ref_b[], c[], s[]
+        end
+    end
+end
+
+## rotm
+for (fname, fname_64, elty) in ((:cublasSrotm_v2, :cublasSrotm_v2_64, :Float32),
+                                (:cublasDrotm_v2, :cublasDrotm_v2_64, :Float64),
+                               )
+    @eval begin
+        function rotm!(n::Integer,
+                       x::StridedCuVecOrDenseMat{$elty},
+                       y::StridedCuVecOrDenseMat{$elty},
+                       param::AbstractVector{$elty})
+            if CUBLAS.version() >= v"12.0"
+                $fname_64(handle(), n, x, stride(x, 1), y, stride(y, 1), param)
+            else
+                $fname(handle(), n, x, stride(x, 1), y, stride(y, 1), param)
+            end
+            x, y
+        end
+    end
+end
+
+## rotmg
+for (fname, elty) in ((:cublasSrotmg_v2, :Float32),
+                      (:cublasDrotmg_v2, :Float64))
+    @eval begin
+        function rotmg!(d1::$elty,
+                        d2::$elty,
+                        x::$elty,
+                        y::$elty,
+                        param::AbstractVector{$elty})
+            ref_d1 = Ref(d1)
+            ref_d2 = Ref(d2)
+            ref_x  = Ref(x)
+            ref_y  = Ref(y)
+            $fname(handle(), ref_d1, ref_d2, ref_x, ref_y, param)
+            ref_d1[], ref_d2[], ref_x[], ref_y[], param 
+        end
+    end
+end
+
 ## swap
 for (fname, fname_64, elty) in ((:cublasSswap_v2, :cublasSswap_v2_64, :Float32),
                                 (:cublasDswap_v2, :cublasDswap_v2_64, :Float64),

test/libraries/cublas/level1.jl

@@ -1,5 +1,4 @@
 using CUDA.CUBLAS
-
 using LinearAlgebra
 
 using BFloat16s
@@ -49,6 +48,77 @@ k = 13
             @test testf(reflect!, rand(T, m), rand(T, m), rand(real(T)), rand(real(T)))
             @test testf(reflect!, rand(T, m), rand(T, m), rand(real(T)), rand(T))
         end
+        
+        @testset "rotg!" begin
+            a = rand(T)
+            b = rand(T)
+            a_copy = copy(a)
+            b_copy = copy(b)
+            a, b, c, s = CUBLAS.rotg!(a, b)
+            rot = [c s; -conj(s) c] * [a_copy; b_copy]
+            @test rot ≈ [a; 0]
+            if T <: Real
+                @test a^2 ≈ a_copy^2 + b_copy^2
+            end
+            @test c^2 + abs2(s) ≈ one(T)
+        end
+        
+        if T <: Real
+            H = rand(T, 2, 2)
+            @testset "flag $flag" for (flag, flag_H) in ((T(-2), [one(T) zero(T); zero(T) one(T)]),
+                                                         (-one(T), H),
+                                                         (zero(T), [one(T) H[1,2]; H[2, 1] one(T)]),
+                                                         (one(T), [H[1,1] one(T); -one(T) H[2, 2]]),
+                                                        )
+                @testset "rotm!" begin
+                    rot_n = 2
+                    x = rand(T, rot_n)
+                    y = rand(T, rot_n)
+                    dx = CuArray(x)
+                    dy = CuArray(y)
+                    dx, dy = CUBLAS.rotm!(rot_n, dx, dy, vcat(flag, H...))
+                    h_x = collect(dx)
+                    h_y = collect(dy)
+                    @test h_x ≈ [x[1] * flag_H[1,1] + y[1] * flag_H[1,2]; x[2] * flag_H[1, 1] + y[2] * flag_H[1, 2]]
+                    @test h_y ≈ [x[1] * flag_H[2,1] + y[1] * flag_H[2,2]; x[2] * flag_H[2, 1] + y[2] * flag_H[2, 2]]
+                end
+            end
+            @testset "rotmg!" begin
+                param = zeros(T, 5)
+                x1 = rand(T)
+                y1 = rand(T)
+                d1 = zero(T)
+                d2 = zero(T)
+                x1_copy = copy(x1)
+                y1_copy = copy(y1)
+                d1, d2, x1, y1, param = CUBLAS.rotmg!(d1, d2, x1, y1, param)
+                flag = param[1]
+                H = zeros(T, 2, 2)
+                if flag == -2
+                    H[1, 1] = one(T) 
+                    H[1, 2] = zero(T)
+                    H[2, 1] = zero(T) 
+                    H[2, 2] = one(T)
+                elseif flag == -1
+                    H[1, 1] = param[2]
+                    H[1, 2] = param[3]
+                    H[2, 1] = param[4]
+                    H[2, 2] = param[5]
+                elseif iszero(flag)
+                    H[1, 1] = one(T) 
+                    H[1, 2] = param[3]
+                    H[2, 1] = param[4]
+                    H[2, 2] = one(T)
+                elseif flag == 1
+                    H[1, 1] = param[2]
+                    H[1, 2] = one(T)
+                    H[2, 1] = -one(T)
+                    H[2, 2] = param[5]
+                end
+                out = H * [(√d1) * x1_copy; (√d2) * y1_copy]
+                @test out[2] ≈ zero(T)
+            end
+        end
 
         @testset "swap!" begin
             # swap is an extension