Skip to content

Commit

Permalink
[flang][hlfir] lower hlfir.get_extent to FIR
Browse files Browse the repository at this point in the history 
By the ConvertToFIR pass, the hlfir.get_shape operation will have been
lowered into a fir.shape operation (during the HFLIR bufferization pass)
and so, lowering get_extent is as simple as fetching the extent from the
shape operation.

Depends on: D146833

Differential Revision: https://reviews.llvm.org/D148222
  • Loading branch information
@tblah
tblah committed Apr 17, 2023
1 parent 683a6e1 commit c2c49f4
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 24 deletions.
28 changes: 26 additions & 2 deletions flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -633,6 +633,30 @@ class NullOpConversion : public mlir::OpRewritePattern<hlfir::NullOp> {
}
};

class GetExtentOpConversion
: public mlir::OpRewritePattern<hlfir::GetExtentOp> {
public:
using mlir::OpRewritePattern<hlfir::GetExtentOp>::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<fir::ShapeOp>(shapeOp)) {
fir::ShapeType shapeTy = shape.getType().cast<fir::ShapeType>();
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<ConvertHLFIRtoFIR> {
public:
Expand All @@ -646,8 +670,8 @@ class ConvertHLFIRtoFIR
mlir::RewritePatternSet patterns(context);
patterns.insert<AssignOpConversion, CopyInOpConversion, CopyOutOpConversion,
DeclareOpConversion, DesignateOpConversion,
NoReassocOpConversion, NullOpConversion,
ParentComponentOpConversion>(context);
GetExtentOpConversion, NoReassocOpConversion,
NullOpConversion, ParentComponentOpConversion>(context);
mlir::ConversionTarget target(*context);
target.addIllegalDialect<hlfir::hlfirDialect>();
target.markUnknownOpDynamicallyLegal(
Expand Down
61 changes: 39 additions & 22 deletions flang/test/HLFIR/extents-of-shape-of.f90
View file
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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.array<?x?xf32>> {fir.bindc_name = "a"}
! CHECK: %[[B_ARG:.*]]: !fir.box<!fir.array<?x?xf32>> {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<?x?xf32>) -> !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<?x?xf32>, index, index) -> f32
! CHECK-NEXT: %[[ASSOC:.*]]:3 = hlfir.associate %[[ELE]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, 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.array<?x?xf32>> {fir.bindc_name = "a"}
! CHECK-ALL: %[[B_ARG:.*]]: !fir.box<!fir.array<?x?xf32>> {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<?x?xf32>) -> !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<?x?xf32>, index, index) -> f32
! CHECK-HLFIR-NEXT: %[[ASSOC:.*]]:3 = hlfir.associate %[[ELE]] {uniq_name = "adapt.valuebyref"} : (f32) -> (!fir.ref<f32>, !fir.ref<f32>, 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<!fir.box<!fir.heap<!fir.array<?x?xf32>>>>
! 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: }

0 comments on commit c2c49f4

Please sign in to comment.