[flang][hlfir] lower hlfir.get_extent to FIR

[tblah]

Summary

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

[tblah]

Fix header diff plus poke CI into re-trying patch application.

[jeanPerier]

Looks good. Beware of the patch application issues.

flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp

646

tothe -> to the

[tblah]

Fix typo

flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp

public:
   mlir::LogicalResult
   matchAndRewrite(hlfir::NullOp nullop,
                   mlir::PatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<fir::ZeroOp>(nullop, nullop.getType());
     return mlir::success();
   }
 };
 
+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:
   void runOnOperation() override {
     // TODO: like "bufferize-hlfir" pass, runtime signature may be added
     // by this pass. This requires the pass to run on the ModuleOp. It would
     // probably be more optimal to have it run on FuncOp and find a way to
     // generate the signatures in a thread safe way.
     auto module = this->getOperation();
     auto *context = &getContext();
     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(
         [](mlir::Operation *) { return true; });
     if (mlir::failed(mlir::applyPartialConversion(module, target,
                                                   std::move(patterns)))) {
       mlir::emitError(mlir::UnknownLoc::get(context),
                       "failure in HLFIR to FIR conversion pass");

flang/test/HLFIR/extents-of-shape-of.f90

-! 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
     elemental subroutine elem_sub(x)
       real, intent(in) :: x
     end subroutine
   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:    }