diff --git a/flang/lib/Optimizer/Builder/HLFIRTools.cpp b/flang/lib/Optimizer/Builder/HLFIRTools.cpp index 9b9174eea80ae..f4481734e9a50 100644 --- a/flang/lib/Optimizer/Builder/HLFIRTools.cpp +++ b/flang/lib/Optimizer/Builder/HLFIRTools.cpp @@ -36,11 +36,22 @@ getExplicitExtentsFromShape(mlir::Value shape, fir::FirOpBuilder &builder) { } else if (mlir::dyn_cast_or_null(shapeOp)) { return {}; } else if (auto s = mlir::dyn_cast_or_null(shapeOp)) { + hlfir::ExprType expr = s.getExpr().getType().cast(); + llvm::ArrayRef exprShape = expr.getShape(); + mlir::Type indexTy = builder.getIndexType(); fir::ShapeType shapeTy = shape.getType().cast(); result.reserve(shapeTy.getRank()); for (unsigned i = 0; i < shapeTy.getRank(); ++i) { - auto op = builder.create(shape.getLoc(), shape, i); - result.emplace_back(op.getResult()); + int64_t extent = exprShape[i]; + mlir::Value extentVal; + if (extent == expr.getUnknownExtent()) { + auto op = builder.create(shape.getLoc(), shape, i); + extentVal = op.getResult(); + } else { + extentVal = + builder.createIntegerConstant(shape.getLoc(), indexTy, extent); + } + result.emplace_back(extentVal); } } else { TODO(shape.getLoc(), "read fir.shape to get extents"); diff --git a/flang/test/HLFIR/extents-of-shape-of.f90 b/flang/test/HLFIR/extents-of-shape-of.f90 index 85479ebb407f6..8fc7133985ffc 100644 --- a/flang/test/HLFIR/extents-of-shape-of.f90 +++ b/flang/test/HLFIR/extents-of-shape-of.f90 @@ -1,7 +1,7 @@ ! RUN: bbc -emit-fir -hlfir %s -o - | FileCheck --check-prefix CHECK-ALL --check-prefix CHECK-HLFIR %s ! RUN: bbc -emit-fir -hlfir %s -o - | fir-opt --lower-hlfir-intrinsics | fir-opt --bufferize-hlfir | fir-opt --convert-hlfir-to-fir | FileCheck --check-prefix CHECK-ALL --check-prefix CHECK-FIR %s subroutine foo(a, b) - real :: a(:, :), b(:, :) + real :: a(2, 2), b(:, :) interface elemental subroutine elem_sub(x) real, intent(in) :: x @@ -10,19 +10,19 @@ elemental subroutine elem_sub(x) call elem_sub(matmul(a, b)) end subroutine ! CHECK-ALL-LABEL: func.func @_QPfoo -! CHECK-ALL: %[[A_ARG:.*]]: !fir.box> {fir.bindc_name = "a"} +! CHECK-ALL: %[[A_ARG:.*]]: !fir.ref> {fir.bindc_name = "a"} ! CHECK-ALL: %[[B_ARG:.*]]: !fir.box> {fir.bindc_name = "b"} ! CHECK-HLFIR-DAG: %[[A_VAR:.*]]:2 = hlfir.declare %[[A_ARG]] ! CHECK-HLFIR-DAG: %[[B_VAR:.*]]:2 = hlfir.declare %[[B_ARG]] ! CHECK-HLFIR-NEXT: %[[MUL:.*]] = hlfir.matmul %[[A_VAR]]#0 %[[B_VAR]]#0 -! CHECK-HLFIR-NEXT: %[[SHAPE:.*]] = hlfir.shape_of %[[MUL]] : (!hlfir.expr) -> !fir.shape<2> -! CHECK-HLFIR-NEXT: %[[EXT0:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 0 : index} : (!fir.shape<2>) -> index +! CHECK-HLFIR-NEXT: %[[SHAPE:.*]] = hlfir.shape_of %[[MUL]] : (!hlfir.expr<2x?xf32>) -> !fir.shape<2> +! CHECK-HLFIR-NEXT: %[[EXT0:.*]] = arith.constant 2 : index ! CHECK-HLFIR-NEXT: %[[EXT1:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 1 : index} : (!fir.shape<2>) -> index ! CHECK-HLFIR-NEXT: %[[C1:.*]] = arith.constant 1 : index ! CHECK-HLFIR-NEXT: fir.do_loop %[[ARG2:.*]] = %[[C1]] to %[[EXT1]] step %[[C1]] { ! CHECK-HLFIR-NEXT: fir.do_loop %[[ARG3:.*]] = %[[C1]] to %[[EXT0]] step %[[C1]] { -! CHECK-HLFIR-NEXT: %[[ELE:.*]] = hlfir.apply %[[MUL]], %[[ARG3]], %[[ARG2]] : (!hlfir.expr, index, index) -> f32 +! CHECK-HLFIR-NEXT: %[[ELE:.*]] = hlfir.apply %[[MUL]], %[[ARG3]], %[[ARG2]] : (!hlfir.expr<2x?xf32>, index, index) -> f32 ! CHECK-HLFIR-NEXT: %[[ASSOC:.*]]:3 = hlfir.associate %[[ELE]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref, !fir.ref, i1) ! CHECK-HLFIR-NEXT: fir.call ! CHECK-HLFIR-NEXT: hlfir.end_associate @@ -39,9 +39,10 @@ elemental subroutine elem_sub(x) ! CHECK-FIR-NEXT: %[[DIMS1:.*]]:3 = fir.box_dims %[[MUL]], %[[C1]] ! ... ! CHECK-FIR: %[[SHAPE:.*]] = fir.shape %[[DIMS0]]#1, %[[DIMS1]]#1 +! CHECK-FIR-NEXT: %[[C2:.*]] = arith.constant 2 : index ! CHECK-FIR-NEXT: %[[C1_1:.*]] = arith.constant 1 : index ! CHECK-FIR-NEXT: fir.do_loop %[[ARG2:.*]] = %[[C1_1]] to %[[DIMS1]]#1 step %[[C1_1]] { -! CHECK-FIR-NEXT: fir.do_loop %[[ARG3:.*]] = %[[C1_1]] to %[[DIMS0]]#1 step %[[C1_1]] { +! CHECK-FIR-NEXT: fir.do_loop %[[ARG3:.*]] = %[[C1_1]] to %[[C2]] step %[[C1_1]] { ! ... ! CHECK-ALL: return