Index: lib/CodeGen/ManagedMemoryRewrite.cpp =================================================================== --- lib/CodeGen/ManagedMemoryRewrite.cpp +++ lib/CodeGen/ManagedMemoryRewrite.cpp @@ -317,6 +317,32 @@ } } +// Replace all uses of `Old` with `New`, even inside `ConstantExpr`. +// `replaceAllUsesWith` does replace values in `ConstantExpr`. This function +// actuall does replace it in `ConstantExpr`. The caveat is that if there is +// a use that is *outside* a function (say, at global declarations), we fail. +// So, this is meant to be used on values which we know will only be used +// within functions. +// +// This process works by looking through the uses of `Old`. If it finds a +// `ConstantExpr`, it recursively looks for the owning instruction. +// Then, it expands all the `ConstantExpr` to instructions and replaces +// `Old` with `New` in the expanded instructions. +static void replaceAllUsesAndConstantUses(Value *Old, Value *New, + PollyIRBuilder &Builder) { + SmallVector UserInstructions; + // Get all instructions that use array. We need to do this weird thing + // because `Constant`s that contain this array neeed to be expanded into + // instructions so that we can replace their parameters. `Constant`s cannot + // be edited easily, so we choose to convert all `Constant`s to + // `Instruction`s and handle all of the uses of `Array` uniformly. + for (Use &ArrayUse : Old->uses()) + getInstructionUsersOfValue(ArrayUse.getUser(), UserInstructions); + + for (Instruction *I : UserInstructions) + rewriteOldValToNew(I, Old, New, Builder); +} + class ManagedMemoryRewritePass : public ModulePass { public: static char ID; @@ -330,18 +356,22 @@ Function *Malloc = M.getFunction("malloc"); if (Malloc) { + PollyIRBuilder Builder(M.getContext()); Function *PollyMallocManaged = getOrCreatePollyMallocManaged(M); assert(PollyMallocManaged && "unable to create polly_mallocManaged"); - Malloc->replaceAllUsesWith(PollyMallocManaged); + + replaceAllUsesAndConstantUses(Malloc, PollyMallocManaged, Builder); Malloc->eraseFromParent(); } Function *Free = M.getFunction("free"); if (Free) { + PollyIRBuilder Builder(M.getContext()); Function *PollyFreeManaged = getOrCreatePollyFreeManaged(M); assert(PollyFreeManaged && "unable to create polly_freeManaged"); - Free->replaceAllUsesWith(PollyFreeManaged); + + replaceAllUsesAndConstantUses(Free, PollyFreeManaged, Builder); Free->eraseFromParent(); } Index: test/GPGPU/managed-memory-rewrite-malloc-free-inside-constexpr.ll =================================================================== --- /dev/null +++ test/GPGPU/managed-memory-rewrite-malloc-free-inside-constexpr.ll @@ -0,0 +1,95 @@ +; RUN: opt %loadPolly -polly-scops \ +; RUN: -analyze < %s | FileCheck %s --check-prefix=SCOP + +; RUN: opt %loadPolly -polly-codegen-ppcg \ +; RUN: -S -polly-acc-codegen-managed-memory \ +; RUN: -polly-acc-rewrite-managed-memory < %s | FileCheck %s --check-prefix=HOST-IR +; +; REQUIRES: pollyacc +; +; Check that we can correctly rewrite `malloc` to `polly_mallocManaged`, and +; `free` to `polly_freeManaged` with the `polly-acc-rewrite-managed-memory` +; pass, even inside `constantExpr`. This is necessary because a cookie cutter +; Inst->replaceUsesOfWith(...) call does not actually work, because this does +; not replace the instruction within a ConstantExpr. +; +; #include +; +; static const int N = 100; +; int* f(int *ToFree) { +; free(ToFree); +; int *A = (int *)malloc(sizeof(int) * N); +; for(int i = 0; i < N; i++) { +; A[i] = 42; +; } +; return A; +; +; } + +; SCOP: Function: f +; SCOP-NEXT: Region: %for.body---%for.end +; SCOP-NEXT: Max Loop Depth: 1 + +; SCOP: Arrays { +; SCOP-NEXT: i32 MemRef_tmp[*]; // Element size 4 +; SCOP-NEXT: } + +; // Check that polly_mallocManaged is declared and used correctly. +; HOST-IR: %1 = bitcast i8* (i64)* @polly_mallocManaged to i32* (i64)* +; HOST-IR: declare i8* @polly_mallocManaged(i64) + +; // Check that polly_freeManaged is declared and used correctly. +; HOST-IR call void @polly_freeManaged(i8* %toFree) +; HOST-IR: declare void @polly_freeManaged(i8*) + +; // Check that we remove the original malloc,free +; HOST-IR-NOT: declare i8* @malloc(i64) +; HOST-IR-NOT: declare void @free(i8*) + +target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128" +target triple = "x86_64-apple-macosx10.12.0" + +define i32* @f(i32 *%toFree) { +entry: + ; Free inside bitcast + call void bitcast (void (i8*)* @free to void (i32 *)*) (i32 * %toFree) + br label %entry.split + +entry.split: ; preds = %entry + ; malloc inside bitcast. + %tmp = call i32* bitcast (i8* (i64)* @malloc to i32* (i64)*) (i64 400) + br label %for.body + +for.body: ; preds = %entry.split, %for.body + %indvars.iv1 = phi i64 [ 0, %entry.split ], [ %indvars.iv.next, %for.body ] + %arrayidx = getelementptr inbounds i32, i32* %tmp, i64 %indvars.iv1 + store i32 42, i32* %arrayidx, align 4, !tbaa !3 + %indvars.iv.next = add nuw nsw i64 %indvars.iv1, 1 + %exitcond = icmp eq i64 %indvars.iv.next, 100 + br i1 %exitcond, label %for.end, label %for.body + +for.end: ; preds = %for.body + ret i32* %tmp +} + +; Function Attrs: argmemonly nounwind +declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #0 + +declare i8* @malloc(i64) +declare void @free(i8*) + +; Function Attrs: argmemonly nounwind +declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #0 + +attributes #0 = { argmemonly nounwind } + +!llvm.module.flags = !{!0, !1} +!llvm.ident = !{!2} + +!0 = !{i32 1, !"wchar_size", i32 4} +!1 = !{i32 7, !"PIC Level", i32 2} +!2 = !{!"clang version 6.0.0 (http://llvm.org/git/clang.git 6660f0d30ef23b3142a6b08f9f41aad3d47c084f) (http://llvm.org/git/llvm.git 052dd78cb30f77a05dc8bb06b851402c4b6c6587)"} +!3 = !{!4, !4, i64 0} +!4 = !{!"int", !5, i64 0} +!5 = !{!"omnipotent char", !6, i64 0} +!6 = !{!"Simple C/C++ TBAA"}