Work consistently in slowness

docker/README.md

@@ -81,19 +81,19 @@ docker build --network=host --file docker/Dockerfile.stride --tag stride .
 And to build the GPU image with `openacc` offloading and the `nvc` compiler, simply run:
 
 ```bash
-docker build --build-arg base=devitocodes/base:nvidia-nvc --network=host --file docker/Dockerfile.stride --tag stride .
+docker build --build-arg base=devitocodes/bases:nvidia-nvc --network=host --file docker/Dockerfile.stride --tag stride .
 ```
 
 or if you wish to use the `llvm-15` (clang) compiler with `openmp` offlaoding:
 
 ```bash
-docker build --build-arg base=devitocodes/base:nvidia-clang --network=host --file docker/Dockerfile.stride --tag stride .
+docker build --build-arg base=devitocodes/bases:nvidia-clang --network=host --file docker/Dockerfile.stride --tag stride .
 ```
 
 and finally for AMD architectures:
 
 ```bash
-docker build --build-arg base=devitocodes/base:amd --network=host --file docker/Dockerfile.stride --tag stride .
+docker build --build-arg base=devitocodes/bases:amd --network=host --file docker/Dockerfile.stride --tag stride .
 ```
 
 

mosaic/cli/mrun.py

@@ -146,7 +146,7 @@ def start_runtime(*args, **extra_kwargs):
     _runtime.init_file(runtime_config)
 
     def run_head():
-        process = cmd_subprocess.run(cmd,
+        process = cmd_subprocess.run(' '.join(cmd), shell=True,
                                      stdout=_stdout,
                                      stderr=_stderr)
 

mosaic/comms/serialisation.py

@@ -35,7 +35,7 @@ def serialise(data):
     """
     try:
         return pickle5_dumps(data)
-    except pickle.PicklingError:
+    except (pickle.PicklingError, AttributeError):
         return cloudpickle.dumps(data), []
 
 

stride/__init__.py

@@ -108,7 +108,7 @@ async def forward(problem, pde, *args, **kwargs):
     published_args = await asyncio.gather(*published_args)
 
     platform = kwargs.get('platform', 'cpu')
-    using_gpu = platform in ['nvidia-acc', 'gpu']
+    using_gpu = 'nvidia' in platform or 'gpu' in platform
     if using_gpu:
         devices = kwargs.pop('devices', None)
         num_gpus = gpu_count() if devices is None else len(devices)
@@ -251,7 +251,7 @@ async def adjoint(problem, pde, loss, optimisation_loop, optimiser, *args, **kwa
     keep_residual = isinstance(step_size, LineSearch)
 
     platform = kwargs.get('platform', 'cpu')
-    using_gpu = platform in ['nvidia-acc', 'gpu']
+    using_gpu = 'nvidia' in platform or 'gpu' in platform
     if using_gpu:
         devices = kwargs.pop('devices', None)
         num_gpus = gpu_count() if devices is None else len(devices)

stride/core.py

@@ -288,6 +288,7 @@ def __init__(self, *args, **kwargs):
 
         self.grad = None
         self.prec = None
+        self.transform = kwargs.pop('transform', None)
 
         self.graph = Graph()
         self.prev_op = None
@@ -384,6 +385,7 @@ def detach(self, *args, **kwargs):
         """
         kwargs['name'] = kwargs.pop('name', self._init_name)
         kwargs['needs_grad'] = kwargs.pop('needs_grad', self.needs_grad)
+        kwargs['transform'] = kwargs.pop('transform', self.transform)
 
         if hasattr(self, 'has_tessera') and self.has_tessera:
             cpy = self.__class__.parameter(*args, **kwargs)
@@ -411,6 +413,7 @@ def as_parameter(self, *args, **kwargs):
         """
         kwargs['name'] = kwargs.pop('name', self._init_name)
         kwargs['needs_grad'] = kwargs.pop('needs_grad', self.needs_grad)
+        kwargs['transform'] = kwargs.pop('transform', self.transform)
 
         cpy = self.__class__.parameter(*args, **kwargs)
 
@@ -437,6 +440,7 @@ def copy(self, *args, **kwargs):
         """
         kwargs['name'] = kwargs.pop('name', self._init_name)
         kwargs['needs_grad'] = kwargs.pop('needs_grad', self.needs_grad)
+        kwargs['transform'] = kwargs.pop('transform', self.transform)
 
         propagate_tessera = kwargs.pop('propagate_tessera', True)
 

stride/optimisation/loss/l2_distance.py

@@ -24,7 +24,7 @@ def __init__(self, **kwargs):
 
         self.residual = None
 
-    async def forward(self, modelled, observed, **kwargs):
+    def forward(self, modelled, observed, **kwargs):
         problem = kwargs.pop('problem', None)
         shot_id = problem.shot_id if problem is not None \
             else kwargs.pop('shot_id', 0)
@@ -38,7 +38,7 @@ async def forward(self, modelled, observed, **kwargs):
 
         return fun
 
-    async def adjoint(self, d_fun, modelled, observed, **kwargs):
+    def adjoint(self, d_fun, modelled, observed, **kwargs):
         grad_modelled = None
         if modelled.needs_grad:
             grad_modelled = +np.asarray(d_fun) * self.residual.copy(name='modelledresidual')

stride/optimisation/optimisers/gradient_descent.py

@@ -30,6 +30,6 @@ async def pre_process(self, grad=None, processed_grad=None, **kwargs):
                                                    **kwargs)
         return processed_grad
 
-    async def update_variable(self, step_size, direction):
-        self.variable.data[:] -= step_size * direction.data
-        return self.variable
+    def update_variable(self, step_size, variable, direction):
+        variable.data[:] -= step_size * direction.data
+        return variable

stride/optimisation/optimisers/optimiser.py

@@ -96,6 +96,7 @@ async def pre_process(self, grad=None, processed_grad=None, **kwargs):
                 if dump_prec and self.variable.grad.prec is not None and problem is not None:
                     self.variable.grad.prec.dump(path=problem.output_folder,
                                                  project_name=problem.name,
+                                                 parameter='raw_%s' % self.variable.grad.prec.name,
                                                  version=iteration.abs_id+1)
 
                 grad = self.variable.process_grad(**kwargs)
@@ -105,6 +106,11 @@ async def pre_process(self, grad=None, processed_grad=None, **kwargs):
                               project_name=problem.name,
                               version=iteration.abs_id+1)
 
+                if dump_prec and self.variable.grad.prec is not None and problem is not None:
+                    self.variable.grad.prec.dump(path=problem.output_folder,
+                                                 project_name=problem.name,
+                                                 version=iteration.abs_id+1)
+
             min_dir = np.min(grad.data)
             max_dir = np.max(grad.data)
 
@@ -171,7 +177,7 @@ async def step(self, step_size=None, grad=None, processed_grad=None, **kwargs):
         step_size = self.step_size if step_size is None else step_size
         step_loop = kwargs.pop('step_loop', None)
         if isinstance(step_size, LineSearch):
-            await step_size.init_search(
+            step_size.init_search(
                 variable=self.variable,
                 direction=direction,
                 **kwargs
@@ -185,7 +191,7 @@ async def step(self, step_size=None, grad=None, processed_grad=None, **kwargs):
                     next_step = 1.
                     done_search = True
                 else:
-                    next_step, done_search = await step_size.next_step(
+                    next_step, done_search = step_size.next_step(
                         variable=self.variable,
                         direction=direction,
                         **kwargs
@@ -216,7 +222,14 @@ async def step(self, step_size=None, grad=None, processed_grad=None, **kwargs):
             self.variable.data[:] = variable_before.data.copy()
 
             # update variable
-            await self.update_variable(next_step, direction)
+            if self.variable.transform is not None:
+                variable = self.variable.transform(self.variable)
+            else:
+                variable = self.variable
+            upd_variable = self.update_variable(next_step, variable, direction)
+            if self.variable.transform is not None:
+                upd_variable = self.variable.transform(upd_variable)
+            self.variable.data[:] = upd_variable.data.copy()
 
             # post-process variable after update
             await self.post_process(**kwargs)
@@ -262,13 +275,15 @@ async def post_process(self, **kwargs):
         self.variable.release_grad()
 
     @abstractmethod
-    async def update_variable(self, step_size, direction):
+    def update_variable(self, step_size, variable, direction):
         """
 
         Parameters
         ----------
         step_size : float
             Step size to use for updating the variable.
+        variable : Data
+            Variable to update.
         direction : Data
             Direction in which to update the variable.
 

stride/optimisation/step_length/line_search.py

@@ -8,12 +8,12 @@
 class LineSearch:
 
     @abstractmethod
-    async def init_search(self, variable, direction, **kwargs):
+    def init_search(self, variable, direction, **kwargs):
         pass
 
     @abstractmethod
-    async def next_step(self, variable, direction, **kwargs):
+    def next_step(self, variable, direction, **kwargs):
         pass
 
-    async def forward_test(self, variable, direction, **kwargs):
+    def forward_test(self, variable, direction, **kwargs):
         pass

stride/physics/common/devito.py

@@ -906,30 +906,30 @@ def set_operator(self, op, **kwargs):
             default_config = {
                 'name': self.name,
                 'subs': subs,
-                'opt': 'advanced-fsg',
+                'opt': 'advanced',
             }
 
         elif platform == 'cpu-icc':
             default_config = {
                 'name': self.name,
                 'subs': subs,
-                'opt': 'advanced-fsg',
+                'opt': 'advanced',
                 'compiler': 'icc',
             }
 
         elif platform == 'cpu-nvc':
             default_config = {
                 'name': self.name,
                 'subs': subs,
-                'opt': 'advanced-fsg',
+                'opt': 'advanced',
                 'compiler': 'nvc',
             }
 
         elif platform == 'nvidia-acc' or platform == 'nvidia-nvc':
             default_config = {
                 'name': self.name,
                 'subs': subs,
-                'opt': 'advanced-fsg',
+                'opt': 'advanced',
                 'autotuning': 'off',
                 'compiler': 'nvc',
                 'language': 'openacc',
@@ -940,7 +940,7 @@ def set_operator(self, op, **kwargs):
             default_config = {
                 'name': self.name,
                 'subs': subs,
-                'opt': 'advanced-fsg',
+                'opt': 'advanced',
                 'compiler': 'cuda',
                 'language': 'cuda',
                 'platform': 'nvidiaX',
@@ -950,7 +950,7 @@ def set_operator(self, op, **kwargs):
             default_config = {
                 'name': self.name,
                 'subs': subs,
-                'opt': 'advanced-fsg',
+                'opt': 'advanced',
                 'compiler': 'hip',
                 'language': 'hip',
                 'platform': 'amdgpuX',
@@ -964,7 +964,7 @@ def set_operator(self, op, **kwargs):
         context = {'log-level': 'DEBUG' if log_level in ['perf', 'debug'] else 'INFO'}
         compiler_config = {}
         for key, value in default_config.items():
-            if key in devito.configuration:
+            if key in devito.configuration and key != 'opt':
                 context[key] = value
             else:
                 compiler_config[key] = value