implement select intrinsic

Author: spall

clang/include/clang/Basic/Builtins.td

@@ -4751,6 +4751,12 @@ def HLSLSaturate : LangBuiltin<"HLSL_LANG"> {
   let Prototype = "void(...)";
 }
 
+def HLSLSelect : LangBuiltin<"HLSL_LANG"> {
+  let Spellings = ["__builtin_hlsl_select"];
+  let Attributes = [NoThrow, Const];
+  let Prototype = "void(...)";
+}
+
 // Builtins for XRay.
 def XRayCustomEvent : Builtin {
   let Spellings = ["__xray_customevent"];

clang/lib/CodeGen/CGBuiltin.cpp

@@ -18695,6 +18695,47 @@ case Builtin::BI__builtin_hlsl_elementwise_isinf: {
         CGM.getHLSLRuntime().getSaturateIntrinsic(), ArrayRef<Value *>{Op0},
         nullptr, "hlsl.saturate");
   }
+  case Builtin::BI__builtin_hlsl_select: {
+    Value *OpCond = EmitScalarExpr(E->getArg(0));
+    Value *OpTrue = EmitScalarExpr(E->getArg(1));
+    Value *OpFalse = EmitScalarExpr(E->getArg(2));
+    llvm::Type *TCond = OpCond->getType();
+
+    // if cond is a bool emit a select instruction
+    if (TCond->isIntegerTy(1))
+      return Builder.CreateSelect(OpCond, OpTrue, OpFalse);
+
+    // if cond is a vector of bools lower to a shufflevector
+    // todo check if that true and false are vectors
+    // todo check that the size of true and false and cond are the same
+    if (TCond->isVectorTy() &&
+	E->getArg(0)->getType()->getAs<VectorType>()->isBooleanType()) {
+      assert(OpTrue->getType()->isVectorTy() && OpFalse->getType()->isVectorTy() &&
+	     "Select's second and third operands must be vectors if first operand is a vector.");
+
+      auto *VecTyTrue = E->getArg(1)->getType()->getAs<VectorType>();
+      auto *VecTyFalse = E->getArg(2)->getType()->getAs<VectorType>();
+
+      assert(VecTyTrue->getElementType() == VecTyFalse->getElementType() &&
+	     "Select's second and third vectors need the same element types.");
+
+      const unsigned N = VecTyTrue->getNumElements();
+      assert(N == VecTyFalse->getNumElements() &&
+	     N == E->getArg(0)->getType()->getAs<VectorType>()->getNumElements() &&
+	     "Select requires vectors to be of the same size.");
+      
+      llvm::SmallVector<Value *> Mask;
+      for (unsigned I = 0; I < N; I++) {
+	Value *Index = ConstantInt::get(IntTy, I);
+	Value *IndexBool = Builder.CreateExtractElement(OpCond, Index);
+	Mask.push_back(Builder.CreateSelect(IndexBool, Index, ConstantInt::get(IntTy, I + N)));
+      }
+      
+      return Builder.CreateShuffleVector(OpTrue, OpFalse, BuildVector(Mask));
+    }
+    
+    llvm_unreachable("Select requires a bool or vector of bools as its first operand.");
+  }
   case Builtin::BI__builtin_hlsl_wave_get_lane_index: {
     return EmitRuntimeCall(CGM.CreateRuntimeFunction(
         llvm::FunctionType::get(IntTy, {}, false), "__hlsl_wave_get_lane_index",

clang/lib/Headers/hlsl/hlsl_intrinsics.h

@@ -1603,6 +1603,30 @@ double3 saturate(double3);
 _HLSL_BUILTIN_ALIAS(__builtin_hlsl_elementwise_saturate)
 double4 saturate(double4);
 
+//===----------------------------------------------------------------------===//
+// select builtins
+//===----------------------------------------------------------------------===//
+
+/// \fn T select(bool Cond, T TrueVal, T FalseVal)
+/// \brief ternary operator.
+/// \param Cond The Condition input value.
+/// \param TrueVal The Value returned if Cond is true.
+/// \param FalseVal The Value returned if Cond is false.
+
+_HLSL_BUILTIN_ALIAS(__builtin_hlsl_select)
+template<typename T>
+T select(bool, T, T);
+  
+/// \fn vector<T,Sz> select(vector<bool,Sz> Conds, vector<T,Sz> TrueVals, vector<T,Sz>, FalseVals)
+/// \brief ternary operator for vectors. All vectors must be the same size.
+/// \param Conds The Condition input values.
+/// \param TrueVals The vector values are chosen from when conditions are true.
+/// \param FalseVals The vector values are chosen from when conditions are false.
+
+_HLSL_BUILTIN_ALIAS(__builtin_hlsl_select)
+template<typename T, int Sz>
+vector<T,Sz> select(vector<bool,Sz>, vector<T,Sz>, vector<T,Sz>);
+  
 //===----------------------------------------------------------------------===//
 // sin builtins
 //===----------------------------------------------------------------------===//

clang/lib/Sema/SemaHLSL.cpp

@@ -1013,6 +1013,66 @@ void SetElementTypeAsReturnType(Sema *S, CallExpr *TheCall,
   TheCall->setType(ReturnType);
 }
 
+bool CheckBoolSelect(Sema *S, CallExpr *TheCall) {
+  assert(TheCall->getNumArgs() == 3);
+  Expr *Arg1 = TheCall->getArg(1);
+  Expr *Arg2 = TheCall->getArg(2);
+  if(!S->Context.hasSameUnqualifiedType(Arg1->getType(),
+					Arg2->getType())) {
+    S->Diag(TheCall->getBeginLoc(),
+	    diag::err_typecheck_call_different_arg_types)
+      << Arg1->getType() << Arg2->getType()
+      << Arg1->getSourceRange() << Arg2->getSourceRange();
+    return true;
+  }
+
+  TheCall->setType(Arg1->getType());
+  return false;
+}
+
+bool CheckVectorSelect(Sema *S, CallExpr *TheCall) {
+  assert(TheCall->getNumArgs() == 3);
+  Expr *Arg1 = TheCall->getArg(1);
+  Expr *Arg2 = TheCall->getArg(2);
+  if(!Arg1->getType()->isVectorType()) {
+    S->Diag(Arg1->getBeginLoc(),
+	    diag::err_builtin_non_vector_type)
+      << "Second" << "__builtin_hlsl_select" << Arg1->getType()
+      << Arg1->getSourceRange();
+    return true;
+  }
+  
+  if(!Arg2->getType()->isVectorType()) {
+    S->Diag(Arg2->getBeginLoc(),
+	    diag::err_builtin_non_vector_type)
+      << "Third" << "__builtin_hlsl_select" << Arg2->getType()
+      << Arg2->getSourceRange();
+    return true;
+  }
+
+  if(!S->Context.hasSameUnqualifiedType(Arg1->getType(),
+					Arg2->getType())) {
+    S->Diag(TheCall->getBeginLoc(),
+	    diag::err_typecheck_call_different_arg_types)
+      << Arg1->getType() << Arg2->getType()
+      << Arg1->getSourceRange() << Arg2->getSourceRange();
+    return true;
+  }
+
+  // caller has checked that Arg0 is a vector.
+  // check all three args have the same length.
+  if(TheCall->getArg(0)->getType()->getAs<VectorType>()->getNumElements() !=
+     Arg1->getType()->getAs<VectorType>()->getNumElements()) {
+    S->Diag(TheCall->getBeginLoc(),
+	    diag::err_typecheck_vector_lengths_not_equal)
+      << TheCall->getArg(0)->getType() << Arg1->getType()
+      << TheCall->getArg(0)->getSourceRange() << Arg1->getSourceRange();
+    return true;
+  }
+
+  return false;
+}
+
 // Note: returning true in this case results in CheckBuiltinFunctionCall
 // returning an ExprError
 bool SemaHLSL::CheckBuiltinFunctionCall(unsigned BuiltinID, CallExpr *TheCall) {
@@ -1046,6 +1106,30 @@ bool SemaHLSL::CheckBuiltinFunctionCall(unsigned BuiltinID, CallExpr *TheCall) {
     break;
   }
   case Builtin::BI__builtin_hlsl_elementwise_saturate:
+  case Builtin::BI__builtin_hlsl_select: {
+    if (SemaRef.checkArgCount(TheCall, 3))
+      return true;
+    QualType ArgTy = TheCall->getArg(0)->getType();
+    if (ArgTy->isBooleanType()) {
+      if (CheckBoolSelect(&SemaRef, TheCall))
+	return true;
+    } else if (ArgTy->isVectorType() &&
+	       ArgTy->getAs<VectorType>()->getElementType()->isBooleanType()) {
+      if (CheckVectorSelect(&SemaRef, TheCall))
+	return true;
+    } else { // first operand is not a bool or a vector of bools.
+      SemaRef.Diag(TheCall->getArg(0)->getBeginLoc(),
+		   diag::err_typecheck_convert_incompatible)
+	<< TheCall->getArg(0)->getType() << SemaRef.Context.getBOOLType()
+	<< 1 << 0 << 0;
+      SemaRef.Diag(TheCall->getArg(0)->getBeginLoc(),
+		   diag::err_builtin_non_vector_type)
+	<< "First" << "__builtin_hlsl_select" << TheCall->getArg(0)->getType()
+	<< TheCall->getArg(0)->getSourceRange();
+      return true;
+    }
+    break;
+  }
   case Builtin::BI__builtin_hlsl_elementwise_rcp: {
     if (CheckAllArgsHaveFloatRepresentation(&SemaRef, TheCall))
       return true;