agd.Eikonal.HFM_CUDA.inf_convolution

 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
 7import numbers
 8
 9from . import cupy_module_helper
10from .cupy_module_helper import SetModuleConstant
11from ... import AutomaticDifferentiation as ad
12
13def dtype_sup(dtype):
14	dtype=np.dtype(dtype)
15	if dtype.kind in ('i','u'): return np.iinfo(dtype).max
16	elif dtype.kind=='f': return dtype(np.inf)
17	else: raise ValueError("Unsupported dtype")
18def dtype_inf(dtype):
19	dtype=np.dtype(dtype)
20	if dtype.kind in ('i','u'): return np.iinfo(dtype).min
21	elif dtype.kind=='f': return dtype(-np.inf)
22	else: raise ValueError("Unsupported dtype")
23
24def distance_kernel(radius,ndim,dtype=np.float32,ord=2,mult=1):
25	rg = range(-radius,radius+1)
26	axes = (rg,)*ndim
27	X = np.meshgrid(*axes)
28	dist = mult*ad.Optimization.norm(X,axis=0,ord=ord)
29	if np.dtype(dtype).kind in ('i','u'): dist = np.round(dist)
30	return dist.astype(dtype)
31
32def inf_convolution(arr,kernel,out=None,niter=1,periodic=False,
33	upper_saturation=None, lower_saturation=None, mix_is_min=True,
34	overwrite=False,block_size=1024):
35	"""
36	Perform an inf convolution of an input with a given kernel, on the GPU.
37	- arr : the input array
38	- kernel : the convolution kernel. A centered kernel will be used.
39	- niter (optional) : number of iterations of the convolution.
40	- periodic (optional, bool or tuple of bool): axes using periodic boundary conditions.
41	- mix_is_min : if false, use sup_convolution instead
42	"""
43	if niter>=2 and not overwrite: arr=arr.copy()
44	arr = cp.ascontiguousarray(arr)
45
46	conv_t = arr.dtype.type
47	int_t = np.int32
48
49	traits = {
50		'T':conv_t,
51		'shape_c':kernel.shape,
52		'mix_is_min_macro':int(mix_is_min),
53		'ndim':arr.ndim,
54		}
55
56	if upper_saturation is None: upper_saturation = dtype_sup(conv_t)
57	else: traits['upper_saturation_macro']=1
58	traits['T_Sup']=(upper_saturation,conv_t)
59	if lower_saturation is None: lower_saturation = dtype_inf(conv_t)
60	else: traits['lower_saturation_macro']=1
61	traits['T_Inf']=(lower_saturation,conv_t)
62
63	if isinstance(periodic,numbers.Number): periodic = (periodic,)*arr.ndim
64	if any(periodic): 
65		traits['periodic_macro'] = 1
66		traits['periodic_axes'] = periodic
67
68	cuda_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),"cuda")
69	date_modified = cupy_module_helper.getmtime_max(cuda_path)
70	source = cupy_module_helper.traits_header(traits,size_of_shape=True)
71
72	source += [
73	'#include "Kernel_InfConvolution.h"',
74	f"// Date cuda code last modified : {date_modified}"]
75	cuoptions = ("-default-device", f"-I {cuda_path}") 
76
77	source="\n".join(source)
78	module = cupy_module_helper.GetModule(source,cuoptions)
79	SetModuleConstant(module,'kernel_c',kernel,conv_t)
80	SetModuleConstant(module,'shape_tot',arr.shape,int_t)
81	SetModuleConstant(module,'size_tot',arr.size,int_t)
82
83	cupy_kernel = module.get_function("InfConvolution")
84
85	if out is None: out = np.empty_like(arr)
86	else: assert out.dtype==arr.dtype and out.size==arr.size and out.flags['C_CONTIGUOUS']
87	grid_size = int(np.ceil(arr.size/block_size))
88
89	for i in range(niter):
90		cupy_kernel((grid_size,),(block_size,),(arr,out))
91		arr,out = out,arr
92
93	return arr
94