agd.Eikonal.HFM_CUDA.VoronoiDecomposition

  1# Copyright 2020 Jean-Marie Mirebeau, University Paris-Sud, CNRS, University Paris-Saclay
  2# Distributed WITHOUT ANY WARRANTY. Licensed under the Apache License, Version 2.0, see http://www.apache.org/licenses/LICENSE-2.0
  3
  4import numpy as np
  5import cupy as cp
  6import os
  7
  8from . import cupy_module_helper
  9from ...Metrics.misc import flatten_symmetric_matrix,expand_symmetric_matrix
 10from ... import LinearParallel as lp
 11
 12def Reconstruct(coefs,offsets):
 13     return lp.mult(coefs,lp.outer_self(offsets)).sum(2)
 14
 15def VoronoiDecomposition(m,offset_t=np.int32,
 16	flattened=False,blockDim=None,traits=None,steps="Both",retry64_tol=2e-5):
 17	"""
 18	Returns the Voronoi decomposition of a family of quadratic forms. 
 19	- m : the (symmetric) matrices of the quadratic forms to decompose.
 20	- offset_t : the type of offsets to be returned. 
 21	- flattened : wether the input matrices are flattened
 22	- retry64_tol (optional) : retries decomposition using 64bit floats if this error 
 23	 is exceeded relative to matrix trace. (Set retry64_tol = 0 to use double precision.)
 24	"""
 25
 26	# Prepare the inputs and outputs
 27	if flattened: m_exp = None
 28	else: m_exp = m; m = flatten_symmetric_matrix(m)
 29	symdim = len(m)
 30	ndim = int(np.sqrt(2*symdim))
 31	assert symdim==ndim*(ndim+1)/2
 32	shape = m.shape[1:]
 33	size = m.size/symdim
 34
 35	if ndim==1: 
 36		from ... import Selling
 37		λ,e = Selling.Decomposition(m[None])
 38		return λ,e.astype(offset_t)
 39	if not (2<=ndim and ndim<=6): 
 40		raise ValueError(f"Voronoi decomposition not implemented in dimension {ndim}")
 41	decompdim = [0,1,3,6,12,15,21][ndim]
 42
 43	float_t = np.float32 if retry64_tol else np.float64
 44	tol = {np.float32:1e-5, np.float64:2e-14}[float_t]
 45	int_t = np.int32
 46	m = cp.ascontiguousarray(m,dtype=float_t)
 47	weights = cp.empty((decompdim,*shape),dtype=float_t)
 48	offsets = cp.empty((ndim,decompdim,*shape),dtype=offset_t)
 49
 50	weights,offsets=map(cp.ascontiguousarray,(weights,offsets))
 51
 52	# Setup the GPU kernel
 53	if traits is None: traits = {}
 54	traits.update({
 55		'ndim_macro':ndim,
 56		'OffsetT':offset_t,
 57		'Scalar':float_t,
 58		'Int':np.int32,
 59		'SIMPLEX_TOL_macro':tol,
 60		})
 61
 62	source = cupy_module_helper.traits_header(traits)
 63	if ndim==5 and float_t==np.float64: source.append('#define DOUBLE')#for linprog solver
 64	cuda_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),"cuda")
 65	date_modified = cupy_module_helper.getmtime_max(cuda_path)
 66
 67	source += [
 68	'#include "Kernel_VoronoiDecomposition.h"',
 69	f"// Date cuda code last modified : {date_modified}"]
 70	cuoptions = ("-default-device", f"-I {cuda_path}") 
 71
 72	source="\n".join(source)
 73	
 74	module = cupy_module_helper.GetModule(source,cuoptions)
 75	cupy_module_helper.SetModuleConstant(module,'size_tot',size,int_t)
 76
 77	if blockDim is None: blockDim = [128,128,128,128,128,32,32][ndim]
 78	gridDim = int(np.ceil(size/blockDim))
 79
 80	def retry64():
 81		if retry64_tol==0 or ndim<5: return
 82		nonlocal m_exp
 83		if m_exp is None: m_exp = expand_symmetric_matrix(m)
 84		mrec = Reconstruct(weights,offsets)
 85		error = np.sum(np.abs(m_exp-mrec),axis=(0,1))
 86		bad = np.logical_not(error < (retry64_tol * lp.trace(m_exp)))
 87		if np.any(bad):
 88#				print(f"nrecomputed {np.sum(bad)}")
 89			out64 = VoronoiDecomposition(m[:,bad],offset_t=offset_t,
 90				flattened=True,traits=traits,retry64_tol=0.,steps=steps)
 91			for a,b in zip(out,out64): a[...,bad]=b
 92
 93	if steps=="Both":
 94		cupy_kernel = module.get_function("VoronoiDecomposition")
 95		cupy_kernel((gridDim,),(blockDim,),(m,weights,offsets))
 96		out = weights,offsets
 97		retry64()
 98		return out
 99
100	# Two steps approach. First minimize
101	a = cp.empty((ndim,ndim,*shape),dtype=float_t)
102	vertex = cp.empty(shape,dtype=int_t)
103	objective = cp.empty(shape,dtype=float_t)
104	a,vertex,objective = map(cp.ascontiguousarray,(a,vertex,objective))
105
106	cupy_kernel = module.get_function("VoronoiMinimization")
107	cupy_kernel((gridDim,),(blockDim,),(m,a,vertex,objective))
108
109	if steps=="Single": 
110		out = a,vertex,objective
111		retry64()
112		return out
113
114	cupy_kernel = module.get_function("VoronoiKKT")
115	cupy_kernel((gridDim,),(blockDim,),(m,a,vertex,objective,weights,offsets))
116	if shape==(): vertex,objective = (e.reshape(()) for e in (vertex,objective))
117
118	out = a,vertex,objective,weights,offsets
119	retry64()
120	return out