[SM6.9] Lower vector any/all calls to vector op (microsoft#7753)

Fixes microsoft#7687

Implement vector reduction lowing for calls to `any` and `all`

---------

Co-authored-by: Helena Kotas <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: 3 people

docs/DXIL.rst

@@ -2423,8 +2423,8 @@ ID  Name                                                  Description
 306 MatVecMulAdd                                          multiplies a MxK dimension matrix and a K sized input vector and adds an M-sized bias vector
 307 OuterProductAccumulate                                Computes the outer product between column vectors and an MxN matrix is accumulated component-wise atomically (with device scope) in memory
 308 VectorAccumulate                                      Accumulates the components of a vector component-wise atomically (with device scope) to the corresponding elements of an array in memory
-309 ReservedD0                                            reserved
-310 ReservedD1                                            reserved
+309 VectorReduceAnd                                       Bitwise AND reduction of the vector returning a scalar
+310 VectorReduceOr                                        Bitwise OR reduction of the vector returning a scalar
 311 FDot                                                  computes the n-dimensional vector dot-product
 === ===================================================== =======================================================================================================================================================================================================================
 

include/dxc/DXIL/DxilConstants.h

@@ -524,8 +524,6 @@ enum class OpCode : unsigned {
   ReservedC7 = 300,  // reserved
   ReservedC8 = 301,  // reserved
   ReservedC9 = 302,  // reserved
-  ReservedD0 = 309,  // reserved
-  ReservedD1 = 310,  // reserved
 
   // Amplification shader instructions
   DispatchMesh = 173, // Amplification shader intrinsic DispatchMesh
@@ -1037,6 +1035,11 @@ enum class OpCode : unsigned {
   Unpack4x8 = 219, // unpacks 4 8-bit signed or unsigned values into int32 or
                    // int16 vector
 
+  // Vector reduce to scalar
+  VectorReduceAnd =
+      309, // Bitwise AND reduction of the vector returning a scalar
+  VectorReduceOr = 310, // Bitwise OR reduction of the vector returning a scalar
+
   // Wave
   WaveActiveAllEqual = 115, // returns 1 if all the lanes have the same value
   WaveActiveBallot = 116, // returns a struct with a bit set for each lane where
@@ -1381,6 +1384,9 @@ enum class OpCodeClass : unsigned {
   // Unpacking intrinsics
   Unpack4x8,
 
+  // Vector reduce to scalar
+  VectorReduce,
+
   // Wave
   WaveActiveAllEqual,
   WaveActiveBallot,
@@ -1417,7 +1423,7 @@ enum class OpCodeClass : unsigned {
   NumOpClasses_Dxil_1_7 = 153,
   NumOpClasses_Dxil_1_8 = 174,
 
-  NumOpClasses = 195 // exclusive last value of enumeration
+  NumOpClasses = 196 // exclusive last value of enumeration
 };
 // OPCODECLASS-ENUM:END
 

include/dxc/DXIL/DxilInstructions.h

@@ -10149,6 +10149,60 @@ struct DxilInst_VectorAccumulate {
   void set_arrayOffset(llvm::Value *val) { Instr->setOperand(3, val); }
 };
 
+/// This instruction Bitwise AND reduction of the vector returning a scalar
+struct DxilInst_VectorReduceAnd {
+  llvm::Instruction *Instr;
+  // Construction and identification
+  DxilInst_VectorReduceAnd(llvm::Instruction *pInstr) : Instr(pInstr) {}
+  operator bool() const {
+    return hlsl::OP::IsDxilOpFuncCallInst(Instr,
+                                          hlsl::OP::OpCode::VectorReduceAnd);
+  }
+  // Validation support
+  bool isAllowed() const { return true; }
+  bool isArgumentListValid() const {
+    if (2 != llvm::dyn_cast<llvm::CallInst>(Instr)->getNumArgOperands())
+      return false;
+    return true;
+  }
+  // Metadata
+  bool requiresUniformInputs() const { return false; }
+  // Operand indexes
+  enum OperandIdx {
+    arg_a = 1,
+  };
+  // Accessors
+  llvm::Value *get_a() const { return Instr->getOperand(1); }
+  void set_a(llvm::Value *val) { Instr->setOperand(1, val); }
+};
+
+/// This instruction Bitwise OR reduction of the vector returning a scalar
+struct DxilInst_VectorReduceOr {
+  llvm::Instruction *Instr;
+  // Construction and identification
+  DxilInst_VectorReduceOr(llvm::Instruction *pInstr) : Instr(pInstr) {}
+  operator bool() const {
+    return hlsl::OP::IsDxilOpFuncCallInst(Instr,
+                                          hlsl::OP::OpCode::VectorReduceOr);
+  }
+  // Validation support
+  bool isAllowed() const { return true; }
+  bool isArgumentListValid() const {
+    if (2 != llvm::dyn_cast<llvm::CallInst>(Instr)->getNumArgOperands())
+      return false;
+    return true;
+  }
+  // Metadata
+  bool requiresUniformInputs() const { return false; }
+  // Operand indexes
+  enum OperandIdx {
+    arg_a = 1,
+  };
+  // Accessors
+  llvm::Value *get_a() const { return Instr->getOperand(1); }
+  void set_a(llvm::Value *val) { Instr->setOperand(1, val); }
+};
+
 /// This instruction computes the n-dimensional vector dot-product
 struct DxilInst_FDot {
   llvm::Instruction *Instr;

lib/DXIL/DxilOperations.cpp

@@ -2687,22 +2687,23 @@ const OP::OpCodeProperty OP::m_OpCodeProps[(unsigned)OP::OpCode::NumOpCodes] = {
      {{0x400}},
      {{0x63}}}, // Overloads: <hfwi
 
-    {OC::ReservedD0,
-     "ReservedD0",
-     OCC::Reserved,
-     "reserved",
-     Attribute::None,
-     0,
-     {},
-     {}}, // Overloads: v
-    {OC::ReservedD1,
-     "ReservedD1",
-     OCC::Reserved,
-     "reserved",
-     Attribute::None,
-     0,
-     {},
-     {}}, // Overloads: v
+    // Vector reduce to scalar
+    {OC::VectorReduceAnd,
+     "VectorReduceAnd",
+     OCC::VectorReduce,
+     "vectorReduce",
+     Attribute::ReadNone,
+     1,
+     {{0x400}},
+     {{0xf8}}}, // Overloads: <18wil
+    {OC::VectorReduceOr,
+     "VectorReduceOr",
+     OCC::VectorReduce,
+     "vectorReduce",
+     Attribute::ReadNone,
+     1,
+     {{0x400}},
+     {{0xf8}}}, // Overloads: <18wil
 
     // Dot
     {OC::FDot,
@@ -6018,14 +6019,16 @@ Function *OP::GetOpFunc(OpCode opCode, Type *pOverloadType) {
     A(pI32);
     break;
 
-    //
-  case OpCode::ReservedD0:
-    A(pV);
+    // Vector reduce to scalar
+  case OpCode::VectorReduceAnd:
+    A(pVecElt);
     A(pI32);
+    A(pETy);
     break;
-  case OpCode::ReservedD1:
-    A(pV);
+  case OpCode::VectorReduceOr:
+    A(pVecElt);
     A(pI32);
+    A(pETy);
     break;
 
     // Dot
@@ -6207,6 +6210,8 @@ llvm::Type *OP::GetOverloadType(OpCode opCode, llvm::Function *F) {
   case OpCode::CreateHandleForLib:
   case OpCode::WaveMatch:
   case OpCode::VectorAccumulate:
+  case OpCode::VectorReduceAnd:
+  case OpCode::VectorReduceOr:
   case OpCode::FDot:
     if (FT->getNumParams() <= 1)
       return nullptr;
@@ -6324,8 +6329,6 @@ llvm::Type *OP::GetOverloadType(OpCode opCode, llvm::Function *F) {
   case OpCode::ReservedC7:
   case OpCode::ReservedC8:
   case OpCode::ReservedC9:
-  case OpCode::ReservedD0:
-  case OpCode::ReservedD1:
     return Type::getVoidTy(Ctx);
   case OpCode::CheckAccessFullyMapped:
   case OpCode::SampleIndex:

lib/HLSL/DxilScalarizeVectorIntrinsics.cpp

@@ -9,11 +9,14 @@
 //                                                                           //
 ///////////////////////////////////////////////////////////////////////////////
 
+#include "dxc/DXIL/DxilConstants.h"
 #include "dxc/DXIL/DxilInstructions.h"
 #include "dxc/DXIL/DxilModule.h"
 #include "dxc/HLSL/DxilGenerationPass.h"
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/IR/Constant.h"
+#include "llvm/IR/Constants.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
@@ -29,6 +32,7 @@ static void scalarizeVectorLoad(hlsl::OP *HlslOP, const DataLayout &DL,
 static void scalarizeVectorStore(hlsl::OP *HlslOP, const DataLayout &DL,
                                  CallInst *CI);
 static void scalarizeVectorIntrinsic(hlsl::OP *HlslOP, CallInst *CI);
+static void scalarizeVectorReduce(hlsl::OP *HlslOP, CallInst *CI);
 
 class DxilScalarizeVectorIntrinsics : public ModulePass {
 public:
@@ -53,24 +57,31 @@ class DxilScalarizeVectorIntrinsics : public ModulePass {
     for (auto F = M.functions().begin(); F != M.functions().end();) {
       Function *Func = &*(F++);
       DXIL::OpCodeClass OpClass;
-      if (HlslOP->GetOpCodeClass(Func, OpClass)) {
+      if (!HlslOP->GetOpCodeClass(Func, OpClass))
+        continue;
+
+      bool NeedsRewrite = (Func->getReturnType()->isVectorTy() ||
+                           OpClass == DXIL::OpCodeClass::RawBufferVectorLoad ||
+                           OpClass == DXIL::OpCodeClass::RawBufferVectorStore ||
+                           OpClass == DXIL::OpCodeClass::VectorReduce);
+      if (!NeedsRewrite)
+        continue;
+
+      for (auto U = Func->user_begin(), UE = Func->user_end(); U != UE;) {
+        CallInst *CI = cast<CallInst>(*(U++));
+
         if (OpClass == DXIL::OpCodeClass::RawBufferVectorLoad)
-          for (auto U = Func->user_begin(), UE = Func->user_end(); U != UE;) {
-            CallInst *CI = cast<CallInst>(*(U++));
-            scalarizeVectorLoad(HlslOP, M.getDataLayout(), CI);
-            Changed = true;
-          }
+          scalarizeVectorLoad(HlslOP, M.getDataLayout(), CI);
         else if (OpClass == DXIL::OpCodeClass::RawBufferVectorStore)
-          for (auto U = Func->user_begin(), UE = Func->user_end(); U != UE;) {
-            CallInst *CI = cast<CallInst>(*(U++));
-            scalarizeVectorStore(HlslOP, M.getDataLayout(), CI);
-            Changed = true;
-          }
+          scalarizeVectorStore(HlslOP, M.getDataLayout(), CI);
+        else if (OpClass == DXIL::OpCodeClass::VectorReduce)
+          scalarizeVectorReduce(HlslOP, CI);
         else if (Func->getReturnType()->isVectorTy())
-          for (auto U = Func->user_begin(), UE = Func->user_end(); U != UE;) {
-            CallInst *CI = cast<CallInst>(*(U++));
-            scalarizeVectorIntrinsic(HlslOP, CI);
-          }
+          scalarizeVectorIntrinsic(HlslOP, CI);
+        else
+          continue;
+
+        Changed = true;
       }
     }
     return Changed;
@@ -226,6 +237,33 @@ static void scalarizeVectorStore(hlsl::OP *HlslOP, const DataLayout &DL,
   CI->eraseFromParent();
 }
 
+static void scalarizeVectorReduce(hlsl::OP *HlslOP, CallInst *CI) {
+  IRBuilder<> Builder(CI);
+
+  OP::OpCode ReduceOp = OP::getOpCode(CI);
+
+  Value *VecArg = CI->getArgOperand(1);
+  Type *VecTy = VecArg->getType();
+
+  Value *Result = Builder.CreateExtractElement(VecArg, (uint64_t)0);
+  for (unsigned I = 1; I < VecTy->getVectorNumElements(); I++) {
+    Value *Elt = Builder.CreateExtractElement(VecArg, I);
+
+    switch (ReduceOp) {
+    case OP::OpCode::VectorReduceAnd:
+      Result = Builder.CreateAnd(Result, Elt);
+      break;
+    case OP::OpCode::VectorReduceOr:
+      Result = Builder.CreateOr(Result, Elt);
+      break;
+    default:
+      assert(false && "Unexpected VectorReduce OpCode");
+    }
+  }
+
+  CI->replaceAllUsesWith(Result);
+}
+
 // Scalarize native vector operation represented by `CI`, generating
 // scalar calls for each element of the its vector parameters.
 // Use `HlslOP` to retrieve the associated scalar op function.