gemm_gpu update

Author: t81dev

include/t81/linalg/gemm_gpu.hpp

@@ -57,6 +57,10 @@ inline constexpr bool backend_available(Backend backend) noexcept {
     return false;
 }
 
+#if !defined(T81LIB_DOXYGEN)
+Backend get_current_backend() noexcept;
+#endif
+
 #if T81LIB_USE_METAL
 bool metal_available() noexcept;
 #endif
@@ -105,6 +109,15 @@ void addcmul(const TensorMetadata &input,
              TensorMetadata &out,
              Backend backend = Backend::Auto);
 
+#if !defined(T81LIB_DOXYGEN)
+void gemm_ternary(const TensorMetadata &A,
+                  const TensorMetadata &B,
+                  TensorMetadata &C,
+                  float alpha,
+                  float beta,
+                  Backend backend = Backend::Auto);
+#endif
+
 #if T81LIB_USE_CUDA
 void cuda_gemm_ternary(std::span<const core::limb> A,
                        std::span<const core::limb> B,
@@ -154,7 +167,7 @@ inline void gemm_ternary_dispatch(std::span<const core::limb> A,
                                   int K,
                                   float alpha,
                                   float beta,
-                                  Backend backend) {
+                                  Backend backend = Backend::Auto) {
     if (M < 0 || N < 0 || K < 0) {
         throw std::invalid_argument("gemm_ternary dimensions must be non-negative");
     }

python/bindings.cpp

@@ -756,14 +756,19 @@ PYBIND11_MODULE(t81lib, module) {
             const int K_limbs = K / core::limb::TRITS;
             const auto a_handle = handle_for_packed_object(A_obj, M, K_limbs);
             const auto b_handle = handle_for_packed_object(B_obj, K_limbs, N);
-            const auto c_handle = extract_tensor_handle(C_obj);
+            auto c_handle = extract_tensor_handle(C_obj);
             if (static_cast<int>(c_handle.meta.sizes.size()) != 2 ||
                 static_cast<int>(c_handle.meta.sizes[0]) != M ||
                 static_cast<int>(c_handle.meta.sizes[1]) != N) {
                 throw py::value_error("C tensor shape must match (M, N)");
             }
             t81::linalg::detail::gemm_ternary(
-                a_handle.meta, b_handle.meta, c_handle.meta, alpha, beta, Backend::Auto);
+                a_handle.meta,
+                b_handle.meta,
+                c_handle.meta,
+                alpha,
+                beta,
+                t81::linalg::Backend::Auto);
         },
         py::arg("A"),
         py::arg("B"),