[llvm] b4a7559 - [OpenMP][FIX] Replace ICVs only with values valid at the getter position

[Johannes Doerfert via llvm-commits]

Author: Johannes Doerfert
Date: 2022-01-19T18:40:13-06:00
New Revision: b4a7559844d06928234d576335316bf82d792c08

URL: https://github.com/llvm/llvm-project/commit/b4a7559844d06928234d576335316bf82d792c08
DIFF: https://github.com/llvm/llvm-project/commit/b4a7559844d06928234d576335316bf82d792c08.diff

LOG: [OpenMP][FIX] Replace ICVs only with values valid at the getter position

While we might know the value if an ICV at a getter position it is not
always clear that we can simply use it. Verify the value is valid first
to avoid invalid IR.

Fixes #53300.

Added: 
    llvm/test/Transforms/OpenMP/icv_tracking_out_of_scope.ll

Modified: 
    llvm/lib/Transforms/IPO/OpenMPOpt.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index 444d638d6b6d0..416c3b8b01d40 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -2166,7 +2166,7 @@ struct AAICVTrackerFunction : public AAICVTracker {
       };
 
       auto CallCheck = [&](Instruction &I) {
-        Optional<Value *> ReplVal = getValueForCall(A, &I, ICV);
+        Optional<Value *> ReplVal = getValueForCall(A, I, ICV);
         if (ReplVal.hasValue() &&
             ValuesMap.insert(std::make_pair(&I, *ReplVal)).second)
           HasChanged = ChangeStatus::CHANGED;
@@ -2192,12 +2192,12 @@ struct AAICVTrackerFunction : public AAICVTracker {
     return HasChanged;
   }
 
-  /// Hepler to check if \p I is a call and get the value for it if it is
+  /// Helper to check if \p I is a call and get the value for it if it is
   /// unique.
-  Optional<Value *> getValueForCall(Attributor &A, const Instruction *I,
+  Optional<Value *> getValueForCall(Attributor &A, const Instruction &I,
                                     InternalControlVar &ICV) const {
 
-    const auto *CB = dyn_cast<CallBase>(I);
+    const auto *CB = dyn_cast<CallBase>(&I);
     if (!CB || CB->hasFnAttr("no_openmp") ||
         CB->hasFnAttr("no_openmp_routines"))
       return None;
@@ -2213,8 +2213,8 @@ struct AAICVTrackerFunction : public AAICVTracker {
     if (CalledFunction == GetterRFI.Declaration)
       return None;
     if (CalledFunction == SetterRFI.Declaration) {
-      if (ICVReplacementValuesMap[ICV].count(I))
-        return ICVReplacementValuesMap[ICV].lookup(I);
+      if (ICVReplacementValuesMap[ICV].count(&I))
+        return ICVReplacementValuesMap[ICV].lookup(&I);
 
       return nullptr;
     }
@@ -2226,8 +2226,11 @@ struct AAICVTrackerFunction : public AAICVTracker {
     const auto &ICVTrackingAA = A.getAAFor<AAICVTracker>(
         *this, IRPosition::callsite_returned(*CB), DepClassTy::REQUIRED);
 
-    if (ICVTrackingAA.isAssumedTracked())
-      return ICVTrackingAA.getUniqueReplacementValue(ICV);
+    if (ICVTrackingAA.isAssumedTracked()) {
+      Optional<Value *> URV = ICVTrackingAA.getUniqueReplacementValue(ICV);
+      if (!URV || (*URV && AA::isValidAtPosition(**URV, I, OMPInfoCache)))
+        return URV;
+    }
 
     // If we don't know, assume it changes.
     return nullptr;
@@ -2279,7 +2282,7 @@ struct AAICVTrackerFunction : public AAICVTracker {
           break;
         }
 
-        Optional<Value *> NewReplVal = getValueForCall(A, CurrInst, ICV);
+        Optional<Value *> NewReplVal = getValueForCall(A, *CurrInst, ICV);
         if (!NewReplVal.hasValue())
           continue;
 

diff  --git a/llvm/test/Transforms/OpenMP/icv_tracking_out_of_scope.ll b/llvm/test/Transforms/OpenMP/icv_tracking_out_of_scope.ll
new file mode 100644
index 0000000000000..00e4574dbba8e
--- /dev/null
+++ b/llvm/test/Transforms/OpenMP/icv_tracking_out_of_scope.ll
@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature
+; RUN: opt -S -openmp-opt-cgscc < %s | FileCheck %s
+; RUN: opt -S -passes=openmp-opt-cgscc < %s | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+ at inofth = external global i32, align 4
+
+define void @set() {
+; CHECK-LABEL: define {{[^@]+}}@set() {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* @inofth, align 4
+; CHECK-NEXT:    call void @omp_set_num_threads(i32 noundef [[TMP0]])
+; CHECK-NEXT:    ret void
+;
+entry:
+  %0 = load i32, i32* @inofth, align 4
+  call void @omp_set_num_threads(i32 noundef %0)
+  ret void
+}
+
+define i32 @test(i32 %a) {
+; CHECK-LABEL: define {{[^@]+}}@test
+; CHECK-SAME: (i32 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    call void @set()
+; CHECK-NEXT:    [[CALL:%.*]] = call i32 @omp_get_max_threads()
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[CALL]], [[A]]
+; CHECK-NEXT:    br i1 [[CMP]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
+; CHECK:       if.then:
+; CHECK-NEXT:    ret i32 0
+; CHECK:       if.end:
+; CHECK-NEXT:    ret i32 1
+;
+entry:
+  call void @set()
+  %call = call i32 @omp_get_max_threads()
+  %cmp = icmp ne i32 %call, %a
+  br i1 %cmp, label %if.then, label %if.end
+
+if.then:                                          ; preds = %entry
+  ret i32 0
+
+if.end:                                           ; preds = %if.then, %entry
+  ret i32 1
+}
+
+declare void @omp_set_num_threads(i32 noundef)
+declare i32 @omp_get_max_threads()
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 7, !"openmp", i32 50}