[flang][hlfir] lower hlfir.shape_of

[tblah]

Summary

If possible the shape is gotten from the definition of the hlfir.expr
argument. Otherwise one is synthesized from the extent information in
the hlfir::ExprType.

The simple cases should already have been resolved during lowering. This
is mostly intended for cases where shape information is added in between
lowering and the end of bufferization (for example transformational
intrinsics).

Depends on: D146832

[jeanPerier]

flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp

177

I think you can get away with a much simpler approach that will cover all cases: the ShapeOfOp operand was translated to an hlfir::Entity that is a variable: you can "simply" call genShape on it (it knows of to retrieve the fir.shape from variables, or to emit fir.box_dim as needed):

hlfir::Entity bufferizedExpr{getBufferizedExprStorage(adaptor.getExpr())};
assert(bufferizedExpr.isVariable() && "array hlfir.expr are bufferized in memory");
mlir::Value shape = hlfir::genShape(loc, builder, bufferizedExpr);
rewriter.replaceOp(shapeOf, shape);
return mlir::success();

[tblah]

Updating to use the simpler implementation suggested by Jean.

[tblah]

flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp

177

Thanks for this, that's a much more elegant solution.

[jeanPerier]

Looks great, thanks. Beware that there is also some patch application issue though.

flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp

194

Did you hit the TODO, or is this here to be safe?

I am not expecting that bufferizedExpr.isVariable() should ever return false here (shape_of must be used on arrays, arrays must be placed in memory by this pass, genShape must succeed for variables), but I am ok with having a TODO to be safe for now.

[tblah]

flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp

194

No I haven't hit this TODO in practice. It is just there to be safe. I'll change it to return mlir::failure

[tblah]

Change from using a TODO to mlir::emitError to make it clearer that this
error condition is not known to occur in real code - it is just there for
safety.

flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp

 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
 #include "flang/Optimizer/HLFIR/HLFIRDialect.h"
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Optimizer/HLFIR/Passes.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
-#include <mlir/Support/LogicalResult.h>
+#include "llvm/ADT/TypeSwitch.h"
 
 namespace hlfir {
 #define GEN_PASS_DEF_BUFFERIZEHLFIR
 #include "flang/Optimizer/HLFIR/Passes.h.inc"
 } // namespace hlfir
 
 namespace {
 
matchAndRewrite(hlfir::AsExprOp asExpr, OpAdaptor adaptor,
     builder.create<hlfir::AssignOp>(loc, source, temp);
     mlir::Value bufferizedExpr =
         packageBufferizedExpr(loc, builder, temp, cleanup);
     rewriter.replaceOp(asExpr, bufferizedExpr);
     return mlir::success();
   }
 };
 
+struct ShapeOfOpConversion
+    : public mlir::OpConversionPattern<hlfir::ShapeOfOp> {
+  using mlir::OpConversionPattern<hlfir::ShapeOfOp>::OpConversionPattern;
+
+  mlir::LogicalResult
+  matchAndRewrite(hlfir::ShapeOfOp shapeOf, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const override {
+    mlir::Location loc = shapeOf.getLoc();
+    mlir::ModuleOp mod = shapeOf->getParentOfType<mlir::ModuleOp>();
+    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(mod));
+
+    mlir::Value shape;
+    hlfir::Entity bufferizedExpr{getBufferizedExprStorage(adaptor.getExpr())};
+    if (bufferizedExpr.isVariable()) {
+      shape = hlfir::genShape(loc, builder, bufferizedExpr);
+    } else {
+      // everything else failed so try to create a shape from static type info
+      hlfir::ExprType exprTy =
+          adaptor.getExpr().getType().dyn_cast_or_null<hlfir::ExprType>();
+      if (exprTy)
+        shape = hlfir::genExprShape(builder, loc, exprTy);
+    }
+    // expected to never happen
+    if (!shape)
+      return emitError(loc,
+                       "Unresolvable hlfir.shape_of where extents are unknown");
+
+    rewriter.replaceOp(shapeOf, shape);
+    return mlir::success();
+  }
+};
+
 struct ApplyOpConversion : public mlir::OpConversionPattern<hlfir::ApplyOp> {
   using mlir::OpConversionPattern<hlfir::ApplyOp>::OpConversionPattern;
   explicit ApplyOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::ApplyOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::ApplyOp apply, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = apply->getLoc();
void runOnOperation() override {
     // intrinsics calls creation. This may create race conflict if the pass is
     // scheduled on FuncOp. A solution could be to provide an optional mutex
     // when building a FirOpBuilder and locking around FuncOp and GlobalOp
     // creation, but this needs a bit more thinking, so at this point the pass
     // is scheduled on the moduleOp.
     auto module = this->getOperation();
     auto *context = &getContext();
     mlir::RewritePatternSet patterns(context);
-    patterns
-        .insert<ApplyOpConversion, AsExprOpConversion, AssignOpConversion,
-                AssociateOpConversion, ConcatOpConversion, DestroyOpConversion,
-                ElementalOpConversion, EndAssociateOpConversion,
-                NoReassocOpConversion, SetLengthOpConversion>(context);
+    patterns.insert<ApplyOpConversion, AsExprOpConversion, AssignOpConversion,
+                    AssociateOpConversion, ConcatOpConversion,
+                    DestroyOpConversion, ElementalOpConversion,
+                    EndAssociateOpConversion, NoReassocOpConversion,
+                    SetLengthOpConversion, ShapeOfOpConversion>(context);
     mlir::ConversionTarget target(*context);
     target.addIllegalOp<hlfir::ApplyOp, hlfir::AssociateOp, hlfir::ElementalOp,
                         hlfir::EndAssociateOp, hlfir::SetLengthOp,
                         hlfir::YieldElementOp>();
     target.markUnknownOpDynamicallyLegal([](mlir::Operation *op) {
       return llvm::all_of(
                  op->getResultTypes(),
                  [](mlir::Type ty) { return !ty.isa<hlfir::ExprType>(); }) &&

flang/test/HLFIR/shapeof-lowering.fir

+// Test hlfir.shape_of lowering
+// RUN: fir-opt %s -bufferize-hlfir | FileCheck %s
+
+func.func @shapeof_asexpr(%arg0: !fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.shape<1> {
+  %c0 = arith.constant 0 : index
+  %59:3 = fir.box_dims %arg0, %c0 : (!fir.box<!fir.heap<!fir.array<?xf32>>>, index) -> (index, index, index)
+  %60 = fir.box_addr %arg0 : (!fir.box<!fir.heap<!fir.array<?xf32>>>) -> !fir.heap<!fir.array<?xf32>>
+  %61 = fir.shape_shift %59#0, %59#1 : (index, index) -> !fir.shapeshift<1>
+  %62:2 = hlfir.declare %60(%61) {uniq_name = ".tmp.intrinsic_result"} : (!fir.heap<!fir.array<?xf32>>, !fir.shapeshift<1>) -> (!fir.box<!fir.array<?xf32>>, !fir.heap<!fir.array<?xf32>>)
+  %true = arith.constant true
+  %63 = hlfir.as_expr %62#0 move %true : (!fir.box<!fir.array<?xf32>>, i1) -> !hlfir.expr<?xf32>
+  %64 = hlfir.shape_of %63 : (!hlfir.expr<?xf32>) -> !fir.shape<1>
+  return %64 : !fir.shape<1>
+}
+// CHECK-LABEL: @shapeof_asexpr
+// CHECK:           %[[ARG0:.*]]: !fir.box<!fir.heap<!fir.array<?xf32>>>
+// CHECK-NEXT:    %[[C0:.*]] = arith.constant 0
+// CHECK-NEXT:    %[[BOX_DIMS:.*]]:3 = fir.box_dims %[[ARG0]], %[[C0]]
+// CHECK-NEXT:    %[[BOX_ADDR:.*]] = fir.box_addr %[[ARG0]]
+// CHECK-NEXT:    %[[SHPE_SHFT:.*]] = fir.shape_shift %[[BOX_DIMS]]#0, %[[BOX_DIMS]]#1
+// CHECK-NEXT:    %[[VAR:.*]]:2 = hlfir.declare %[[BOX_ADDR]](%[[SHPE_SHFT]])
+// CHECK-NEXT:    %[[TRUE:.*]] = arith.constant true
+// CHECK-NEXT:    %[[TUPLE0:.*]] = fir.undefined tuple
+// CHECK-NEXT:    %[[TUPLE1:.*]] = fir.insert_value %[[TUPLE0]], %[[TRUE]]
+// CHECK-NEXT:    %[[TUPLE2:.*]] = fir.insert_value %[[TUPLE1]], %[[VAR]]#0
+// CHECK-NEXT:    %[[SHAPE:.*]] = fir.shape %[[BOX_DIMS]]#1
+// CHECK-NEXT:    return %[[SHAPE]]
+
+func.func @shapeof_elemental() -> !fir.shape<1> {
+  %c1 = arith.constant 1 : index
+  %0 = fir.shape %c1 : (index) -> !fir.shape<1>
+  %1 = hlfir.elemental %0 : (!fir.shape<1>) -> !hlfir.expr<?xindex> {
+  ^bb0(%arg3: index):
+    hlfir.yield_element %arg3 : index
+  }
+  %2 = hlfir.shape_of %1 : (!hlfir.expr<?xindex>) -> !fir.shape<1>
+  return %2 : !fir.shape<1>
+}
+// CHECK-LABEL: @shapeof_elemental
+// CHECK:         %[[C1:.*]] = arith.constant 1 : index
+// CHECK-NEXT:    %[[SHAPE:.*]] = fir.shape %[[C1]]
+// CHECK:         fir.do_loop %{{.*}} = %{{.*}} to %[[C1:.*]]
+// CHECK:         return %[[SHAPE]]
+
+func.func @shapeof_fallback(%arg0: !hlfir.expr<1x2x3xi32>) -> !fir.shape<3> {
+  %shape = hlfir.shape_of %arg0 : (!hlfir.expr<1x2x3xi32>) -> !fir.shape<3>
+  return %shape : !fir.shape<3>
+}
+// CHECK-LABEL: @shapeof_fallback
+// CHECK:           %[[EXPR:.*]]: !hlfir.expr<1x2x3xi32>
+// CHECK-NEXT:    %[[C1:.*]] = arith.constant 1 : index
+// CHECK-NEXT:    %[[C2:.*]] = arith.constant 2 : index
+// CHECK-NEXT:    %[[C3:.*]] = arith.constant 3 : index
+// CHECK-NEXT:    %[[SHAPE:.*]] = fir.shape %[[C1]], %[[C2]], %[[C3]] :
+// CHECK-NEXT:    return %[[SHAPE]]