From: Martin Kronbichler Date: Sun, 15 Mar 2020 15:09:20 +0000 (+0100) Subject: Manually unroll a loop in matrix-free tensor product kernel X-Git-Tag: v9.2.0-rc1~420^2 X-Git-Url: https://gitweb.dealii.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=64350c91b7d837d285a57e1f35408a2a9e39a7d1;p=dealii.git Manually unroll a loop in matrix-free tensor product kernel --- diff --git a/include/deal.II/matrix_free/tensor_product_kernels.h b/include/deal.II/matrix_free/tensor_product_kernels.h index 761f08188f..d639ccee40 100644 --- a/include/deal.II/matrix_free/tensor_product_kernels.h +++ b/include/deal.II/matrix_free/tensor_product_kernels.h @@ -675,47 +675,119 @@ namespace internal Utilities::fixed_power(n_rows); Assert(n_rows <= 128, ExcNotImplemented()); - for (int i2 = 0; i2 < n_blocks2; ++i2) + // specialization for n_rows = 2 that manually unrolls the innermost loop + // to make the operation perform better (not completely as good as the + // templated one, but much better than the generic version down below, + // because the loop over col can be more effectively unrolled by the + // compiler) + if (contract_over_rows && n_rows == 2) { - for (int i1 = 0; i1 < n_blocks1; ++i1) + const Number2 *shape_data_1 = shape_data + n_columns; + for (int i2 = 0; i2 < n_blocks2; ++i2) { - Number x[129]; - for (int i = 0; i < mm; ++i) - x[i] = in[stride * i]; - for (int col = 0; col < nn; ++col) + for (int i1 = 0; i1 < n_blocks1; ++i1) { - Number2 val0; - if (contract_over_rows == true) - val0 = shape_data[col]; - else - val0 = shape_data[col * n_columns]; - Number res0 = val0 * x[0]; - for (int i = 1; i < mm; ++i) + const Number x0 = in[0], x1 = in[stride]; + for (int col = 0; col < nn; ++col) { - if (contract_over_rows == true) - val0 = shape_data[i * n_columns + col]; + const Number result = + shape_data[col] * x0 + shape_data_1[col] * x1; + if (add == false) + out[stride * col] = result; else - val0 = shape_data[col * n_columns + i]; - res0 += val0 * x[i]; + out[stride * col] += result; } - if (add == false) - out[stride * col] = res0; - else - out[stride * col] += res0; - } + if (one_line == false) + { + ++in; + ++out; + } + } if (one_line == false) { - ++in; - ++out; + in += stride * (mm - 1); + out += stride * (nn - 1); } } - if (one_line == false) + } + // specialization for n = 3 + else if (contract_over_rows && n_rows == 3) + { + const Number2 *shape_data_1 = shape_data + n_columns; + const Number2 *shape_data_2 = shape_data + 2 * n_columns; + for (int i2 = 0; i2 < n_blocks2; ++i2) { - in += stride * (mm - 1); - out += stride * (nn - 1); + for (int i1 = 0; i1 < n_blocks1; ++i1) + { + const Number x0 = in[0], x1 = in[stride], x2 = in[2 * stride]; + for (int col = 0; col < nn; ++col) + { + const Number result = shape_data[col] * x0 + + shape_data_1[col] * x1 + + shape_data_2[col] * x2; + if (add == false) + out[stride * col] = result; + else + out[stride * col] += result; + } + + if (one_line == false) + { + ++in; + ++out; + } + } + if (one_line == false) + { + in += stride * (mm - 1); + out += stride * (nn - 1); + } } } + // general loop for all other cases + else + for (int i2 = 0; i2 < n_blocks2; ++i2) + { + for (int i1 = 0; i1 < n_blocks1; ++i1) + { + Number x[129]; + for (int i = 0; i < mm; ++i) + x[i] = in[stride * i]; + for (int col = 0; col < nn; ++col) + { + Number2 val0; + if (contract_over_rows == true) + val0 = shape_data[col]; + else + val0 = shape_data[col * n_columns]; + Number res0 = val0 * x[0]; + for (int i = 1; i < mm; ++i) + { + if (contract_over_rows == true) + val0 = shape_data[i * n_columns + col]; + else + val0 = shape_data[col * n_columns + i]; + res0 += val0 * x[i]; + } + if (add == false) + out[stride * col] = res0; + else + out[stride * col] += res0; + } + + if (one_line == false) + { + ++in; + ++out; + } + } + if (one_line == false) + { + in += stride * (mm - 1); + out += stride * (nn - 1); + } + } }