Fixed vectorisation tests

lbmpy/lbstep.py

@@ -22,7 +22,7 @@ class LatticeBoltzmannStep:
                  velocity_data_name=None, density_data_name=None, density_data_index=None,
                  compute_velocity_in_every_step=False, compute_density_in_every_step=False,
                  velocity_input_array_name=None, time_step_order='stream_collide', flag_interface=None,
-                 alignment_if_vectorized=64, fixed_loop_sizes=True, fixed_relaxation_rates=True,
+                 alignment_if_vectorized=64, fixed_loop_sizes=True,
                  timeloop_creation_function=TimeLoop,
                  lbm_config=None, lbm_optimisation=None, config=None, **method_parameters):
 

lbmpy_tests/test_vectorization.py

 import numpy as np
 import pytest
+from dataclasses import replace
+
+import pystencils as ps
+
+from lbmpy import LBMOptimisation
 
 from pystencils.backends.simd_instruction_sets import get_supported_instruction_sets
 from lbmpy.scenarios import create_lid_driven_cavity
@@ -43,26 +48,29 @@ def test_lbm_vectorization(instruction_set, aligned_and_padding, nontemporal, do
     size2 = (666, 34)
     relaxation_rate = 1.8
 
+    if instruction_set == "neon" and fixed_loop_sizes is False:
+        pytest.skip("Has do be aligned for neon")
+
     print("Computing reference solutions")
     ldc1_ref = create_lid_driven_cavity(size1, relaxation_rate=relaxation_rate)
     ldc1_ref.run(time_steps)
     ldc2_ref = create_lid_driven_cavity(size2, relaxation_rate=relaxation_rate)
     ldc2_ref.run(time_steps)
 
-    optimization = {'double_precision': double_precision,
-                    'vectorization': vectorization_options,
-                    'cse_global': True,
-                    }
-    print("Vectorization test, double precision {}, vectorization {}, fixed loop sizes {}".format(
-        double_precision, vectorization_options, fixed_loop_sizes))
-    ldc1 = create_lid_driven_cavity(size1, relaxation_rate=relaxation_rate, optimization=optimization,
-                                    fixed_loop_sizes=fixed_loop_sizes)
+    data_type = "float64" if double_precision else "float32"
+    config = ps.CreateKernelConfig(data_type=data_type, cpu_vectorize_info=vectorization_options)
+    lbm_optimisation = LBMOptimisation(cse_global=True)
+
+    print(f"Vectorization test, double precision {double_precision}, vectorization {vectorization_options}, "
+          f"fixed loop sizes {fixed_loop_sizes}")
+    ldc1 = create_lid_driven_cavity(size1, relaxation_rate=relaxation_rate, fixed_loop_sizes=fixed_loop_sizes,
+                                    config=config, lbm_optimisation=lbm_optimisation)
     ldc1.run(time_steps)
     np.testing.assert_almost_equal(ldc1_ref.velocity[:, :], ldc1.velocity[:, :])
 
-    optimization['split'] = True
-    ldc2 = create_lid_driven_cavity(size2, relaxation_rate=relaxation_rate, optimization=optimization,
-                                    fixed_loop_sizes=fixed_loop_sizes)
+    lbm_optimisation = replace(lbm_optimisation, split=True)
+    ldc2 = create_lid_driven_cavity(size2, relaxation_rate=relaxation_rate, fixed_loop_sizes=fixed_loop_sizes,
+                                    config=config, lbm_optimisation=lbm_optimisation)
     ldc2.run(time_steps)
     np.testing.assert_almost_equal(ldc2_ref.velocity[:, :], ldc2.velocity[:, :])