Repository: google/jax

Issue:
slow linalg.solve
Dear jax team,

I'd like to perform a lot (1e5-1e7) `np.linalg.solve` operations on small (3x3) matrices. jax is a lot slower than the numpy version: 17s in jax vs. 382ms in numpy. 

Is 3x3 too small to be effective on GPU? Or is the `solve` method just not as efficiently implemented for such a task as the numpy version?

```python
import jax
import jax.numpy as np
import numpy as onp

def to_gpu(arr):
    return jax.device_put(arr.astype(np.float32))

num_matrices = int(1e6)
sidelen = 3

A = to_gpu(onp.random.normal(size=(num_matrices, sidelen, sidelen)))
b = to_gpu(onp.random.normal(size=(num_matrices, sidelen)))

solve_jit = jax.jit(np.linalg.solve)
solve_vmap = jax.vmap(np.linalg.solve)
solve_vmap_jit = jax.jit(jax.vmap(np.linalg.solve))

# run jit'ed versions once
_ = solve_jit(A, b)
_ = solve_vmap_jit(A, b)
                                     
%time np.linalg.solve(A, b).block_until_ready()   # 16.5 s
%time solve_jit(A, b).block_until_ready()         # 17.0 s
%time solve_vmap(A, b).block_until_ready()        # 17.0 s
%time solve_vmap_jit(A, b).block_until_ready()    # 15.3 s

%time onp.linalg.solve(A, b)  # 383 ms
```


Discussion:
We didn't spend much time optimizing batched operations yet, especially on GPU.

I suspect we should be calling `cublasSgetrfBatched` in this case rather than `cusolverDnSgetrf`, and similarly we should be calling a batched triangular solve.
I have change that adds support for batched LU decompositions and triangular solves on GPU to JAX. I now see on a GTX 1080:
```
In [16]: %time solve_jit(A, b).block_until_ready()         # 17.0 s
CPU times: user 21.4 ms, sys: 3.13 ms, total: 24.5 ms
Wall time: 23.5 ms
```
which is presumably more like what you wanted! I'll clean it up and get it submitted.
Thank you very much for the explanation and the incredibly fast help and fix!

Is there a under-the-hood difference in using the "builtin" batching of `solve` vs. an explicit `vmap`? Also, is `jit`'ing a `vmap` generally useful/advisable?

Resolving patch:
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -29,7 +29,7 @@
 from jax.interpreters import xla
 from jax.interpreters import ad
 from jax.interpreters import batching
-from jax.util import partial
+from jax.util import partial, prod
 from jax.abstract_arrays import ShapedArray
 from jax.core import Primitive
 from jax.lax import (standard_primitive, standard_unop, binop_dtype_rule,
@@ -371,7 +371,7 @@ def triangular_solve_batching_rule(batched_args, batch_dims, left_side,
 batching.primitive_batchers[triangular_solve_p] = triangular_solve_batching_rule
 
 
-def triangular_solve_cpu_translation_rule(
+def _triangular_solve_cpu_translation_rule(
     c, a, b, left_side, lower, transpose_a, conjugate_a, unit_diagonal):
   shape = c.GetShape(a)
   dtype = shape.element_type().type
@@ -389,8 +389,30 @@ def triangular_solve_cpu_translation_rule(
     return c.TriangularSolve(a, b, left_side, lower, transpose_a, conjugate_a,
                              unit_diagonal)
 
-xla.backend_specific_translations['cpu'][triangular_solve_p] = triangular_solve_cpu_translation_rule
+xla.backend_specific_translations['cpu'][triangular_solve_p] = \
+  _triangular_solve_cpu_translation_rule
 
+def _triangular_solve_gpu_translation_rule(
+    c, a, b, left_side, lower, transpose_a, conjugate_a, unit_diagonal):
+  shape = c.GetShape(a)
+  dtype = shape.element_type().type
+  dims = shape.dimensions()
+  m, n = dims[-2:]
+  batch = prod(dims[:-2])
+  if batch > 1 and m <= 32 and n <= 32:
+    if conjugate_a and not transpose_a:
+      a = c.Conj(a)
+      conjugate_a = False
+    return cusolver.trsm(
+      c, a, b, left_side, lower, transpose_a, conjugate_a, unit_diagonal)
+  else:
+    # Use the XLA implementation for unbatched triangular_solve.
+    return c.TriangularSolve(a, b, left_side, lower, transpose_a, conjugate_a,
+                             unit_diagonal)
+
+if cusolver:
+  xla.backend_specific_translations['gpu'][triangular_solve_p] = \
+      _triangular_solve_gpu_translation_rule
 
 # LU decomposition
 
diff --git a/jax/numpy/linalg.py b/jax/numpy/linalg.py
--- a/jax/numpy/linalg.py
+++ b/jax/numpy/linalg.py
@@ -228,6 +228,7 @@ def qr(a, mode="reduced"):
 
 
 @_wraps(onp.linalg.solve)
+@jit
 def solve(a, b):
   a, b = _promote_arg_dtypes(np.asarray(a), np.asarray(b))
   a_shape = np.shape(a)
diff --git a/jaxlib/cusolver.py b/jaxlib/cusolver.py
--- a/jaxlib/cusolver.py
+++ b/jaxlib/cusolver.py
@@ -16,10 +16,20 @@
 from __future__ import division
 from __future__ import print_function
 
+import operator
+
 import numpy as np
+from six.moves import reduce
 
 from jaxlib import xla_client
 
+try:
+  from jaxlib import cublas_kernels
+  for _name, _value in cublas_kernels.registrations().items():
+    xla_client.register_custom_call_target(_name, _value, platform="gpu")
+except ImportError:
+  pass
+
 try:
   from jaxlib import cusolver_kernels
   for _name, _value in cusolver_kernels.registrations().items():
@@ -27,6 +37,7 @@
 except ImportError:
   pass
 
+
 _Shape = xla_client.Shape
 
 
@@ -43,6 +54,50 @@ def _real_type(dtype):
   else:
     raise NotImplementedError("Unsupported dtype {}".format(dtype))
 
+_prod = lambda xs: reduce(operator.mul, xs, 1)
+
+def trsm(c, a, b, left_side=False, lower=False, trans_a=False, conj_a=False,
+         diag=False):
+  """Batched triangular solve.
+
+  XLA implements unbatched triangular solve directly, so we need only implement
+  the batched case."""
+  b_shape = c.GetShape(b)
+  dtype = b_shape.element_type()
+  dims = b_shape.dimensions()
+  assert len(dims) >= 2
+  m, n = dims[-2:]
+  batch_dims = tuple(dims[:-2])
+  num_bd = len(batch_dims)
+  batch = _prod(batch_dims)
+  k = m if left_side else n
+
+  a_shape = c.GetShape(a)
+  if (batch_dims + (k, k) != a_shape.dimensions() or
+      a_shape.element_type() != dtype):
+    raise ValueError("Argument mismatch for trsm, got {} and {}".format(
+      a_shape, b_shape))
+
+  if conj_a and not trans_a:
+    raise NotImplementedError("Conjugation without transposition not supported")
+
+  lwork, opaque = cublas_kernels.build_trsm_batched_descriptor(
+    np.dtype(dtype), batch, m, n, left_side, lower, trans_a, conj_a, diag)
+  layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
+  out = c.CustomCall(
+      b"cublas_trsm_batched",
+      operands=(a, b),
+      shape_with_layout=_Shape.tuple_shape((
+          _Shape.array_shape(dtype, b_shape.dimensions(), layout),
+          _Shape.array_shape(np.dtype(np.int8), (lwork,), (0,)),
+          _Shape.array_shape(np.dtype(np.int8), (lwork,), (0,)))),
+      operand_shapes_with_layout=(
+          _Shape.array_shape(dtype, a_shape.dimensions(), layout),
+          _Shape.array_shape(dtype, b_shape.dimensions(), layout),
+      ),
+      opaque=opaque)
+  return c.GetTupleElement(out, 0)
+
 
 def getrf(c, a):
   """LU decomposition."""
@@ -53,32 +108,39 @@ def getrf(c, a):
   m, n = dims[-2:]
   batch_dims = tuple(dims[:-2])
   num_bd = len(batch_dims)
-  b = 1
-  for d in batch_dims:
-    b *= d
+  batch = _prod(batch_dims)
+
+  if batch > 1 and m == n and m // batch <= 128:
+    lwork, opaque = cublas_kernels.build_getrf_batched_descriptor(
+      np.dtype(dtype), batch, m)
+    workspace = _Shape.array_shape(np.dtype(np.int8), (lwork,), (0,))
+    kernel = b"cublas_getrf_batched"
+  else:
+    lwork, opaque = cusolver_kernels.build_getrf_descriptor(
+        np.dtype(dtype), batch, m, n)
+    workspace = _Shape.array_shape(dtype, (lwork,), (0,))
+    kernel = b"cusolver_getrf"
 
-  lwork, opaque = cusolver_kernels.build_getrf_descriptor(
-      np.dtype(dtype), b, m, n)
   out = c.CustomCall(
-      b"cusolver_getrf",
+      kernel,
       operands=(a,),
       shape_with_layout=_Shape.tuple_shape((
           _Shape.array_shape(
               dtype, batch_dims + (m, n),
               (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))),
-          _Shape.array_shape(dtype, (lwork,), (0,)),
           _Shape.array_shape(
               np.dtype(np.int32), batch_dims + (min(m, n),),
               tuple(range(num_bd, -1, -1))),
           _Shape.array_shape(
               np.dtype(np.int32), batch_dims, tuple(range(num_bd - 1, -1, -1))),
+          workspace,
       )),
       operand_shapes_with_layout=(_Shape.array_shape(
           dtype, batch_dims + (m, n),
           (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))),),
       opaque=opaque)
-  return (c.GetTupleElement(out, 0), c.GetTupleElement(out, 2),
-          c.GetTupleElement(out, 3))
+  return (c.GetTupleElement(out, 0), c.GetTupleElement(out, 1),
+          c.GetTupleElement(out, 2))
 
 
 def syevd(c, a, lower=False):
@@ -92,19 +154,17 @@ def syevd(c, a, lower=False):
   assert m == n
   batch_dims = tuple(dims[:-2])
   num_bd = len(batch_dims)
-  b = 1
-  for d in batch_dims:
-    b *= d
+  batch = _prod(batch_dims)
   layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
 
   if n <= 32:
     kernel = b"cusolver_syevj"
     lwork, opaque = cusolver_kernels.build_syevj_descriptor(
-        np.dtype(dtype), lower, b, n)
+        np.dtype(dtype), lower, batch, n)
   else:
     kernel = b"cusolver_syevd"
     lwork, opaque = cusolver_kernels.build_syevd_descriptor(
-        np.dtype(dtype), lower, b, n)
+        np.dtype(dtype), lower, batch, n)
   eigvals_type = _real_type(dtype)
 
   out = c.CustomCall(


Test patch: