[flang] Set low probability for array repacking code. (#144830)

This allows LLVM to place the most probably cold blocks
that do the repacking out of the line of the potentially hot code.

Author: vzakhari

flang/include/flang/Optimizer/Dialect/FIROps.td

@@ -2381,6 +2381,16 @@ def fir_IfOp
     static constexpr llvm::StringRef getWeightsAttrAssemblyName() {
       return "weights";
     }
+
+    /// Sets WeightedRegionBranchOpInterface weights to indicate
+    /// that either THEN or ELSE branch is unlikely.
+    /// By default, THEN branch is set to be unlikely.
+    void setUnlikelyIfWeights(bool unlikelyElse = false) {
+      if (unlikelyElse)
+        setWeights({1, 0});
+      else
+        setWeights({0, 1});
+    }
   }];
 }
 

flang/lib/Optimizer/CodeGen/LowerRepackArrays.cpp

@@ -250,6 +250,8 @@ PackArrayConversion::genRepackedBox(fir::FirOpBuilder &builder,
 
   fir::IfOp ifOp =
       builder.create<fir::IfOp>(loc, boxType, doPack, /*withElseRegion=*/true);
+  // Assume that the repacking is unlikely.
+  ifOp.setUnlikelyIfWeights();
 
   // Return original box.
   builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
@@ -322,20 +324,24 @@ UnpackArrayConversion::matchAndRewrite(fir::UnpackArrayOp op,
 
     auto isNotSame = builder.genPtrCompare(loc, mlir::arith::CmpIPredicate::ne,
                                            tempAddr, originalAddr);
-    builder.genIfThen(loc, isNotSame).genThen([&]() {});
-    // Copy from temporary to the original.
-    if (!op.getNoCopy())
-      fir::runtime::genShallowCopy(builder, loc, originalBox, tempBox,
-                                   /*resultIsAllocated=*/true);
-
-    // Deallocate, if it was allocated in heap.
-    // Note that the stack attribute does not always mean
-    // that the allocation was actually done in stack memory.
-    // There are currently cases where we delegate the allocation
-    // to the runtime that uses heap memory, even when the stack
-    // attribute is set on fir.pack_array.
-    if (!op.getStack() || !canAllocateTempOnStack(originalBox))
-      builder.create<fir::FreeMemOp>(loc, tempAddr);
+    builder.genIfThen(loc, isNotSame)
+        .genThen([&]() {
+          // Copy from temporary to the original.
+          if (!op.getNoCopy())
+            fir::runtime::genShallowCopy(builder, loc, originalBox, tempBox,
+                                         /*resultIsAllocated=*/true);
+
+          // Deallocate, if it was allocated in heap.
+          // Note that the stack attribute does not always mean
+          // that the allocation was actually done in stack memory.
+          // There are currently cases where we delegate the allocation
+          // to the runtime that uses heap memory, even when the stack
+          // attribute is set on fir.pack_array.
+          if (!op.getStack() || !canAllocateTempOnStack(originalBox))
+            builder.create<fir::FreeMemOp>(loc, tempAddr);
+        })
+        .getIfOp()
+        .setUnlikelyIfWeights();
   });
   rewriter.eraseOp(op);
   return mlir::success();

flang/test/Integration/cold_array_repacking.f90

@@ -0,0 +1,30 @@
+! Check that the branch weights used by the array repacking
+! are propagated all the way to LLVM IR:
+! RUN: %flang_fc1 -frepack-arrays -emit-llvm %s -o - | FileCheck %s
+
+! CHECK-LABEL: define void @test_(
+! CHECK-SAME:      ptr [[TMP0:%.*]]) {
+! CHECK:    [[TMP4:%.*]] = ptrtoint ptr [[TMP0]] to i64
+! CHECK:    [[TMP5:%.*]] = icmp ne i64 [[TMP4]], 0
+! CHECK:    br i1 [[TMP5]], label %[[BB6:.*]], label %[[BB46:.*]]
+! CHECK:  [[BB6]]:
+! CHECK:    [[TMP7:%.*]] = call i1 @_FortranAIsContiguous(ptr [[TMP0]])
+! CHECK:    [[TMP8:%.*]] = icmp eq i1 [[TMP7]], false
+! CHECK:    [[TMP13:%.*]] = and i1 [[TMP8]], [[TMP12:.*]]
+! CHECK:    br i1 [[TMP13]], label %[[BB14:.*]], label %[[BB46]], !prof [[PROF2:![0-9]+]]
+! CHECK:  [[BB14]]:
+! CHECK:    call void @_FortranAShallowCopyDirect
+! CHECK:    br label %[[BB46]]
+! CHECK:  [[BB46]]:
+! CHECK:    br i1 [[TMP5]], label %[[BB48:.*]], label %[[BB57:.*]]
+! CHECK:  [[BB48]]:
+! CHECK:    br i1 [[TMP55:.*]], label %[[BB56:.*]], label %[[BB57]], !prof [[PROF2]]
+! CHECK:  [[BB56]]:
+! CHECK:    call void @_FortranAShallowCopyDirect
+! CHECK:    br label %[[BB57]]
+! CHECK:  [[BB57]]:
+! CHECK:    ret void
+! CHECK: [[PROF2]] = !{!"branch_weights", i32 0, i32 1}
+subroutine test(x)
+  real :: x(:)
+end subroutine test