def __init__(self, slvr_cfg): """ RimeSolver Constructor Parameters: slvr_cfg : SolverConfiguration Solver Configuration variables """ # Call the parent constructor super(RimeSolver, self).__init__(slvr_cfg) self.register_default_dimensions() # Configure the dimensions of the beam cube self.register_dimension('beam_lw', slvr_cfg[Options.E_BEAM_WIDTH], description='E cube l width') self.register_dimension('beam_mh', slvr_cfg[Options.E_BEAM_HEIGHT], description='E cube m height') self.register_dimension('beam_nud', slvr_cfg[Options.E_BEAM_DEPTH], description='E cube nu depth') self.rime_e_beam = RimeEBeam() self.rime_b_sqrt = RimeBSqrt() self.rime_ekb_sqrt = RimeEKBSqrt() self.rime_sum = RimeSumCoherencies() self.rime_reduce = RimeReduction() from montblanc.impl.rime.v4.ant_pairs import monkey_patch_antenna_pairs monkey_patch_antenna_pairs(self) # Create # (1) A stream that this solver will asynchronously # operate on # (2) An event indicating when an iteration of # the kernels above have finished executing with self.context: self.stream = cuda.Stream() self.kernels_done = cuda.Event() # Create constant data for transfer to GPU self._const_data = mbu.create_rime_const_data(self) # Indicate these variables have not been set self._dev_mem_pool = None self._pinned_mem_pool = None self._pool_lock = None
class RimeSolver(MontblancCUDASolver): """ RIME Solver Implementation """ def __init__(self, slvr_cfg): """ RimeSolver Constructor Parameters: slvr_cfg : SolverConfiguration Solver Configuration variables """ # Call the parent constructor super(RimeSolver, self).__init__(slvr_cfg) self.register_default_dimensions() # Configure the dimensions of the beam cube self.register_dimension('beam_lw', slvr_cfg[Options.E_BEAM_WIDTH], description='E cube l width') self.register_dimension('beam_mh', slvr_cfg[Options.E_BEAM_HEIGHT], description='E cube m height') self.register_dimension('beam_nud', slvr_cfg[Options.E_BEAM_DEPTH], description='E cube nu depth') self.rime_e_beam = RimeEBeam() self.rime_b_sqrt = RimeBSqrt() self.rime_ekb_sqrt = RimeEKBSqrt() self.rime_sum = RimeSumCoherencies() self.rime_reduce = RimeReduction() from montblanc.impl.rime.v4.ant_pairs import monkey_patch_antenna_pairs monkey_patch_antenna_pairs(self) # Create # (1) A stream that this solver will asynchronously # operate on # (2) An event indicating when an iteration of # the kernels above have finished executing with self.context: self.stream = cuda.Stream() self.kernels_done = cuda.Event() # Create constant data for transfer to GPU self._const_data = mbu.create_rime_const_data(self) # Indicate these variables have not been set self._dev_mem_pool = None self._pinned_mem_pool = None self._pool_lock = None def const_data(self): return self._const_data def update_dimension(self, **update_dict): """ Override update_dimension on HyperCube to also update rime_const_data. """ # Defer to parent method on the base solver super(RimeSolver, self).update_dimension(**update_dict) # Update constant data, updating nsrc with sum of source counts self._const_data.update(self, sum_nsrc=True) @property def dev_mem_pool(self): return self._dev_mem_pool @dev_mem_pool.setter def dev_mem_pool(self, pool): self._dev_mem_pool = pool @property def pinned_mem_pool(self): return self._pinned_mem_pool @pinned_mem_pool.setter def pinned_mem_pool(self, pool): self._pinned_mem_pool = pool @property def pool_lock(self): return self._pool_lock @pool_lock.setter def pool_lock(self, lock): self._pool_lock = lock def initialise(self): with self.context: self.rime_e_beam.initialise(self) self.rime_b_sqrt.initialise(self) self.rime_ekb_sqrt.initialise(self) self.rime_sum.initialise(self) self.rime_reduce.initialise(self) def solve(self): with self.context: self.rime_e_beam.execute(self) self.rime_b_sqrt.execute(self) self.rime_ekb_sqrt.execute(self) self.rime_sum.execute(self) self.rime_reduce.execute(self) def shutdown(self): with self.context: self.rime_e_beam.shutdown(self) self.rime_b_sqrt.shutdown(self) self.rime_ekb_sqrt.shutdown(self) self.rime_sum.shutdown(self) self.rime_reduce.shutdown(self)