Skip to content

Commit

Permalink
[flang][OpenMP] Handle usage of array elements in loop-control expres…
Browse files Browse the repository at this point in the history
…sions

Extends the fix-up logic for `trip_count` calculation in `target`
regions. Previously, if an array element was used to compute any of the
loop bounds, the trip-count calculation ops would extract arra elements
from the mapped declaration of the array inside the target region. This
commit hanles that situation.
  • Loading branch information
ergawy committed Aug 3, 2024
1 parent 6eb4e22 commit b50e1d1
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 2 deletions.
17 changes: 15 additions & 2 deletions flang/lib/Lower/OpenMP/OpenMP.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,8 @@ class HostClausesInsertionGuard {
assert(targetParentRegion != nullptr &&
"Expected omp.target op to be nested in a parent region.");

llvm::DenseSet<mlir::Operation *> visitedOps;

// Walk the parent region in pre-order to make sure we visit `targetOp`
// before its nested ops.
targetParentRegion->walk<mlir::WalkOrder::PreOrder>(
Expand All @@ -149,6 +151,10 @@ class HostClausesInsertionGuard {
if (operandDefiningOp == nullptr)
continue;

if (visitedOps.contains(operandDefiningOp))
continue;

visitedOps.insert(operandDefiningOp);
auto parentTargetOp =
operandDefiningOp->getParentOfType<mlir::omp::TargetOp>();

Expand Down Expand Up @@ -179,7 +185,9 @@ class HostClausesInsertionGuard {

mlir::IRMapping mapper;
builder.setInsertionPoint(escapingOperand->getOwner());

mlir::Operation *lastSliceOp = nullptr;
llvm::SetVector<mlir::Operation *> opsToClone;

for (auto *op : backwardSlice) {
// DeclareOps need special handling by searching for the corresponding ops
Expand All @@ -188,12 +196,17 @@ class HostClausesInsertionGuard {
//
// TODO this might need a more elaborate handling in the future but for
// now this seems sufficient for our purposes.
if (llvm::isa<hlfir::DeclareOp>(op))
if (llvm::isa<hlfir::DeclareOp>(op)) {
opsToClone.clear();
break;
}

lastSliceOp = builder.clone(*op, mapper);
opsToClone.insert(op);
}

for (mlir::Operation *op : opsToClone)
lastSliceOp = builder.clone(*op, mapper);

builder.restoreInsertionPoint(ip);
return lastSliceOp;
}
Expand Down
31 changes: 31 additions & 0 deletions flang/test/Lower/OpenMP/target-do-loop-control-exprs.f90
Original file line number Diff line number Diff line change
Expand Up @@ -62,3 +62,34 @@ subroutine foo_with_dummy_arg(nodes)
! CHECK: omp.target

! CHECK: }


subroutine bounds_expr_in_loop_control(array)
real, intent(out) :: array(:,:)
integer :: bounds(2), i, j
bounds = shape(array)

!$omp target teams distribute parallel do simd collapse(2)
do j = 1,bounds(2)
do i = 1,bounds(1)
array(i,j) = 0.
enddo
enddo
end subroutine bounds_expr_in_loop_control


! CHECK: func.func @_QPbounds_expr_in_loop_control(%[[FUNC_ARG:.*]]: {{.*}}) {

! CHECK: %[[BOUNDS_DECL:.*]]:2 = hlfir.declare %{{.*}}(%{{.*}}) {uniq_name = "{{.*}}Ebounds"} : (!fir.ref<!fir.array<2xi32>>, !fir.shape<1>) -> ({{.*}})

! Verify that the host declaration of `bounds` (i.e. not the target/mapped one)
! is used for the trip count calculation. Trip count is calculation ops are emitted
! directly before the `omp.target` op and after all `omp.map.info` op; hence the
! `CHECK-NOT: ...` line.

! CHECK: hlfir.designate %[[BOUNDS_DECL:.*]]#0 (%c2{{.*}})
! CHECK: hlfir.designate %[[BOUNDS_DECL:.*]]#0 (%c1{{.*}})
! CHECK-NOT: omp.map.info
! CHECK: omp.target

! CHECK: }

0 comments on commit b50e1d1

Please sign in to comment.