numo-tiny_linalg 0.0.1

data/ext/numo/tiny_linalg/dot_sub.hpp

@@ -0,0 +1,85 @@
+namespace TinyLinalg {
+
+struct ZDotuSub {
+  void call(const int n, const double* x, const int incx, const double* y, const int incy, double* ret) {
+    cblas_zdotu_sub(n, x, incx, y, incy, ret);
+  }
+};
+
+struct CDotuSub {
+  void call(const int n, const float* x, const int incx, const float* y, const int incy, float* ret) {
+    cblas_cdotu_sub(n, x, incx, y, incy, ret);
+  }
+};
+
+template <int nary_dtype_id, typename dtype, class BlasFn>
+class DotSub {
+public:
+  static void define_module_function(VALUE mBlas, const char* modfn_name) {
+    rb_define_module_function(mBlas, modfn_name, RUBY_METHOD_FUNC(tiny_linalg_dot_sub), 2);
+  };
+
+private:
+  static void iter_dot_sub(na_loop_t* const lp) {
+    dtype* x = (dtype*)NDL_PTR(lp, 0);
+    dtype* y = (dtype*)NDL_PTR(lp, 1);
+    dtype* d = (dtype*)NDL_PTR(lp, 2);
+    const size_t n = NDL_SHAPE(lp, 0)[0];
+    BlasFn().call(n, x, 1, y, 1, d);
+  };
+
+  static VALUE tiny_linalg_dot_sub(VALUE self, VALUE x, VALUE y) {
+    VALUE nary_dtype = NaryTypes[nary_dtype_id];
+
+    if (CLASS_OF(x) != nary_dtype) {
+      x = rb_funcall(nary_dtype, rb_intern("cast"), 1, x);
+    }
+    if (!RTEST(nary_check_contiguous(x))) {
+      x = nary_dup(x);
+    }
+    if (CLASS_OF(y) != nary_dtype) {
+      y = rb_funcall(nary_dtype, rb_intern("cast"), 1, y);
+    }
+    if (!RTEST(nary_check_contiguous(y))) {
+      y = nary_dup(y);
+    }
+
+    narray_t* x_nary = NULL;
+    GetNArray(x, x_nary);
+    narray_t* y_nary = NULL;
+    GetNArray(y, y_nary);
+
+    if (NA_NDIM(x_nary) != 1) {
+      rb_raise(rb_eArgError, "x must be 1-dimensional");
+      return Qnil;
+    }
+    if (NA_NDIM(y_nary) != 1) {
+      rb_raise(rb_eArgError, "y must be 1-dimensional");
+      return Qnil;
+    }
+    if (NA_SIZE(x_nary) == 0) {
+      rb_raise(rb_eArgError, "x must not be empty");
+      return Qnil;
+    }
+    if (NA_SIZE(y_nary) == 0) {
+      rb_raise(rb_eArgError, "x must not be empty");
+      return Qnil;
+    }
+    if (NA_SIZE(x_nary) != NA_SIZE(y_nary)) {
+      rb_raise(rb_eArgError, "x and y must have same size");
+      return Qnil;
+    }
+
+    ndfunc_arg_in_t ain[2] = { { nary_dtype, 1 }, { nary_dtype, 1 } };
+    size_t shape_out[1] = { 1 };
+    ndfunc_arg_out_t aout[1] = { { nary_dtype, 0, shape_out } };
+    ndfunc_t ndf = { iter_dot_sub, NO_LOOP | NDF_EXTRACT, 2, 1, ain, aout };
+    VALUE ret = na_ndloop(&ndf, 2, x, y);
+
+    RB_GC_GUARD(x);
+    RB_GC_GUARD(y);
+    return ret;
+  };
+};
+
+} // namespace TinyLinalg

data/ext/numo/tiny_linalg/extconf.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+require 'mkmf'
+require 'numo/narray'
+
+$LOAD_PATH.each do |lp|
+  if File.exist?(File.join(lp, 'numo/numo/narray.h'))
+    $INCFLAGS = "-I#{lp}/numo #{$INCFLAGS}"
+    break
+  end
+end
+
+abort 'numo/narray.h is not found' unless have_header('numo/narray.h')
+
+if RUBY_PLATFORM.match?(/mswin|cygwin|mingw/)
+  $LOAD_PATH.each do |lp|
+    if File.exist?(File.join(lp, 'numo/libnarray.a'))
+      $LDFLAGS = "-L#{lp}/numo #{$LDFLAGS}"
+      break
+    end
+  end
+  abort 'libnarray.a is not found' unless have_library('narray', 'nary_new')
+end
+
+if RUBY_PLATFORM.include?('darwin') && Gem::Version.new('3.1.0') <= Gem::Version.new(RUBY_VERSION) &&
+   try_link('int main(void){return 0;}', '-Wl,-undefined,dynamic_lookup')
+  $LDFLAGS << ' -Wl,-undefined,dynamic_lookup'
+end
+
+use_accelerate = false
+# NOTE: Accelerate framework does not support LAPACKE.
+# if RUBY_PLATFORM.include?('darwin') && have_framework('Accelerate')
+#   $CFLAGS << ' -DTINYLINALG_USE_ACCELERATE'
+#   use_accelerate = true
+# end
+
+unless use_accelerate
+  if have_library('openblas')
+    $CFLAGS << ' -DTINYLINALG_USE_OPENBLAS'
+  else
+    abort 'libblas is not found' unless have_library('blas')
+    $CFLAGS << ' -DTINYLINALG_USE_BLAS'
+  end
+
+  abort 'liblapacke is not found' if !have_func('LAPACKE_dsyevr') && !have_library('lapacke')
+  abort 'cblas.h is not found' unless have_header('cblas.h')
+  abort 'lapacke.h is not found' unless have_header('lapacke.h')
+end
+
+abort 'libstdc++ is not found.' unless have_library('stdc++')
+
+$CXXFLAGS << ' -std=c++11'
+
+create_makefile('numo/tiny_linalg/tiny_linalg')

data/ext/numo/tiny_linalg/gemm.hpp

@@ -0,0 +1,223 @@
+namespace TinyLinalg {
+
+struct DGemm {
+  void call(const enum CBLAS_ORDER order, const enum CBLAS_TRANSPOSE transa, const enum CBLAS_TRANSPOSE transb,
+            const blasint m, const blasint n, const blasint k,
+            const double alpha, const double* a, const blasint lda, const double* b, const blasint ldb, const double beta,
+            double* c, const blasint ldc) {
+    cblas_dgemm(order, transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
+  }
+};
+
+struct SGemm {
+  void call(const enum CBLAS_ORDER order, const enum CBLAS_TRANSPOSE transa, const enum CBLAS_TRANSPOSE transb,
+            const blasint m, const blasint n, const blasint k,
+            const float alpha, const float* a, const blasint lda, const float* b, const blasint ldb, const float beta,
+            float* c, const blasint ldc) {
+    cblas_sgemm(order, transa, transb, m, n, k, alpha, a, lda, b, ldb, beta, c, ldc);
+  }
+};
+
+struct ZGemm {
+  void call(const enum CBLAS_ORDER order, const enum CBLAS_TRANSPOSE transa, const enum CBLAS_TRANSPOSE transb,
+            const blasint m, const blasint n, const blasint k,
+            const dcomplex alpha, const dcomplex* a, const blasint lda, const dcomplex* b, const blasint ldb, const dcomplex beta,
+            dcomplex* c, const blasint ldc) {
+    cblas_zgemm(order, transa, transb, m, n, k, &alpha, a, lda, b, ldb, &beta, c, ldc);
+  }
+};
+
+struct CGemm {
+  void call(const enum CBLAS_ORDER order, const enum CBLAS_TRANSPOSE transa, const enum CBLAS_TRANSPOSE transb,
+            const blasint m, const blasint n, const blasint k,
+            const scomplex alpha, const scomplex* a, const blasint lda, const scomplex* b, const blasint ldb, const scomplex beta,
+            scomplex* c, const blasint ldc) {
+    cblas_cgemm(order, transa, transb, m, n, k, &alpha, a, lda, b, ldb, &beta, c, ldc);
+  }
+};
+
+template <int nary_dtype_id, typename dtype, class BlasFn, class Converter>
+class Gemm {
+public:
+  static void define_module_function(VALUE mBlas, const char* modfn_name) {
+    rb_define_module_function(mBlas, modfn_name, RUBY_METHOD_FUNC(tiny_linalg_gemm), -1);
+  };
+
+private:
+  struct options {
+    dtype alpha;
+    dtype beta;
+    enum CBLAS_ORDER order;
+    enum CBLAS_TRANSPOSE transa;
+    enum CBLAS_TRANSPOSE transb;
+    blasint m;
+    blasint n;
+    blasint k;
+  };
+
+  static void iter_gemm(na_loop_t* const lp) {
+    const dtype* a = (dtype*)NDL_PTR(lp, 0);
+    const dtype* b = (dtype*)NDL_PTR(lp, 1);
+    dtype* c = (dtype*)NDL_PTR(lp, 2);
+    const options* opt = (options*)(lp->opt_ptr);
+    const blasint lda = opt->transa == CblasNoTrans ? opt->k : opt->m;
+    const blasint ldb = opt->transb == CblasNoTrans ? opt->n : opt->k;
+    const blasint ldc = opt->n;
+    BlasFn().call(opt->order, opt->transa, opt->transb, opt->m, opt->n, opt->k, opt->alpha, a, lda, b, ldb, opt->beta, c, ldc);
+  };
+
+  static VALUE tiny_linalg_gemm(int argc, VALUE* argv, VALUE self) {
+    VALUE nary_dtype = NaryTypes[nary_dtype_id];
+
+    VALUE a = Qnil;
+    VALUE b = Qnil;
+    VALUE c = Qnil;
+    VALUE kw_args = Qnil;
+    rb_scan_args(argc, argv, "21:", &a, &b, &c, &kw_args);
+
+    ID kw_table[5] = { rb_intern("alpha"), rb_intern("beta"), rb_intern("order"), rb_intern("transa"), rb_intern("transb") };
+    VALUE kw_values[5] = { Qundef, Qundef, Qundef, Qundef, Qundef };
+    rb_get_kwargs(kw_args, kw_table, 0, 5, kw_values);
+
+    if (CLASS_OF(a) != nary_dtype) {
+      a = rb_funcall(nary_dtype, rb_intern("cast"), 1, a);
+    }
+    if (!RTEST(nary_check_contiguous(a))) {
+      a = nary_dup(a);
+    }
+    if (CLASS_OF(b) != nary_dtype) {
+      b = rb_funcall(nary_dtype, rb_intern("cast"), 1, b);
+    }
+    if (!RTEST(nary_check_contiguous(b))) {
+      b = nary_dup(b);
+    }
+    if (!NIL_P(c)) {
+      if (CLASS_OF(c) != nary_dtype) {
+        c = rb_funcall(nary_dtype, rb_intern("cast"), 1, c);
+      }
+      if (!RTEST(nary_check_contiguous(c))) {
+        c = nary_dup(c);
+      }
+    }
+
+    dtype alpha = kw_values[0] != Qundef ? Converter().to_dtype(kw_values[0]) : Converter().one();
+    dtype beta = kw_values[1] != Qundef ? Converter().to_dtype(kw_values[1]) : Converter().zero();
+    enum CBLAS_ORDER order = kw_values[2] != Qundef ? get_cblas_order(kw_values[2]) : CblasRowMajor;
+    enum CBLAS_TRANSPOSE transa = kw_values[3] != Qundef ? get_cblas_trans(kw_values[3]) : CblasNoTrans;
+    enum CBLAS_TRANSPOSE transb = kw_values[4] != Qundef ? get_cblas_trans(kw_values[4]) : CblasNoTrans;
+
+    narray_t* a_nary = NULL;
+    GetNArray(a, a_nary);
+    narray_t* b_nary = NULL;
+    GetNArray(b, b_nary);
+
+    if (NA_NDIM(a_nary) != 2) {
+      rb_raise(rb_eArgError, "a must be 2-dimensional");
+      return Qnil;
+    }
+    if (NA_NDIM(b_nary) != 2) {
+      rb_raise(rb_eArgError, "b must be 2-dimensional");
+      return Qnil;
+    }
+    if (NA_SIZE(a_nary) == 0) {
+      rb_raise(rb_eArgError, "a must not be empty");
+      return Qnil;
+    }
+    if (NA_SIZE(b_nary) == 0) {
+      rb_raise(rb_eArgError, "b must not be empty");
+      return Qnil;
+    }
+
+    const blasint ma = NA_SHAPE(a_nary)[0];
+    const blasint ka = NA_SHAPE(a_nary)[1];
+    const blasint kb = NA_SHAPE(b_nary)[0];
+    const blasint nb = NA_SHAPE(b_nary)[1];
+    const blasint m = transa == CblasNoTrans ? ma : ka;
+    const blasint n = transb == CblasNoTrans ? nb : kb;
+    const blasint k = transa == CblasNoTrans ? ka : ma;
+    const blasint l = transb == CblasNoTrans ? kb : nb;
+
+    if (k != l) {
+      rb_raise(nary_eShapeError, "shape1[1](=%d) != shape2[0](=%d)", k, l);
+      return Qnil;
+    }
+
+    options opt = { alpha, beta, order, transa, transb, m, n, k };
+    size_t shape_out[2] = { static_cast<size_t>(m), static_cast<size_t>(n) };
+    ndfunc_arg_out_t aout[1] = { { nary_dtype, 2, shape_out } };
+    VALUE ret = Qnil;
+
+    if (!NIL_P(c)) {
+      narray_t* c_nary = NULL;
+      GetNArray(c, c_nary);
+      blasint nc = NA_SHAPE(c_nary)[0];
+      if (m > nc) {
+        rb_raise(nary_eShapeError, "shape3[0](=%d) >= shape1[0]=%d", nc, m);
+        return Qnil;
+      }
+      ndfunc_arg_in_t ain[3] = { { nary_dtype, 2 }, { nary_dtype, 2 }, { OVERWRITE, 2 } };
+      ndfunc_t ndf = { iter_gemm, NO_LOOP, 3, 0, ain, aout };
+      na_ndloop3(&ndf, &opt, 3, a, b, c);
+      ret = c;
+    } else {
+      c = INT2NUM(0);
+      ndfunc_arg_in_t ain[3] = { { nary_dtype, 2 }, { nary_dtype, 2 }, { sym_init, 0 } };
+      ndfunc_t ndf = { iter_gemm, NO_LOOP, 3, 1, ain, aout };
+      ret = na_ndloop3(&ndf, &opt, 3, a, b, c);
+    }
+
+    RB_GC_GUARD(a);
+    RB_GC_GUARD(b);
+    RB_GC_GUARD(c);
+
+    return ret;
+  };
+
+  static enum CBLAS_TRANSPOSE get_cblas_trans(VALUE val) {
+    const char* option_str = StringValueCStr(val);
+    enum CBLAS_TRANSPOSE res = CblasNoTrans;
+
+    if (std::strlen(option_str) > 0) {
+      switch (option_str[0]) {
+      case 'n':
+      case 'N':
+        res = CblasNoTrans;
+        break;
+      case 't':
+      case 'T':
+        res = CblasTrans;
+        break;
+      case 'c':
+      case 'C':
+        res = CblasConjTrans;
+        break;
+      }
+    }
+
+    RB_GC_GUARD(val);
+
+    return res;
+  }
+
+  static enum CBLAS_ORDER get_cblas_order(VALUE val) {
+    const char* option_str = StringValueCStr(val);
+
+    if (std::strlen(option_str) > 0) {
+      switch (option_str[0]) {
+      case 'r':
+      case 'R':
+        break;
+      case 'c':
+      case 'C':
+        rb_warn("Numo::TinyLinalg::BLAS.gemm does not support column major.");
+        break;
+      }
+    }
+
+    RB_GC_GUARD(val);
+
+    return CblasRowMajor;
+  }
+};
+
+} // namespace TinyLinalg

data/ext/numo/tiny_linalg/gemv.hpp

@@ -0,0 +1,211 @@
+namespace TinyLinalg {
+
+struct DGemv {
+  void call(const enum CBLAS_ORDER order, const enum CBLAS_TRANSPOSE trans, const blasint m, const blasint n,
+            const double alpha, const double* a, const blasint lda,
+            const double* x, const blasint incx, const double beta, double* y, const blasint incy) {
+    cblas_dgemv(order, trans, m, n, alpha, a, lda, x, incx, beta, y, incy);
+  }
+};
+
+struct SGemv {
+  void call(const enum CBLAS_ORDER order, const enum CBLAS_TRANSPOSE trans, const blasint m, const blasint n,
+            const float alpha, const float* a, const blasint lda,
+            const float* x, const blasint incx, const float beta, float* y, const blasint incy) {
+    cblas_sgemv(order, trans, m, n, alpha, a, lda, x, incx, beta, y, incy);
+  }
+};
+
+struct ZGemv {
+  void call(const enum CBLAS_ORDER order, const enum CBLAS_TRANSPOSE trans, const blasint m, const blasint n,
+            const dcomplex alpha, const dcomplex* a, const blasint lda,
+            const dcomplex* x, const blasint incx, const dcomplex beta, dcomplex* y, const blasint incy) {
+    cblas_zgemv(order, trans, m, n, &alpha, a, lda, x, incx, &beta, y, incy);
+  }
+};
+
+struct CGemv {
+  void call(const enum CBLAS_ORDER order, const enum CBLAS_TRANSPOSE trans, const blasint m, const blasint n,
+            const scomplex alpha, const scomplex* a, const blasint lda,
+            const scomplex* x, const blasint incx, const scomplex beta, scomplex* y, const blasint incy) {
+    cblas_cgemv(order, trans, m, n, &alpha, a, lda, x, incx, &beta, y, incy);
+  }
+};
+
+template <int nary_dtype_id, typename dtype, class BlasFn, class Converter>
+class Gemv {
+public:
+  static void define_module_function(VALUE mBlas, const char* modfn_name) {
+    rb_define_module_function(mBlas, modfn_name, RUBY_METHOD_FUNC(tiny_linalg_gemv), -1);
+  };
+
+private:
+  struct options {
+    dtype alpha;
+    dtype beta;
+    enum CBLAS_ORDER order;
+    enum CBLAS_TRANSPOSE trans;
+    blasint m;
+    blasint n;
+  };
+
+  static void iter_gemv(na_loop_t* const lp) {
+    const dtype* a = (dtype*)NDL_PTR(lp, 0);
+    const dtype* x = (dtype*)NDL_PTR(lp, 1);
+    dtype* y = (dtype*)NDL_PTR(lp, 2);
+    const options* opt = (options*)(lp->opt_ptr);
+    const blasint lda = opt->n;
+    BlasFn().call(opt->order, opt->trans, opt->m, opt->n, opt->alpha, a, lda, x, 1, opt->beta, y, 1);
+  };
+
+  static VALUE tiny_linalg_gemv(int argc, VALUE* argv, VALUE self) {
+    VALUE nary_dtype = NaryTypes[nary_dtype_id];
+
+    VALUE a = Qnil;
+    VALUE x = Qnil;
+    VALUE y = Qnil;
+    VALUE kw_args = Qnil;
+    rb_scan_args(argc, argv, "21:", &a, &x, &y, &kw_args);
+
+    ID kw_table[4] = { rb_intern("alpha"), rb_intern("beta"), rb_intern("order"), rb_intern("trans") };
+    VALUE kw_values[4] = { Qundef, Qundef, Qundef, Qundef };
+    rb_get_kwargs(kw_args, kw_table, 0, 4, kw_values);
+
+    if (CLASS_OF(a) != nary_dtype) {
+      a = rb_funcall(nary_dtype, rb_intern("cast"), 1, a);
+    }
+    if (!RTEST(nary_check_contiguous(a))) {
+      a = nary_dup(a);
+    }
+    if (CLASS_OF(x) != nary_dtype) {
+      x = rb_funcall(nary_dtype, rb_intern("cast"), 1, x);
+    }
+    if (!RTEST(nary_check_contiguous(x))) {
+      x = nary_dup(x);
+    }
+    if (!NIL_P(y)) {
+      if (CLASS_OF(y) != nary_dtype) {
+        y = rb_funcall(nary_dtype, rb_intern("cast"), 1, y);
+      }
+      if (!RTEST(nary_check_contiguous(y))) {
+        y = nary_dup(y);
+      }
+    }
+
+    dtype alpha = kw_values[0] != Qundef ? Converter().to_dtype(kw_values[0]) : Converter().one();
+    dtype beta = kw_values[1] != Qundef ? Converter().to_dtype(kw_values[1]) : Converter().zero();
+    enum CBLAS_ORDER order = kw_values[2] != Qundef ? get_cblas_order(kw_values[2]) : CblasRowMajor;
+    enum CBLAS_TRANSPOSE trans = kw_values[3] != Qundef ? get_cblas_trans(kw_values[3]) : CblasNoTrans;
+
+    narray_t* a_nary = NULL;
+    GetNArray(a, a_nary);
+    narray_t* x_nary = NULL;
+    GetNArray(x, x_nary);
+
+    if (NA_NDIM(a_nary) != 2) {
+      rb_raise(rb_eArgError, "a must be 2-dimensional");
+      return Qnil;
+    }
+    if (NA_NDIM(x_nary) != 1) {
+      rb_raise(rb_eArgError, "x must be 1-dimensional");
+      return Qnil;
+    }
+    if (NA_SIZE(a_nary) == 0) {
+      rb_raise(rb_eArgError, "a must not be empty");
+      return Qnil;
+    }
+    if (NA_SIZE(x_nary) == 0) {
+      rb_raise(rb_eArgError, "x must not be empty");
+      return Qnil;
+    }
+
+    const blasint ma = NA_SHAPE(a_nary)[0];
+    const blasint na = NA_SHAPE(a_nary)[1];
+    const blasint mx = NA_SHAPE(x_nary)[0];
+    const blasint m = trans == CblasNoTrans ? ma : na;
+    const blasint n = trans == CblasNoTrans ? na : ma;
+
+    if (n != mx) {
+      rb_raise(nary_eShapeError, "shape1[1](=%d) != shape2[0](=%d)", n, mx);
+      return Qnil;
+    }
+
+    options opt = { alpha, beta, order, trans, ma, na };
+    size_t shape_out[1] = { static_cast<size_t>(m) };
+    ndfunc_arg_out_t aout[1] = { { nary_dtype, 1, shape_out } };
+    VALUE ret = Qnil;
+
+    if (!NIL_P(y)) {
+      narray_t* y_nary = NULL;
+      GetNArray(y, y_nary);
+      blasint my = NA_SHAPE(y_nary)[0];
+      if (m > my) {
+        rb_raise(nary_eShapeError, "shape3[0](=%d) >= shape1[0]=%d", my, m);
+        return Qnil;
+      }
+      ndfunc_arg_in_t ain[3] = { { nary_dtype, 2 }, { nary_dtype, 1 }, { OVERWRITE, 1 } };
+      ndfunc_t ndf = { iter_gemv, NO_LOOP, 3, 0, ain, aout };
+      na_ndloop3(&ndf, &opt, 3, a, x, y);
+      ret = y;
+    } else {
+      y = INT2NUM(0);
+      ndfunc_arg_in_t ain[3] = { { nary_dtype, 2 }, { nary_dtype, 1 }, { sym_init, 0 } };
+      ndfunc_t ndf = { iter_gemv, NO_LOOP, 3, 1, ain, aout };
+      ret = na_ndloop3(&ndf, &opt, 3, a, x, y);
+    }
+
+    RB_GC_GUARD(a);
+    RB_GC_GUARD(x);
+    RB_GC_GUARD(y);
+
+    return ret;
+  }
+
+  static enum CBLAS_TRANSPOSE get_cblas_trans(VALUE val) {
+    const char* option_str = StringValueCStr(val);
+    enum CBLAS_TRANSPOSE res = CblasNoTrans;
+
+    if (std::strlen(option_str) > 0) {
+      switch (option_str[0]) {
+      case 'n':
+      case 'N':
+        res = CblasNoTrans;
+        break;
+      case 't':
+      case 'T':
+        res = CblasTrans;
+        break;
+      case 'c':
+      case 'C':
+        res = CblasConjTrans;
+        break;
+      }
+    }
+
+    RB_GC_GUARD(val);
+
+    return res;
+  }
+
+  static enum CBLAS_ORDER get_cblas_order(VALUE val) {
+    const char* option_str = StringValueCStr(val);
+
+    if (std::strlen(option_str) > 0) {
+      switch (option_str[0]) {
+      case 'r':
+      case 'R':
+        break;
+      case 'c':
+      case 'C':
+        rb_warn("Numo::TinyLinalg::BLAS.gemm does not support column major.");
+        break;
+      }
+    }
+
+    RB_GC_GUARD(val);
+
+    return CblasRowMajor;
+  }
+};
+
+} // namespace TinyLinalg

data/ext/numo/tiny_linalg/gesdd.hpp

@@ -0,0 +1,134 @@
+namespace TinyLinalg {
+
+struct DGESDD {
+  lapack_int call(int matrix_order, char jobz, lapack_int m, lapack_int n,
+                  double* a, lapack_int lda, double* s, double* u, lapack_int ldu, double* vt, lapack_int ldvt) {
+    return LAPACKE_dgesdd(matrix_order, jobz, m, n, a, lda, s, u, ldu, vt, ldvt);
+  };
+};
+
+struct SGESDD {
+  lapack_int call(int matrix_order, char jobz, lapack_int m, lapack_int n,
+                  float* a, lapack_int lda, float* s, float* u, lapack_int ldu, float* vt, lapack_int ldvt) {
+    return LAPACKE_sgesdd(matrix_order, jobz, m, n, a, lda, s, u, ldu, vt, ldvt);
+  };
+};
+
+struct ZGESDD {
+  lapack_int call(int matrix_order, char jobz, lapack_int m, lapack_int n,
+                  lapack_complex_double* a, lapack_int lda, double* s, lapack_complex_double* u, lapack_int ldu, lapack_complex_double* vt, lapack_int ldvt) {
+    return LAPACKE_zgesdd(matrix_order, jobz, m, n, a, lda, s, u, ldu, vt, ldvt);
+  };
+};
+
+struct CGESDD {
+  lapack_int call(int matrix_order, char jobz, lapack_int m, lapack_int n,
+                  lapack_complex_float* a, lapack_int lda, float* s, lapack_complex_float* u, lapack_int ldu, lapack_complex_float* vt, lapack_int ldvt) {
+    return LAPACKE_cgesdd(matrix_order, jobz, m, n, a, lda, s, u, ldu, vt, ldvt);
+  };
+};
+
+template <int nary_dtype_id, int nary_rtype_id, typename DType, typename RType, typename FncType>
+class GESDD {
+public:
+  static void define_module_function(VALUE mLapack, const char* mf_name) {
+    rb_define_module_function(mLapack, mf_name, RUBY_METHOD_FUNC(tiny_linalg_gesdd), -1);
+  };
+
+private:
+  struct gesdd_opt {
+    int matrix_order;
+    char jobz;
+  };
+
+  static void iter_gesdd(na_loop_t* const lp) {
+    DType* a = (DType*)NDL_PTR(lp, 0);
+    RType* s = (RType*)NDL_PTR(lp, 1);
+    DType* u = (DType*)NDL_PTR(lp, 2);
+    DType* vt = (DType*)NDL_PTR(lp, 3);
+    int* info = (int*)NDL_PTR(lp, 4);
+    gesdd_opt* opt = (gesdd_opt*)(lp->opt_ptr);
+
+    const size_t m = opt->matrix_order == LAPACK_ROW_MAJOR ? NDL_SHAPE(lp, 0)[0] : NDL_SHAPE(lp, 0)[1];
+    const size_t n = opt->matrix_order == LAPACK_ROW_MAJOR ? NDL_SHAPE(lp, 0)[1] : NDL_SHAPE(lp, 0)[0];
+    const size_t min_mn = m < n ? m : n;
+    const lapack_int lda = n;
+    const lapack_int ldu = opt->jobz == 'S' ? min_mn : m;
+    const lapack_int ldvt = opt->jobz == 'S' ? min_mn : n;
+
+    lapack_int i = FncType().call(opt->matrix_order, opt->jobz, m, n, a, lda, s, u, ldu, vt, ldvt);
+    *info = static_cast<int>(i);
+  };
+
+  static VALUE tiny_linalg_gesdd(int argc, VALUE* argv, VALUE self) {
+    VALUE nary_dtype = NaryTypes[nary_dtype_id];
+    VALUE nary_rtype = NaryTypes[nary_rtype_id];
+    VALUE a_vnary = Qnil;
+    VALUE kw_args = Qnil;
+
+    rb_scan_args(argc, argv, "1:", &a_vnary, &kw_args);
+
+    ID kw_table[2] = { rb_intern("jobz"), rb_intern("order") };
+    VALUE kw_values[2] = { Qundef, Qundef };
+
+    rb_get_kwargs(kw_args, kw_table, 0, 2, kw_values);
+
+    const char jobz = kw_values[0] == Qundef ? 'A' : StringValueCStr(kw_values[0])[0];
+    const char order = kw_values[1] == Qundef ? 'R' : StringValueCStr(kw_values[1])[0];
+
+    if (CLASS_OF(a_vnary) != nary_dtype) {
+      rb_raise(rb_eTypeError, "type of input array is invalid for overwriting");
+      return Qnil;
+    }
+    if (!RTEST(nary_check_contiguous(a_vnary))) {
+      a_vnary = nary_dup(a_vnary);
+    }
+
+    narray_t* a_nary = NULL;
+    GetNArray(a_vnary, a_nary);
+    const int n_dims = NA_NDIM(a_nary);
+    if (n_dims != 2) {
+      rb_raise(rb_eArgError, "input array must be 2-dimensional");
+      return Qnil;
+    }
+
+    const int matrix_order = order == 'C' ? LAPACK_COL_MAJOR : LAPACK_ROW_MAJOR;
+    const size_t m = matrix_order == LAPACK_ROW_MAJOR ? NA_SHAPE(a_nary)[0] : NA_SHAPE(a_nary)[1];
+    const size_t n = matrix_order == LAPACK_ROW_MAJOR ? NA_SHAPE(a_nary)[1] : NA_SHAPE(a_nary)[0];
+
+    const size_t min_mn = m < n ? m : n;
+    size_t shape_s[1] = { min_mn };
+    size_t shape_u[2] = { m, m };
+    size_t shape_vt[2] = { n, n };
+
+    ndfunc_arg_in_t ain[1] = { { OVERWRITE, 2 } };
+    ndfunc_arg_out_t aout[4] = { { nary_rtype, 1, shape_s }, { nary_dtype, 2, shape_u }, { nary_dtype, 2, shape_vt }, { numo_cInt32, 0 } };
+
+    switch (jobz) {
+    case 'A':
+      break;
+    case 'S':
+      shape_u[matrix_order == LAPACK_ROW_MAJOR ? 1 : 0] = min_mn;
+      shape_vt[matrix_order == LAPACK_ROW_MAJOR ? 0 : 1] = min_mn;
+      break;
+    case 'O':
+      break;
+    case 'N':
+      aout[1].dim = 0;
+      aout[2].dim = 0;
+      break;
+    default:
+      rb_raise(rb_eArgError, "jobz must be one of 'A', 'S', 'O', or 'N'");
+      return Qnil;
+    }
+
+    ndfunc_t ndf = { iter_gesdd, NO_LOOP | NDF_EXTRACT, 1, 4, ain, aout };
+    gesdd_opt opt = { matrix_order, jobz };
+    VALUE ret = na_ndloop3(&ndf, &opt, 1, a_vnary);
+
+    RB_GC_GUARD(a_vnary);
+    return ret;
+  };
+};
+
+} // namespace TinyLinalg