[OpenMP] Use assumptions during ICV tracking

[jdoerfert]

Summary

The OpenMP 5.1 assumptions no_openmp and no_openmp_routines allow us
to ignore calls that would otherwise prevent ICV tracking.

Once we track more ICVs we might need to distinguish the ones that could
be impacted even with no_openmp_routines.

[sstefan1]

Cool :)

LGTM

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

struct AAICVTrackerFunction : public AAICVTracker {
   }
 
   /// Hepler 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,
                                     InternalControlVar &ICV) const {
 
     const auto *CB = dyn_cast<CallBase>(I);
-    if (!CB)
+    if (!CB || CB->hasFnAttr("no_openmp") ||
+        CB->hasFnAttr("no_openmp_routines"))
       return None;
 
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
     auto &GetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Getter];
     auto &SetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Setter];
     Function *CalledFunction = CB->getCalledFunction();
 
     // Indirect call, assume ICV changes.

llvm/test/Transforms/OpenMP/icv_tracking.ll

 ; CHECK-NEXT:    [[TMP2:%.*]] = add nsw i32 [[TMP0]], 1
 ; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
   tail call void @use(i32 %0)
   %2 = add nsw i32 %0, 1
   ret i32 %2
 }
 
+define i32 @ok_use_assume(i32 %0) {
+; CHECK-LABEL: define {{[^@]+}}@ok_use_assume
+; CHECK-SAME: (i32 [[TMP0:%.*]]) {
+; CHECK-NEXT:    call void @use(i32 [[TMP0]]) [[ATTR1:#.*]]
+; CHECK-NEXT:    call void @use(i32 [[TMP0]]) [[ATTR2:#.*]]
+; CHECK-NEXT:    call void @no_openmp_use(i32 [[TMP0]])
+; CHECK-NEXT:    [[TMP2:%.*]] = add nsw i32 [[TMP0]], 1
+; CHECK-NEXT:    ret i32 [[TMP2]]
+;
+  call void @use(i32 %0) "no_openmp"
+  call void @use(i32 %0) "no_openmp_routines"
+  call void @no_openmp_use(i32 %0)
+  %2 = add nsw i32 %0, 1
+  ret i32 %2
+}
+
 define void @indirect_call(void ()* %0) {
 ; CHECK-LABEL: define {{[^@]+}}@indirect_call
 ; CHECK-SAME: (void ()* [[TMP0:%.*]])
 ; CHECK-NEXT:    call void @omp_set_num_threads(i32 4)
 ; CHECK-NEXT:    tail call void [[TMP0]]()
 ; CHECK-NEXT:    [[TMP2:%.*]] = tail call i32 @omp_get_max_threads()
 ; CHECK-NEXT:    tail call void @use(i32 [[TMP2]])
 ; CHECK-NEXT:    ret void
; FIXME: this value should be tracked and the rest of the getters deduplicated and replaced with it.
   ret i32 0
 }
 
 declare dso_local void @omp_set_num_threads(i32)
 
 declare dso_local i32 @omp_get_max_threads()
 
 declare dso_local void @use(i32)
+declare dso_local void @no_openmp_use(i32) "no_openmp"
 
 define internal void @.omp_outlined.(i32* %0, i32* %1) {
 ; CHECK-LABEL: define {{[^@]+}}@.omp_outlined.
 ; CHECK-SAME: (i32* [[TMP0:%.*]], i32* [[TMP1:%.*]])
 ; CHECK-NEXT:    [[TMP3:%.*]] = tail call i32 @omp_get_max_threads()
 ; CHECK-NEXT:    [[TMP4:%.*]] = tail call i32 @omp_get_max_threads()
 ; CHECK-NEXT:    tail call void @use(i32 [[TMP4]])
 ; CHECK-NEXT:    tail call void @omp_set_num_threads(i32 10)
3:                                                ; preds = %1
   br label %5
 
 5:                                                ; preds = %3, %1
   %6 = call i32 @omp_get_max_threads()
   call void @use(i32 %6)
   ret void
 }
 
+define void @ok_use_assume_test(i1 %0) {
+; CHECK-LABEL: define {{[^@]+}}@ok_use_assume_test
+; CHECK-SAME: (i1 [[TMP0:%.*]]) {
+; CHECK-NEXT:    call void @omp_set_num_threads(i32 2)
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i1 [[TMP0]], false
+; CHECK-NEXT:    br i1 [[TMP2]], label [[TMP5:%.*]], label [[TMP3:%.*]]
+; CHECK:       3:
+; CHECK-NEXT:    [[TMP4:%.*]] = call i32 @ok_use_assume(i32 10)
+; CHECK-NEXT:    br label [[TMP5]]
+; CHECK:       5:
+; CHECK-NEXT:    call void @use(i32 2)
+; CHECK-NEXT:    ret void
+;
+  call void @omp_set_num_threads(i32 2)
+  %2 = icmp eq i1 %0, 0
+  br i1 %2, label %5, label %3
+
+3:                                                ; preds = %1
+  %4 = call i32 @ok_use_assume(i32 10)
+  br label %5
+
+5:                                                ; preds = %3, %1
+  %6 = call i32 @omp_get_max_threads()
+  call void @use(i32 %6)
+  ret void
+}
+
 define weak void @weak_known_unique_icv(i1 %0) {
 ; CHECK-LABEL: define {{[^@]+}}@weak_known_unique_icv
 ; CHECK-SAME: (i1 [[TMP0:%.*]])
 ; CHECK-NEXT:    call void @omp_set_num_threads(i32 2)
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i1 [[TMP0]], false
 ; CHECK-NEXT:    br i1 [[TMP2]], label [[TMP5:%.*]], label [[TMP3:%.*]]
 ; CHECK:       3:
 ; CHECK-NEXT:    [[TMP4:%.*]] = call i32 @icv_free_use(i32 10)