diff --git a/src/pystencils/gpu/kernelcreation.py b/src/pystencils/gpu/kernelcreation.py index c7670c2d1311567d43881283bd445f3fb98eef85..2feb8883a8df644ec2f4dca5d6452831bea619f3 100644 --- a/src/pystencils/gpu/kernelcreation.py +++ b/src/pystencils/gpu/kernelcreation.py @@ -72,11 +72,11 @@ def create_cuda_kernel(assignments: NodeCollection, config: CreateKernelConfig): if len(indexed_elements) > 0: common_indexed_element = get_common_indexed_element(indexed_elements) - index = list(common_indexed_element.indices[0].atoms(TypedSymbol)) + index = common_indexed_element.indices[0].atoms(TypedSymbol) assert len(index) == 1, "index expressions must only contain one symbol representing the index" indexing = indexing_creator(iteration_space=(slice(0, common_indexed_element.shape[0], 1), *iteration_space), data_layout=common_field.layout) - extended_ctrs = [index[0], *loop_counter_symbols] + extended_ctrs = [index.pop(), *loop_counter_symbols] loop_counter_assignments = indexing.get_loop_ctr_assignments(extended_ctrs) else: indexing = indexing_creator(iteration_space=iteration_space, data_layout=common_field.layout) diff --git a/src/pystencils/transformations.py b/src/pystencils/transformations.py index 628da58670b4a4f10b0b44abf631c3a2bc97dc56..5374806a9ff1ccf7508f11a4d742d2bd2eaa676d 100644 --- a/src/pystencils/transformations.py +++ b/src/pystencils/transformations.py @@ -276,10 +276,10 @@ def add_outer_loop_over_indexed_elements(loop_node: ast.Block) -> ast.Block: if len(indexed_elements) == 0: return loop_node reference_element = get_common_indexed_element(indexed_elements) - index = list(reference_element.indices[0].atoms(TypedSymbol)) + index = reference_element.indices[0].atoms(TypedSymbol) assert len(index) == 1, "index expressions must only contain one symbol representing the index" new_loop = ast.LoopOverCoordinate(loop_node, 0, 0, - reference_element.shape[0], 1, custom_loop_ctr=index[0]) + reference_element.shape[0], 1, custom_loop_ctr=index.pop()) return ast.Block([new_loop])