Implement support for @llvm.sadd.sat intrinsic. (#1524)

* Implement support for @llvm.sadd.sat intrinsic.
* Rename LowerSaddWithOverflow to LowerSaddIntrinsics.

Author: jcranmer-intel

include/LLVMSPIRVLib.h

@@ -57,7 +57,7 @@ void initializeSPIRVLowerBoolLegacyPass(PassRegistry &);
 void initializeSPIRVLowerConstExprLegacyPass(PassRegistry &);
 void initializeSPIRVLowerOCLBlocksLegacyPass(PassRegistry &);
 void initializeSPIRVLowerMemmoveLegacyPass(PassRegistry &);
-void initializeSPIRVLowerSaddWithOverflowLegacyPass(PassRegistry &);
+void initializeSPIRVLowerSaddIntrinsicsLegacyPass(PassRegistry &);
 void initializeSPIRVRegularizeLLVMLegacyPass(PassRegistry &);
 void initializeSPIRVToOCL12LegacyPass(PassRegistry &);
 void initializeSPIRVToOCL20LegacyPass(PassRegistry &);
@@ -183,7 +183,7 @@ ModulePass *createSPIRVLowerOCLBlocksLegacy();
 ModulePass *createSPIRVLowerMemmoveLegacy();
 
 /// Create a pass for lowering llvm.sadd.with.overflow
-ModulePass *createSPIRVLowerSaddWithOverflowLegacy();
+ModulePass *createSPIRVLowerSaddIntrinsicsLegacy();
 
 /// Create a pass for regularize LLVM module to be translated to SPIR-V.
 ModulePass *createSPIRVRegularizeLLVMLegacy();

lib/SPIRV/CMakeLists.txt

@@ -14,7 +14,7 @@ add_llvm_library(LLVMSPIRVLib
   SPIRVLowerConstExpr.cpp
   SPIRVLowerMemmove.cpp
   SPIRVLowerOCLBlocks.cpp
-  SPIRVLowerSaddWithOverflow.cpp
+  SPIRVLowerSaddIntrinsics.cpp
   SPIRVReader.cpp
   SPIRVRegularizeLLVM.cpp
   SPIRVToLLVMDbgTran.cpp

lib/SPIRV/SPIRVLowerSaddIntrinsics.cpp

@@ -0,0 +1,192 @@
+//===- SPIRVLowerSaddIntrinsics.cpp - Lower llvm.sadd.* -------------------===//
+//
+//                     The LLVM/SPIRV Translator
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+// Copyright (c) 2020 Intel Corporation. All rights reserved.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal with the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// Redistributions of source code must retain the above copyright notice,
+// this list of conditions and the following disclaimers.
+// Redistributions in binary form must reproduce the above copyright notice,
+// this list of conditions and the following disclaimers in the documentation
+// and/or other materials provided with the distribution.
+// Neither the names of Intel Corporation, nor the names of its
+// contributors may be used to endorse or promote products derived from this
+// Software without specific prior written permission.
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS WITH
+// THE SOFTWARE.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements lowering of llvm.sadd.* into basic LLVM
+// operations. Probably, in the future this pass can be generalized for other
+// function calls
+//
+//===----------------------------------------------------------------------===//
+#define DEBUG_TYPE "spv-lower-llvm_sadd_intrinsics"
+
+#include "SPIRVLowerSaddIntrinsics.h"
+#include "LLVMSaddWithOverflow.h"
+
+#include "LLVMSPIRVLib.h"
+#include "SPIRVError.h"
+#include "libSPIRV/SPIRVDebug.h"
+
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IRReader/IRReader.h"
+#include "llvm/Linker/Linker.h"
+#include "llvm/Support/SourceMgr.h"
+
+using namespace llvm;
+using namespace SPIRV;
+
+namespace SPIRV {
+
+void SPIRVLowerSaddIntrinsicsBase::replaceSaddOverflow(Function &F) {
+  assert(F.getIntrinsicID() == Intrinsic::sadd_with_overflow);
+
+  StringRef IntrinsicName = F.getName();
+  std::string FuncName = "llvm_sadd_with_overflow_i";
+  if (IntrinsicName.endswith(".i16"))
+    FuncName += "16";
+  else if (IntrinsicName.endswith(".i32"))
+    FuncName += "32";
+  else if (IntrinsicName.endswith(".i64"))
+    FuncName += "64";
+  else {
+    assert(false &&
+           "Unsupported overloading of llvm.sadd.with.overflow intrinsic");
+    return;
+  }
+
+  // Redirect @llvm.sadd.with.overflow.* call to the function we have in
+  // the loaded module @llvm_sadd_with_overflow_*
+  Function *ReplacementFunc = Mod->getFunction(FuncName);
+  if (!ReplacementFunc) { // This function needs linking.
+    Mod->getOrInsertFunction(FuncName, F.getFunctionType());
+    // Read LLVM IR with the intrinsic's implementation
+    SMDiagnostic Err;
+    auto MB = MemoryBuffer::getMemBuffer(LLVMSaddWithOverflow);
+    auto SaddWithOverflowModule =
+        parseIR(MB->getMemBufferRef(), Err, *Context,
+                [&](StringRef) { return Mod->getDataLayoutStr(); });
+    if (!SaddWithOverflowModule) {
+      std::string ErrMsg;
+      raw_string_ostream ErrStream(ErrMsg);
+      Err.print("", ErrStream);
+      SPIRVErrorLog EL;
+      EL.checkError(false, SPIRVEC_InvalidLlvmModule, ErrMsg);
+      return;
+    }
+
+    // Link in the intrinsic's implementation.
+    if (!Linker::linkModules(*Mod, std::move(SaddWithOverflowModule),
+                             Linker::LinkOnlyNeeded))
+      TheModuleIsModified = true;
+
+    ReplacementFunc = Mod->getFunction(FuncName);
+    assert(ReplacementFunc && "How did we not link in the necessary function?");
+  }
+
+  F.replaceAllUsesWith(ReplacementFunc);
+}
+
+void SPIRVLowerSaddIntrinsicsBase::replaceSaddSat(Function &F) {
+  assert(F.getIntrinsicID() == Intrinsic::sadd_sat);
+
+  SmallVector<IntrinsicInst *, 4> Intrinsics;
+  for (User *U : F.users()) {
+    if (auto *II = dyn_cast<IntrinsicInst>(U))
+      Intrinsics.push_back(II);
+  }
+
+  // Get the corresponding sadd_with_overflow intrinsic for the sadd_sat.
+  Type *IntTy = F.getFunctionType()->getReturnType();
+  Function *SaddO =
+      Intrinsic::getDeclaration(Mod, Intrinsic::sadd_with_overflow, IntTy);
+
+  // Replace all uses of the intrinsic with equivalent code relying on
+  // sadd_with_overflow
+  IRBuilder<> Builder(F.getContext());
+  unsigned BitWidth = IntTy->getIntegerBitWidth();
+  Value *IntMin = Builder.getInt(APInt::getSignedMinValue(BitWidth));
+  Value *ShiftWidth = Builder.getIntN(BitWidth, BitWidth - 1);
+
+  for (IntrinsicInst *II : Intrinsics) {
+    Builder.SetInsertPoint(II);
+    // {res, overflow} = @llvm.sadd_with_overflow(a, b)
+    // sadd_sat(a, b) => overflow ? (res >> bitwidth) ^ intmin : res;
+    Value *StructRes =
+        Builder.CreateCall(SaddO, {II->getArgOperand(0), II->getArgOperand(1)});
+    Value *Sum = Builder.CreateExtractValue(StructRes, 0);
+    Value *Overflow = Builder.CreateExtractValue(StructRes, 1);
+    Value *OverflowedRes =
+        Builder.CreateXor(Builder.CreateAShr(Sum, ShiftWidth), IntMin);
+    Value *Result = Builder.CreateSelect(Overflow, OverflowedRes, Sum);
+    II->replaceAllUsesWith(Result);
+    II->eraseFromParent();
+  }
+
+  // Now replace the sadd_with_overflow intrinsic itself.
+  replaceSaddOverflow(*SaddO);
+}
+
+bool SPIRVLowerSaddIntrinsicsBase::runLowerSaddIntrinsics(Module &M) {
+  Context = &M.getContext();
+  Mod = &M;
+  for (Function &F : M) {
+    Intrinsic::ID IntrinId = F.getIntrinsicID();
+    if (IntrinId == Intrinsic::sadd_with_overflow)
+      replaceSaddOverflow(F);
+    else if (IntrinId == Intrinsic::sadd_sat)
+      replaceSaddSat(F);
+  }
+
+  verifyRegularizationPass(M, "SPIRVLowerSaddIntrinsics");
+  return TheModuleIsModified;
+}
+
+llvm::PreservedAnalyses
+SPIRVLowerSaddIntrinsicsPass::run(llvm::Module &M,
+                                  llvm::ModuleAnalysisManager &MAM) {
+  return runLowerSaddIntrinsics(M) ? llvm::PreservedAnalyses::none()
+                                   : llvm::PreservedAnalyses::all();
+}
+
+SPIRVLowerSaddIntrinsicsLegacy::SPIRVLowerSaddIntrinsicsLegacy()
+    : ModulePass(ID) {
+  initializeSPIRVLowerSaddIntrinsicsLegacyPass(
+      *PassRegistry::getPassRegistry());
+}
+
+bool SPIRVLowerSaddIntrinsicsLegacy::runOnModule(Module &M) {
+  return runLowerSaddIntrinsics(M);
+}
+
+char SPIRVLowerSaddIntrinsicsLegacy::ID = 0;
+
+} // namespace SPIRV
+
+INITIALIZE_PASS(SPIRVLowerSaddIntrinsicsLegacy,
+                "spv-lower-llvm_sadd_intrinsics",
+                "Lower llvm.sadd.* intrinsics", false, false)
+
+ModulePass *llvm::createSPIRVLowerSaddIntrinsicsLegacy() {
+  return new SPIRVLowerSaddIntrinsicsLegacy();
+}

lib/SPIRV/SPIRVLowerSaddWithOverflow.h renamed to lib/SPIRV/SPIRVLowerSaddIntrinsics.h

@@ -1,4 +1,4 @@
-//===- SPIRVLowerSaddWithOverflowBase.h - sadd lowering  --------*- C++ -*-===//
+//===- SPIRVLowerSaddIntrinsics.h - sadd lowering  --------------*- C++ -*-===//
 //
 //                     The LLVM/SPIR-V Translator
 //
@@ -32,42 +32,41 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef SPIRV_SPIRVLOWERSADDWITHOVERFLOW_H
-#define SPIRV_SPIRVLOWERSADDWITHOVERFLOW_H
+#ifndef SPIRV_SPIRVLOWERSADDINTRINSICS_H
+#define SPIRV_SPIRVLOWERSADDINTRINSICS_H
 
-#include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Pass.h"
 
 namespace SPIRV {
 
-class SPIRVLowerSaddWithOverflowBase
-    : public llvm::InstVisitor<SPIRVLowerSaddWithOverflowBase> {
+class SPIRVLowerSaddIntrinsicsBase {
 public:
-  SPIRVLowerSaddWithOverflowBase() : Context(nullptr), Mod(nullptr) {}
-  virtual ~SPIRVLowerSaddWithOverflowBase() {}
-  virtual void visitIntrinsicInst(llvm::CallInst &I);
+  SPIRVLowerSaddIntrinsicsBase() : Context(nullptr), Mod(nullptr) {}
 
-  bool runLowerSaddWithOverflow(llvm::Module &M);
+  bool runLowerSaddIntrinsics(llvm::Module &M);
 
 private:
+  void replaceSaddOverflow(llvm::Function &F);
+  void replaceSaddSat(llvm::Function &F);
+
   llvm::LLVMContext *Context;
   llvm::Module *Mod;
   bool TheModuleIsModified = false;
 };
 
-class SPIRVLowerSaddWithOverflowPass
-    : public llvm::PassInfoMixin<SPIRVLowerSaddWithOverflowPass>,
-      public SPIRVLowerSaddWithOverflowBase {
+class SPIRVLowerSaddIntrinsicsPass
+    : public llvm::PassInfoMixin<SPIRVLowerSaddIntrinsicsPass>,
+      public SPIRVLowerSaddIntrinsicsBase {
 public:
   llvm::PreservedAnalyses run(llvm::Module &M,
                               llvm::ModuleAnalysisManager &MAM);
 };
 
-class SPIRVLowerSaddWithOverflowLegacy : public llvm::ModulePass,
-                                         public SPIRVLowerSaddWithOverflowBase {
+class SPIRVLowerSaddIntrinsicsLegacy : public llvm::ModulePass,
+                                       public SPIRVLowerSaddIntrinsicsBase {
 public:
-  SPIRVLowerSaddWithOverflowLegacy();
+  SPIRVLowerSaddIntrinsicsLegacy();
 
   bool runOnModule(llvm::Module &M) override;
 
@@ -76,4 +75,4 @@ class SPIRVLowerSaddWithOverflowLegacy : public llvm::ModulePass,
 
 } // namespace SPIRV
 
-#endif // SPIRV_SPIRVLOWERSADDWITHOVERFLOW_H
+#endif // SPIRV_SPIRVLOWERSADDINTRINSICS_H