[flang][hlfir] add hlfir.get_extent

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

Differential Revision: https://reviews.llvm.org/D148220

diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
index c942d97..79bba3d 100644 (file)
--- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
+++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
@@ -793,4 +793,25 @@ def hlfir_ShapeOfOp : hlfir_Op<"shape_of", [Pure]> {
   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

diff --git a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
index 88acb2b..a479d64 100644 (file)
--- a/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
+++ b/flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
@@ -16,10 +16,12 @@
 #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>
@@ -937,5 +939,26 @@ hlfir::ShapeOfOp::canonicalize(ShapeOfOp 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"

diff --git a/flang/test/HLFIR/getextent.fir b/flang/test/HLFIR/getextent.fir
new file mode 100644 (file)
index 0000000..3c50d1a
--- /dev/null
+++ b/flang/test/HLFIR/getextent.fir
@@ -0,0 +1,16 @@
+// 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:  }

diff --git a/flang/test/HLFIR/invalid.fir b/flang/test/HLFIR/invalid.fir
index 7ecb7cd..9b80cbd 100644 (file)
--- a/flang/test/HLFIR/invalid.fir
+++ b/flang/test/HLFIR/invalid.fir
@@ -512,3 +512,9 @@ 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
+}