Merge pull request #2294 from devitocodes/example-tweaks

example: small cleanup of tti for easier reuse

devito/arch/archinfo.py

@@ -23,7 +23,7 @@
            'INTEL64', 'SNB', 'IVB', 'HSW', 'BDW', 'KNL', 'KNL7210',
            'SKX', 'KLX', 'CLX', 'CLK', 'SPR',
            # ARM CPUs
-           'AMD', 'ARM', 'M1', 'GRAVITON',
+           'AMD', 'ARM', 'AppleArm', 'M1', 'M2', 'M3', 'GRAVITON',
            # Other legacy CPUs
            'POWER8', 'POWER9',
            # Generic GPUs
@@ -709,6 +709,15 @@ class Arm(Cpu64):
     known_isas = ('fp', 'asimd', 'asimdrdm')
 
 
+class AppleArm(Arm):
+
+    @cached_property
+    def march(self):
+        sysinfo = run(["sysctl", "-n", "machdep.cpu.brand_string"],
+                      stdout=PIPE, stderr=DEVNULL).stdout.decode("utf-8")
+        return sysinfo.split(' ')[1].lower()
+
+
 class Amd(Cpu64):
 
     known_isas = ('cpp', 'sse', 'avx', 'avx2')
@@ -839,7 +848,9 @@ def march(cls):
 
 ARM = Arm('arm')
 GRAVITON = Arm('graviton')
-M1 = Arm('m1')
+M1 = AppleArm('m1')
+M2 = AppleArm('m2')
+M3 = AppleArm('m3')
 
 AMD = Amd('amd')
 

devito/arch/compiler.py

@@ -14,7 +14,7 @@
 from codepy.toolchain import (GCCToolchain,
                               call_capture_output as _call_capture_output)
 
-from devito.arch import (AMDGPUX, Cpu64, M1, NVIDIAX, POWER8, POWER9, GRAVITON,
+from devito.arch import (AMDGPUX, Cpu64, AppleArm, NVIDIAX, POWER8, POWER9, GRAVITON,
                          IntelDevice, get_nvidia_cc, check_cuda_runtime,
                          get_m1_llvm_path)
 from devito.exceptions import CompilationError
@@ -486,14 +486,16 @@ def __init_finalize__(self, **kwargs):
                                  '-fopenmp-targets=amdgcn-amd-amdhsa',
                                  '-Xopenmp-target=amdgcn-amd-amdhsa']
                 self.ldflags += ['-march=%s' % platform.march]
-        elif platform is M1:
+        elif isinstance(platform, AppleArm):
             # NOTE:
-            # Apple M1 supports OpenMP through Apple's LLVM compiler.
+            # Apple Mx supports OpenMP through Apple's LLVM compiler.
             # The compiler can be installed with Homebrew or can be built from scratch.
             # Check if installed and set compiler flags accordingly
             llvmm1 = get_m1_llvm_path(language)
             if llvmm1 and language == 'openmp':
-                self.ldflags += ['-mcpu=apple-m1', '-fopenmp', '-L%s' % llvmm1['libs']]
+                mx = platform.march
+                self.ldflags += ['-mcpu=apple-%s' % mx,
+                                 '-fopenmp', '-L%s' % llvmm1['libs']]
                 self.cflags += ['-Xclang', '-I%s' % llvmm1['include']]
         else:
             if platform in [POWER8, POWER9]:
@@ -895,7 +897,7 @@ def __new__(cls, *args, **kwargs):
         platform = kwargs.pop('platform', configuration['platform'])
         language = kwargs.pop('language', configuration['language'])
 
-        if platform is M1:
+        if isinstance(platform, AppleArm):
             _base = ClangCompiler
         elif isinstance(platform, IntelDevice):
             _base = OneapiCompiler

devito/builtins/utils.py

@@ -27,7 +27,7 @@ def __init__(self, *functions, op=dv.mpi.MPI.SUM, dtype=None):
         else:
             dtype = {f.dtype for f in functions}
             if len(dtype) == 1:
-                self.dtype = dtype.pop()
+                self.dtype = np.result_type(dtype.pop(), np.float32).type
             else:
                 raise ValueError("Illegal mixed data types")
         self.v = None
@@ -37,7 +37,7 @@ def __enter__(self):
         i = dv.Dimension(name='mri',)
         self.n = dv.Function(name='n', shape=(1,), dimensions=(i,),
                              grid=self.grid, dtype=self.dtype, space='host')
-        self.n.data[0] = 0
+        self.n.data[:] = 0
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):

devito/types/dimension.py

@@ -922,7 +922,7 @@ def _arg_defaults(self, _min=None, size=None, alias=None):
         except AttributeError:
             factor = dim._factor
 
-        defaults[dim.parent.max_name] = range(1, factor*size - 1)
+        defaults[dim.parent.max_name] = range(0, factor*size - 1)
 
         return defaults
 

examples/seismic/preset_models.py

@@ -222,12 +222,12 @@ def demo_model(preset, **kwargs):
         for i in range(1, nlayers):
             v[..., i*int(shape[-1] / nlayers):] = vp_i[i]  # Bottom velocity
 
-        epsilon = .3*(v - 1.5)
-        delta = .2*(v - 1.5)
-        theta = .5*(v - 1.5)
+        epsilon = .1*(v - vp_top)
+        delta = .05*(v - vp_top)
+        theta = .5*(v - vp_top)
         phi = None
         if len(shape) > 2 and preset.lower() not in ['layers-tti-noazimuth']:
-            phi = .25*(v - 1.5)
+            phi = .25*(v - vp_top)
 
         if density:
             kwargs['b'] = Gardners(v)