diff --git a/loopy/kernel/tools.py b/loopy/kernel/tools.py index 58915650f7b6aedcf03484d12f4db39cb1123a53..2c11d4a60f57b058c71aa9b4c2e14d8a19a6193f 100644 --- a/loopy/kernel/tools.py +++ b/loopy/kernel/tools.py @@ -296,6 +296,10 @@ class DomainChanger: """ def __init__(self, kernel, inames): + """ + :arg inames: a non-mutable iterable + """ + self.kernel = kernel if inames: ldi = kernel.get_leaf_domain_indices(inames) @@ -311,6 +315,9 @@ class DomainChanger: self.domain = kernel.combine_domains(()) self.leaf_domain_index = None + def get_original_domain(self): + return self.kernel.domains[self.leaf_domain_index] + def get_domains_with(self, replacement): result = self.kernel.domains[:] if self.leaf_domain_index is not None: @@ -320,6 +327,14 @@ class DomainChanger: return result + def get_kernel_with(self, replacement): + return self.kernel.copy( + domains=self.get_domains_with(replacement), + + # Changing the domain might look like it wants to change grid + # sizes. Not true. + get_grid_sizes=self.kernel.get_grid_sizes) + # }}}