def decorate(nodes): processed = [] for node in nodes: mapper = {} for tree in retrieve_iteration_tree(node): vector_iterations = [i for i in tree if i.is_Vectorizable] for i in vector_iterations: handle = FindSymbols('symbolics').visit(i) try: aligned = [ j for j in handle if j.is_Tensor and j.shape[-1] % get_simd_items(j.dtype) == 0 ] except KeyError: aligned = [] if aligned: simd = omplang['simd-for-aligned'] simd = as_tuple( simd(','.join([j.name for j in aligned]), simdinfo[get_simd_flag()])) else: simd = as_tuple(omplang['simd-for']) mapper[i] = i._rebuild(pragmas=i.pragmas + ignore_deps + simd) processed.append(Transformer(mapper).visit(node)) return processed
def _simdize(self, nodes, state): """ Add compiler-specific or, if not available, OpenMP pragmas to the Iteration/Expression tree to emit SIMD-friendly code. """ ignore_deps = as_tuple(self._compiler_decoration('ignore-deps')) mapper = {} for tree in retrieve_iteration_tree(nodes): vector_iterations = [i for i in tree if i.is_Vectorizable] for i in vector_iterations: handle = FindSymbols('symbolics').visit(i) try: aligned = [ j for j in handle if j.is_Tensor and j.shape[-1] % get_simd_items(j.dtype) == 0 ] except KeyError: aligned = [] if aligned: simd = Ompizer.lang['simd-for-aligned'] simd = as_tuple( simd(','.join([j.name for j in aligned]), simdinfo[get_simd_flag()])) else: simd = as_tuple(Ompizer.lang['simd-for']) mapper[i] = i._rebuild(pragmas=i.pragmas + ignore_deps + simd) processed = Transformer(mapper).visit(nodes) return processed, {}