[Clang][OpenMP] Fix the issue that `llvm.lifetime.end` is emitted too early for variables captured in linear clause

[tianshilei1992]

Summary

Currently if an OpenMP program uses linear clause, and is compiled with
optimization, llvm.lifetime.end for variables listed in linear clause are
emitted too early such that there could still be uses after that. Let's take the
following code as example:

// loop.c
int j;
int *u;

void loop(int n) {
  int i;
  for (i = 0; i < n; ++i) {
    ++j;
    u = &j;
  }
}

We compile using the command:

clang -cc1 -fopenmp-simd -O3 -x c -triple x86_64-apple-darwin10 -emit-llvm loop.c -o loop.ll

The following IR (simplified) will be generated:

@j = local_unnamed_addr global i32 0, align 4
@u = local_unnamed_addr global ptr null, align 8

define void @loop(i32 noundef %n) local_unnamed_addr {
entry:
  %j = alloca i32, align 4
  %cmp = icmp sgt i32 %n, 0
  br i1 %cmp, label %simd.if.then, label %simd.if.end

simd.if.then:                                     ; preds = %entry
  call void @llvm.lifetime.start.p0(i64 4, ptr nonnull %j)
  store ptr %j, ptr @u, align 8
  call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %j)
  %0 = load i32, ptr %j, align 4
  store i32 %0, ptr @j, align 4
  br label %simd.if.end

simd.if.end:                                      ; preds = %simd.if.then, %entry
  ret void
}

The most important part is:

call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %j)
%0 = load i32, ptr %j, align 4
store i32 %0, ptr @j, align 4

%j is still loaded after @llvm.lifetime.end.p0(i64 4, ptr nonnull %j). This
could cause the backend incorrectly optimizes the code and further generates
incorrect code. The root cause is, when we emit a construct that could have
linear clause, it usually has the following pattern:

EmitOMPLinearClauseInit(S)
{
  OMPPrivateScope LoopScope(*this);
  ...
  EmitOMPLinearClause(S, LoopScope);
  ...
  (void)LoopScope.Privatize();
  ...
}
EmitOMPLinearClauseFinal(S, [](CodeGenFunction &) { return nullptr; });

Variables that need to be privatized are added into LoopScope, which also
serves as a RAII object. When LoopScope is destructed and if optimization is
enabled, a @llvm.lifetime.end is also emitted for each privatized variable.
However, the writing back to original variables in linear clause happens after
the scope in EmitOMPLinearClauseFinal, causing the issue we see above.

A quick "fix" seems to be, moving EmitOMPLinearClauseFinal inside the scope.
However, it doesn't work. That's because the local variable map has been updated
by LoopScope such that a variable declaration is mapped to the privatized
variable, instead of the actual one. In that way, the following code will be
generated:

%0 = load i32, ptr %j, align 4
store i32 %0, ptr %j, align 4
call void @llvm.lifetime.end.p0(i64 4, ptr nonnull %j)

Well, now the life time is correct, but apparently the writing back is broken.

In this patch, a new function OMPPrivateScope::restoreMap is added and called
before calling EmitOMPLinearClauseFinal. This can make sure that
EmitOMPLinearClauseFinal can find the orignal varaibls to write back.

Fixes #56913.

[ABataev]

clang/lib/CodeGen/CodeGenFunction.h

1070

Better to make it a property of MappedVars. Also, are there any issues if just restore it several times without checking?

[tianshilei1992]

clang/lib/CodeGen/CodeGenFunction.h

1070

Better to make it a property of MappedVars.

Will do.

Also, are there any issues if just restore it several times without checking?

TBH, I don't know clearly, but I feel it should not be a problem. We will not have two scopes at the same time, will we? But I will add an assertion that if the mapped variables have been restored, addPrivate should not be called anymore.

[tianshilei1992]

rebase and fix comments

[tianshilei1992]

clang/lib/CodeGen/CodeGenFunction.h

1070

We don't have to worry about restore for multiple times as every time it is called, the container is cleared.

[ABataev]

LG

clang/lib/CodeGen/CGStmtOpenMP.cpp

emitCommonSimdLoop(
         });
     CGF.EmitOMPSimdFinal(S, [](CodeGenFunction &) { return nullptr; });
     // Emit final copy of the lastprivate variables at the end of loops.
     if (HasLastprivateClause)
       CGF.EmitOMPLastprivateClauseFinal(S, /*NoFinals=*/true);
     CGF.EmitOMPReductionClauseFinal(S, /*ReductionKind=*/OMPD_simd);
     emitPostUpdateForReductionClause(CGF, S,
                                      [](CodeGenFunction &) { return nullptr; });
-  }
+    LoopScope.restoreMap();
     CGF.EmitOMPLinearClauseFinal(S, [](CodeGenFunction &) { return nullptr; });
+  }
   // Emit: if (PreCond) - end.
   if (ContBlock) {
     CGF.EmitBranch(ContBlock);
     CGF.EmitBlock(ContBlock, true);
   }
 }
 
 static bool isSupportedByOpenMPIRBuilder(const OMPSimdDirective &S) {
// Emit 'then' code.
             return CGF.Builder.CreateIsNotNull(
                 CGF.EmitLoadOfScalar(IL, S.getBeginLoc()));
           });
       // Emit final copy of the lastprivate variables if IsLastIter != 0.
       if (HasLastprivateClause)
         EmitOMPLastprivateClauseFinal(
             S, isOpenMPSimdDirective(S.getDirectiveKind()),
             Builder.CreateIsNotNull(EmitLoadOfScalar(IL, S.getBeginLoc())));
-    }
+      LoopScope.restoreMap();
       EmitOMPLinearClauseFinal(S, [IL, &S](CodeGenFunction &CGF) {
         return CGF.Builder.CreateIsNotNull(
             CGF.EmitLoadOfScalar(IL, S.getBeginLoc()));
       });
+    }
     DoacrossCleanupScope.ForceCleanup();
     // We're now done with the loop, so jump to the continuation block.
     if (ContBlock) {
       EmitBranch(ContBlock);
       EmitBlock(ContBlock, /*IsFinished=*/true);
     }
   }
   return HasLastprivateClause;
auto &&BodyGen = [CS, &S](CodeGenFunction &CGF, PrePostActionTy &) {
     // Emit final copy of the lastprivate variables if IsLastIter != 0.
     if (HasLastprivateClause) {
       CGF.EmitOMPLastprivateClauseFinal(
           S, isOpenMPSimdDirective(S.getDirectiveKind()),
           CGF.Builder.CreateIsNotNull(CGF.EmitLoadOfScalar(
               CGF.GetAddrOfLocalVar(*LIP), /*Volatile=*/false,
               (*LIP)->getType(), S.getBeginLoc())));
     }
+    LoopScope.restoreMap();
     CGF.EmitOMPLinearClauseFinal(S, [LIP, &S](CodeGenFunction &CGF) {
       return CGF.Builder.CreateIsNotNull(
           CGF.EmitLoadOfScalar(CGF.GetAddrOfLocalVar(*LIP), /*Volatile=*/false,
                                (*LIP)->getType(), S.getBeginLoc()));
     });
   };
   auto &&TaskGen = [&S, SharedsTy, CapturedStruct,
                     IfCond](CodeGenFunction &CGF, llvm::Function *OutlinedFn,

clang/lib/CodeGen/CodeGenFunction.h

public:
 
   /// The scope used to remap some variables as private in the OpenMP loop body
   /// (or other captured region emitted without outlining), and to restore old
   /// vars back on exit.
   class OMPPrivateScope : public RunCleanupsScope {
     OMPMapVars MappedVars;
     OMPPrivateScope(const OMPPrivateScope &) = delete;
     void operator=(const OMPPrivateScope &) = delete;
 
   public:
     /// Enter a new OpenMP private scope.
     explicit OMPPrivateScope(CodeGenFunction &CGF) : RunCleanupsScope(CGF) {}
 
     /// Registers \p LocalVD variable as a private with \p Addr as the address
     /// of the corresponding private variable. \p
     /// PrivateGen is the address of the generated private variable.
     /// \return true if the variable is registered as private, false if it has
public:
     /// variable initialized by a constructor that references other private
     /// variables. But at initialization original variables must be used, not
     /// private copies.
     /// \return true if at least one variable was privatized, false otherwise.
     bool Privatize() { return MappedVars.apply(CGF); }
 
     void ForceCleanup() {
       RunCleanupsScope::ForceCleanup();
-      MappedVars.restore(CGF);
+      restoreMap();
     }
 
     /// Exit scope - all the mapped variables are restored.
     ~OMPPrivateScope() {
       if (PerformCleanup)
         ForceCleanup();
     }
 
     /// Checks if the global variable is captured in current function.
     bool isGlobalVarCaptured(const VarDecl *VD) const {
       VD = VD->getCanonicalDecl();
       return !VD->isLocalVarDeclOrParm() && CGF.LocalDeclMap.count(VD) > 0;
     }
+
+    /// Restore all mapped variables w/o clean up. This is usefully when we want
+    /// to reference the original variables but don't want the clean up because
+    /// that could emit lifetime end too early, causing backend issue #56913.
+    void restoreMap() { MappedVars.restore(CGF); }
   };
 
   /// Save/restore original map of previously emitted local vars in case when we
   /// need to duplicate emission of the same code several times in the same
   /// function for OpenMP code.
   class OMPLocalDeclMapRAII {
     CodeGenFunction &CGF;
     DeclMapTy SavedMap;

clang/test/OpenMP/bug56913.c

+// NOTE: Assertions have been autogenerated by utils/update_cc_test_checks.py UTC_ARGS: --function-signature --include-generated-funcs --prefix-filecheck-ir-name _
+// RUN: %clang_cc1 -fopenmp-simd -O1 -x c -triple x86_64-apple-darwin10 -emit-llvm %s -o - | FileCheck %s --check-prefix=CHECK
+
+int j;
+int *u;
+
+void loop(int n) {
+  int i;
+#pragma omp parallel master taskloop simd linear(j)
+  for (i = 0; i < n; ++i) {
+    ++j;
+    u = &j;
+  }
+}
+// CHECK-LABEL: define {{[^@]+}}@loop
+// CHECK-SAME: (i32 noundef [[N:%.*]]) local_unnamed_addr #[[ATTR0:[0-9]+]] {
+// CHECK-NEXT:  entry:
+// CHECK-NEXT:    [[J:%.*]] = alloca i32, align 4
+// CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i32 [[N]], 0
+// CHECK-NEXT:    br i1 [[CMP]], label [[SIMD_IF_THEN:%.*]], label [[SIMD_IF_END:%.*]]
+// CHECK:       simd.if.then:
+// CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr @j, align 4, !tbaa [[TBAA2:![0-9]+]]
+// CHECK-NEXT:    call void @llvm.lifetime.start.p0(i64 4, ptr nonnull [[J]]) #[[ATTR2:[0-9]+]]
+// CHECK-NEXT:    store ptr [[J]], ptr @u, align 8, !tbaa [[TBAA6:![0-9]+]], !llvm.access.group [[ACC_GRP8:![0-9]+]]
+// CHECK-NEXT:    [[INC_LE:%.*]] = add i32 [[TMP0]], [[N]]
+// CHECK-NEXT:    store i32 [[INC_LE]], ptr [[J]], align 4, !tbaa [[TBAA2]]
+// CHECK-NEXT:    store i32 [[INC_LE]], ptr @j, align 4, !tbaa [[TBAA2]]
+// CHECK-NEXT:    call void @llvm.lifetime.end.p0(i64 4, ptr nonnull [[J]]) #[[ATTR2]]
+// CHECK-NEXT:    br label [[SIMD_IF_END]]
+// CHECK:       simd.if.end:
+// CHECK-NEXT:    ret void
+//

clang/test/OpenMP/for_linear_codegen.cpp

 // CHECK1-NEXT:    [[TMP18:%.*]] = load i32, i32* [[DOTOMP_IV]], align 4
 // CHECK1-NEXT:    [[ADD9:%.*]] = add nsw i32 [[TMP18]], 1
 // CHECK1-NEXT:    store i32 [[ADD9]], i32* [[DOTOMP_IV]], align 4
 // CHECK1-NEXT:    br label [[OMP_INNER_FOR_COND]]
 // CHECK1:       omp.inner.for.end:
 // CHECK1-NEXT:    br label [[OMP_LOOP_EXIT:%.*]]
 // CHECK1:       omp.loop.exit:
 // CHECK1-NEXT:    call void @__kmpc_for_static_fini(%struct.ident_t* @[[GLOB2]], i32 [[TMP5]])
-// CHECK1-NEXT:    [[TMP19:%.*]] = bitcast i64* [[DOTLVAR__ADDR]] to i8*
-// CHECK1-NEXT:    call void @__kmpc_free(i32 [[TMP5]], i8* [[TMP19]], i8* inttoptr (i64 5 to i8*))
-// CHECK1-NEXT:    [[TMP20:%.*]] = load i32, i32* [[DOTOMP_IS_LAST]], align 4
-// CHECK1-NEXT:    [[TMP21:%.*]] = icmp ne i32 [[TMP20]], 0
-// CHECK1-NEXT:    br i1 [[TMP21]], label [[DOTOMP_LINEAR_PU:%.*]], label [[DOTOMP_LINEAR_PU_DONE:%.*]]
+// CHECK1-NEXT:    [[TMP19:%.*]] = load i32, i32* [[DOTOMP_IS_LAST]], align 4
+// CHECK1-NEXT:    [[TMP20:%.*]] = icmp ne i32 [[TMP19]], 0
+// CHECK1-NEXT:    br i1 [[TMP20]], label [[DOTOMP_LINEAR_PU:%.*]], label [[DOTOMP_LINEAR_PU_DONE:%.*]]
 // CHECK1:       .omp.linear.pu:
-// CHECK1-NEXT:    [[TMP22:%.*]] = load float*, float** [[PVAR2]], align 8
-// CHECK1-NEXT:    store float* [[TMP22]], float** [[TMP0]], align 8
-// CHECK1-NEXT:    [[TMP23:%.*]] = load i64, i64* [[DOTLVAR__ADDR]], align 8
-// CHECK1-NEXT:    store i64 [[TMP23]], i64* [[TMP1]], align 8
+// CHECK1-NEXT:    [[TMP21:%.*]] = load float*, float** [[PVAR2]], align 8
+// CHECK1-NEXT:    store float* [[TMP21]], float** [[TMP0]], align 8
+// CHECK1-NEXT:    [[TMP22:%.*]] = load i64, i64* [[DOTLVAR__ADDR]], align 8
+// CHECK1-NEXT:    store i64 [[TMP22]], i64* [[TMP1]], align 8
 // CHECK1-NEXT:    br label [[DOTOMP_LINEAR_PU_DONE]]
 // CHECK1:       .omp.linear.pu.done:
+// CHECK1-NEXT:    [[TMP23:%.*]] = bitcast i64* [[DOTLVAR__ADDR]] to i8*
+// CHECK1-NEXT:    call void @__kmpc_free(i32 [[TMP5]], i8* [[TMP23]], i8* inttoptr (i64 5 to i8*))
 // CHECK1-NEXT:    call void @__kmpc_barrier(%struct.ident_t* @[[GLOB1]], i32 [[TMP5]])
 // CHECK1-NEXT:    ret void
 //
 //
 // CHECK1-LABEL: define {{[^@]+}}@_Z5tmainIiET_v
 // CHECK1-SAME: () #[[ATTR5:[0-9]+]] {
 // CHECK1-NEXT:  entry:
 // CHECK1-NEXT:    [[TEST:%.*]] = alloca [[STRUCT_S_0:%.*]], align 4

clang/test/OpenMP/parallel_master_taskloop_simd_lastprivate_codegen.cpp

 // CHECK6-NEXT:    [[ADD6_I:%.*]] = add nsw i32 [[TMP36]], 1
 // CHECK6-NEXT:    store i32 [[ADD6_I]], i32* [[DOTOMP_IV_I]], align 4, !noalias !14, !llvm.access.group [[ACC_GRP15]]
 // CHECK6-NEXT:    br label [[OMP_INNER_FOR_COND_I]], !llvm.loop [[LOOP16:![0-9]+]]
 // CHECK6:       omp.inner.for.end.i:
 // CHECK6-NEXT:    [[TMP37:%.*]] = load i32, i32* [[DOTLITER__ADDR_I]], align 4, !noalias !14
 // CHECK6-NEXT:    [[TMP38:%.*]] = icmp ne i32 [[TMP37]], 0
 // CHECK6-NEXT:    br i1 [[TMP38]], label [[DOTOMP_LINEAR_PU_I:%.*]], label [[DOTOMP_OUTLINED__1_EXIT:%.*]]
 // CHECK6:       .omp.linear.pu.i:
-// CHECK6-NEXT:    [[TMP39:%.*]] = load i32, i32* [[J_I]], align 4, !noalias !14
-// CHECK6-NEXT:    store i32 [[TMP39]], i32* [[J_I]], align 4, !noalias !14
+// CHECK6-NEXT:    [[TMP39:%.*]] = getelementptr inbounds [[STRUCT_ANON]], %struct.anon* [[TMP20]], i32 0, i32 0
+// CHECK6-NEXT:    [[TMP40:%.*]] = load i32*, i32** [[TMP39]], align 8
+// CHECK6-NEXT:    [[TMP41:%.*]] = getelementptr inbounds [[STRUCT_ANON]], %struct.anon* [[TMP20]], i32 0, i32 1
+// CHECK6-NEXT:    [[TMP42:%.*]] = load i32*, i32** [[TMP41]], align 8
+// CHECK6-NEXT:    [[TMP43:%.*]] = load i32, i32* [[J_I]], align 4, !noalias !14
+// CHECK6-NEXT:    store i32 [[TMP43]], i32* [[TMP42]], align 4
 // CHECK6-NEXT:    br label [[DOTOMP_OUTLINED__1_EXIT]]
 // CHECK6:       .omp_outlined..1.exit:
 // CHECK6-NEXT:    ret i32 0
 //
 //
 // CHECK7-LABEL: define {{[^@]+}}@main
 // CHECK7-SAME: () #[[ATTR0:[0-9]+]] {
 // CHECK7-NEXT:  entry: