From 0ac8cb1b3df724f549a62f6b277745af3d50b23a Mon Sep 17 00:00:00 2001 From: Slava Zakharin Date: Mon, 31 Mar 2025 11:41:43 -0700 Subject: [PATCH] [flang] Recognize fir.pack_array in LoopVersioning. (#133191) This change enables LoopVersioning when `fir.pack_array` is met in the def-use chain. It fixes a couple of huge performance regressions caused by enabling `-frepack-arrays`. --- flang/docs/ArrayRepacking.md | 2 + .../Optimizer/Transforms/LoopVersioning.cpp | 31 ++++++++++---- .../loop-versioning-with-repack-arrays.fir | 40 +++++++++++++++++++ 3 files changed, 65 insertions(+), 8 deletions(-) create mode 100644 flang/test/Transforms/loop-versioning-with-repack-arrays.fir diff --git a/flang/docs/ArrayRepacking.md b/flang/docs/ArrayRepacking.md index f22e26ce4973..87cfc5d1bb4b 100755 --- a/flang/docs/ArrayRepacking.md +++ b/flang/docs/ArrayRepacking.md @@ -432,6 +432,8 @@ There is an existing optimization pass (controlled via `-f[no-]version-loops-for The array repacking is targeting better data cache utilization, and is not intended to enable more unit-strided vectorization for the assumed-shape arrays. At the same time, combining array repacking with the loop versioning may provide better performance for programs where the actual array arguments are non-contiguous, but then their repacked copies can be accessed using unit strides. +It is suggested that the LoopVersioning pass is run before the lowering of `fir.pack_array` and `fir.unpack_array` operations, and recognizes `fir.pack_array` on the path from `fir.declare` to the function entry block argument. The pass generates the dynamic contiguity checks, and multiversions the loops. In case the repacking actually happens, the most optimal versions of the loops are executed. + In cases where `fir.pack_array` is statically known to produce a copy that is contiguous in the innermost dimension, the loop versioning pass can skip the generation of the dynamic checks. ### Driver: user options diff --git a/flang/lib/Optimizer/Transforms/LoopVersioning.cpp b/flang/lib/Optimizer/Transforms/LoopVersioning.cpp index 1f3495569c9d..42e149bb3dba 100644 --- a/flang/lib/Optimizer/Transforms/LoopVersioning.cpp +++ b/flang/lib/Optimizer/Transforms/LoopVersioning.cpp @@ -184,13 +184,28 @@ getRankAndElementSize(const fir::KindMapping &kindMap, return {0, 0}; } -/// if a value comes from a fir.declare, follow it to the original source, -/// otherwise return the value -static mlir::Value unwrapFirDeclare(mlir::Value val) { - // fir.declare is for source code variables. We don't have declares of - // declares - if (fir::DeclareOp declare = val.getDefiningOp()) - return declare.getMemref(); +/// If a value comes from a fir.declare of fir.pack_array, +/// follow it to the original source, otherwise return the value. +static mlir::Value unwrapPassThroughOps(mlir::Value val) { + // Instead of unwrapping fir.declare, we may try to start + // the analysis in this pass from fir.declare's instead + // of the function entry block arguments. This way the loop + // versioning would work even after FIR inlining. + while (true) { + if (fir::DeclareOp declare = val.getDefiningOp()) { + val = declare.getMemref(); + continue; + } + // fir.pack_array might be met before fir.declare - this is how + // it is orifinally generated. + // It might also be met after fir.declare - after the optimization + // passes that sink fir.pack_array closer to the uses. + if (auto packArray = val.getDefiningOp()) { + val = packArray.getArray(); + continue; + } + break; + } return val; } @@ -242,7 +257,7 @@ static mlir::Value unwrapReboxOp(mlir::Value val) { /// normalize a value (removing fir.declare and fir.rebox) so that we can /// more conveniently spot values which came from function arguments static mlir::Value normaliseVal(mlir::Value val) { - return unwrapFirDeclare(unwrapReboxOp(val)); + return unwrapPassThroughOps(unwrapReboxOp(val)); } /// some FIR operations accept a fir.shape, a fir.shift or a fir.shapeshift. diff --git a/flang/test/Transforms/loop-versioning-with-repack-arrays.fir b/flang/test/Transforms/loop-versioning-with-repack-arrays.fir new file mode 100644 index 000000000000..7a2cac4b56a2 --- /dev/null +++ b/flang/test/Transforms/loop-versioning-with-repack-arrays.fir @@ -0,0 +1,40 @@ +// RUN: fir-opt --loop-versioning %s | FileCheck %s + +// Check that LoopVersioning kicks in when there is fir.pack_array +// in between fir.declare and the block argument. + +module attributes {dlti.dl_spec = #dlti.dl_spec<>} { + func.func @_QPtest(%arg0: !fir.box> {fir.bindc_name = "x"}) { + %c1 = arith.constant 1 : index + %c0 = arith.constant 0 : index + %cst = arith.constant 1.000000e+00 : f32 + %0 = fir.dummy_scope : !fir.dscope + %1 = fir.pack_array %arg0 stack whole : (!fir.box>) -> !fir.box> + %2 = fir.declare %1 dummy_scope %0 {uniq_name = "_QFtestEx"} : (!fir.box>, !fir.dscope) -> !fir.box> + %3 = fir.rebox %2 : (!fir.box>) -> !fir.box> + %4:3 = fir.box_dims %3, %c0 : (!fir.box>, index) -> (index, index, index) + fir.do_loop %arg1 = %c1 to %4#1 step %c1 unordered { + %5 = fir.array_coor %2 %arg1 : (!fir.box>, index) -> !fir.ref + fir.store %cst to %5 : !fir.ref + } + fir.unpack_array %1 to %arg0 stack : !fir.box> + return + } +} +// CHECK-LABEL: func.func @_QPtest( +// CHECK-SAME: %[[VAL_0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !fir.box> {fir.bindc_name = "x"}) { +// CHECK: %[[VAL_5:.*]] = fir.pack_array %[[VAL_0]] stack whole : (!fir.box>) -> !fir.box> +// CHECK: %[[VAL_6:.*]] = fir.declare %[[VAL_5]] dummy_scope %{{.*}} {uniq_name = "_QFtestEx"} : (!fir.box>, !fir.dscope) -> !fir.box> +// CHECK: %[[VAL_10:.*]]:3 = fir.box_dims %[[VAL_6]], %{{.*}} : (!fir.box>, index) -> (index, index, index) +// CHECK: %[[VAL_11:.*]] = arith.constant 4 : index +// CHECK: %[[VAL_12:.*]] = arith.cmpi eq, %[[VAL_10]]#2, %[[VAL_11]] : index +// CHECK: fir.if %[[VAL_12]] { +// CHECK: fir.do_loop {{.*}} { +// CHECK: fir.coordinate_of {{.*}} : (!fir.ref>, index) -> !fir.ref +// CHECK: } +// CHECK: } else { +// CHECK: fir.do_loop {{.*}} { +// CHECK: fir.array_coor {{.*}} : (!fir.box>, index) -> !fir.ref +// CHECK: } +// CHECK: } +// CHECK: fir.unpack_array %[[VAL_5]] to %[[VAL_0]] stack : !fir.box>