Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

api: Introduce complex numbers support (np.complex64/128) #2375

Merged
merged 74 commits into from
Mar 19, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
74 commits
Select commit Hold shift + click to select a range
9b0420e
api: add support for complex dtype
mloubout Aug 2, 2023
5398c68
api: fix printer for complex dtype
mloubout May 22, 2024
2b4f7b0
compiler: fix alias dtype with complex numbers
mloubout May 22, 2024
0cf3f78
api: move complex ctype to dtype lowering
mloubout May 22, 2024
386d5a2
compiler: generate std:complex for cpp compilers
mloubout May 28, 2024
b5bf79c
compiler: add std::complex arithmetic defs for unsupported types
mloubout May 30, 2024
ac28372
compiler: fix alias dtype with complex numbers
mloubout May 30, 2024
89b6fea
compiler: fix internal language specific types and cast
mloubout May 31, 2024
d5a2542
compiler: rework dtype lowering
mloubout Jun 20, 2024
64801cb
compiler: switch to c++14 for complex_literals
mloubout Jun 27, 2024
2f0cacf
compiler: subdtype numpy for dtype lowering
mloubout Jul 8, 2024
7ad523b
compiler: use structs to pass complex arguments
enwask Jul 9, 2024
e8e51f1
compiler: add Dereference scalar case
enwask Jul 11, 2024
d0c55f9
compiler: implement float16 support
enwask Jul 11, 2024
9081b2e
symbolics: fix printer for half precision
enwask Jul 11, 2024
89c8bf5
misc: fix formatting
enwask Jul 11, 2024
43b82db
compiler: refactor float16 and lower_dtypes
enwask Jul 11, 2024
91ab1e9
compiler: add dtype_alloc_ctype helper for allocation size
enwask Jul 11, 2024
d907e8d
misc: more float16 refactoring/formatting fixes
enwask Jul 15, 2024
425d568
Remove dtypes lowering from IET layer
enwask Jul 16, 2024
8db3f96
compiler: reimplement float16/complex lowering
enwask Jul 26, 2024
53462a3
misc: cleanup, docs and typing for half support
enwask Jul 29, 2024
8a8f24f
compiler: FindSymbols 'scalars' -> 'abstractsymbols'
enwask Jul 29, 2024
b0c9ee9
test: include scalar parameters in complex tests
enwask Jul 30, 2024
344c435
test: add test_dtypes with initial tests for float16 + complex
enwask Jul 30, 2024
4f5ee36
misc: more lower_dtypes cleanup + type hints
enwask Jul 30, 2024
49a9bec
api: use grid dtype for extent and origin, add test_grid
enwask Jul 31, 2024
e8c7fc0
test: clean up and add more half/complex tests
enwask Jul 31, 2024
12b936a
test: fix test_grid_objs, add test_grid_dtypes
enwask Jul 31, 2024
1981465
api: allow side for cross derivatives, fixes #2442
mloubout Aug 13, 2024
da2ea2a
compiler: process dtypes through printer
mloubout Jan 15, 2025
f76dbad
symbolics: specialize sizeof
mloubout Jan 16, 2025
d55931f
compiler: move dtype pass to top level operator iet pass
mloubout Jan 16, 2025
cb6fdc8
symbolics: fix SizeOf rebuild
mloubout Jan 16, 2025
11d6b32
symbolics: use std namespace for c++
mloubout Jan 16, 2025
ed5dbe6
compiler: fix std math func names
mloubout Jan 16, 2025
63e589d
symbolics: move printers rogether through registry
mloubout Jan 17, 2025
f781aa4
symbolics: rework Cast
mloubout Jan 17, 2025
05d4e3f
compiler: fix complex headers
mloubout Jan 17, 2025
6fc54e3
api: remove un-needed dtype reconstruction mode
mloubout Jan 17, 2025
0f17026
compiler: fix dtype for mpi routines
mloubout Jan 17, 2025
2ce9817
compiler: fix missing algorithm include for min/max
mloubout Jan 18, 2025
9307bf0
arch: switch sycl error to warning for no-compile codegen
mloubout Jan 18, 2025
2b2848b
symbolics: rework cast/sizeof for pickling
mloubout Jan 22, 2025
3e9e931
api: fix c_datatype hack
mloubout Jan 22, 2025
58e6310
compiler: make visitor language parametric
mloubout Jan 23, 2025
f1a082d
compiler: make sure complex ctype is handled properly for typedata
mloubout Jan 23, 2025
0c914f2
symbolics: cleaner repr of Cast
mloubout Jan 23, 2025
1c9ab2e
test: improve dtype tests log
mloubout Jan 24, 2025
461fd43
compiler: make sure cpp is used for c++ compilers
mloubout Jan 26, 2025
c515253
compiler: make printer part of the target and differentiate C and CXX
mloubout Jan 27, 2025
4899f11
compiler: add all cxx target to operator registry
mloubout Jan 27, 2025
f968a89
compiler: cleanup operator class names
mloubout Jan 28, 2025
d385957
compiler: switch cxx backend to static_cast
mloubout Jan 28, 2025
b419050
compiler: add switch for static_cast vs reinterpret_cast
mloubout Jan 28, 2025
ad7271f
compiler: handle plain text header
mloubout Jan 30, 2025
91f2018
compiler: convert all in visitors to f-string
mloubout Jan 30, 2025
049f17a
compiler: convert printer to f-string
mloubout Jan 31, 2025
8f4f221
arch: add intel gpu basic gpu_info support
mloubout Feb 13, 2025
d24e6e1
compiler: fix header order
mloubout Feb 20, 2025
8e0a2d3
compiler: add scalar type option
mloubout Feb 25, 2025
65f1b37
compiler: fix real dtype
mloubout Mar 2, 2025
1da795e
compiler: more robust safeinv
mloubout Mar 5, 2025
e9fa6ec
compiler: mssing substraction cxx def
mloubout Mar 6, 2025
83737b0
compiler: make dtype lowering more flexible
mloubout Mar 6, 2025
6a563a2
examples: add on the fly dft tutorial
mloubout Mar 6, 2025
601d99d
api: fix norm with complex numbers
mloubout Mar 6, 2025
c1a31a9
api: dix sympy assumptions for complex valued objects
mloubout Mar 6, 2025
43da5a2
misc: f-string formatting
mloubout Mar 11, 2025
8185a5e
compiler: cleanup default includes/header/namespaces
mloubout Mar 14, 2025
124b9ba
misc: fix typos and formatting
mloubout Mar 17, 2025
24faf28
compiler: rename lang option to langbb for clarity
mloubout Mar 17, 2025
5f89aed
api: enforce pow_to_mul to be un-evaluable
mloubout Mar 18, 2025
be87c17
compiler: rename lang to langbb throughout for clarity
mloubout Mar 18, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .github/workflows/docker-bases.yml
Original file line number Diff line number Diff line change
Expand Up @@ -255,4 +255,6 @@ jobs:
file: './docker/Dockerfile.amd'
push: true
target: 'hip'
build-args: |
ROCM_VERSION=6.3.4
tags: devitocodes/bases:amd-hip
5 changes: 3 additions & 2 deletions devito/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,8 @@ def reinit_compiler(val):
"""
Re-initialize the Compiler.
"""
configuration['compiler'].__init__(suffix=configuration['compiler'].suffix,
configuration['compiler'].__init__(name=configuration['compiler'].name,
suffix=configuration['compiler'].suffix,
mpi=configuration['mpi'])
return val

Expand All @@ -65,7 +66,7 @@ def reinit_compiler(val):
configuration.add('platform', 'cpu64', list(platform_registry),
callback=lambda i: platform_registry[i]())
configuration.add('compiler', 'custom', compiler_registry,
callback=lambda i: compiler_registry[i]())
callback=lambda i: compiler_registry[i](name=i))

# Setup language for shared-memory parallelism
preprocessor = lambda i: {0: 'C', 1: 'openmp'}.get(i, i) # Handles DEVITO_OPENMP deprec
Expand Down
65 changes: 64 additions & 1 deletion devito/arch/archinfo.py
Original file line number Diff line number Diff line change
Expand Up @@ -330,6 +330,69 @@ def cbk(deviceid=0):
except OSError:
pass

# *** Third try: `sycl-ls`, clearly only works with Intel cards
try:
gpu_infos = {}

# sycl-ls sometimes finds gpu twice with opencl and without so
# we need to make sure we don't get duplicates
selected_platform = None
platform_block = ""

proc = Popen(["sycl-ls", "--verbose"], stdout=PIPE, stderr=DEVNULL, text=True)
sycl_output, _ = proc.communicate()

# Extract platform blocks
platforms = re.findall(r"Platform \[#(\d+)\]:([\s\S]*?)(?=Platform \[#\d+\]:|$)",
sycl_output)

# Select Level-Zero if available, otherwise use OpenCL
for platform_id, platform_content in platforms:
if "Intel(R) Level-Zero" in platform_content:
selected_platform = platform_id
platform_block = platform_content
break
elif "Intel(R) OpenCL Graphics" in platform_content and \
selected_platform is None:
selected_platform = platform_id
platform_block = platform_content

# Extract GPU devices from the selected platform
devices = re.findall(r"Device \[#(\d+)\]:([\s\S]*?)(?=Device \[#\d+\]:|$)",
platform_block)

for device_id, device_block in devices:
if re.search(r"^\s*Type\s*:\s*gpu", device_block, re.MULTILINE):
name_match = re.search(r"^\s*Name\s*:\s*(.+)", device_block, re.MULTILINE)

if name_match:
name = name_match.group(1).strip()

# Store GPU info with correct physical ID
gpu_infos[device_id] = {
"physicalid": device_id,
"product": name
}

gpu_info = homogenise_gpus(list(gpu_infos.values()))

# Also attach callbacks to retrieve instantaneous memory info
# Now this should be done using xpu-smi but for some reason
# it throws a lot of weird errors in docker so skipping for now
for i in ['total', 'free', 'used']:
def make_cbk(i):
def cbk(deviceid=0):
return None
return cbk

gpu_info['mem.%s' % i] = make_cbk(i)

gpu_infos['architecture'] = 'Intel'
return gpu_info

except OSError:
pass

# *** Second try: `lshw`
try:
info_cmd = ['lshw', '-C', 'video']
Expand Down Expand Up @@ -391,7 +454,7 @@ def parse_product_arch():
gpu_infos = []
for line in lines:
# Graphics cards are listed as VGA or 3D controllers in lspci
if 'VGA' in line or '3D' in line:
if any(i in line for i in ('VGA', '3D', 'Display')):
gpu_info = {}
# Lines produced by lspci command are of the form:
# xxxx:xx:xx.x Device Type: Name
Expand Down
74 changes: 43 additions & 31 deletions devito/arch/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,12 +180,21 @@ def __init__(self):
"""

fields = {'cc', 'ld'}
_cpp = False
_default_cpp = False
_cxxstd = 'c++14'
_cstd = 'c99'

def __init__(self, **kwargs):
maybe_name = kwargs.pop('name', self.__class__.__name__)
if isinstance(maybe_name, Compiler):
self._name = maybe_name.name
else:
self._name = maybe_name

super().__init__(**kwargs)

self.__lookup_cmds__()
self._cpp = kwargs.get('cpp', self._default_cpp)

self.suffix = kwargs.get('suffix')
if not kwargs.get('mpi'):
Expand All @@ -195,7 +204,7 @@ def __init__(self, **kwargs):
self.cc = self.MPICC if self._cpp is False else self.MPICXX
self.ld = self.cc # Wanted by the superclass

self.cflags = ['-O3', '-g', '-fPIC', '-Wall', '-std=c99']
self.cflags = ['-O3', '-g', '-fPIC', '-Wall', f'-std={self.std}']
self.ldflags = ['-shared']

self.include_dirs = []
Expand Down Expand Up @@ -225,13 +234,13 @@ def __new_with__(self, **kwargs):
Create a new Compiler from an existing one, inherenting from it
the flags that are not specified via ``kwargs``.
"""
return self.__class__(suffix=kwargs.pop('suffix', self.suffix),
return self.__class__(name=self.name, suffix=kwargs.pop('suffix', self.suffix),
mpi=kwargs.pop('mpi', configuration['mpi']),
**kwargs)

@property
def name(self):
return self.__class__.__name__
return self._name

@property
def version(self):
Expand All @@ -247,6 +256,10 @@ def version(self):

return version

@property
def std(self):
return self._cxxstd if self._cpp else self._cstd

def get_version(self):
result, stdout, stderr = call_capture_output((self.cc, "--version"))
if result != 0:
Expand Down Expand Up @@ -482,15 +495,15 @@ def __init_finalize__(self, **kwargs):
platform = kwargs.pop('platform', configuration['platform'])

if isinstance(platform, NvidiaDevice):
self.cflags.remove('-std=c99')
self.cflags.remove(f'-std={self.std}')
# Add flags for OpenMP offloading
if language in ['C', 'openmp']:
cc = get_nvidia_cc()
if cc:
self.cflags += ['-Xopenmp-target', f'-march=sm_{cc}']
self.ldflags += ['-fopenmp', '-fopenmp-targets=nvptx64-nvidia-cuda']
elif platform is AMDGPUX:
self.cflags.remove('-std=c99')
self.cflags.remove(f'-std={self.std}')
# Add flags for OpenMP offloading
if language in ['C', 'openmp']:
self.ldflags += ['-target', 'x86_64-pc-linux-gnu']
Expand Down Expand Up @@ -553,9 +566,9 @@ def __init_finalize__(self, **kwargs):
self.ldflags += ['-fopenmp']

if isinstance(platform, NvidiaDevice):
self.cflags.remove('-std=c99')
self.cflags.remove(f'-std={self.std}')
elif platform is AMDGPUX:
self.cflags.remove('-std=c99')
self.cflags.remove(f'-std={self.std}')
# Add flags for OpenMP offloading
if language in ['C', 'openmp']:
self.ldflags += ['-target', 'x86_64-pc-linux-gnu']
Expand Down Expand Up @@ -590,16 +603,13 @@ def __lookup_cmds__(self):

class PGICompiler(Compiler):

_cpp = True
_default_cpp = True

def __init_finalize__(self, **kwargs):

self.cflags.remove('-std=c99')
self.cflags.remove('-O3')
self.cflags.remove('-Wall')

self.cflags.append('-std=c++11')

language = kwargs.pop('language', configuration['language'])
platform = kwargs.pop('platform', configuration['platform'])

Expand Down Expand Up @@ -643,14 +653,13 @@ def __lookup_cmds__(self):

class CudaCompiler(Compiler):

_cpp = True
_default_cpp = True

def __init_finalize__(self, **kwargs):

self.cflags.remove('-std=c99')
self.cflags.remove('-Wall')
self.cflags.remove('-fPIC')
self.cflags.extend(['-std=c++14', '-Xcompiler', '-fPIC'])
self.cflags.extend(['-Xcompiler', '-fPIC'])

if configuration['mpi']:
# We rather use `nvcc` to compile MPI, but for this we have to
Expand Down Expand Up @@ -717,15 +726,10 @@ def __lookup_cmds__(self):

class HipCompiler(Compiler):

_cpp = True
_default_cpp = True

def __init_finalize__(self, **kwargs):

self.cflags.remove('-std=c99')
self.cflags.remove('-Wall')
self.cflags.remove('-fPIC')
self.cflags.extend(['-std=c++14', '-fPIC'])

if configuration['mpi']:
# We rather use `hipcc` to compile MPI, but for this we have to
# explicitly pass the flags that an `mpicc` would implicitly use
Expand Down Expand Up @@ -831,7 +835,7 @@ def __init_finalize__(self, **kwargs):
language = kwargs.pop('language', configuration['language'])

if language == 'sycl':
raise ValueError("Use SyclCompiler to jit-compile sycl")
warning(f"Use SyclCompiler (`sycl`) to jit-compile sycl, not {self.name}")

elif language == 'openmp':
# Earlier versions to OneAPI 2023.2.0 (clang17 underneath), have an
Expand Down Expand Up @@ -878,7 +882,7 @@ def __lookup_cmds__(self):

class SyclCompiler(OneapiCompiler):

_cpp = True
_default_cpp = True

def __init_finalize__(self, **kwargs):
IntelCompiler.__init_finalize__(self, **kwargs)
Expand All @@ -887,9 +891,9 @@ def __init_finalize__(self, **kwargs):
language = kwargs.pop('language', configuration['language'])

if language != 'sycl':
raise ValueError("Expected language sycl with SyclCompiler")
warning(f"Expected language sycl with SyclCompiler, not {language}")

self.cflags.remove('-std=c99')
self.cflags.remove(f'-std={self.std}')
self.cflags.append('-fsycl')

self.cflags.remove('-g') # -g disables some optimizations in IGC
Expand All @@ -903,7 +907,7 @@ def __init_finalize__(self, **kwargs):
elif isinstance(platform, IntelDevice):
self.cflags.append('-fsycl-targets=spir64')
else:
raise NotImplementedError(f"Unsupported platform {platform}")
warning(f"Unsupported platform {platform}")


class CustomCompiler(Compiler):
Expand Down Expand Up @@ -945,7 +949,6 @@ def __new__(cls, *args, **kwargs):
obj = super().__new__(cls)
# Keep base to initialize accordingly
obj._base = kwargs.pop('base', _base)
obj._cpp = obj._base._cpp

return obj

Expand Down Expand Up @@ -976,6 +979,10 @@ def __lookup_cmds__(self):
def __new_with__(self, **kwargs):
return super().__new_with__(base=self._base, **kwargs)

@property
def _default_cpp(self):
return self._base._default_cpp


class CompilerRegistry(dict):
"""
Expand All @@ -984,15 +991,19 @@ class CompilerRegistry(dict):
"""

def __getitem__(self, key):
if isinstance(key, Compiler):
key = key.name

if key.startswith('gcc-'):
i = key.split('-')[1]
return partial(GNUCompiler, suffix=i)

return super().__getitem__(key)

def __contains__(self, k):
if isinstance(k, Compiler):
k = k.name
return k in self.keys() or k.startswith('gcc-')
def __contains__(self, key):
if isinstance(key, Compiler):
key = key.name
return key in self.keys() or key.startswith('gcc-')


_compiler_registry = {
Expand All @@ -1011,6 +1022,7 @@ def __contains__(self, k):
'nvc++': NvidiaCompiler,
'nvidia': NvidiaCompiler,
'cuda': CudaCompiler,
'nvcc': CudaCompiler,
'osx': ClangCompiler,
'intel': OneapiCompiler,
'icx': OneapiCompiler,
Expand Down
4 changes: 2 additions & 2 deletions devito/builtins/arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,13 +32,13 @@ def norm(f, order=2):
s = dv.types.Symbol(name='sum', dtype=n.dtype)

op = dv.Operator([dv.Eq(s, 0.0)] + eqns +
[dv.Inc(s, dv.Abs(Pow(p, order))), dv.Eq(n[0], s)],
[dv.Inc(s, Pow(dv.Abs(p), order)), dv.Eq(n[0], s)],
name='norm%d' % order)
op.apply(**kwargs)

v = np.power(n.data[0], 1/order)

return f.dtype(v)
return np.real(f.dtype(v))


@dv.switchconfig(log_level='ERROR')
Expand Down
5 changes: 4 additions & 1 deletion devito/builtins/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,10 @@
# NOTE: np.float128 isn't really a thing, see for example
# https://github.com/numpy/numpy/issues/10288
# https://docs.oracle.com/cd/E19957-01/806-3568/ncg_goldberg.html#1070
np.float64: np.float64
np.float64: np.float64,
# ComplexX accumulates on Complex2X
np.complex64: np.complex128,
np.complex128: np.complex128,
}


Expand Down
Loading
Loading