Let backends decide if they support kernel stacking for specific inst…

…ructions

Summary: Pull Request resolved: #2876

Differential Revision: D15285608

Pulled By: opti-mix

fbshipit-source-id: cefe45707d2332eb7e150c2587a7efdfd4a433f4

include/glow/LLVMIRCodeGen/LLVMIRGen.h

@@ -330,6 +330,14 @@ class LLVMIRGen {
   /// \returns true if a global symbol \p GV needs to be preserved in the module
   /// and not interalized during optimizations.
   virtual bool preserveSymbol(const llvm::GlobalValue &GV);
+  /// \returns true if an instruction \p I can be part of a data parallel
+  /// kernel. This gives backends a possibility to provide a custom logic to
+  /// decide on a per-instruction basis what can be part of data parallel
+  /// kernels. Typically an instruction which is isDataParallel() can be part of
+  /// a data parallel kernel. But a backend may decide that a specific
+  /// instruction \p I cannot be part of data-parallel kernels, because there is
+  /// no support for this functionality in this backend yet.
+  virtual bool canBePartOfDataParallelKernel(const glow::Instruction *I) const;
 };
 
 } // namespace glow

lib/Backends/CPU/CPULLVMIRGen.cpp

@@ -36,7 +36,7 @@ void CPULLVMIRGen::generateLLVMIRForModule(llvm::IRBuilder<> &builder) {
 void CPULLVMIRGen::generateLLVMIRForInstr(llvm::IRBuilder<> &builder,
                                           const glow::Instruction *I) {
   setCurrentDebugLocation(builder, I);
-  assert(!I->isDataParallel() &&
+  assert(!canBePartOfDataParallelKernel(I) &&
          "data parallel instructions are not handled here");
   // Perform any backend-specific code generation here and delegate everything
   // else to LLVMIRGen.
@@ -121,7 +121,8 @@ void CPULLVMIRGen::generateLLVMIRForDataParallelInstr(
     llvm::Function *kernel, llvm::DenseMap<Value *, int> &bufferToArgNum,
     llvm::Value *loopCount) {
   setCurrentDebugLocation(builder, I);
-  assert(I->isDataParallel() && "Expected a data parallel instruction");
+  assert(canBePartOfDataParallelKernel(I) &&
+         "Expected a data parallel instruction");
   // Perform any backend-specific code generation here and delegate everything
   // else to LLVMIRGen.
   switch (I->getKind()) {

lib/LLVMIRCodeGen/LLVMIRGen.cpp

@@ -676,7 +676,8 @@ void LLVMIRGen::emitDataParallelKernelImpl(
   // instruction.
   for (auto &BI : bundle) {
     // Name of the stacked operation to be invoked.
-    assert(BI->isDataParallel() && "Data parallel operation is expected");
+    assert(canBePartOfDataParallelKernel(BI) &&
+           "Data parallel operation is expected");
     generateLLVMIRForDataParallelInstr(kernelBuilder, BI, kernelFunc,
                                        bufferToArgNum, kernelLoopIdx);
   }
@@ -786,7 +787,7 @@ void LLVMIRGen::generateLLVMIRForModule(llvm::IRBuilder<> &builder) {
   // instructions and emit them.
   llvm::SmallVector<const Instruction *, 32> bundle;
   for (auto &I : instrs) {
-    if (!I.isDataParallel()) {
+    if (!canBePartOfDataParallelKernel(&I)) {
       // Ignore memory management instructions as they are handled by the
       // MemoryManager and are NOPs for a JIT.
       if (isa<AllocActivationInst>(&I) || isa<DeallocActivationInst>(&I) ||
@@ -845,7 +846,8 @@ void LLVMIRGen::generateLLVMIRForDataParallelInstr(
     llvm::Function *kernel, llvm::DenseMap<Value *, int> &bufferToArgNum,
     llvm::Value *loopCount) {
   setCurrentDebugLocation(builder, I);
-  assert(I->isDataParallel() && "Expected a data parallel instruction");
+  assert(canBePartOfDataParallelKernel(I) &&
+         "Instruction cannot be part of a data parallel kernel");
   switch (I->getKind()) {
 
 #define ARITHMETIC_UNARY_OP_WITH_IMM_CASE(INST_NAME_, FUN_NAME_, VALUE_)       \
@@ -1341,7 +1343,7 @@ void LLVMIRGen::generateLLVMIRForDataParallelInstr(
 void LLVMIRGen::generateLLVMIRForInstr(llvm::IRBuilder<> &builder,
                                        const glow::Instruction *I) {
   setCurrentDebugLocation(builder, I);
-  assert(!I->isDataParallel() &&
+  assert((!canBePartOfDataParallelKernel(I)) &&
          "data parallel instructions are not handled here");
   switch (I->getKind()) {
   case Kinded::Kind::MatMulInstKind: {
@@ -2392,3 +2394,8 @@ unsigned LLVMIRGen::getLibjitSizeTWidth() const {
 bool LLVMIRGen::isEligibleForSpecialization(const llvm::CallInst *call) {
   return true;
 }
+
+bool LLVMIRGen::canBePartOfDataParallelKernel(
+    const glow::Instruction *I) const {
+  return I->isDataParallel();
+}