[Local][SimplifyCFG][GVN] Handle !nontemporal in combineMetadata

[dwoodwor-intel]

Summary

SimplifyCFG currently drops !nontemporal metadata when sinking common instructions. With this change, SimplifyCFG and similar transforms will preserve !nontemporal metadata as long as it is set on both original instructions.

[StephenFan]

It seems the change in patchReplacementInstruction is not tested. I think it's better to make it tested.

[dwoodwor-intel]

I've added a new GVN test to exercise the patchReplacementInstruction change, and I also added a GVNSink test to make sure this change works properly in that pass too even though it uses the same code path as SimplifyCFG. Are there other passes I should add tests for/update to better combine !nontemporal metadata?

[nikic]

LGTM

[dwoodwor-intel]

I don't have commit access; can you commit this for me? You can credit my name/email as Daniel Woodworth <[e-mail]>.

llvm/lib/Transforms/Utils/Local.cpp

switch (Kind) {
       case LLVMContext::MD_preserve_access_index:
         // Preserve !preserve.access.index in K.
         break;
       case LLVMContext::MD_noundef:
         // If K does move, keep noundef if it is present in both instructions.
         if (DoesKMove)
           K->setMetadata(Kind, JMD);
         break;
+      case LLVMContext::MD_nontemporal:
+        // Preserve !nontemporal if it is present on both instructions.
+        K->setMetadata(Kind, JMD);
+        break;
     }
   }
   // Set !invariant.group from J if J has it. If both instructions have it
   // then we will just pick it from J - even when they are different.
   // Also make sure that K is load or store - f.e. combining bitcast with load
   // could produce bitcast with invariant.group metadata, which is invalid.
   // FIXME: we should try to preserve both invariant.group md if they are
   // different, but right now instruction can only have one invariant.group.
   if (auto *JMD = J->getMetadata(LLVMContext::MD_invariant_group))
     if (isa<LoadInst>(K) || isa<StoreInst>(K))
       K->setMetadata(LLVMContext::MD_invariant_group, JMD);
 }
 
 void llvm::combineMetadataForCSE(Instruction *K, const Instruction *J,
                                  bool KDominatesJ) {
   unsigned KnownIDs[] = {
       LLVMContext::MD_tbaa,            LLVMContext::MD_alias_scope,
       LLVMContext::MD_noalias,         LLVMContext::MD_range,
       LLVMContext::MD_invariant_load,  LLVMContext::MD_nonnull,
       LLVMContext::MD_invariant_group, LLVMContext::MD_align,
       LLVMContext::MD_dereferenceable,
       LLVMContext::MD_dereferenceable_or_null,
-      LLVMContext::MD_access_group,    LLVMContext::MD_preserve_access_index};
+      LLVMContext::MD_access_group,    LLVMContext::MD_preserve_access_index,
+      LLVMContext::MD_nontemporal};
   combineMetadata(K, J, KnownIDs, KDominatesJ);
 }
 
 void llvm::copyMetadataForLoad(LoadInst &Dest, const LoadInst &Source) {
   SmallVector<std::pair<unsigned, MDNode *>, 8> MD;
   Source.getAllMetadata(MD);
   MDBuilder MDB(Dest.getContext());
   Type *NewType = Dest.getType();
void llvm::patchReplacementInstruction(Instruction *I, Value *Repl) {
   // regions, and so we need a conservative combination of the noalias
   // scopes.
   static const unsigned KnownIDs[] = {
       LLVMContext::MD_tbaa,            LLVMContext::MD_alias_scope,
       LLVMContext::MD_noalias,         LLVMContext::MD_range,
       LLVMContext::MD_fpmath,          LLVMContext::MD_invariant_load,
       LLVMContext::MD_invariant_group, LLVMContext::MD_nonnull,
       LLVMContext::MD_access_group,    LLVMContext::MD_preserve_access_index,
-      LLVMContext::MD_noundef};
+      LLVMContext::MD_noundef,         LLVMContext::MD_nontemporal};
   combineMetadata(ReplInst, I, KnownIDs, false);
 }
 
 template <typename RootType, typename DominatesFn>
 static unsigned replaceDominatedUsesWith(Value *From, Value *To,
                                          const RootType &Root,
                                          const DominatesFn &Dominates) {
   assert(From->getType() == To->getType());

llvm/test/Transforms/GVN/nontemporal.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=gvn -S < %s | FileCheck %s
+
+; Check that !nontemporal metadata is preserved only if both original
+; instructions are marked.
+define i64 @one_nontemporal(ptr %p) {
+; CHECK-LABEL: @one_nontemporal(
+; CHECK-NEXT:    [[A:%.*]] = load i64, ptr [[P:%.*]], align 4
+; CHECK-NEXT:    [[C:%.*]] = add i64 [[A]], [[A]]
+; CHECK-NEXT:    ret i64 [[C]]
+;
+  %a = load i64, ptr %p
+  %b = load i64, ptr %p, !nontemporal !0
+  %c = add i64 %a, %b
+  ret i64 %c
+}
+
+define i64 @both_nontemporal(ptr %p) {
+; CHECK-LABEL: @both_nontemporal(
+; CHECK-NEXT:    [[A:%.*]] = load i64, ptr [[P:%.*]], align 4, !nontemporal !0
+; CHECK-NEXT:    [[C:%.*]] = add i64 [[A]], [[A]]
+; CHECK-NEXT:    ret i64 [[C]]
+;
+  %a = load i64, ptr %p, !nontemporal !0
+  %b = load i64, ptr %p, !nontemporal !0
+  %c = add i64 %a, %b
+  ret i64 %c
+}
+
+!0 = !{i32 1}

llvm/test/Transforms/GVNSink/sink-combine-metadata.ll

if.else:
   %bb= load ptr, ptr %b, !nonnull !0
   br label %if.end
 
 if.end:
   %c = phi ptr [ %aa, %if.then ], [ %bb, %if.else ]
   ret ptr %c
 }
 
+; Check that nontemporal markings are propagated if both original stores are
+; marked.
+; CHECK-LABEL: @nontemporal(
+; CHECK-LABEL: if.end:
+; CHECK: !nontemporal
+; CHECK: ret void
+define void @nontemporal(i1 zeroext %flag, ptr %p) {
+entry:
+  br i1 %flag, label %if.then, label %if.else
+
+if.then:
+  %a = load ptr, ptr %p
+  store ptr %p, ptr %a, align 8, !nontemporal !1
+  br label %if.end
+
+if.else:
+  %b = load ptr, ptr %p
+  store ptr %p, ptr %b, align 8, !nontemporal !1
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+; CHECK-LABEL: @nontemporal_mismatch(
+; CHECK-NOT: !nontemporal
+define void @nontemporal_mismatch(i1 zeroext %flag, ptr %p) {
+entry:
+  br i1 %flag, label %if.then, label %if.else
+
+if.then:
+  %a = load ptr, ptr %p
+  store ptr %p, ptr %a, align 8
+  br label %if.end
+
+if.else:
+  %b = load ptr, ptr %p
+  store ptr %p, ptr %b, align 8, !nontemporal !1
+  br label %if.end
+
+if.end:
+  ret void
+}
 
 !0 = !{}
+!1 = !{i32 1}

llvm/test/Transforms/SimplifyCFG/X86/sink-common-code.ll

 
 uncond_pred1:
   call void @direct_callee()
   br label %end
 
 end:
   ret void
 }
+
+define void @nontemporal(ptr %ptr, i1 %cond) {
+; CHECK-LABEL: @nontemporal(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i64 0, ptr [[PTR:%.*]], align 8, !nontemporal !7
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 %cond, label %if.then, label %if.else
+
+if.then:
+  store i64 0, ptr %ptr, align 8, !nontemporal !12
+  br label %if.end
+
+if.else:
+  store i64 0, ptr %ptr, align 8, !nontemporal !12
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+define void @nontemporal_mismatch(ptr %ptr, i1 %cond) {
+; CHECK-LABEL: @nontemporal_mismatch(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i64 0, ptr [[PTR:%.*]], align 8
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 %cond, label %if.then, label %if.else
+
+if.then:
+  store i64 0, ptr %ptr, align 8, !nontemporal !12
+  br label %if.end
+
+if.else:
+  store i64 0, ptr %ptr, align 8
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+!12 = !{i32 1}