Support more vector instructions in InstRebuilder.

Author: dmitryryintel

IGC/VectorCompiler/include/vc/Utils/General/InstRebuilder.h

@@ -13,6 +13,7 @@ SPDX-License-Identifier: MIT
 
 #include <llvm/ADT/ArrayRef.h>
 #include <llvm/IR/Instructions.h>
+#include <llvm/IR/IntrinsicInst.h>
 #include <llvm/IR/Use.h>
 #include <llvm/Support/Casting.h>
 
@@ -246,12 +247,20 @@ class InstructionRebuilder {
     return Replace;
   }
 
+  static std::vector<llvm::Value *>
+  getOrigOperands(llvm::Instruction &OrigInst) {
+    if (llvm::isa<llvm::IntrinsicInst>(OrigInst)) {
+      auto &OrigIntr = llvm::cast<llvm::IntrinsicInst>(OrigInst);
+      return {OrigIntr.arg_begin(), OrigIntr.arg_end()};
+    }
+    return {OrigInst.value_op_begin(), OrigInst.value_op_end()};
+  }
+
   // Takes arguments of the original instruction (OrigInst.User) and rewrites
   // the required ones with new values according to info in \p OrigInst
   static std::vector<llvm::Value *>
   createNewOperands(const InstToRebuild &OrigInst) {
-    std::vector<llvm::Value *> NewOperands{OrigInst.User->value_op_begin(),
-                                           OrigInst.User->value_op_end()};
+    auto NewOperands = getOrigOperands(*OrigInst.User);
     for (auto &&OpReplacement :
          zip(OrigInst.OperandNos, OrigInst.NewOperands)) {
       int OperandNo = std::get<0>(OpReplacement);

IGC/VectorCompiler/include/vc/Utils/General/Types.h

@@ -11,6 +11,8 @@ SPDX-License-Identifier: MIT
 
 #include <llvm/IR/DerivedTypes.h>
 
+#include "llvmWrapper/IR/DerivedTypes.h"
+
 namespace vc {
 
 // Returns potentially new pointer type with the provided \p AddrSpace
@@ -20,6 +22,16 @@ inline llvm::PointerType *changeAddrSpace(llvm::PointerType *OrigTy,
   return llvm::PointerType::get(OrigTy->getElementType(), AddrSpace);
 }
 
+// Changes addrspace inside a vector of pointers type.
+IGCLLVM::FixedVectorType *changeAddrSpace(IGCLLVM::FixedVectorType *OrigTy,
+                                          int AddrSpace);
+
+// Change addrspace of a pointer or a vector of pointers type.
+llvm::Type *changeAddrSpace(llvm::Type *OrigTy, int AddrSpace);
+
+// Get addrspace of a pointer or a vector of pointers type.
+int getAddrSpace(llvm::Type *PtrOrPtrVec);
+
 } // namespace vc
 
 #endif // VC_UTILS_GENERAL_TYPES_H

IGC/VectorCompiler/lib/GenXOpts/CMTrans/CMABI.cpp

@@ -58,6 +58,7 @@ SPDX-License-Identifier: MIT
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/InstVisitor.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Module.h"
 #include "llvm/InitializePasses.h"
@@ -541,8 +542,13 @@ bool CMABI::runOnSCC(CallGraphSCC &SCC) {
 // change will stop.
 static bool isRebuildTerminal(const Instruction &Inst) {
   // Result of a load inst is no longer a pointer so here propogation will stop.
-  return isa<LoadInst>(Inst) || isa<AddrSpaceCastInst>(Inst) ||
-         isa<StoreInst>(Inst);
+  if (isa<LoadInst>(Inst) || isa<AddrSpaceCastInst>(Inst) ||
+      isa<StoreInst>(Inst))
+    return true;
+  if (!isa<IntrinsicInst>(Inst))
+    return false;
+  auto IID = cast<IntrinsicInst>(Inst).getIntrinsicID();
+  return IID == Intrinsic::masked_gather || IID == Intrinsic::masked_scatter;
 }
 
 // Replaces uses of global variables with the corresponding allocas inside a

IGC/VectorCompiler/lib/Utils/General/CMakeLists.txt

@@ -10,6 +10,7 @@ set(GENERAL_UTILS_SOURCES
   BiF.cpp
   InstRebuilder.cpp
   BreakConst.cpp
+  Types.cpp
   )
 
 add_library(VCGeneralUtils ${GENERAL_UTILS_SOURCES})

IGC/VectorCompiler/lib/Utils/General/InstRebuilder.cpp

@@ -12,11 +12,52 @@ SPDX-License-Identifier: MIT
 #include "Probe/Assertion.h"
 #include "llvmWrapper/Support/Alignment.h"
 
+#include <llvm/ADT/ArrayRef.h>
 #include <llvm/IR/InstVisitor.h>
+#include <llvm/IR/IntrinsicInst.h>
+
+#include <algorithm>
+#include <iterator>
 
 using namespace llvm;
 using namespace vc;
 
+// Define intrinsic return type based on its arguments types.
+static Type *getIntrinsicRetTypeBasedOnArgs(Intrinsic::ID IID,
+                                            ArrayRef<Type *> ArgTys,
+                                            LLVMContext &C) {
+  switch (IID) {
+  case Intrinsic::masked_gather:
+    // "Pass through" operand.
+    return ArgTys[3];
+  case Intrinsic::masked_scatter:
+    return Type::getVoidTy(C);
+  default:
+    IGC_ASSERT_MESSAGE(0, "unsupported intrinsic");
+    return nullptr;
+  }
+}
+
+// Emits list of overloaded types for the provided intrinsic. Takes all the
+// types that may take part in overloading: return type, types of arguments.
+static std::vector<Type *>
+getIntrinsicOverloadedTypes(Intrinsic::ID IID, Type *RetTy,
+                            ArrayRef<Type *> ArgTys) {
+  IGC_ASSERT_MESSAGE(RetTy, "wrong argument");
+  // FIXME: generalize for any ID using IntrinsicInfoTable.
+  switch (IID) {
+  case Intrinsic::masked_gather:
+    // Loaded value type and pointer operand type.
+    return {RetTy, ArgTys[0]};
+  case Intrinsic::masked_scatter:
+    // Value and pointer types.
+    return {ArgTys[0], ArgTys[1]};
+  default:
+    IGC_ASSERT_MESSAGE(0, "unsupported intrinsic");
+    return {};
+  }
+}
+
 namespace {
 class cloneInstWithNewOpsImpl
     : public InstVisitor<cloneInstWithNewOpsImpl, Instruction *> {
@@ -61,18 +102,17 @@ class cloneInstWithNewOpsImpl
   // type addrspace corresponds with this operand.
   CastInst *visitBitCastInst(BitCastInst &OrigCast) {
     Value &NewOp = getSingleNewOperand();
-    if (isa<PointerType>(OrigCast.getType()))
-      return visitPointerBitCastInst(OrigCast);
+    // If the operand changed addrspace the bitcast type should change it too.
+    if (OrigCast.getType()->isPtrOrPtrVectorTy())
+      return visitPtrOrPtrVectorBitCastInst(OrigCast);
     return new BitCastInst{&NewOp, OrigCast.getType()};
   }
 
-  CastInst *visitPointerBitCastInst(BitCastInst &OrigCast) {
+  CastInst *visitPtrOrPtrVectorBitCastInst(BitCastInst &OrigCast) {
     Value &NewOp = getSingleNewOperand();
-    auto NewOpAS = cast<PointerType>(NewOp.getType())->getAddressSpace();
-    // If the operand changed addrspace the bitcast type should change it too.
-    return new BitCastInst{
-        &NewOp,
-        changeAddrSpace(cast<PointerType>(OrigCast.getType()), NewOpAS)};
+    auto NewOpAS = getAddrSpace(NewOp.getType());
+    return new BitCastInst{&NewOp,
+                           changeAddrSpace(OrigCast.getType(), NewOpAS)};
   }
 
   CastInst *visitAddrSpaceCastInst(AddrSpaceCastInst &OrigCast) {
@@ -89,6 +129,24 @@ class cloneInstWithNewOpsImpl
     return SelectInst::Create(NewOperands[0], NewOperands[1], NewOperands[2]);
   }
 
+  IntrinsicInst *visitIntrinsicInst(IntrinsicInst &OrigIntrinsic) {
+    auto IID = OrigIntrinsic.getIntrinsicID();
+    if (IID != Intrinsic::masked_gather && IID != Intrinsic::masked_scatter) {
+      IGC_ASSERT_MESSAGE(0, "yet unsupported instruction");
+      return nullptr;
+    }
+    std::vector<Type *> ArgTys;
+    std::transform(NewOperands.begin(), NewOperands.end(),
+                   std::back_inserter(ArgTys),
+                   [](Value *Operand) { return Operand->getType(); });
+    auto *RetTy =
+        getIntrinsicRetTypeBasedOnArgs(IID, ArgTys, OrigIntrinsic.getContext());
+    auto OverloadedTys = getIntrinsicOverloadedTypes(IID, RetTy, ArgTys);
+    auto *Decl = Intrinsic::getDeclaration(OrigIntrinsic.getModule(), IID,
+                                           OverloadedTys);
+    return cast<IntrinsicInst>(CallInst::Create(Decl, NewOperands));
+  }
+
 private:
   Value &getSingleNewOperand() {
     IGC_ASSERT_MESSAGE(

IGC/VectorCompiler/lib/Utils/General/Types.cpp

@@ -0,0 +1,44 @@
+/*========================== begin_copyright_notice ============================
+
+Copyright (C) 2021 Intel Corporation
+
+SPDX-License-Identifier: MIT
+
+============================= end_copyright_notice ===========================*/
+
+#include "vc/Utils/General/Types.h"
+
+#include "Probe/Assertion.h"
+
+#include <llvm/Support/Casting.h>
+
+using namespace llvm;
+
+IGCLLVM::FixedVectorType *vc::changeAddrSpace(IGCLLVM::FixedVectorType *OrigTy,
+                                              int AddrSpace) {
+  IGC_ASSERT_MESSAGE(OrigTy, "wrong argument");
+  auto *PointeeTy = OrigTy->getElementType()->getPointerElementType();
+  auto EC = OrigTy->getElementCount();
+  return IGCLLVM::FixedVectorType::get(
+      llvm::PointerType::get(PointeeTy, AddrSpace), EC);
+}
+
+Type *vc::changeAddrSpace(Type *OrigTy, int AddrSpace) {
+  IGC_ASSERT_MESSAGE(OrigTy, "wrong argument");
+  IGC_ASSERT_MESSAGE(
+      OrigTy->isPtrOrPtrVectorTy(),
+      "wrong argument: pointer or vector of pointers type is expected");
+  if (OrigTy->isPointerTy())
+    return changeAddrSpace(cast<PointerType>(OrigTy), AddrSpace);
+  return changeAddrSpace(cast<IGCLLVM::FixedVectorType>(OrigTy), AddrSpace);
+}
+
+int vc::getAddrSpace(Type *PtrOrPtrVec) {
+  IGC_ASSERT_MESSAGE(PtrOrPtrVec, "wrong argument");
+  IGC_ASSERT_MESSAGE(
+      PtrOrPtrVec->isPtrOrPtrVectorTy(),
+      "wrong argument: pointer or vector of pointers type is expected");
+  if (PtrOrPtrVec->isPointerTy())
+    return PtrOrPtrVec->getPointerAddressSpace();
+  return PtrOrPtrVec->getVectorElementType()->getPointerAddressSpace();
+}