[MSAN] Implement experiemental vector reduction intrinsics

[guiand]

Summary

Implement llvm.experiemental.vector.{add,mul,or,and,...}.
An IR test is included but no C test for lack of good way to
get the compiler to emit these.

These vectorized shadow operations for the and/or case
were shown mathematically to be equivalent to the binary
operations already present in MSAN.

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3164

This is not correct for v2 intrinsics, they have an extra scalar argument.

[guiand]

I removed the v2 intrinsic handling

[eugenis]

LGTM

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

   void handleVectorCompareScalarIntrinsic(IntrinsicInst &I) {
     IRBuilder<> IRB(&I);
     Value *S0 = IRB.CreateOr(getShadow(&I, 0), getShadow(&I, 1));
     Value *S = LowerElementShadowExtend(IRB, S0, getShadowTy(&I));
     setShadow(&I, S);
     setOriginForNaryOp(I);
   }
 
+  // Instrument generic vector reduction intrinsics
+  // by ORing together all their fields.
+  void handleVectorReduceIntrinsic(IntrinsicInst &I) {
+    IRBuilder<> IRB(&I);
+    Value *S = IRB.CreateOrReduce(getShadow(&I, 0));
+    setShadow(&I, S);
+    setOrigin(&I, getOrigin(&I, 0));
+  }
+
+  // Instrument experimental.vector.reduce.or intrinsic.
+  // Valid (non-poisoned) set bits in the operand pull low the
+  // corresponding shadow bits.
+  void handleVectorReduceOrIntrinsic(IntrinsicInst &I) {
+    IRBuilder<> IRB(&I);
+    Value *OperandShadow = getShadow(&I, 0);
+    Value *OperandUnsetBits = IRB.CreateNot(I.getOperand(0));
+    Value *OperandUnsetOrPoison = IRB.CreateOr(OperandUnsetBits, OperandShadow);
+    // Bit N is clean if any field's bit N is 1 and unpoison
+    Value *OutShadowMask = IRB.CreateAndReduce(OperandUnsetOrPoison);
+    // Otherwise, it is clean if every field's bit N is unpoison
+    Value *OrShadow = IRB.CreateOrReduce(OperandShadow);
+    Value *S = IRB.CreateAnd(OutShadowMask, OrShadow);
+
+    setShadow(&I, S);
+    setOrigin(&I, getOrigin(&I, 0));
+  }
+
+  // Instrument experimental.vector.reduce.or intrinsic.
+  // Valid (non-poisoned) unset bits in the operand pull down the
+  // corresponding shadow bits.
+  void handleVectorReduceAndIntrinsic(IntrinsicInst &I) {
+    IRBuilder<> IRB(&I);
+    Value *OperandShadow = getShadow(&I, 0);
+    Value *OperandSetOrPoison = IRB.CreateOr(I.getOperand(0), OperandShadow);
+    // Bit N is clean if any field's bit N is 0 and unpoison
+    Value *OutShadowMask = IRB.CreateAndReduce(OperandSetOrPoison);
+    // Otherwise, it is clean if every field's bit N is unpoison
+    Value *OrShadow = IRB.CreateOrReduce(OperandShadow);
+    Value *S = IRB.CreateAnd(OutShadowMask, OrShadow);
+
+    setShadow(&I, S);
+    setOrigin(&I, getOrigin(&I, 0));
+  }
+
   void handleStmxcsr(IntrinsicInst &I) {
     IRBuilder<> IRB(&I);
     Value* Addr = I.getArgOperand(0);
     Type *Ty = IRB.getInt32Ty();
     Value *ShadowPtr =
         getShadowOriginPtr(Addr, IRB, Ty, Align(1), /*isStore*/ true).first;
 
     IRB.CreateStore(getCleanShadow(Ty),
case Intrinsic::bswap:
       handleBswap(I);
       break;
     case Intrinsic::masked_store:
       handleMaskedStore(I);
       break;
     case Intrinsic::masked_load:
       handleMaskedLoad(I);
       break;
+    case Intrinsic::experimental_vector_reduce_and:
+      handleVectorReduceAndIntrinsic(I);
+      break;
+    case Intrinsic::experimental_vector_reduce_or:
+      handleVectorReduceOrIntrinsic(I);
+      break;
+    case Intrinsic::experimental_vector_reduce_add:
+    case Intrinsic::experimental_vector_reduce_xor:
+    case Intrinsic::experimental_vector_reduce_mul:
+      handleVectorReduceIntrinsic(I);
+      break;
     case Intrinsic::x86_sse_stmxcsr:
       handleStmxcsr(I);
       break;
     case Intrinsic::x86_sse_ldmxcsr:
       handleLdmxcsr(I);
       break;
     case Intrinsic::x86_avx512_vcvtsd2usi64:
     case Intrinsic::x86_avx512_vcvtsd2usi32:

llvm/test/Instrumentation/MemorySanitizer/experimental-reduce.ll

+
+; RUN: opt < %s -msan-check-access-address=0 -msan-track-origins=1 -S -passes='module(msan-module),function(msan)' 2>&1 | \
+; RUN:   FileCheck -allow-deprecated-dag-overlap -check-prefixes=CHECK,CHECK-ORIGINS %s
+
+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"
+
+declare i32 @llvm.experimental.vector.reduce.add(<3 x i32>)
+declare i32 @llvm.experimental.vector.reduce.and(<3 x i32>)
+declare i32 @llvm.experimental.vector.reduce.or(<3 x i32>)
+
+; CHECK-LABEL: @reduce_add
+define i32 @reduce_add() sanitize_memory {
+; CHECK: [[P:%.*]] = inttoptr i64 0 to <3 x i32>*
+  %p = inttoptr i64 0 to <3 x i32> *
+; CHECK: [[O:%.*]] = load <3 x i32>, <3 x i32>* [[P]]
+  %o = load <3 x i32>, <3 x i32> *%p
+; CHECK: [[O_SHADOW:%.*]] = load <3 x i32>, <3 x i32>*
+; CHECK: [[O_ORIGIN:%.*]] = load i32, i32*
+; CHECK: [[R_SHADOW:%.*]] = call i32 @llvm.experimental.vector.reduce.or.v3i32(<3 x i32> [[O_SHADOW]])
+; CHECK: [[R:%.*]] = call i32 @llvm.experimental.vector.reduce.add.v3i32(<3 x i32> [[O]])
+  %r = call i32 @llvm.experimental.vector.reduce.add(<3 x i32> %o)
+; CHECK: store i32 [[R_SHADOW]], {{.*}} @__msan_retval_tls
+; CHECK: store i32 [[O_ORIGIN]], {{.*}} @__msan_retval_origin_tls
+; CHECK: ret i32 [[R]]
+  ret i32 %r
+}
+
+; CHECK-LABEL: @reduce_and
+define i32 @reduce_and() sanitize_memory {
+; CHECK: [[P:%.*]] = inttoptr i64 0 to <3 x i32>*
+  %p = inttoptr i64 0 to <3 x i32> *
+; CHECK: [[O:%.*]] = load <3 x i32>, <3 x i32>* [[P]]
+  %o = load <3 x i32>, <3 x i32> *%p
+; CHECK: [[O_SHADOW:%.*]] = load <3 x i32>, <3 x i32>*
+; CHECK: [[O_ORIGIN:%.*]] = load i32, i32*
+; CHECK: [[O_SHADOW_1:%.*]] = or <3 x i32> [[O]], [[O_SHADOW]]
+; CHECK: [[O_SHADOW_2:%.*]] = call i32 @llvm.experimental.vector.reduce.and.v3i32(<3 x i32> [[O_SHADOW_1]]
+; CHECK: [[O_SHADOW_3:%.*]] = call i32 @llvm.experimental.vector.reduce.or.v3i32(<3 x i32> [[O_SHADOW]])
+; CHECK: [[R_SHADOW:%.*]] = and i32 [[O_SHADOW_2]], [[O_SHADOW_3]]
+; CHECK: [[R:%.*]] = call i32 @llvm.experimental.vector.reduce.and.v3i32(<3 x i32> [[O]])
+  %r = call i32 @llvm.experimental.vector.reduce.and(<3 x i32> %o)
+; CHECK: store i32 [[R_SHADOW]], {{.*}} @__msan_retval_tls
+; CHECK: store i32 [[O_ORIGIN]], {{.*}} @__msan_retval_origin_tls
+; CHECK: ret i32 [[R]]
+  ret i32 %r
+}
+
+; CHECK-LABEL: @reduce_or
+define i32 @reduce_or() sanitize_memory {
+; CHECK: [[P:%.*]] = inttoptr i64 0 to <3 x i32>*
+  %p = inttoptr i64 0 to <3 x i32> *
+; CHECK: [[O:%.*]] = load <3 x i32>, <3 x i32>* [[P]]
+  %o = load <3 x i32>, <3 x i32> *%p
+; CHECK: [[O_SHADOW:%.*]] = load <3 x i32>, <3 x i32>*
+; CHECK: [[O_ORIGIN:%.*]] = load i32, i32*
+; CHECK: [[NOT_O:%.*]] = xor <3 x i32> [[O]], <i32 -1, i32 -1, i32 -1>
+; CHECK: [[O_SHADOW_1:%.*]] = or <3 x i32> [[NOT_O]], [[O_SHADOW]]
+; CHECK: [[O_SHADOW_2:%.*]] = call i32 @llvm.experimental.vector.reduce.and.v3i32(<3 x i32> [[O_SHADOW_1]]
+; CHECK: [[O_SHADOW_3:%.*]] = call i32 @llvm.experimental.vector.reduce.or.v3i32(<3 x i32> [[O_SHADOW]])
+; CHECK: [[R_SHADOW:%.*]] = and i32 [[O_SHADOW_2]], [[O_SHADOW_3]]
+; CHECK: [[R:%.*]] = call i32 @llvm.experimental.vector.reduce.or.v3i32(<3 x i32> [[O]])
+  %r = call i32 @llvm.experimental.vector.reduce.or(<3 x i32> %o)
+; CHECK: store i32 [[R_SHADOW]], {{.*}} @__msan_retval_tls
+; CHECK: store i32 [[O_ORIGIN]], {{.*}} @__msan_retval_origin_tls
+; CHECK: ret i32 [[R]]
+  ret i32 %r
+}