Add physical register counting functions

llvm-svn: 118397

llvm/lib/Target/PTX/CMakeLists.txt

@@ -15,6 +15,7 @@ add_llvm_target(PTXCodeGen
   PTXInstrInfo.cpp
   PTXMCAsmInfo.cpp
   PTXMCAsmStreamer.cpp
+  PTXMFInfoExtract.cpp
   PTXRegisterInfo.cpp
   PTXSubtarget.cpp
   PTXTargetMachine.cpp

llvm/lib/Target/PTX/PTX.h

@@ -24,6 +24,9 @@ namespace llvm {
   FunctionPass *createPTXISelDag(PTXTargetMachine &TM,
                                  CodeGenOpt::Level OptLevel);
 
+  FunctionPass *createPTXMFInfoExtract(PTXTargetMachine &TM,
+                                       CodeGenOpt::Level OptLevel);
+
   extern Target ThePTXTarget;
 } // namespace llvm;
 

llvm/lib/Target/PTX/PTXISelLowering.cpp

@@ -13,6 +13,7 @@
 
 #include "PTX.h"
 #include "PTXISelLowering.h"
+#include "PTXMachineFunctionInfo.h"
 #include "PTXRegisterInfo.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -107,6 +108,9 @@ SDValue PTXTargetLowering::
                        SmallVectorImpl<SDValue> &InVals) const {
   if (isVarArg) llvm_unreachable("PTX does not support varargs");
 
+  MachineFunction &MF = DAG.getMachineFunction();
+  PTXMachineFunctionInfo *MFI = MF.getInfo<PTXMachineFunctionInfo>();
+
   lower_argument_func lower_argument;
 
   switch (CallConv) {
@@ -114,9 +118,11 @@ SDValue PTXTargetLowering::
       llvm_unreachable("Unsupported calling convention");
       break;
     case CallingConv::PTX_Kernel:
+      MFI->setKernel();
       lower_argument = lower_kernel_argument;
       break;
     case CallingConv::PTX_Device:
+      MFI->setKernel(false);
       lower_argument = lower_device_argument;
       break;
   }
@@ -137,8 +143,15 @@ SDValue PTXTargetLowering::
     unsigned reg;
     SDValue arg = lower_argument(i, Chain, dl, VT, entry, DAG, &reg);
     InVals.push_back(arg);
+
+    if (!MFI->isDoneAddArg())
+      MFI->addArgReg(reg);
   }
 
+  // Make sure we don't add argument registers twice
+  if (!MFI->isDoneAddArg())
+    MFI->doneAddArg();
+
   return Chain;
 }
 
@@ -174,6 +187,10 @@ SDValue PTXTargetLowering::
   SDValue Flag;
   unsigned reg = PTX::R0;
 
+  MachineFunction &MF = DAG.getMachineFunction();
+  PTXMachineFunctionInfo *MFI = MF.getInfo<PTXMachineFunctionInfo>();
+  MFI->setRetReg(reg);
+
   // If this is the first return lowered for this function, add the regs to the
   // liveout set for the function
   if (DAG.getMachineFunction().getRegInfo().liveout_empty())

llvm/lib/Target/PTX/PTXMFInfoExtract.cpp

@@ -0,0 +1,76 @@
+//===-- PTXMFInfoExtract.cpp - Extract PTX machine function info ----------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines an information extractor for PTX machine functions.
+//
+//===----------------------------------------------------------------------===//
+
+#define DEBUG_TYPE "ptx-mf-info-extract"
+
+#include "PTX.h"
+#include "PTXTargetMachine.h"
+#include "PTXMachineFunctionInfo.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace llvm {
+  /// PTXMFInfoExtract - PTX specific code to extract of PTX machine
+  /// function information for PTXAsmPrinter
+  ///
+  class PTXMFInfoExtract : public MachineFunctionPass {
+    private:
+      static char ID;
+
+    public:
+      PTXMFInfoExtract(PTXTargetMachine &TM, CodeGenOpt::Level OptLevel)
+        : MachineFunctionPass(ID) {}
+
+      virtual bool runOnMachineFunction(MachineFunction &MF);
+
+      virtual const char *getPassName() const {
+        return "PTX Machine Function Info Extractor";
+      }
+  }; // class PTXMFInfoExtract
+} // namespace llvm
+
+using namespace llvm;
+
+char PTXMFInfoExtract::ID = 0;
+
+bool PTXMFInfoExtract::runOnMachineFunction(MachineFunction &MF) {
+  PTXMachineFunctionInfo *MFI = MF.getInfo<PTXMachineFunctionInfo>();
+  MachineRegisterInfo &MRI = MF.getRegInfo();
+
+  DEBUG(dbgs() << "****** PTX FUNCTION LOCAL VAR REG DEF ******\n");
+
+  unsigned reg_ret = MFI->retReg();
+
+  // FIXME: This is a slow linear scanning
+  for (unsigned reg = PTX::NoRegister + 1; reg < PTX::NUM_TARGET_REGS; ++reg)
+    if (MRI.isPhysRegUsed(reg) && reg != reg_ret && !MFI->isArgReg(reg))
+      MFI->addLocalVarReg(reg);
+
+  // Notify MachineFunctionInfo that I've done adding local var reg
+  MFI->doneAddLocalVar();
+
+  DEBUG(for (PTXMachineFunctionInfo::reg_iterator
+             i = MFI->localVarRegBegin(), e = MFI->localVarRegEnd();
+	     i != e; ++i)
+        dbgs() << "Used Reg: " << *i << "\n";);
+
+  return false;
+}
+
+FunctionPass *llvm::createPTXMFInfoExtract(PTXTargetMachine &TM,
+                                           CodeGenOpt::Level OptLevel) {
+  return new PTXMFInfoExtract(TM, OptLevel);
+}

llvm/lib/Target/PTX/PTXMachineFunctionInfo.h

@@ -0,0 +1,78 @@
+//===- PTXMachineFuctionInfo.h - PTX machine function info -------*- C++ -*-==//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares PTX-specific per-machine-function information.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PTX_MACHINE_FUNCTION_INFO_H
+#define PTX_MACHINE_FUNCTION_INFO_H
+
+#include "PTX.h"
+#include "llvm/CodeGen/MachineFunction.h"
+
+namespace llvm {
+/// PTXMachineFunctionInfo - This class is derived from MachineFunction and
+/// contains private PTX target-specific information for each MachineFunction.
+///
+class PTXMachineFunctionInfo : public MachineFunctionInfo {
+private:
+  bool is_kernel;
+  std::vector<unsigned> reg_arg, reg_local_var;
+  unsigned reg_ret;
+  bool _isDoneAddArg;
+
+public:
+  PTXMachineFunctionInfo(MachineFunction &MF)
+    : is_kernel(false), reg_ret(PTX::NoRegister), _isDoneAddArg(false) {
+      reg_arg.reserve(32);
+      reg_local_var.reserve(64);
+    }
+
+  void setKernel(bool _is_kernel=true) { is_kernel = _is_kernel; }
+
+  void addArgReg(unsigned reg) { reg_arg.push_back(reg); }
+  void addLocalVarReg(unsigned reg) { reg_local_var.push_back(reg); }
+  void setRetReg(unsigned reg) { reg_ret = reg; }
+
+  void doneAddArg(void) {
+    std::sort(reg_arg.begin(), reg_arg.end());
+    _isDoneAddArg = true;
+  }
+  void doneAddLocalVar(void) {
+    std::sort(reg_local_var.begin(), reg_local_var.end());
+  }
+
+  bool isDoneAddArg(void) { return _isDoneAddArg; }
+
+  bool isKernel() const { return is_kernel; }
+
+  typedef std::vector<unsigned>::const_iterator reg_iterator;
+
+  bool argRegEmpty() const { return reg_arg.empty(); }
+  reg_iterator argRegBegin() const { return reg_arg.begin(); }
+  reg_iterator argRegEnd()   const { return reg_arg.end(); }
+
+  bool localVarRegEmpty() const { return reg_local_var.empty(); }
+  reg_iterator localVarRegBegin() const { return reg_local_var.begin(); }
+  reg_iterator localVarRegEnd()   const { return reg_local_var.end(); }
+
+  unsigned retReg() const { return reg_ret; }
+
+  bool isArgReg(unsigned reg) const {
+    return std::binary_search(reg_arg.begin(), reg_arg.end(), reg);
+  }
+
+  bool isLocalVarReg(unsigned reg) const {
+    return std::binary_search(reg_local_var.begin(), reg_local_var.end(), reg);
+  }
+}; // class PTXMachineFunctionInfo
+} // namespace llvm
+
+#endif // PTX_MACHINE_FUNCTION_INFO_H

llvm/lib/Target/PTX/PTXTargetMachine.cpp

@@ -41,5 +41,6 @@ PTXTargetMachine::PTXTargetMachine(const Target &T,
 bool PTXTargetMachine::addInstSelector(PassManagerBase &PM,
                                        CodeGenOpt::Level OptLevel) {
   PM.add(createPTXISelDag(*this, OptLevel));
+  PM.add(createPTXMFInfoExtract(*this, OptLevel));
   return false;
 }