diff --git a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp b/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp index d07bf51e50d01..fbf0a4f962d4e 100644 --- a/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp +++ b/flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp @@ -633,6 +633,30 @@ class NullOpConversion : public mlir::OpRewritePattern { } }; +class GetExtentOpConversion + : public mlir::OpRewritePattern { +public: + using mlir::OpRewritePattern::OpRewritePattern; + + mlir::LogicalResult + matchAndRewrite(hlfir::GetExtentOp getExtentOp, + mlir::PatternRewriter &rewriter) const override { + mlir::Value shape = getExtentOp.getShape(); + mlir::Operation *shapeOp = shape.getDefiningOp(); + // the hlfir.shape_of operation which led to the creation of this get_extent + // operation should now have been lowered to a fir.shape operation + if (auto s = mlir::dyn_cast_or_null(shapeOp)) { + fir::ShapeType shapeTy = shape.getType().cast(); + llvm::APInt dim = getExtentOp.getDim(); + uint64_t dimVal = dim.getLimitedValue(shapeTy.getRank()); + mlir::Value extent = s.getExtents()[dimVal]; + rewriter.replaceOp(getExtentOp, extent); + return mlir::success(); + } + return mlir::failure(); + } +}; + class ConvertHLFIRtoFIR : public hlfir::impl::ConvertHLFIRtoFIRBase { public: @@ -646,8 +670,8 @@ class ConvertHLFIRtoFIR mlir::RewritePatternSet patterns(context); patterns.insert(context); + GetExtentOpConversion, NoReassocOpConversion, + NullOpConversion, ParentComponentOpConversion>(context); mlir::ConversionTarget target(*context); target.addIllegalDialect(); target.markUnknownOpDynamicallyLegal( diff --git a/flang/test/HLFIR/extents-of-shape-of.f90 b/flang/test/HLFIR/extents-of-shape-of.f90 index ff1a657dc0ea5..85479ebb407f6 100644 --- a/flang/test/HLFIR/extents-of-shape-of.f90 +++ b/flang/test/HLFIR/extents-of-shape-of.f90 @@ -1,4 +1,5 @@ -! RUN: bbc -emit-fir -hlfir %s -o - | FileCheck %s +! 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(:, :) interface @@ -8,24 +9,40 @@ elemental subroutine elem_sub(x) end interface call elem_sub(matmul(a, b)) end subroutine -! CHECK-LABEL: func.func @_QPfoo -! CHECK: %[[A_ARG:.*]]: !fir.box> {fir.bindc_name = "a"} -! CHECK: %[[B_ARG:.*]]: !fir.box> {fir.bindc_name = "b"} -! CHECK-DAG: %[[A_VAR:.*]]:2 = hlfir.declare %[[A_ARG]] -! CHECK-DAG: %[[B_VAR:.*]]:2 = hlfir.declare %[[B_ARG]] -! CHECK-NEXT: %[[MUL:.*]] = hlfir.matmul %[[A_VAR]]#0 %[[B_VAR]]#0 -! CHECK-NEXT: %[[SHAPE:.*]] = hlfir.shape_of %[[MUL]] : (!hlfir.expr) -> !fir.shape<2> -! CHECK-NEXT: %[[EXT0:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 0 : index} : (!fir.shape<2>) -> index -! CHECK-NEXT: %[[EXT1:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 1 : index} : (!fir.shape<2>) -> index -! CHECK-NEXT: %[[C1:.*]] = arith.constant 1 : index -! CHECK-NEXT: fir.do_loop %[[ARG2:.*]] = %[[C1]] to %[[EXT1]] step %[[C1]] { -! CHECK-NEXT: fir.do_loop %[[ARG3:.*]] = %[[C1]] to %[[EXT0]] step %[[C1]] { -! CHECK-NEXT: %[[ELE:.*]] = hlfir.apply %[[MUL]], %[[ARG3]], %[[ARG2]] : (!hlfir.expr, index, index) -> f32 -! CHECK-NEXT: %[[ASSOC:.*]]:3 = hlfir.associate %[[ELE]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref, !fir.ref, i1) -! CHECK-NEXT: fir.call -! CHECK-NEXT: hlfir.end_associate -! CHECK-NEXT: } -! CHECK-NEXT: } -! CHECK-NEXT: hlfir.destroy %[[MUL]] -! CHECK-NEXT: return -! CHECK-NEXT: } +! CHECK-ALL-LABEL: func.func @_QPfoo +! CHECK-ALL: %[[A_ARG:.*]]: !fir.box> {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: %[[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: %[[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 +! CHECK-HLFIR-NEXT: } +! CHECK-HLFIR-NEXT: } +! CHECK-HLFIR-NEXT: hlfir.destroy %[[MUL]] + +! ... +! CHECK-FIR: fir.call @_FortranAMatmul +! CHECK-FIR-NEXT: %[[MUL:.*]] = fir.load %[[MUL_BOX:.*]] : !fir.ref>>> +! CHECK-FIR-NEXT: %[[C0:.*]] = arith.constant 0 : index +! CHECK-FIR-NEXT: %[[DIMS0:.*]]:3 = fir.box_dims %[[MUL]], %[[C0]] +! CHECK-FIR-NEXT: %[[C1:.*]] = arith.constant 1 : index +! CHECK-FIR-NEXT: %[[DIMS1:.*]]:3 = fir.box_dims %[[MUL]], %[[C1]] +! ... +! CHECK-FIR: %[[SHAPE:.*]] = fir.shape %[[DIMS0]]#1, %[[DIMS1]]#1 +! 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-ALL: return +! CHECK-ALL-NEXT: }