from __future__ import division from __future__ import absolute_import from pycuda.sparse.inner import AsyncInnerProduct from pytools import memoize_method import pycuda.driver as drv import pycuda.gpuarray as gpuarray import numpy as np class ConvergenceError(RuntimeError): pass class CGStateContainer: def __init__(self, operator, precon=None, pagelocked_allocator=None): if precon is None: from pycuda.sparse.operator import IdentityOperator precon = IdentityOperator(operator.dtype, operator.shape[0]) self.operator = operator self.precon = precon self.pagelocked_allocator = pagelocked_allocator @memoize_method def make_lc2_kernel(self, dtype, a_is_gpu, b_is_gpu): from pycuda.elementwise import get_linear_combination_kernel return get_linear_combination_kernel( ((a_is_gpu, dtype, dtype), (b_is_gpu, dtype, dtype)), dtype ) def lc2(self, a, x, b, y, out=None): if out is None: out = gpuarray.empty(x.shape, dtype=x.dtype, allocator=x.allocator) assert x.dtype == y.dtype == out.dtype a_is_gpu = isinstance(a, gpuarray.GPUArray) b_is_gpu = isinstance(b, gpuarray.GPUArray) assert x.shape == y.shape == out.shape kernel, texrefs = self.make_lc2_kernel(x.dtype, a_is_gpu, b_is_gpu) texrefs = texrefs[:] args = [] if a_is_gpu: assert a.dtype == x.dtype assert a.shape == () a.bind_to_texref_ext(texrefs.pop(0), allow_double_hack=True) else: args.append(a) args.append(x.gpudata) if b_is_gpu: assert b.dtype == y.dtype assert b.shape == () b.bind_to_texref_ext(texrefs.pop(0), allow_double_hack=True) else: args.append(b) args.append(y.gpudata) args.append(out.gpudata) args.append(x.mem_size) kernel.prepared_call(x._grid, x._block, *args) return out @memoize_method def guarded_div_kernel(self, dtype_x, dtype_y, dtype_z): from pycuda.elementwise import get_elwise_kernel from pycuda.tools import dtype_to_ctype return get_elwise_kernel( "%(tp_x)s *x, %(tp_y)s *y, %(tp_z)s *z" % { "tp_x": dtype_to_ctype(dtype_x), "tp_y": dtype_to_ctype(dtype_y), "tp_z": dtype_to_ctype(dtype_z), }, "z[i] = y[i] == 0 ? 0 : (x[i] / y[i])", "divide", ) def guarded_div(self, a, b): from pycuda.gpuarray import _get_common_dtype result = a._new_like_me(_get_common_dtype(a, b)) assert a.shape == b.shape func = self.guarded_div_kernel(a.dtype, b.dtype, result.dtype) func.prepared_async_call( a._grid, a._block, None, a.gpudata, b.gpudata, result.gpudata, a.mem_size ) return result def reset(self, rhs, x=None): self.rhs = rhs if x is None: x = np.zeros((self.operator.shape[0],)) self.x = x self.residual = rhs - self.operator(x) self.d = self.precon(self.residual) # grows at the end delta = AsyncInnerProduct(self.residual, self.d, self.pagelocked_allocator) self.real_delta_queue = [delta] self.delta = delta.gpu_result def one_iteration(self, compute_real_residual=False): # typed up from J.R. Shewchuk, # An Introduction to the Conjugate Gradient Method # Without the Agonizing Pain, Edition 1 1/4 [8/1994] # Appendix B3 q = self.operator(self.d) myip = gpuarray.dot(self.d, q) alpha = self.guarded_div(self.delta, myip) self.lc2(1, self.x, alpha, self.d, out=self.x) if compute_real_residual: self.residual = self.lc2(1, self.rhs, -1, self.operator(self.x)) else: self.lc2(1, self.residual, -alpha, q, out=self.residual) s = self.precon(self.residual) delta_old = self.delta delta = AsyncInnerProduct(self.residual, s, self.pagelocked_allocator) self.delta = delta.gpu_result beta = self.guarded_div(self.delta, delta_old) self.lc2(1, s, beta, self.d, out=self.d) if compute_real_residual: self.real_delta_queue.append(delta) def run(self, max_iterations=None, tol=1e-7, debug_callback=None): check_interval = 20 if max_iterations is None: max_iterations = max(3 * check_interval + 1, 10 * self.operator.shape[0]) real_resid_interval = min(self.operator.shape[0], 50) iterations = 0 delta_0 = None while iterations < max_iterations: compute_real_residual = iterations % real_resid_interval == 0 self.one_iteration(compute_real_residual=compute_real_residual) if debug_callback is not None: if compute_real_residual: what = "it+residual" else: what = "it" debug_callback( what, iterations, self.x, self.residual, self.d, self.delta ) # do often enough to allow AsyncInnerProduct # to progress through (polled) event chain rdq = self.real_delta_queue if iterations % check_interval == 0: if delta_0 is None: delta_0 = rdq[0].get_host_result() if delta_0 is not None: rdq.pop(0) if delta_0 is not None: i = 0 while i < len(rdq): delta = rdq[i].get_host_result() if delta is not None: if abs(delta) < tol * tol * abs(delta_0): if debug_callback is not None: debug_callback( "end", iterations, self.x, self.residual, self.d, self.delta, ) return self.x rdq.pop(i) else: i += 1 iterations += 1 raise ConvergenceError("cg failed to converge") def solve_pkt_with_cg( pkt_spmv, b, precon=None, x=None, tol=1e-7, max_iterations=None, debug=False, pagelocked_allocator=None, ): if x is None: x = gpuarray.zeros( pkt_spmv.shape[0], dtype=pkt_spmv.dtype, allocator=b.allocator ) else: x = pkt_spmv.permute(x) if pagelocked_allocator is None: pagelocked_allocator = drv.pagelocked_empty cg = CGStateContainer(pkt_spmv, precon, pagelocked_allocator=pagelocked_allocator) cg.reset(pkt_spmv.permute(b), x) it_count = [0] res_count = [0] def debug_callback(what, it_number, x, resid, d, delta): if what == "it": it_count[0] += 1 elif what == "it+residual": res_count[0] += 1 it_count[0] += 1 result = cg.run(max_iterations, tol, debug_callback=debug_callback) return pkt_spmv.unpermute(result), it_count[0], res_count[0]