[flang][hlfir] add hlfir.get_extent

[tblah]

Summary

This operation fetches an extent value from a fir.shape. The operation
could just as easily live in the fir namespace, but is only needed for
hlfir lowering so I put it here.

This operation is required to allow one to defer getting the extents of a shape
generated by hlfir.get_shape until after that shape has been resolved
(after bufferization of the hlfir.expr).

This operation will be lowered to FIR as an arith.constant created using
the definition of the fir.shape argument.

Depends On: D146830

[jeanPerier]

Thanks a lot for filling the HLFIR gaps here! (the dependency was not set properly for the bot, I think you need to write "Depends on" instead of "Depends On:" in the message to get it set automatically. I added it manually and relaunched the build).

[jeanPerier]

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

796

While I am here, I started adding side effect in my sandbox so that we can leverage some of MLIR analysis. You probably want to add Pure here (it has no memory side effect and is speculatable (can be executed as soon as its operand is ready)).

[tblah]

Add the Pure attribute to the operation definition

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

def hlfir_ShapeOfOp : hlfir_Op<"shape_of", [Pure]> {
 
   let assemblyFormat = [{
     $expr attr-dict `:` functional-type(operands, results)
   }];
 
   let builders = [OpBuilder<(ins "mlir::Value":$expr)>];
 }
 
+def hlfir_GetExtentOp : hlfir_Op<"get_extent", [Pure]> {
+  let summary = "Get an extent value from a fir.shape";
+  let description = [{
+    Gets an extent value from a fir.shape. The dimension argument uses C style
+    indexing and so should be between 0 and 1 less than the rank of the shape
+  }];
+
+  let arguments = (ins fir_ShapeType:$shape,
+                       IndexAttr:$dim);
+
+  let results = (outs Index);
+
+  let hasVerifier = 1;
+
+  let assemblyFormat = [{
+    $shape attr-dict `:` functional-type(operands, results)
+  }];
+
+  let builders = [OpBuilder<(ins "mlir::Value":$shape, "unsigned":$dim)>];
+}
+
 #endif // FORTRAN_DIALECT_HLFIR_OPS

flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

 //===----------------------------------------------------------------------===//
 
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
 #include "flang/Optimizer/HLFIR/HLFIRDialect.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
+#include "llvm/ADT/APInt.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include <iterator>
 #include <optional>
 #include <tuple>
 
 //===----------------------------------------------------------------------===//
 // DeclareOp
 //===----------------------------------------------------------------------===//
if (!shape)
     // shape information is not available at compile time
     return mlir::LogicalResult::failure();
 
   rewriter.replaceAllUsesWith(shapeOf.getResult(), shape);
   rewriter.eraseOp(shapeOf);
   return mlir::LogicalResult::success();
 }
 
+//===----------------------------------------------------------------------===//
+// GetExtent
+//===----------------------------------------------------------------------===//
+
+void hlfir::GetExtentOp::build(mlir::OpBuilder &builder,
+                               mlir::OperationState &result, mlir::Value shape,
+                               unsigned dim) {
+  mlir::Type indexTy = builder.getIndexType();
+  mlir::IntegerAttr dimAttr = mlir::IntegerAttr::get(indexTy, dim);
+  build(builder, result, indexTy, shape, dimAttr);
+}
+
+mlir::LogicalResult hlfir::GetExtentOp::verify() {
+  fir::ShapeType shapeTy = getShape().getType().cast<fir::ShapeType>();
+  std::uint64_t rank = shapeTy.getRank();
+  llvm::APInt dim = getDim();
+  if (dim.sge(rank))
+    return emitOpError("dimension index out of bounds");
+  return mlir::success();
+}
+
 #define GET_OP_CLASSES
 #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc"

flang/test/HLFIR/getextent.fir

+// Test hlfir.get_extent operaiton parse, verify (no errors), and unparse
+// RUN: fir-opt %s | fir-opt | FileCheck %s
+
+func.func @getextent(%arg0: !fir.shape<3>) {
+  %0 = hlfir.get_extent %arg0 {dim = 0 : index} : (!fir.shape<3>) -> index
+  %1 = hlfir.get_extent %arg0 {dim = 1 : index} : (!fir.shape<3>) -> index
+  %2 = hlfir.get_extent %arg0 {dim = 2 : index} : (!fir.shape<3>) -> index
+  return
+}
+// CHECK-LABEL: func.func @getextent
+// CHECK:           %[[SHAPE:.*]]: !fir.shape<3>
+// CHECK-NEXT:    %[[EXT0:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 0 : index} : (!fir.shape<3>) -> index
+// CHECK-NEXT:    %[[EXT1:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 1 : index} : (!fir.shape<3>) -> index
+// CHECK-NEXT:    %[[EXT2:.*]] = hlfir.get_extent %[[SHAPE]] {dim = 2 : index} : (!fir.shape<3>) -> index
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }

flang/test/HLFIR/invalid.fir

func.func @bad_shapeof(%arg0: !hlfir.expr<!fir.char<1,10>>) {
   %0 = hlfir.shape_of %arg0 : (!hlfir.expr<!fir.char<1,10>>) -> !fir.shape<1>
 }
 
 // -----
 func.func @bad_shapeof2(%arg0: !hlfir.expr<10xi32>) {
   // expected-error@+1 {{'hlfir.shape_of' op result rank and expr rank do not match}}
   %0 = hlfir.shape_of %arg0 : (!hlfir.expr<10xi32>) -> !fir.shape<42>
 }
+
+// -----
+func.func @bad_getextent(%arg0: !fir.shape<1>) {
+  // expected-error@+1 {{'hlfir.get_extent' op dimension index out of bounds}}
+  %0 = hlfir.get_extent %arg0 {dim = 1 : index} : (!fir.shape<1>) -> index
+}