[dfsan] Propagate origins at non-memory/phi/call instructions

[stephan.yichao.zhao]

Summary

This is a part of https://reviews.llvm.org/D95835.

[morehouse]

Please fix lints.

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

563

Zero appears unused.

[stephan.yichao.zhao]

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

563

We will use it when propagating labels at load.

[stephan.yichao.zhao]

fixed lint

llvm/lib/Transforms/Instrumentation/DataFlowSanitizer.cpp

struct DFSanFunction {
   Value *getArgTLS(Type *T, unsigned ArgOffset, IRBuilder<> &IRB);
 
   /// Computes the shadow address for a return value.
   Value *getRetvalTLS(Type *T, IRBuilder<> &IRB);
 
   /// Computes the origin address for a given function argument.
   ///
   /// Origin = ArgOriginTLS[ArgNo].
-  // Value *getArgOriginTLS(unsigned ArgNo, IRBuilder<> &IRB);
+  Value *getArgOriginTLS(unsigned ArgNo, IRBuilder<> &IRB);
 
   /// Computes the origin address for a return value.
-  // Value *getRetvalOriginTLS();
+  Value *getRetvalOriginTLS();
+
+  Value *getOrigin(Value *V);
+  void setOrigin(Instruction *I, Value *Origin);
+  /// Generates IR to compute the origin of the last operand with a taint label.
+  Value *combineOperandOrigins(Instruction *Inst);
+  /// Before the instruction Pos, generates IR to compute the last origin with a
+  /// taint label. Labels and origins are from vectors Shadows and Origins
+  /// correspondingly. The generated IR is like
+  ///   Sn-1 != Zero ? On-1: ... S2 != Zero ? O2: S1 != Zero ? O1: O0
+  /// When Zero is nullptr, it uses ZeroPrimitiveShadow. Otherwise it can be
+  /// zeros with other bitwidths.
+  Value *combineOrigins(const std::vector<Value *> &Shadows,
+                        const std::vector<Value *> &Origins, Instruction *Pos,
+                        ConstantInt *Zero = nullptr);
 
-  // Value *getOrigin(Value *V);
-  // void setOrigin(Instruction *I, Value *Origin);
   Value *getShadow(Value *V);
   void setShadow(Instruction *I, Value *Shadow);
   /// Generates IR to compute the union of the two given shadows, inserting it
   /// before Pos. The combined value is with primitive type.
   Value *combineShadows(Value *V1, Value *V2, Instruction *Pos);
   /// Combines the shadow values of V1 and V2, then converts the combined value
   /// with primitive type into a shadow value with the original type T.
   Value *combineShadowsThenConvert(Type *T, Value *V1, Value *V2,
public:
   DFSanFunction &DFSF;
 
   DFSanVisitor(DFSanFunction &DFSF) : DFSF(DFSF) {}
 
   const DataLayout &getDataLayout() const {
     return DFSF.F->getParent()->getDataLayout();
   }
 
-  // Combines shadow values for all of I's operands. Returns the combined shadow
-  // value.
-  Value *visitOperandShadowInst(Instruction &I);
+  // Combines shadow values and origins for all of I's operands.
+  void visitInstOperands(Instruction &I);
 
   void visitUnaryOperator(UnaryOperator &UO);
   void visitBinaryOperator(BinaryOperator &BO);
   void visitCastInst(CastInst &CI);
   void visitCmpInst(CmpInst &CI);
   void visitGetElementPtrInst(GetElementPtrInst &GEPI);
   void visitLoadInst(LoadInst &LI);
   void visitStoreInst(StoreInst &SI);
   void visitReturnInst(ReturnInst &RI);
   void visitCallBase(CallBase &CB);
   void visitPHINode(PHINode &PN);
   void visitExtractElementInst(ExtractElementInst &I);
   void visitInsertElementInst(InsertElementInst &I);
   void visitShuffleVectorInst(ShuffleVectorInst &I);
   void visitExtractValueInst(ExtractValueInst &I);
   void visitInsertValueInst(InsertValueInst &I);
   void visitAllocaInst(AllocaInst &I);
   void visitSelectInst(SelectInst &I);
   void visitMemSetInst(MemSetInst &I);
   void visitMemTransferInst(MemTransferInst &I);
 
 private:
   // Returns false when this is an invoke of a custom function.
   bool visitWrappedCallBase(Function &F, CallBase &CB);
+
+  // Combines origins for all of I's operands.
+  void visitInstOperandOrigins(Instruction &I);
 };
 
 } // end anonymous namespace
 
 DataFlowSanitizer::DataFlowSanitizer(
     const std::vector<std::string> &ABIListFiles) {
   std::vector<std::string> AllABIListFiles(std::move(ABIListFiles));
   llvm::append_range(AllABIListFiles, ClABIListFiles);
Value *DFSanFunction::getArgTLS(Type *T, unsigned ArgOffset, IRBuilder<> &IRB) {
   return IRB.CreateIntToPtr(Base, PointerType::get(DFS.getShadowTy(T), 0),
                             "_dfsarg");
 }
 
 Value *DFSanFunction::getRetvalTLS(Type *T, IRBuilder<> &IRB) {
   return IRB.CreatePointerCast(
       DFS.RetvalTLS, PointerType::get(DFS.getShadowTy(T), 0), "_dfsret");
 }
-/*
+
 Value *DFSanFunction::getRetvalOriginTLS() { return DFS.RetvalOriginTLS; }
 
 Value *DFSanFunction::getArgOriginTLS(unsigned ArgNo, IRBuilder<> &IRB) {
   return IRB.CreateConstGEP2_64(DFS.ArgOriginTLSTy, DFS.ArgOriginTLS, 0, ArgNo,
                                 "_dfsarg_o");
 }
 
 Value *DFSanFunction::getOrigin(Value *V) {

 
 void DFSanFunction::setOrigin(Instruction *I, Value *Origin) {
   if (!DFS.shouldTrackOrigins())
     return;
   assert(!ValOriginMap.count(I));
   assert(Origin->getType() == DFS.OriginTy);
   ValOriginMap[I] = Origin;
 }
-*/
+
 Value *DFSanFunction::getShadowForTLSArgument(Argument *A) {
   unsigned ArgOffset = 0;
   const DataLayout &DL = F->getParent()->getDataLayout();
   for (auto &FArg : F->args()) {
     if (!FArg.getType()->isSized()) {
       if (A == &FArg)
         break;
       continue;
Value *DFSanFunction::combineOperandShadows(Instruction *Inst) {
 
   Value *Shadow = getShadow(Inst->getOperand(0));
   for (unsigned i = 1, n = Inst->getNumOperands(); i != n; ++i) {
     Shadow = combineShadows(Shadow, getShadow(Inst->getOperand(i)), Inst);
   }
   return expandFromPrimitiveShadow(Inst->getType(), Shadow, Inst);
 }
 
-Value *DFSanVisitor::visitOperandShadowInst(Instruction &I) {
+void DFSanVisitor::visitInstOperands(Instruction &I) {
   Value *CombinedShadow = DFSF.combineOperandShadows(&I);
   DFSF.setShadow(&I, CombinedShadow);
-  return CombinedShadow;
+  visitInstOperandOrigins(I);
+}
+
+Value *DFSanFunction::combineOrigins(const std::vector<Value *> &Shadows,
+                                     const std::vector<Value *> &Origins,
+                                     Instruction *Pos, ConstantInt *Zero) {
+  assert(Shadows.size() == Origins.size());
+  size_t Size = Origins.size();
+  if (Size == 0)
+    return DFS.ZeroOrigin;
+  Value *Origin = nullptr;
+  if (!Zero)
+    Zero = DFS.ZeroPrimitiveShadow;
+  for (size_t i = 0; i != Size; ++i) {
+    Value *OpOrigin = Origins[i];
+    Constant *ConstOpOrigin = dyn_cast<Constant>(OpOrigin);
+    if (ConstOpOrigin && ConstOpOrigin->isNullValue())
+      continue;
+    if (!Origin) {
+      Origin = OpOrigin;
+      continue;
+    }
+    Value *OpShadow = Shadows[i];
+    Value *primitiveShadow = collapseToPrimitiveShadow(OpShadow, Pos);
+    IRBuilder<> IRB(Pos);
+    Value *Cond = IRB.CreateICmpNE(primitiveShadow, Zero);
+    Origin = IRB.CreateSelect(Cond, OpOrigin, Origin);
+  }
+  return Origin ? Origin : DFS.ZeroOrigin;
+}
+
+Value *DFSanFunction::combineOperandOrigins(Instruction *Inst) {
+  size_t Size = Inst->getNumOperands();
+  std::vector<Value *> Shadows(Size);
+  std::vector<Value *> Origins(Size);
+  for (unsigned I = 0; I != Size; ++I) {
+    Shadows[I] = getShadow(Inst->getOperand(I));
+    Origins[I] = getOrigin(Inst->getOperand(I));
+  }
+  return combineOrigins(Shadows, Origins, Inst);
+}
+
+void DFSanVisitor::visitInstOperandOrigins(Instruction &I) {
+  if (!DFSF.DFS.shouldTrackOrigins())
+    return;
+  Value *CombinedOrigin = DFSF.combineOperandOrigins(&I);
+  DFSF.setOrigin(&I, CombinedOrigin);
 }
 
 Value *DFSanFunction::loadFast16ShadowFast(Value *ShadowAddr, uint64_t Size,
                                            Align ShadowAlign,
                                            Instruction *Pos) {
   // First OR all the WideShadows, then OR individual shadows within the
   // combined WideShadow. This is fewer instructions than ORing shadows
   // individually.
void DFSanVisitor::visitStoreInst(StoreInst &SI) {
   if (ClEventCallbacks) {
     IRBuilder<> IRB(&SI);
     Value *Addr8 = IRB.CreateBitCast(SI.getPointerOperand(), DFSF.DFS.Int8Ptr);
     IRB.CreateCall(DFSF.DFS.DFSanStoreCallbackFn, {PrimitiveShadow, Addr8});
   }
 }
 
 void DFSanVisitor::visitUnaryOperator(UnaryOperator &UO) {
-  visitOperandShadowInst(UO);
+  visitInstOperands(UO);
 }
 
 void DFSanVisitor::visitBinaryOperator(BinaryOperator &BO) {
-  visitOperandShadowInst(BO);
+  visitInstOperands(BO);
 }
 
-void DFSanVisitor::visitCastInst(CastInst &CI) { visitOperandShadowInst(CI); }
+void DFSanVisitor::visitCastInst(CastInst &CI) { visitInstOperands(CI); }
 
 void DFSanVisitor::visitCmpInst(CmpInst &CI) {
-  Value *CombinedShadow = visitOperandShadowInst(CI);
+  visitInstOperands(CI);
   if (ClEventCallbacks) {
     IRBuilder<> IRB(&CI);
+    Value *CombinedShadow = DFSF.getShadow(&CI);
     IRB.CreateCall(DFSF.DFS.DFSanCmpCallbackFn, CombinedShadow);
   }
 }
 
 void DFSanVisitor::visitGetElementPtrInst(GetElementPtrInst &GEPI) {
-  visitOperandShadowInst(GEPI);
+  visitInstOperands(GEPI);
 }
 
 void DFSanVisitor::visitExtractElementInst(ExtractElementInst &I) {
-  visitOperandShadowInst(I);
+  visitInstOperands(I);
 }
 
 void DFSanVisitor::visitInsertElementInst(InsertElementInst &I) {
-  visitOperandShadowInst(I);
+  visitInstOperands(I);
 }
 
 void DFSanVisitor::visitShuffleVectorInst(ShuffleVectorInst &I) {
-  visitOperandShadowInst(I);
+  visitInstOperands(I);
 }
 
 void DFSanVisitor::visitExtractValueInst(ExtractValueInst &I) {
   if (!DFSF.DFS.shouldTrackFieldsAndIndices()) {
-    visitOperandShadowInst(I);
+    visitInstOperands(I);
     return;
   }
 
   IRBuilder<> IRB(&I);
   Value *Agg = I.getAggregateOperand();
   Value *AggShadow = DFSF.getShadow(Agg);
   Value *ResShadow = IRB.CreateExtractValue(AggShadow, I.getIndices());
   DFSF.setShadow(&I, ResShadow);
+  visitInstOperandOrigins(I);
 }
 
 void DFSanVisitor::visitInsertValueInst(InsertValueInst &I) {
   if (!DFSF.DFS.shouldTrackFieldsAndIndices()) {
-    visitOperandShadowInst(I);
+    visitInstOperands(I);
     return;
   }
 
   IRBuilder<> IRB(&I);
   Value *AggShadow = DFSF.getShadow(I.getAggregateOperand());
   Value *InsShadow = DFSF.getShadow(I.getInsertedValueOperand());
   Value *Res = IRB.CreateInsertValue(AggShadow, InsShadow, I.getIndices());
   DFSF.setShadow(&I, Res);
+  visitInstOperandOrigins(I);
 }
 
 void DFSanVisitor::visitAllocaInst(AllocaInst &I) {
   bool AllLoadsStores = true;
   for (User *U : I.users()) {
     if (isa<LoadInst>(U))
       continue;
 
void DFSanVisitor::visitAllocaInst(AllocaInst &I) {
   DFSF.setShadow(&I, DFSF.DFS.ZeroPrimitiveShadow);
 }
 
 void DFSanVisitor::visitSelectInst(SelectInst &I) {
   Value *CondShadow = DFSF.getShadow(I.getCondition());
   Value *TrueShadow = DFSF.getShadow(I.getTrueValue());
   Value *FalseShadow = DFSF.getShadow(I.getFalseValue());
   Value *ShadowSel = nullptr;
+  const bool ShouldTrackOrigins = DFSF.DFS.shouldTrackOrigins();
+  std::vector<Value *> Shadows;
+  std::vector<Value *> Origins;
+  Value *TrueOrigin =
+      ShouldTrackOrigins ? DFSF.getOrigin(I.getTrueValue()) : nullptr;
+  Value *FalseOrigin =
+      ShouldTrackOrigins ? DFSF.getOrigin(I.getFalseValue()) : nullptr;
 
   if (isa<VectorType>(I.getCondition()->getType())) {
     ShadowSel = DFSF.combineShadowsThenConvert(I.getType(), TrueShadow,
                                                FalseShadow, &I);
+    if (ShouldTrackOrigins) {
+      Shadows.push_back(TrueShadow);
+      Shadows.push_back(FalseShadow);
+      Origins.push_back(TrueOrigin);
+      Origins.push_back(FalseOrigin);
+    }
   } else {
     if (TrueShadow == FalseShadow) {
       ShadowSel = TrueShadow;
+      if (ShouldTrackOrigins) {
+        Shadows.push_back(TrueShadow);
+        Origins.push_back(TrueOrigin);
+      }
     } else {
       ShadowSel =
           SelectInst::Create(I.getCondition(), TrueShadow, FalseShadow, "", &I);
+      if (ShouldTrackOrigins) {
+        Shadows.push_back(ShadowSel);
+        Origins.push_back(SelectInst::Create(I.getCondition(), TrueOrigin,
+                                             FalseOrigin, "", &I));
+      }
     }
   }
   DFSF.setShadow(&I, ClTrackSelectControlFlow
                          ? DFSF.combineShadowsThenConvert(
                                I.getType(), CondShadow, ShadowSel, &I)
                          : ShadowSel);
+  if (ShouldTrackOrigins) {
+    if (ClTrackSelectControlFlow) {
+      Shadows.push_back(CondShadow);
+      Origins.push_back(DFSF.getOrigin(I.getCondition()));
+    }
+    DFSF.setOrigin(&I, DFSF.combineOrigins(Shadows, Origins, &I));
+  }
 }
 
 void DFSanVisitor::visitMemSetInst(MemSetInst &I) {
   IRBuilder<> IRB(&I);
   Value *ValShadow = DFSF.getShadow(I.getValue());
   IRB.CreateCall(DFSF.DFS.DFSanSetLabelFn,
                  {ValShadow, IRB.CreateBitCast(I.getDest(), Type::getInt8PtrTy(
                                                                 *DFSF.DFS.Ctx)),
case DataFlowSanitizer::IA_TLS: {
       unsigned Size =
           getDataLayout().getTypeAllocSize(DFSF.DFS.getShadowTy(RT));
       if (Size <= kRetvalTLSSize) {
         // If the size overflows, stores nothing. At callsite, oversized return
         // shadows are set to zero.
         IRB.CreateAlignedStore(S, DFSF.getRetvalTLS(RT, IRB),
                                kShadowTLSAlignment);
       }
+      if (DFSF.DFS.shouldTrackOrigins()) {
+        Value *O = DFSF.getOrigin(RI.getReturnValue());
+        IRB.CreateStore(O, DFSF.getRetvalOriginTLS());
+      }
       break;
     }
     case DataFlowSanitizer::IA_Args: {
       IRBuilder<> IRB(&RI);
       Type *RT = DFSF.F->getFunctionType()->getReturnType();
       Value *InsVal =
           IRB.CreateInsertValue(UndefValue::get(RT), RI.getReturnValue(), 0);
       Value *InsShadow =
case DataFlowSanitizer::WK_Warning:
     DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB));
     return true;
   case DataFlowSanitizer::WK_Discard:
     CB.setCalledFunction(&F);
     DFSF.setShadow(&CB, DFSF.DFS.getZeroShadow(&CB));
     return true;
   case DataFlowSanitizer::WK_Functional:
     CB.setCalledFunction(&F);
-    visitOperandShadowInst(CB);
+    visitInstOperands(CB);
     return true;
   case DataFlowSanitizer::WK_Custom:
     // Don't try to handle invokes of custom functions, it's too complicated.
     // Instead, invoke the dfsw$ wrapper, which will in turn call the __dfsw_
     // wrapper.
     CallInst *CI = dyn_cast<CallInst>(&CB);
     if (!CI)
       return false;
case DataFlowSanitizer::WK_Custom:
     return true;
   }
   return false;
 }
 
 void DFSanVisitor::visitCallBase(CallBase &CB) {
   Function *F = CB.getCalledFunction();
   if ((F && F->isIntrinsic()) || CB.isInlineAsm()) {
-    visitOperandShadowInst(CB);
+    visitInstOperands(CB);
     return;
   }
 
   // Calls to this function are synthesized in wrappers, and we shouldn't
   // instrument them.
   if (F == DFSF.DFS.DFSanVarargWrapperFn.getCallee()->stripPointerCasts())
     return;
 

llvm/test/Instrumentation/DataFlowSanitizer/origin_other_ops.ll

+; RUN: opt < %s -dfsan -dfsan-track-origins=1 -dfsan-fast-16-labels=true -S | FileCheck %s --check-prefix=CHECK
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define float @unop(float %f) {
+  ; CHECK: @"dfs$unop"
+  ; CHECK: [[FO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: store i32 [[FO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  %r = fneg float %f
+  ret float %r
+}
+
+define i1 @binop(i1 %a, i1 %b) {
+  ; CHECK: @"dfs$binop"
+  ; CHECK: [[BO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; CHECK: [[AO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: [[BS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 2) to i16*), align 2
+  ; CHECK: [[NE:%.*]] = icmp ne i16 [[BS]], 0
+  ; CHECK: [[MO:%.*]] = select i1 [[NE]], i32 [[BO]], i32 [[AO]]
+  ; CHECK: store i32 [[MO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  %r = add i1 %a, %b
+  ret i1 %r
+}
+
+define i8 @castop(i32* %p) {
+  ; CHECK: @"dfs$castop"
+  ; CHECK: [[PO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: store i32 [[PO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  %r = ptrtoint i32* %p to i8
+  ret i8 %r
+}
+
+define i1 @cmpop(i1 %a, i1 %b) {
+  ; CHECK: @"dfs$cmpop"
+  ; CHECK: [[BO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; CHECK: [[AO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: [[BS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 2) to i16*), align 2
+  ; CHECK: [[NE:%.*]] = icmp ne i16 [[BS]], 0
+  ; CHECK: [[MO:%.*]] = select i1 [[NE]], i32 [[BO]], i32 [[AO]]
+  ; CHECK: store i32 [[MO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  %r = icmp eq i1 %a, %b
+  ret i1 %r
+}
+
+define i32* @gepop([10 x [20 x i32]]* %p, i32 %a, i32 %b, i32 %c) {
+  ; CHECK: @"dfs$gepop"
+  ; CHECK: [[CO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 3), align 4
+  ; CHECK: [[BO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 2), align 4
+  ; CHECK: [[AO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; CHECK: [[PO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: [[CS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 6) to i16*), align 2
+  ; CHECK: [[BS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 4) to i16*), align 2
+  ; CHECK: [[AS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 2) to i16*), align 2
+  ; CHECK: [[AS_NE:%.*]] = icmp ne i16 [[AS]], 0
+  ; CHECK: [[APO:%.*]] = select i1 [[AS_NE]], i32 [[AO]], i32 [[PO]]
+  ; CHECK: [[BS_NE:%.*]] = icmp ne i16 [[BS]], 0
+  ; CHECK: [[ABPO:%.*]] = select i1 [[BS_NE]], i32 [[BO]], i32 [[APO]]
+  ; CHECK: [[CS_NE:%.*]] = icmp ne i16 [[CS]], 0
+  ; CHECK: [[ABCPO:%.*]] = select i1 [[CS_NE]], i32 [[CO]], i32 [[ABPO]]
+  ; CHECK: store i32 [[ABCPO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  %e = getelementptr [10 x [20 x i32]], [10 x [20 x i32]]* %p, i32 %a, i32 %b, i32 %c
+  ret i32* %e
+}
+
+define i32 @eeop(<4 x i32> %a, i32 %b) {
+  ; CHECK: @"dfs$eeop"
+  ; CHECK: [[BO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; CHECK: [[AO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: [[BS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 2) to i16*), align 2
+  ; CHECK: [[NE:%.*]] = icmp ne i16 [[BS]], 0
+  ; CHECK: [[MO:%.*]] = select i1 [[NE]], i32 [[BO]], i32 [[AO]]
+  ; CHECK: store i32 [[MO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  %e = extractelement <4 x i32> %a, i32 %b
+  ret i32 %e
+}
+
+define <4 x i32> @ieop(<4 x i32> %p, i32 %a, i32 %b) {
+  ; CHECK: @"dfs$ieop"
+  ; CHECK: [[BO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 2), align 4
+  ; CHECK: [[AO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; CHECK: [[PO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: [[BS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 4) to i16*), align 2
+  ; CHECK: [[AS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 2) to i16*), align 2
+  ; CHECK: [[AS_NE:%.*]] = icmp ne i16 [[AS]], 0
+  ; CHECK: [[APO:%.*]] = select i1 [[AS_NE]], i32 [[AO]], i32 [[PO]]
+  ; CHECK: [[BS_NE:%.*]] = icmp ne i16 [[BS]], 0
+  ; CHECK: [[ABPO:%.*]] = select i1 [[BS_NE]], i32 [[BO]], i32 [[APO]]
+  ; CHECK: store i32 [[ABPO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  %e = insertelement <4 x i32> %p, i32 %a, i32 %b
+  ret <4 x i32> %e
+}
+
+define <4 x i32> @svop(<4 x i32> %a, <4 x i32> %b) {
+  ; CHECK: @"dfs$svop"
+  ; CHECK: [[BO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; CHECK: [[AO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: [[BS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 2) to i16*), align 2
+  ; CHECK: [[NE:%.*]] = icmp ne i16 [[BS]], 0
+  ; CHECK: [[MO:%.*]] = select i1 [[NE]], i32 [[BO]], i32 [[AO]]
+  ; CHECK: store i32 [[MO]], i32* @__dfsan_retval_origin_tls, align 4
+  
+  %e = shufflevector <4 x i32> %a, <4 x i32> %b, <4 x i32> <i32 0, i32 4, i32 1, i32 5>
+  ret <4 x i32> %e
+}  
+
+define i32 @evop({i32, float} %a) {
+  ; CHECK: @"dfs$evop"
+  ; CHECK: [[AO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: store i32 [[AO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  %e = extractvalue {i32, float} %a, 0
+  ret i32 %e
+}
+
+define {i32, {float, float}} @ivop({i32, {float, float}} %a, {float, float} %b) {
+  ; CHECK: @"dfs$ivop"
+  ; CHECK: [[BO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; CHECK: [[AO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; CHECK: [[BS:%.*]] = load { i16, i16 }, { i16, i16 }* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 6) to { i16, i16 }*), align 2
+  ; CHECK: [[BS0:%.*]] = extractvalue { i16, i16 } [[BS]], 0
+  ; CHECK: [[BS1:%.*]] = extractvalue { i16, i16 } [[BS]], 1
+  ; CHECK: [[BS01:%.*]] = or i16 [[BS0]], [[BS1]]
+  ; CHECK: [[NE:%.*]] = icmp ne i16 [[BS01]], 0
+  ; CHECK: [[MO:%.*]] = select i1 [[NE]], i32 [[BO]], i32 [[AO]]
+  ; CHECK: store i32 [[MO]], i32* @__dfsan_retval_origin_tls, align 4
+  
+  %e = insertvalue {i32, {float, float}} %a, {float, float} %b, 1
+  ret {i32, {float, float}} %e
+}
 No newline at end of file

llvm/test/Instrumentation/DataFlowSanitizer/origin_select.ll

+; RUN: opt < %s -dfsan -dfsan-track-select-control-flow=1 -dfsan-track-origins=1 -dfsan-fast-16-labels=true -S | FileCheck %s --check-prefix=TRACK_CONTROL_FLOW
+; RUN: opt < %s -dfsan -dfsan-track-select-control-flow=0 -dfsan-track-origins=1 -dfsan-fast-16-labels=true -S | FileCheck %s --check-prefix=NO_TRACK_CONTROL_FLOW
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i8 @select8(i1 %c, i8 %t, i8 %f) {
+  ; TRACK_CONTROL_FLOW: @"dfs$select8"
+  ; TRACK_CONTROL_FLOW: [[CO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; TRACK_CONTROL_FLOW: [[FO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 2), align 4
+  ; TRACK_CONTROL_FLOW: [[TO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; TRACK_CONTROL_FLOW: [[CS:%.*]] = load i16, i16* bitcast ([100 x i64]* @__dfsan_arg_tls to i16*), align 2
+  ; TRACK_CONTROL_FLOW: [[TFO:%.*]] = select i1 %c, i32 [[TO]], i32 [[FO]]
+  ; TRACK_CONTROL_FLOW: [[CS_NE:%.*]] = icmp ne i16 [[CS]], 0
+  ; TRACK_CONTROL_FLOW: [[CTFO:%.*]] = select i1 [[CS_NE]], i32 [[CO]], i32 [[TFO]]
+  ; TRACK_CONTROL_FLOW: store i32 [[CTFO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  ; NO_TRACK_CONTROL_FLOW: @"dfs$select8"
+  ; NO_TRACK_CONTROL_FLOW: [[FO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 2), align 4
+  ; NO_TRACK_CONTROL_FLOW: [[TO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; NO_TRACK_CONTROL_FLOW: [[TFO:%.*]] = select i1 %c, i32 [[TO]], i32 [[FO]]
+  ; NO_TRACK_CONTROL_FLOW: store i32 [[TFO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  %a = select i1 %c, i8 %t, i8 %f
+  ret i8 %a
+}
+
+define i8 @select8e(i1 %c, i8 %tf) {
+  ; TRACK_CONTROL_FLOW: @"dfs$select8e"
+  ; TRACK_CONTROL_FLOW: [[CO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; TRACK_CONTROL_FLOW: [[TFO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; TRACK_CONTROL_FLOW: [[CS:%.*]] = load i16, i16* bitcast ([100 x i64]* @__dfsan_arg_tls to i16*), align 2
+  ; TRACK_CONTROL_FLOW: [[CS_NE:%.*]] = icmp ne i16 [[CS]], 0
+  ; TRACK_CONTROL_FLOW: [[CTFO:%.*]] = select i1 [[CS_NE]], i32 [[CO]], i32 [[TFO]]
+  ; TRACK_CONTROL_FLOW: store i32 [[CTFO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  ; NO_TRACK_CONTROL_FLOW: @"dfs$select8e"
+  ; NO_TRACK_CONTROL_FLOW: [[TFO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; NO_TRACK_CONTROL_FLOW: store i32 [[TFO]], i32* @__dfsan_retval_origin_tls, align 4
+
+%a = select i1 %c, i8 %tf, i8 %tf
+  ret i8 %a
+}
+
+define <4 x i8> @select8v(<4 x i1> %c, <4 x i8> %t, <4 x i8> %f) {
+  ; TRACK_CONTROL_FLOW: @"dfs$select8v"
+  ; TRACK_CONTROL_FLOW: [[CO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 0), align 4
+  ; TRACK_CONTROL_FLOW: [[FO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 2), align 4
+  ; TRACK_CONTROL_FLOW: [[TO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; TRACK_CONTROL_FLOW: [[FS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 4) to i16*), align 2
+  ; TRACK_CONTROL_FLOW: [[CS:%.*]] = load i16, i16* bitcast ([100 x i64]* @__dfsan_arg_tls to i16*), align 2
+  ; TRACK_CONTROL_FLOW: [[FS_NE:%.*]] = icmp ne i16 [[FS]], 0
+  ; TRACK_CONTROL_FLOW: [[FTO:%.*]] = select i1 [[FS_NE]], i32 [[FO]], i32 [[TO]]
+  ; TRACK_CONTROL_FLOW: [[CS_NE:%.*]] = icmp ne i16 [[CS]], 0
+  ; TRACK_CONTROL_FLOW: [[CFTO:%.*]] = select i1 [[CS_NE]], i32 [[CO]], i32 [[FTO]]
+  ; TRACK_CONTROL_FLOW: store i32 [[CFTO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  ; NO_TRACK_CONTROL_FLOW: @"dfs$select8v"
+  ; NO_TRACK_CONTROL_FLOW: [[FO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 2), align 4
+  ; NO_TRACK_CONTROL_FLOW: [[TO:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__dfsan_arg_origin_tls, i64 0, i64 1), align 4
+  ; NO_TRACK_CONTROL_FLOW: [[FS:%.*]] = load i16, i16* inttoptr (i64 add (i64 ptrtoint ([100 x i64]* @__dfsan_arg_tls to i64), i64 4) to i16*), align 2
+  ; NO_TRACK_CONTROL_FLOW: [[FS_NE:%.*]] = icmp ne i16 [[FS]], 0
+  ; NO_TRACK_CONTROL_FLOW: [[FTO:%.*]] = select i1 [[FS_NE]], i32 [[FO]], i32 [[TO]]
+  ; NO_TRACK_CONTROL_FLOW: store i32 [[FTO]], i32* @__dfsan_retval_origin_tls, align 4
+
+  %a = select <4 x i1> %c, <4 x i8> %t, <4 x i8> %f
+  ret <4 x i8> %a
+}
 No newline at end of file