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


@pytest.mark.skipif(not get_supported_instruction_sets(), reason='cannot detect CPU instruction set')
def test_lbm_vectorization_short():
    print("Computing reference solutions")
    size1 = (64, 32)
    relaxation_rate = 1.8

    ldc1_ref = create_lid_driven_cavity(size1, relaxation_rate=relaxation_rate)
    ldc1_ref.run(10)

    ldc1 = create_lid_driven_cavity(size1, relaxation_rate=relaxation_rate,
                                    optimization={
                                        'vectorization': {'instruction_set': get_supported_instruction_sets()[-1],
                                                          'assume_aligned': True,
                                                          'nontemporal': True,
                                                          'assume_inner_stride_one': True,
                                                          'assume_sufficient_line_padding': False,
                                                          }},
                                    fixed_loop_sizes=False)
    ldc1.run(10)


@pytest.mark.parametrize('instruction_set', get_supported_instruction_sets())
@pytest.mark.parametrize('aligned_and_padding', [[False, False], [True, False], [True, True]])
@pytest.mark.parametrize('nontemporal', [False, True])
@pytest.mark.parametrize('double_precision', [False, True])
@pytest.mark.parametrize('fixed_loop_sizes', [False, True])
@pytest.mark.longrun
def test_lbm_vectorization(instruction_set, aligned_and_padding, nontemporal, double_precision, fixed_loop_sizes):
    vectorization_options = {'instruction_set': instruction_set,
                             'assume_aligned': aligned_and_padding[0],
                             'nontemporal': nontemporal,
                             'assume_inner_stride_one': True,
                             'assume_sufficient_line_padding': aligned_and_padding[1]}
    time_steps = 100
    size1 = (64, 32)
    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)

    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[:, :])

    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[:, :])


if __name__ == '__main__':
    test_lbm_vectorization()