diff --git a/libinterp/corefcn/dot.cc b/libinterp/corefcn/dot.cc --- a/libinterp/corefcn/dot.cc +++ b/libinterp/corefcn/dot.cc @@ -227,6 +227,112 @@ but avoids forming a temporary array and %!error dot ([1 2], [1 2], 0) */ +template +void +blkmm_internal (const T& x, const T& y, T& z, + F77_INT m, F77_INT n, F77_INT k, F77_INT np); + +template <> +void +blkmm_internal (const FloatComplexNDArray& x, const FloatComplexNDArray& y, + FloatComplexNDArray& z, + F77_INT m, F77_INT n, F77_INT k, F77_INT np) +{ + F77_XFCN (cmatm3, CMATM3, (m, n, k, np, + F77_CONST_CMPLX_ARG (x.data ()), + F77_CONST_CMPLX_ARG (y.data ()), + F77_CMPLX_ARG (z.fortran_vec ()))); +} + +template <> +void +blkmm_internal (const ComplexNDArray& x, const ComplexNDArray& y, + ComplexNDArray& z, + F77_INT m, F77_INT n, F77_INT k, F77_INT np) +{ + F77_XFCN (zmatm3, ZMATM3, (m, n, k, np, + F77_CONST_DBLE_CMPLX_ARG (x.data ()), + F77_CONST_DBLE_CMPLX_ARG (y.data ()), + F77_DBLE_CMPLX_ARG (z.fortran_vec ()))); +} + +template <> +void +blkmm_internal (const FloatNDArray& x, const FloatNDArray& y, FloatNDArray& z, + F77_INT m, F77_INT n, F77_INT k, F77_INT np) +{ + F77_XFCN (smatm3, SMATM3, (m, n, k, np, + x.data (), y.data (), + z.fortran_vec ())); +} + +template <> +void +blkmm_internal (const NDArray& x, const NDArray& y, NDArray& z, + F77_INT m, F77_INT n, F77_INT k, F77_INT np) +{ + F77_XFCN (dmatm3, DMATM3, (m, n, k, np, + x.data (), y.data (), + z.fortran_vec ())); +} + +static void +get_blkmm_dims (const dim_vector& dimx, const dim_vector& dimy, + F77_INT& m, F77_INT& n, F77_INT& k, F77_INT& np, + dim_vector& dimz) +{ + int nd = dimx.ndims (); + + m = octave::to_f77_int (dimx(0)); + k = octave::to_f77_int (dimx(1)); + n = octave::to_f77_int (dimy(1)); + + octave_idx_type tmp_np = 1; + + bool match = dimy(0) == k && nd == dimy.ndims (); + + dimz = dim_vector::alloc (nd); + + dimz(0) = m; + dimz(1) = n; + for (int i = 2; match && i < nd; i++) + { + match = match && dimx(i) == dimy(i); + dimz(i) = dimx(i); + tmp_np *= dimz(i); + } + + np = octave::to_f77_int (tmp_np); + + if (! match) + error ("blkmm: A and B dimensions don't match: (%s) and (%s)", + dimx.str ().c_str (), dimy.str ().c_str ()); +} + +template +T +blkmm (const T& x, const T& y) +{ + F77_INT m, n, k, np; + dim_vector dimz; + + get_blkmm_dims (x.dims (), y.dims (), m, n, k, np, dimz); + + T z (dimz); + + if (n != 0 && m != 0) + blkmm_internal (x, y, z, m, n, k, np); + + return z; +} + +template +T +blkmm_ov (const octave_value& x, const octave_value& y) +{ + return blkmm (octave_value_extract (x), octave_value_extract (y)); +} + DEFUN (blkmm, args, , doc: /* -*- texinfo -*- @deftypefn {} {} blkmm (@var{A}, @var{B}) @@ -246,89 +352,30 @@ endfor @end example @end deftypefn */) { + octave_value retval; + if (args.length () != 2) print_usage (); - octave_value retval; - octave_value argx = args(0); octave_value argy = args(1); if (! argx.isnumeric () || ! argy.isnumeric ()) error ("blkmm: A and B must be numeric"); - const dim_vector dimx = argx.dims (); - const dim_vector dimy = argy.dims (); - int nd = dimx.ndims (); - F77_INT m = octave::to_f77_int (dimx(0)); - F77_INT k = octave::to_f77_int (dimx(1)); - F77_INT n = octave::to_f77_int (dimy(1)); - octave_idx_type tmp_np = 1; - bool match = dimy(0) == k && nd == dimy.ndims (); - dim_vector dimz = dim_vector::alloc (nd); - dimz(0) = m; - dimz(1) = n; - for (int i = 2; match && i < nd; i++) - { - match = match && dimx(i) == dimy(i); - dimz(i) = dimx(i); - tmp_np *= dimz(i); - } - F77_INT np = octave::to_f77_int (tmp_np); - - if (! match) - error ("blkmm: A and B dimensions don't match: (%s) and (%s)", - dimx.str ().c_str (), dimy.str ().c_str ()); - if (argx.iscomplex () || argy.iscomplex ()) { if (argx.is_single_type () || argy.is_single_type ()) - { - FloatComplexNDArray x = argx.float_complex_array_value (); - FloatComplexNDArray y = argy.float_complex_array_value (); - FloatComplexNDArray z (dimz); - - F77_XFCN (cmatm3, CMATM3, (m, n, k, np, - F77_CONST_CMPLX_ARG (x.data ()), F77_CONST_CMPLX_ARG (y.data ()), - F77_CMPLX_ARG (z.fortran_vec ()))); - retval = z; - } + retval = blkmm_ov (argx, argy); else - { - ComplexNDArray x = argx.complex_array_value (); - ComplexNDArray y = argy.complex_array_value (); - ComplexNDArray z (dimz); - - F77_XFCN (zmatm3, ZMATM3, (m, n, k, np, - F77_CONST_DBLE_CMPLX_ARG (x.data ()), F77_CONST_DBLE_CMPLX_ARG (y.data ()), - F77_DBLE_CMPLX_ARG (z.fortran_vec ()))); - retval = z; - } + retval = blkmm_ov (argx, argy); } else { if (argx.is_single_type () || argy.is_single_type ()) - { - FloatNDArray x = argx.float_array_value (); - FloatNDArray y = argy.float_array_value (); - FloatNDArray z (dimz); - - F77_XFCN (smatm3, SMATM3, (m, n, k, np, - x.data (), y.data (), - z.fortran_vec ())); - retval = z; - } + retval = blkmm_ov (argx, argy); else - { - NDArray x = argx.array_value (); - NDArray y = argy.array_value (); - NDArray z (dimz); - - F77_XFCN (dmatm3, DMATM3, (m, n, k, np, - x.data (), y.data (), - z.fortran_vec ())); - retval = z; - } + retval = blkmm_ov (argx, argy); } return retval;