[flang-commits] [flang] [Flang][HLFIR] Lower PACK(array, .TRUE.) to hlfir.reshape (PR #220860)

[via flang-commits]

================
@@ -3068,6 +3069,71 @@ class DotProductConversion
   }
 };
 
+static std::optional<bool> getLogicalConstant(mlir::Value value) {
+  if (auto convertOp = value.getDefiningOp<fir::ConvertOp>())
+    value = convertOp.getValue();
+  if (auto cst = fir::getIntIfConstant(value))
+    return *cst != 0;
+  return std::nullopt;
+}
+
+class PackAsReshapeConversion : public mlir::OpRewritePattern<hlfir::PackOp> {
+public:
+  using mlir::OpRewritePattern<hlfir::PackOp>::OpRewritePattern;
+
+  llvm::LogicalResult
+  matchAndRewrite(hlfir::PackOp pack,
+                  mlir::PatternRewriter &rewriter) const override {
+    if (pack.getVector())
+      return rewriter.notifyMatchFailure(pack, "PACK with VECTOR");
+    hlfir::Entity mask{pack.getMask()};
+    if (mask.getRank() != 0)
+      return rewriter.notifyMatchFailure(pack, "non-scalar mask");
+    if (!getLogicalConstant(pack.getMask()).value_or(false))
+      return rewriter.notifyMatchFailure(pack, "mask is not .TRUE.");
+    hlfir::Entity array{pack.getArray()};
+    if (!fir::isa_trivial(array.getFortranElementType()) ||
+        array.isPolymorphic())
+      return rewriter.notifyMatchFailure(pack, "unsupported array type");
+
+    mlir::Location loc = pack.getLoc();
+    fir::FirOpBuilder builder{rewriter, pack.getOperation()};
+    builder.setIntegerOverflowFlags(mlir::arith::IntegerOverflowFlags::nuw);
+
+    llvm::SmallVector<mlir::Value, Fortran::common::maxRank> arrayExtents =
+        hlfir::genExtentsVector(loc, builder, array);
+    mlir::Type indexType = builder.getIndexType();
+    mlir::Value totalSize = builder.createIntegerConstant(loc, indexType, 1);
+    for (mlir::Value extent : arrayExtents)
+      totalSize = mlir::arith::MulIOp::create(
+          builder, loc, totalSize,
+          builder.createConvert(loc, indexType, extent));
+
+    mlir::Value one = builder.createIntegerConstant(loc, indexType, 1);
+    mlir::Value unitShape = fir::ShapeOp::create(builder, loc, one);
+    mlir::Type shapeExprType =
+        hlfir::ExprType::get(builder.getContext(), {1}, indexType,
----------------
ejose02 wrote:

@tblah, thanks for the feedback. You're right that the lit test is a different case, and I understand the type-code mapping concern.
@vzakhari, given both jeanPerier and tblah have flagged the awkwardness with index types, do you still think that's the right approach, or should we explore using sized integers instead?

https://github.com/llvm/llvm-project/pull/220860