Fix issues with some hlfir.declare shape arguments

Author: skatrak

flang/lib/Optimizer/OpenMP/FunctionFiltering.cpp

@@ -22,6 +22,7 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 namespace flangomp {
 #define GEN_PASS_DEF_FUNCTIONFILTERINGPASS
@@ -443,16 +444,38 @@ class FunctionFilteringPass
                                                  builder.getI64IntegerAttr(0));
 
       if (auto shape = declareOp.getShape()) {
+        // The pre-cg rewrite pass requires the shape to be defined by one of
+        // fir.shape, fir.shapeshift or fir.shift, so we need to make sure it's
+        // still defined by one of these after this pass.
         Operation *shapeOp = shape.getDefiningOp();
-        unsigned numArgs = shapeOp->getNumOperands();
-        if (isa<fir::ShapeShiftOp>(shapeOp))
-          numArgs /= 2;
-
-        // Since the pre-cg rewrite pass requires the shape to be defined by one
-        // of fir.shape, fir.shapeshift or fir.shift, we need to create one of
-        // these.
-        llvm::SmallVector<Value> extents(numArgs, zero);
-        auto newShape = builder.create<fir::ShapeOp>(shape.getLoc(), extents);
+        llvm::SmallVector<Value> extents(shapeOp->getNumOperands(), zero);
+        Value newShape =
+            llvm::TypeSwitch<Operation *, Value>(shapeOp)
+                .Case([&](fir::ShapeOp op) {
+                  return builder.create<fir::ShapeOp>(op.getLoc(), extents);
+                })
+                .Case([&](fir::ShapeShiftOp op) {
+                  auto type = fir::ShapeShiftType::get(op.getContext(),
+                                                       extents.size() / 2);
+                  return builder.create<fir::ShapeShiftOp>(op.getLoc(), type,
+                                                           extents);
+                })
+                .Case([&](fir::ShiftOp op) {
+                  auto type =
+                      fir::ShiftType::get(op.getContext(), extents.size());
+                  return builder.create<fir::ShiftOp>(op.getLoc(), type,
+                                                      extents);
+                })
+                .Default([](Operation *op) {
+                  op->emitOpError()
+                      << "hlfir.declare shape expected to be one of: "
+                         "fir.shape, fir.shapeshift or fir.shift";
+                  return nullptr;
+                });
+
+        if (!newShape)
+          return failure();
+
         declareOp.getShapeMutable().assign(newShape);
       }
 

flang/test/Transforms/OpenMP/function-filtering-host-ops.mlir

@@ -163,7 +163,7 @@ module attributes {omp.is_target_device = true} {
   // CHECK-SAME: (%[[X:.*]]: [[X_TYPE:[^)]*]])
   func.func @box_ptr(%x: !fir.ref<!fir.box<!fir.ptr<!fir.array<?xi32>>>>) {
     // CHECK-NEXT: %[[ZERO:.*]] = arith.constant 0 : i64
-    // CHECK-NEXT: %[[SHAPE:.*]] = fir.shape %[[ZERO]] : (i64) -> !fir.shape<1>
+    // CHECK-NEXT: %[[SHAPE:.*]] = fir.shape_shift %[[ZERO]], %[[ZERO]] : (i64, i64) -> !fir.shapeshift<1>
     // CHECK-NEXT: %[[PLACEHOLDER_X:.*]] = fir.alloca i1
     // CHECK-NEXT: %[[ALLOCA_X:.*]] = fir.convert %[[PLACEHOLDER_X]] : (!fir.ref<i1>) -> [[X_TYPE]]
     %0 = fir.alloca !fir.box<!fir.ptr<!fir.array<?xi32>>>
@@ -187,7 +187,7 @@ module attributes {omp.is_target_device = true} {
     %9:3 = fir.box_dims %3, %c0_2 : (!fir.box<!fir.ptr<!fir.array<?xi32>>>, index) -> (index, index, index)
     %10 = fir.shape_shift %9#0, %9#1 : (index, index) -> !fir.shapeshift<1>
 
-    // CHECK-NEXT: %[[Y_DECL:.*]]:2 = hlfir.declare %[[ALLOCA_Y]](%[[SHAPE]]) {fortran_attrs = #fir.var_attrs<target>, uniq_name = "y"} : ([[Y_TYPE]], !fir.shape<1>) -> (!fir.box<!fir.array<?xi32>>, [[Y_TYPE]])
+    // CHECK-NEXT: %[[Y_DECL:.*]]:2 = hlfir.declare %[[ALLOCA_Y]](%[[SHAPE]]) {fortran_attrs = #fir.var_attrs<target>, uniq_name = "y"} : ([[Y_TYPE]], !fir.shapeshift<1>) -> (!fir.box<!fir.array<?xi32>>, [[Y_TYPE]])
     %11:2 = hlfir.declare %8(%10) {fortran_attrs = #fir.var_attrs<target>, uniq_name = "y"} : (!fir.ptr<!fir.array<?xi32>>, !fir.shapeshift<1>) -> (!fir.box<!fir.array<?xi32>>, !fir.ptr<!fir.array<?xi32>>)
     %c1_3 = arith.constant 1 : index
     %c0_4 = arith.constant 0 : index