Improve MPI rank/size simplification (rust-lang#484)

wsmoses · vchuravy · web-flow · commit 9e48a51bf7b7 · 2022-02-04T17:43:16.000-05:00
* make MPI wrapper actually useful

* Add test

* Improve julia mpi aliasing

Co-authored-by: Valentin Churavy &lt;v.churavy@gmail.com&gt;
diff --git a/enzyme/Enzyme/Enzyme.cpp b/enzyme/Enzyme/Enzyme.cpp
@@ -230,7 +230,8 @@ static void handleKnownFunctions(llvm::Function &F) {
     }
     F.addParamAttr(6, Attribute::WriteOnly);
   }
-  if (F.getName() == "MPI_Comm_rank") {
+  if (F.getName() == "MPI_Comm_rank" || F.getName() == "PMPI_Comm_rank" ||
+      F.getName() == "MPI_Comm_size" || F.getName() == "PMPI_Comm_size") {
     F.addFnAttr(Attribute::InaccessibleMemOrArgMemOnly);
     F.addFnAttr(Attribute::NoUnwind);
     F.addFnAttr(Attribute::NoRecurse);
@@ -243,8 +244,10 @@ static void handleKnownFunctions(llvm::Function &F) {
       F.addParamAttr(0, Attribute::NoCapture);
       F.addParamAttr(0, Attribute::ReadOnly);
     }
-    F.addParamAttr(1, Attribute::WriteOnly);
-    F.addParamAttr(1, Attribute::NoCapture);
+    if (F.getFunctionType()->getParamType(1)->isPointerTy()) {
+      F.addParamAttr(1, Attribute::WriteOnly);
+      F.addParamAttr(1, Attribute::NoCapture);
+    }
   }
   if (F.getName() == "MPI_Wait") {
     F.addFnAttr(Attribute::NoUnwind);
diff --git a/enzyme/Enzyme/FunctionUtils.cpp b/enzyme/Enzyme/FunctionUtils.cpp
@@ -657,6 +657,12 @@ Function *CreateMPIWrapper(Function *F) {
   Function *W = Function::Create(FT, GlobalVariable::InternalLinkage, name,
                                  F->getParent());
   llvm::Attribute::AttrKind attrs[] = {
+#if LLVM_VERSION_MAJOR >= 9
+    Attribute::WillReturn,
+#endif
+#if LLVM_VERSION_MAJOR >= 12
+    Attribute::MustProgress,
+#endif
     Attribute::ReadOnly,
     Attribute::Speculatable,
     Attribute::NoUnwind,
@@ -684,6 +690,12 @@ Function *CreateMPIWrapper(Function *F) {
   IRBuilder<> B(entry);
   auto alloc = B.CreateAlloca(F->getReturnType());
   Value *args[] = {W->arg_begin(), alloc};
+
+  auto T = F->getFunctionType()->getParamType(1);
+  if (!isa<PointerType>(T)) {
+    assert(isa<IntegerType>(T));
+    args[1] = B.CreatePtrToInt(args[1], T);
+  }
   B.CreateCall(F, args);
 #if LLVM_VERSION_MAJOR > 7
   B.CreateRet(B.CreateLoad(F->getReturnType(), alloc));
@@ -692,7 +704,8 @@ Function *CreateMPIWrapper(Function *F) {
 #endif
   return W;
 }
-static void SimplifyMPIQueries(Function &NewF) {
+static void SimplifyMPIQueries(Function &NewF, FunctionAnalysisManager &FAM) {
+  DominatorTree &DT = FAM.getResult<DominatorTreeAnalysis>(NewF);
   SmallVector<CallInst *, 4> Todo;
   SmallVector<CallInst *, 0> OMPBounds;
   for (auto &BB : NewF) {
@@ -703,10 +716,8 @@ static void SimplifyMPIQueries(Function &NewF) {
           continue;
         if (Fn->getName() == "MPI_Comm_rank" ||
             Fn->getName() == "PMPI_Comm_rank" ||
-            Fn->getName() == "MPI_Comm_size") {
-          if (!CI->use_empty()) {
-            continue;
-          }
+            Fn->getName() == "MPI_Comm_size" ||
+            Fn->getName() == "PMPI_Comm_size") {
           Todo.push_back(CI);
         }
         if (Fn->getName() == "__kmpc_for_static_init_4" ||
@@ -721,9 +732,56 @@ static void SimplifyMPIQueries(Function &NewF) {
   for (auto CI : Todo) {
     IRBuilder<> B(CI);
     Value *arg[] = {CI->getArgOperand(0)};
-    auto res = B.CreateCall(CreateMPIWrapper(CI->getCalledFunction()), arg);
-    B.CreateStore(res, CI->getArgOperand(1));
+    SmallVector<OperandBundleDef, 2> Defs;
+    CI->getOperandBundlesAsDefs(Defs);
+    auto res =
+        B.CreateCall(CreateMPIWrapper(CI->getCalledFunction()), arg, Defs);
+    Value *storePointer = CI->getArgOperand(1);
+
+    // Comm_rank and Comm_size return Err, assume 0 is success
+    CI->replaceAllUsesWith(ConstantInt::get(CI->getType(), 0));
     CI->eraseFromParent();
+
+    while (auto Cast = dyn_cast<CastInst>(storePointer)) {
+      storePointer = Cast->getOperand(0);
+      if (Cast->use_empty())
+        Cast->eraseFromParent();
+    }
+
+    B.SetInsertPoint(res);
+
+    if (auto PT = dyn_cast<PointerType>(storePointer->getType())) {
+      if (PT->getElementType() != res->getType())
+        storePointer = B.CreateBitCast(
+            storePointer,
+            PointerType::get(res->getType(), PT->getAddressSpace()));
+    } else {
+      assert(isa<IntegerType>(storePointer->getType()));
+      storePointer = B.CreateIntToPtr(storePointer,
+                                      PointerType::getUnqual(res->getType()));
+    }
+    if (isa<AllocaInst>(storePointer)) {
+      // If this is only loaded from, immedaitely replace
+      // Immediately replace all dominated stores.
+      SmallVector<LoadInst *, 2> LI;
+      bool nonload = false;
+      for (auto &U : storePointer->uses()) {
+        if (auto L = dyn_cast<LoadInst>(U.getUser())) {
+          LI.push_back(L);
+        } else
+          nonload = true;
+      }
+      if (!nonload) {
+        for (auto L : LI) {
+          if (DT.dominates(res, L)) {
+            L->replaceAllUsesWith(res);
+            L->eraseFromParent();
+          }
+        }
+      }
+    }
+    B.SetInsertPoint(res->getNextNode());
+    B.CreateStore(res, storePointer);
   }
   for (auto Bound : OMPBounds) {
     for (int i = 4; i <= 6; i++) {
@@ -747,6 +805,13 @@ static void SimplifyMPIQueries(Function &NewF) {
       Bound->addParamAttr(i, Attribute::NoCapture);
     }
   }
+  PreservedAnalyses PA;
+  PA.preserve<AssumptionAnalysis>();
+  PA.preserve<TargetLibraryAnalysis>();
+  PA.preserve<LoopAnalysis>();
+  PA.preserve<DominatorTreeAnalysis>();
+  PA.preserve<PostDominatorTreeAnalysis>();
+  FAM.invalidate(NewF, PA);
 }
 
 /// Perform recursive inlinining on NewF up to the given limit
@@ -1104,11 +1169,7 @@ Function *PreProcessCache::preprocessForClone(Function *F,
       ConstantFoldTerminator(BE);
   }
 
-  {
-    SimplifyMPIQueries(*NewF);
-    PreservedAnalyses PA;
-    FAM.invalidate(*NewF, PA);
-  }
+  SimplifyMPIQueries(*NewF, FAM);
 
   if (EnzymeLowerGlobals) {
     std::vector<CallInst *> Calls;
diff --git a/enzyme/test/Enzyme/ReverseMode/mpirep.ll b/enzyme/test/Enzyme/ReverseMode/mpirep.ll
@@ -0,0 +1,102 @@
+; RUN: %opt < %s %loadEnzyme -enzyme -enzyme-preopt=false -mem2reg -instsimplify -simplifycfg -correlated-propagation -adce -S | FileCheck %s
+source_filename = "text"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+
+; Function Attrs: cold noreturn nounwind
+declare void @llvm.trap() #0
+
+declare dso_local i32 @MPI_Comm_rank(i64, i64)
+
+define double @sum(double* %arg, i64 %comm) {
+bb:
+  %alloc = alloca i32, align 8
+  %i5 = ptrtoint i32* %alloc to i64
+  br label %bb11
+
+bb11:                                             ; preds = %bb
+  %idx = phi i64 [ 0, %bb ], [ %inc, %bb22 ]
+  %sum = phi double [ 0.000000e+00, %bb ], [ %add, %bb22 ]
+  %inc = add i64 %idx, 1
+  %i13 = getelementptr inbounds double, double* %arg, i64 %idx
+  %i14 = load double, double* %i13, align 8
+  %i16 = fmul double %i14, %i14
+  %i19 = call i32 @MPI_Comm_rank(i64 %comm, i64 %i5)
+  %ld = load i32, i32* %alloc
+  %cf = uitofp i32 %ld to double
+  %mm = fmul double %i16, %cf
+  %add = fadd double %sum, %mm
+  %i20 = icmp eq i32 %i19, 0
+  br i1 %i20, label %bb22, label %bb21
+
+bb21:                                             ; preds = %bb11, %bb
+  call void @llvm.trap() #1
+  unreachable
+
+bb22:
+  %cmp = icmp eq i64 %idx, 9
+  br i1 %cmp, label %exit, label %bb11
+
+exit:
+  ret double %add
+}
+
+define void @dsum(double* %x, double* %xp, i64 %n) {
+entry:
+  %0 = tail call double (double (double*, i64)*, ...) @__enzyme_autodiff(double (double*, i64)* nonnull @sum, double* %x, double* %xp, i64 %n)
+  ret void
+}
+
+declare double @__enzyme_autodiff(double (double*, i64)*, ...)
+
+attributes #0 = { cold noreturn nounwind }
+attributes #1 = { noreturn }
+
+; CHECK: define internal void @diffesum(double* %arg, double* %"arg'", i64 %comm, double %differeturn) 
+; CHECK-NEXT: bb:
+; CHECK-NEXT:   %0 = alloca i32
+; CHECK-NEXT:   %1 = alloca i32
+; CHECK-NEXT:   br label %bb11
+
+; CHECK: bb11:                                             ; preds = %bb11, %bb
+; CHECK-NEXT:   %iv = phi i64 [ %iv.next, %bb11 ], [ 0, %bb ]
+; CHECK-NEXT:   %iv.next = add nuw nsw i64 %iv, 1
+; CHECK-NEXT:   %2 = bitcast i32* %1 to i8*
+; CHECK-NEXT:   call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %2)
+; CHECK-NEXT:   %3 = ptrtoint i32* %1 to i64
+; CHECK-NEXT:   %4 = call i32 @MPI_Comm_rank(i64 %comm, i64 %3) 
+; CHECK-NEXT:   %5 = bitcast i32* %1 to i8*
+; CHECK-NEXT:   call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %5)
+; CHECK-NEXT:   %cmp = icmp eq i64 %iv, 9
+; CHECK-NEXT:   br i1 %cmp, label %invertbb22, label %bb11
+
+; CHECK: invertbb:                                         ; preds = %invertbb22
+; CHECK-NEXT:   ret void
+
+; CHECK: incinvertbb11:                                    ; preds = %invertbb22
+; CHECK-NEXT:   %6 = add nsw i64 %"iv'ac.0", -1
+; CHECK-NEXT:   br label %invertbb22
+
+; CHECK: invertbb22:                                       ; preds = %bb11, %incinvertbb11
+; CHECK-NEXT:   %"iv'ac.0" = phi i64 [ %6, %incinvertbb11 ], [ 9, %bb11 ]
+; CHECK-NEXT:   %7 = bitcast i32* %0 to i8*
+; CHECK-NEXT:   call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %7)
+; CHECK-NEXT:   %8 = ptrtoint i32* %0 to i64
+; CHECK-NEXT:   %9 = call i32 @MPI_Comm_rank(i64 %comm, i64 %8) 
+; CHECK-NEXT:   %10 = load i32, i32* %0
+; CHECK-NEXT:   %11 = bitcast i32* %0 to i8*
+; CHECK-NEXT:   call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %11)
+; CHECK-NEXT:   %cf_unwrap = uitofp i32 %10 to double
+; CHECK-NEXT:   %m0diffei16 = fmul fast double %differeturn, %cf_unwrap
+; CHECK-NEXT:   %i13_unwrap = getelementptr inbounds double, double* %arg, i64 %"iv'ac.0"
+; CHECK-NEXT:   %i14_unwrap = load double, double* %i13_unwrap, align 8, !invariant.group !0
+; CHECK-NEXT:   %m0diffei14 = fmul fast double %m0diffei16, %i14_unwrap
+; CHECK-NEXT:   %m1diffei14 = fmul fast double %m0diffei16, %i14_unwrap
+; CHECK-NEXT:   %12 = fadd fast double %m0diffei14, %m1diffei14
+; CHECK-NEXT:   %"i13'ipg_unwrap" = getelementptr inbounds double, double* %"arg'", i64 %"iv'ac.0"
+; CHECK-NEXT:   %13 = load double, double* %"i13'ipg_unwrap", align 8
+; CHECK-NEXT:   %14 = fadd fast double %13, %12
+; CHECK-NEXT:   store double %14, double* %"i13'ipg_unwrap", align 8
+; CHECK-NEXT:   %15 = icmp eq i64 %"iv'ac.0", 0
+; CHECK-NEXT:   br i1 %15, label %invertbb, label %incinvertbb11
+; CHECK-NEXT: }

Original file line number	Diff line number	Diff line change
`@@ -230,7 +230,8 @@ static void handleKnownFunctions(llvm::Function &F) {`
`230`	`230`	`}`
`231`	`231`	`F.addParamAttr(6, Attribute::WriteOnly);`
`232`	`232`	`}`
`233`		`- if (F.getName() == "MPI_Comm_rank") {`
	`233`	`+ if (F.getName() == "MPI_Comm_rank" \|\| F.getName() == "PMPI_Comm_rank" \|\|`
	`234`	`+ F.getName() == "MPI_Comm_size" \|\| F.getName() == "PMPI_Comm_size") {`
`234`	`235`	`F.addFnAttr(Attribute::InaccessibleMemOrArgMemOnly);`
`235`	`236`	`F.addFnAttr(Attribute::NoUnwind);`
`236`	`237`	`F.addFnAttr(Attribute::NoRecurse);`
`@@ -243,8 +244,10 @@ static void handleKnownFunctions(llvm::Function &F) {`
`243`	`244`	`F.addParamAttr(0, Attribute::NoCapture);`
`244`	`245`	`F.addParamAttr(0, Attribute::ReadOnly);`
`245`	`246`	`}`
`246`		`- F.addParamAttr(1, Attribute::WriteOnly);`
`247`		`- F.addParamAttr(1, Attribute::NoCapture);`
	`247`	`+ if (F.getFunctionType()->getParamType(1)->isPointerTy()) {`
	`248`	`+ F.addParamAttr(1, Attribute::WriteOnly);`
	`249`	`+ F.addParamAttr(1, Attribute::NoCapture);`
	`250`	`+ }`
`248`	`251`	`}`
`249`	`252`	`if (F.getName() == "MPI_Wait") {`
`250`	`253`	`F.addFnAttr(Attribute::NoUnwind);`