diff --git a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td --- a/flang/include/flang/Optimizer/HLFIR/HLFIROps.td +++ b/flang/include/flang/Optimizer/HLFIR/HLFIROps.td @@ -793,4 +793,25 @@ 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 --- 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 #include @@ -937,5 +939,26 @@ 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(); + 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 --- /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 --- a/flang/test/HLFIR/invalid.fir +++ b/flang/test/HLFIR/invalid.fir @@ -512,3 +512,9 @@ // 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 +}