Cold-Air-Networks/mirror-llvm-project
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 7

[flang][hlfir] Prefer a constant-shape operand for elemental array results (#209791)

Browse source 
Elemental operands are conformable, so their extents are identical and
the result shape may be taken from either operand. This patch prefers an
operand whose shape is a compile-time constant, producing a statically
shaped result. The left operand is checked first to avoid churn in lit
test expectations.
This commit is contained in:
Caroline Newcombe 2026-07-27 10:48:35 -04:00 committed by GitHub
parent f613ead38e
commit 491f16472b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 71 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
flang/lib/Lower/ConvertExprToHLFIR.cpp
View file

@ -1774,10 +1774,22 @@ private:
Fortran::lower::getFIRType(builder.getContext(), R::category, R::kind,
/*params=*/{});
// TODO: "merge" shape, get cst shape from front-end if possible.
// Prefer a compile-time constant shape to get a statically shaped result.
auto hasConstantShape = [](hlfir::Entity entity) -> bool {
if (!entity.isArray())
return false;
auto seqTy =
mlir::dyn_cast<fir::SequenceType>(entity.getElementOrSequenceType());
return seqTy && !fir::sequenceWithNonConstantShape(seqTy);
};
mlir::Value shape;
if (left.isArray()) {
if (hasConstantShape(left))
shape = hlfir::genShape(loc, builder, left);
} else {
else if (hasConstantShape(right))
shape = hlfir::genShape(loc, builder, right);
else if (left.isArray())
shape = hlfir::genShape(loc, builder, left);
else {
assert(right.isArray() && "must have at least one array operand");
shape = hlfir::genShape(loc, builder, right);
}

53
flang/test/Lower/HLFIR/elemental-array-constant-shape.f90 Normal file
View file

@ -0,0 +1,53 @@
! Test that when a binary elemental operation has one operand with a
! compile-time constant shape and another with only a dynamic (descriptor-based)
! shape, lowering infers the constant shape for the operation result.
! Conforming elemental operands have identical extents, so preferring the
! constant shape is safe and yields a statically shaped, more precise result.
! RUN: %flang_fc1 -emit-hlfir -o - %s | FileCheck %s
! Constant-shape operand on the right ("a" is assumed-shape/dynamic, "x" is
! explicit-shape/constant): the result takes x's constant shape.
subroutine const_right_operand(a)
real :: a(:)
real :: x(3)
a = a + x
end subroutine
! CHECK-LABEL: func.func @_QPconst_right_operand(
! CHECK: %[[ELEM:.*]] = hlfir.elemental %{{.*}} unordered : (!fir.shape<1>) -> !hlfir.expr<3xf32>
! CHECK: hlfir.assign %[[ELEM]] to %{{.*}} : !hlfir.expr<3xf32>, !fir.box<!fir.array<?xf32>>
! Constant-shape operand on the left: the result still takes the constant shape.
subroutine const_left_operand(a)
real :: a(:)
real :: x(3)
a = x + a
end subroutine
! CHECK-LABEL: func.func @_QPconst_left_operand(
! CHECK: hlfir.elemental %{{.*}} unordered : (!fir.shape<1>) -> !hlfir.expr<3xf32>
! No constant-shape operand: the result stays dynamically shaped.
subroutine both_dynamic(a, b)
real :: a(:), b(:)
a = a + b
end subroutine
! CHECK-LABEL: func.func @_QPboth_dynamic(
! CHECK: hlfir.elemental %{{.*}} unordered : (!fir.shape<1>) -> !hlfir.expr<?xf32>
! Scalar operand with a constant-shape array: the scalar has no shape, so the
! result takes the array's constant shape.
subroutine scalar_times_const(a, s)
real :: a(:), s
real :: x(3)
a = s * x
end subroutine
! CHECK-LABEL: func.func @_QPscalar_times_const(
! CHECK: hlfir.elemental %{{.*}} unordered : (!fir.shape<1>) -> !hlfir.expr<3xf32>
! Scalar operand with a dynamic array: the result stays dynamically shaped.
subroutine scalar_times_dyn(a, b, s)
real :: a(:), b(:), s
a = s * b
end subroutine
! CHECK-LABEL: func.func @_QPscalar_times_dyn(
! CHECK: hlfir.elemental %{{.*}} unordered : (!fir.shape<1>) -> !hlfir.expr<?xf32>

11
flang/test/Lower/HLFIR/matmul.f90
View file

@ -43,18 +43,15 @@ endsubroutine
! CHECK: fir.call @_QPfill
! CHECK: fir.call @_QPfill
! CHECK-NEXT: %[[B_BOX:.*]] = fir.load %[[B_BOX_DECL]]#0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xi32>>>>
! CHECK-NEXT: %[[C0:.*]] = arith.constant 0 : index
! CHECK-NEXT: %[[B_DIMS_0:.*]]:3 = fir.box_dims %[[B_BOX]], %[[C0]]
! CHECK-NEXT: %[[C1:.*]] = arith.constant 1 : index
! CHECK-NEXT: %[[B_DIMS_1:.*]]:3 = fir.box_dims %[[B_BOX]], %[[C1]]
! CHECK-NEXT: %[[B_SHAPE:.*]] = fir.shape %[[B_DIMS_0]]#1, %[[B_DIMS_1]]#1
! CHECK-NEXT: %[[ELEMENTAL:.*]] = hlfir.elemental %[[B_SHAPE]] unordered : (!fir.shape<2>) -> !hlfir.expr<?x?xi32> {
! The elemental shape is taken from the constant-shape operand "x" (see above),
! so "b - x" is inferred as a 4x4 array rather than a dynamically shaped one.
! CHECK-NEXT: %[[ELEMENTAL:.*]] = hlfir.elemental %{{.*}} unordered : (!fir.shape<2>) -> !hlfir.expr<4x4xi32> {
! CHECK: }
! CHECK-NEXT: %[[A_BOX:.*]] = fir.load %{{.*}} : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xi32>>>>
! The shapes in these types are what is being tested:
! CHECK-NEXT: %[[MATMUL:.*]] = hlfir.matmul %[[A_BOX]] %[[ELEMENTAL]] {{.*}} : (!fir.box<!fir.heap<!fir.array<?x?xi32>>>, !hlfir.expr<?x?xi32>) -> !hlfir.expr<?x4xi32>
! CHECK-NEXT: %[[MATMUL:.*]] = hlfir.matmul %[[A_BOX]] %[[ELEMENTAL]] {{.*}} : (!fir.box<!fir.heap<!fir.array<?x?xi32>>>, !hlfir.expr<4x4xi32>) -> !hlfir.expr<?x4xi32>
subroutine matmul3(lhs, rhs, res)
integer, allocatable :: lhs(:,:), rhs(:,:), res(:,:)