[flang][hlfir] add hlfir.shape_of

[tblah]

Summary

This is an operation which returns the fir.shape for a hlfir.expr.

A hlfir.expr can be defined by:

• A transformational intrinsic (e.g. hlfir.matmul)
• hlfir.as_expr
• hlfir.elemental

hlfir.elemental is easy because there is a compulsory shape operand.
hlfir.as_expr is defined as operating on a variable (defined using a
hlfir.declare). hlfir.declare has an optional shape argument. The
transformational intrinsics do not have an associated shape.

If all extents are known at compile time, the extents for the shape can
be fetched from the hlfir.expr's type. For example, the result of a
hlfir.matmul with arguments who's extents are known at compile time will
have constant extents which can be queried from the type. In this case
the hlfir.shape_of will be canonicalised to a fir.shape operation using
those extents.

If not all extents are known at compile time, shapes have to be read
from boxes after bufferization. In the case of the transformational
intrinsics, the shape read from the result box can be queried from the
hlfir.declare operation for the buffer allocated to that hlfir.expr (via
the hlfir.as_expr).

[jeanPerier]

Other than the circular dependency issue, LGTM.

Some cases are impossible to resolve using hlfir.shape_of. For example,

func.func @bad(%arg0: !hlfir.expr<?xi32>) {

%shape = hlfir.shape_of %arg0 : (!hlfir.expr<?xi32>) -> !fir.shape<1>

}
Here there is no way to get additional information about this
hlfir.expr and its extents are not known at compile time. Lowering must
not create MLIR like this. If this happens, there will be a failure in
the bufferization pass.

I agree this is not a desired pattern, but do not think it is impossible to support if we ever need to. This could translate to func.func @bad(%arg0: tuple<!fir.box<!fir.array<?xi32>>, i1>) { .... (hlfir.expr are translated to a tuple with the storage of the expression value, and a boolean flag indicating if the storage must be deallocated when hitting an hlfir.destroy). So the !fir.box would be here to support the shape. This comment does not change anything to your patch though (my code suggestion in D146833 will put some light on this. Note that we indeed do not currently translate block argument hlfir.expr type as illustrated above, so this would still fail to compile, but we could if we wanted/needed to).

flang/include/flang/Optimizer/HLFIR/HLFIRDialect.h

17

I think genExprShape cannot be here, this will create circular dependencies in shared library builds. Can't genExprShape use an mlir::OpBuilder &builder instead of a fir::FirOpBuilder, I am not sure the FirOpBuilder specific helpers are used by it, and since a FirOpBuilder inherits from mlir::OpBuilder, it can still be passed to genExprShape.

[tblah]

Thanks for review.

Updated to fix the circular dependency.

[jeanPerier]

LGTM

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

766

Can probably also be given the Pure interface.

[tblah]

Add pure attribute to operation definition

flang/include/flang/Optimizer/HLFIR/HLFIRDialect.h

 //===- HLFIRDialect.h - High Level Fortran IR dialect -----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the HLFIR dialect that models Fortran expressions and
 // assignments without requiring storage allocation and manipulations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef FORTRAN_OPTIMIZER_HLFIR_HLFIRDIALECT_H
 #define FORTRAN_OPTIMIZER_HLFIR_HLFIRDIALECT_H
 
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "mlir/IR/Dialect.h"
 
 namespace hlfir {
 /// Is this a type that can be used for an HLFIR variable ?
 bool isFortranVariableType(mlir::Type);
 bool isFortranScalarCharacterType(mlir::Type);
 bool isFortranScalarCharacterExprType(mlir::Type);
 } // namespace hlfir

 bool isFortranNumericalArrayObject(mlir::Type);
 bool isFortranNumericalOrLogicalArrayObject(mlir::Type);
 bool isFortranArrayObject(mlir::Type);
 bool isPassByRefOrIntegerType(mlir::Type);
 bool isI1Type(mlir::Type);
 // scalar i1 or logical, or sequence of logical (via (boxed?) array or expr)
 bool isMaskArgument(mlir::Type);
 
+/// If an expression's extents are known at compile time, generate a fir.shape
+/// for this expression. Otherwise return {}
+mlir::Value genExprShape(mlir::OpBuilder &builder, const mlir::Location &loc,
+                         const hlfir::ExprType &expr);
+
 } // namespace hlfir
 
 #endif // FORTRAN_OPTIMIZER_HLFIR_HLFIRDIALECT_H

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

let arguments = (ins fir_BaseBoxType:$temp,
                        Optional<fir_BaseBoxType>:$var);
 
   let assemblyFormat = [{
     $temp `,` $was_copied (`to` $var^)?
     attr-dict `:` functional-type(operands, results)
   }];
 }
 
+def hlfir_ShapeOfOp : hlfir_Op<"shape_of", [Pure]> {
+  let summary = "Get the shape of a hlfir.expr";
+  let description = [{
+    Gets the runtime shape of a hlfir.expr. In lowering to FIR, the
+    hlfir.shape_of operation will be replaced by an fir.shape.
+    It is not valid to request the shape of a hlfir.expr which has no shape.
+  }];
+
+  let arguments = (ins hlfir_ExprType:$expr);
+
+  let results = (outs fir_ShapeType);
+
+  let hasVerifier = 1;
+
+  // If all extents are known at compile time, the hlfir.shape_of can be
+  // immediately folded into a fir.shape operation. This makes information
+  // available sooner to inform bufferization decisions
+  let hasCanonicalizeMethod = 1;
+
+  let extraClassDeclaration = [{
+    std::size_t getRank();
+  }];
+
+  let assemblyFormat = [{
+    $expr attr-dict `:` functional-type(operands, results)
+  }];
+
+  let builders = [OpBuilder<(ins "mlir::Value":$expr)>];
+}
+
 #endif // FORTRAN_DIALECT_HLFIR_OPS

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

 //===-- HLFIRDialect.cpp --------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
 //
 //===----------------------------------------------------------------------===//
 
 #include "flang/Optimizer/HLFIR/HLFIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 #include "flang/Optimizer/HLFIR/HLFIRDialect.cpp.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "flang/Optimizer/HLFIR/HLFIRTypes.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
bool hlfir::isMaskArgument(mlir::Type type) {
   mlir::Type elementType = getFortranElementType(unwrappedType);
   if (unwrappedType != elementType)
     // input type is an array
     return mlir::isa<fir::LogicalType>(elementType);
 
   // input is a scalar, so allow i1 too
   return mlir::isa<fir::LogicalType>(elementType) || isI1Type(elementType);
 }
+
+mlir::Value hlfir::genExprShape(mlir::OpBuilder &builder,
+                                const mlir::Location &loc,
+                                const hlfir::ExprType &expr) {
+  mlir::IndexType indexTy = builder.getIndexType();
+  llvm::SmallVector<mlir::Value> extents;
+  extents.reserve(expr.getRank());
+
+  for (std::int64_t extent : expr.getShape()) {
+    if (extent == hlfir::ExprType::getUnknownExtent())
+      return {};
+    extents.emplace_back(builder.create<mlir::arith::ConstantOp>(
+        loc, indexTy, builder.getIntegerAttr(indexTy, extent)));
+  }
+
+  fir::ShapeType shapeTy =
+      fir::ShapeType::get(builder.getContext(), expr.getRank());
+  fir::ShapeOp shape = builder.create<fir::ShapeOp>(loc, shapeTy, extents);
+  return shape.getResult();
+}

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

 //===-- HLFIROps.cpp ------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
 //
 //===----------------------------------------------------------------------===//
 
 #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/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include <iterator>

 
 void hlfir::CopyInOp::build(mlir::OpBuilder &builder,
                             mlir::OperationState &odsState, mlir::Value var,
                             mlir::Value var_is_present) {
   return build(builder, odsState, {var.getType(), builder.getI1Type()}, var,
                var_is_present);
 }
 
+//===----------------------------------------------------------------------===//
+// ShapeOfOp
+//===----------------------------------------------------------------------===//
+
+void hlfir::ShapeOfOp::build(mlir::OpBuilder &builder,
+                             mlir::OperationState &result, mlir::Value expr) {
+  hlfir::ExprType exprTy = expr.getType().cast<hlfir::ExprType>();
+  mlir::Type type = fir::ShapeType::get(builder.getContext(), exprTy.getRank());
+  build(builder, result, type, expr);
+}
+
+std::size_t hlfir::ShapeOfOp::getRank() {
+  mlir::Type resTy = getResult().getType();
+  fir::ShapeType shape = resTy.cast<fir::ShapeType>();
+  return shape.getRank();
+}
+
+mlir::LogicalResult hlfir::ShapeOfOp::verify() {
+  mlir::Value expr = getExpr();
+  hlfir::ExprType exprTy = expr.getType().cast<hlfir::ExprType>();
+  std::size_t exprRank = exprTy.getShape().size();
+
+  if (exprRank == 0)
+    return emitOpError("cannot get the shape of a shape-less expression");
+
+  std::size_t shapeRank = getRank();
+  if (shapeRank != exprRank)
+    return emitOpError("result rank and expr rank do not match");
+
+  return mlir::success();
+}
+
+mlir::LogicalResult
+hlfir::ShapeOfOp::canonicalize(ShapeOfOp shapeOf,
+                               mlir::PatternRewriter &rewriter) {
+  // if extent information is available at compile time, immediately fold the
+  // hlfir.shape_of into a fir.shape
+  mlir::Location loc = shapeOf.getLoc();
+  hlfir::ExprType expr = shapeOf.getExpr().getType().cast<hlfir::ExprType>();
+
+  mlir::Value shape = hlfir::genExprShape(rewriter, loc, expr);
+  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();
+}
+
 #define GET_OP_CLASSES
 #include "flang/Optimizer/HLFIR/HLFIROps.cpp.inc"

flang/test/HLFIR/invalid.fir

 // -----
 func.func @bad_parent_comp6(%arg0: !fir.box<!fir.array<10x!fir.type<t2{i:i32,j:i32}>>>) {
   %c10 = arith.constant 10 : index
   %1 = fir.shape %c10 : (index) -> !fir.shape<1>
   // expected-error@+1 {{'hlfir.parent_comp' op result type must be a fir.box if the result is an array or has length parameters}}
   %2 = hlfir.parent_comp %arg0 shape %1 : (!fir.box<!fir.array<10x!fir.type<t2{i:i32,j:i32}>>>, !fir.shape<1>) -> !fir.ref<!fir.array<10x!fir.type<t1{i:i32}>>>
   return
 }
+
+// -----
+func.func @bad_shapeof(%arg0: !hlfir.expr<!fir.char<1,10>>) {
+  // expected-error@+1 {{'hlfir.shape_of' op cannot get the shape of a shape-less expression}}
+  %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>
+}

flang/test/HLFIR/shapeof.fir

+// Test hlfir.shape_of operation parse, verify (no errors), and unparse
+// RUN: fir-opt %s | fir-opt | FileCheck --check-prefix CHECK --check-prefix CHECK-ALL %s
+
+// Test canonicalization
+// RUN: fir-opt %s --canonicalize | FileCheck --check-prefix CHECK-CANON --check-prefix CHECK-ALL %s
+
+func.func @shapeof(%arg0: !hlfir.expr<2x2xi32>) -> !fir.shape<2> {
+  %shape = hlfir.shape_of %arg0 : (!hlfir.expr<2x2xi32>) -> !fir.shape<2>
+  return %shape : !fir.shape<2>
+}
+// CHECK-ALL-LABEL: func.func @shapeof
+// CHECK-ALL:          %[[EXPR:.*]]: !hlfir.expr<2x2xi32>
+
+// CHECK-NEXT:       %[[SHAPE:.*]] = hlfir.shape_of %[[EXPR]] : (!hlfir.expr<2x2xi32>) -> !fir.shape<2>
+
+// CHECK-CANON-NEXT: %[[C2:.*]] = arith.constant 2 : index
+// CHECK-CANON-NEXT: %[[SHAPE:.*]] = fir.shape %[[C2]], %[[C2]] : (index, index) -> !fir.shape<2>
+
+// CHECK-ALL-NEXT:   return %[[SHAPE]]
+
+// no canonicalization of expressions with unknown extents
+func.func @shapeof2(%arg0: !hlfir.expr<?x2xi32>) -> !fir.shape<2> {
+  %shape = hlfir.shape_of %arg0 : (!hlfir.expr<?x2xi32>) -> !fir.shape<2>
+  return %shape : !fir.shape<2>
+}
+// CHECK-ALL-LABEL: func.func @shapeof2
+// CHECK-ALL:          %[[EXPR:.*]]: !hlfir.expr<?x2xi32>
+// CHECK-ALL-NEXT:   %[[SHAPE:.*]] = hlfir.shape_of %[[EXPR]] : (!hlfir.expr<?x2xi32>) -> !fir.shape<2>
+// CHECK-ALL-NEXT:   return %[[SHAPE]]