[flang-commits] [flang] 80e45a7 - [flang][hlfir] add hlfir.get_extent

[Tom Eccles via flang-commits]

Author: Tom Eccles
Date: 2023-04-17T13:25:54Z
New Revision: 80e45a7947fa17a14db575a71aa178747de19b20

URL: https://github.com/llvm/llvm-project/commit/80e45a7947fa17a14db575a71aa178747de19b20
DIFF: https://github.com/llvm/llvm-project/commit/80e45a7947fa17a14db575a71aa178747de19b20.diff

LOG: [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

Added: 
    flang/test/HLFIR/getextent.fir

Modified: 
    flang/include/flang/Optimizer/HLFIR/HLFIROps.td
    flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp
    flang/test/HLFIR/invalid.fir

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td
index c942d97f4218b..79bba3dc39e2b 100644
--- 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 88acb2b4ac790..a479d64f39db8 100644
--- 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
index 0000000000000..3c50d1a64796c
--- /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 7ecb7cd221839..9b80cbdc25c4f 100644
--- 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 at +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 at +1 {{'hlfir.get_extent' op dimension index out of bounds}}
+  %0 = hlfir.get_extent %arg0 {dim = 1 : index} : (!fir.shape<1>) -> index
+}