separate c index array and vectorized index array

python/dune/codegen/blockstructured/vectorization.py

@@ -338,38 +338,67 @@ def find_accumulation_inames(knl):
     return inames
 
 
-def add_iname_array(knl, vec_iname):
-    insns_with_macro_points = lp.find_instructions(knl, Tagged(vec_iname))
+def add_iname_array(knl, iname):
+    insns_with_macro_points = lp.find_instructions(knl, Tagged(iname))
 
     if insns_with_macro_points:
-        array_name = vec_iname + '_arr'
-        vector_name = vec_iname + '_vec'
+        new_iname = knl.get_var_name_generator()(iname)
+
+        new_dom = BasicSet('{{ [{0}] : 0<={0}<{1} }}'.format(new_iname, get_form_option('number_of_blocks')))
+
+        array_name = iname + '_arr'
 
         new_temporaries = dict()
         new_temporaries[array_name] = DuneTemporaryVariable(array_name, managed=True,
                                                             shape=(get_form_option('number_of_blocks'),),
                                                             scope=lp.temp_var_scope.PRIVATE, dtype=np.float64,
-                                                            base_storage=vec_iname + '_buff',
+                                                            base_storage=array_name + '_buff',
                                                             _base_storage_access_may_be_aliasing=True)
+
+        replacemap = dict()
+        replacemap[iname] = prim.Subscript(prim.Variable(array_name), (prim.Variable(iname),))
+
+        new_insns = [lp.Assignment(assignee=prim.Subscript(prim.Variable(array_name), (prim.Variable(new_iname),)),
+                                   expression=prim.Variable(new_iname),
+                                   id='init_{}_buffer'.format(array_name),
+                                   within_inames=frozenset({new_iname}), within_inames_is_final=True)]
+
+        for insn in knl.instructions:
+            if insn in insns_with_macro_points:
+                transformed_insn = insn.with_transformed_expressions(lambda expr: substitute(expr, replacemap))
+                new_insns.append(transformed_insn.copy(depends_on='init_{}_buffer'.format(array_name)))
+            else:
+                new_insns.append(insn)
+
+        knl = knl.copy(instructions=new_insns, domains=knl.domains + [new_dom],
+                       temporary_variables=dict(**knl.temporary_variables, **new_temporaries))
+
+    return knl
+
+
+def add_vcl_iname_array(knl, iname, vec_iname, vcl_size):
+    insns_with_macro_points = lp.find_instructions(knl, And((Tagged(iname), Iname(vec_iname))))
+
+    if insns_with_macro_points:
+        iname_array = iname + '_arr'
+        vector_name = iname + '_vec{}'.format(vcl_size)
+
+        new_temporaries = dict()
         new_temporaries[vector_name] = DuneTemporaryVariable(vector_name, managed=True,
                                                              shape=(get_form_option('number_of_blocks'),),
                                                              scope=lp.temp_var_scope.PRIVATE, dtype=np.float64,
-                                                             base_storage=vec_iname + '_buff',
+                                                             base_storage=iname_array + '_buff',
                                                              _base_storage_access_may_be_aliasing=True)
         silenced_warning = ["read_no_write({})".format(vector_name)]
 
-        new_insns = [lp.Assignment(assignee=prim.Subscript(prim.Variable(array_name), (prim.Variable(vec_iname),)),
-                                   expression=prim.Variable(vec_iname),
-                                   id='init_{}_buffer'.format(vec_iname),
-                                   within_inames=frozenset({vec_iname}), within_inames_is_final=True)]
-
         replacemap = dict()
-        replacemap[vec_iname] = prim.Subscript(prim.Variable(vector_name), (prim.Variable(vec_iname),))
+        replacemap[iname_array] = prim.Variable(vector_name)
 
+        new_insns = []
         for insn in knl.instructions:
             if insn in insns_with_macro_points:
                 transformed_insn = insn.with_transformed_expressions(lambda expr: substitute(expr, replacemap))
-                new_insns.append(transformed_insn.copy(depends_on='init_{}_buffer'.format(vec_iname)))
+                new_insns.append(transformed_insn.copy(depends_on='init_{}_buffer'.format(iname_array)))
             else:
                 new_insns.append(insn)
 
@@ -377,7 +406,8 @@ def add_iname_array(knl, vec_iname):
                        temporary_variables=dict(**knl.temporary_variables, **new_temporaries),
                        silenced_warnings=knl.silenced_warnings + silenced_warning)
 
-        knl = lp.duplicate_inames(knl, [vec_iname], within=Id('init_{}_buffer'.format(vec_iname)))
+        knl = lp.split_array_axis(knl, (vector_name,), 0, vcl_size)
+        knl = lp.tag_array_axes(knl, (vector_name,), ('c', 'vec'))
 
     return knl