[pISA][Translator] Refactor function signature for pISA printing

This is a refactor of intel#500 to populate function signature printing without chaning IR.
bader · Mar 28, 2024 · 2d63681 · 2d63681
1 parent c6f6f44
commit 2d63681
Showing 1 changed file with 31 additions and 11 deletions.
diff --git a/llvm/lib/Target/pISA/pISAAsmPrinter.cpp b/llvm/lib/Target/pISA/pISAAsmPrinter.cpp
@@ -65,6 +65,7 @@ class pISAAsmPrinter : public AsmPrinter {
 
 private:
 
+  bool isFromTranslator() { return PTI; }
   void collectRegDcls(pISA::RegDcls &);
   void collectStackVariableDcls(pISA::StackVariableDcls &);
 
@@ -80,6 +81,8 @@ class pISAAsmPrinter : public AsmPrinter {
   void collectKernelParameters(pISA::FunctionSignature &);
   void collectFunctionDirectiveAndName(pISA::FunctionDirectiveAndName &DN,
                                        const Function &F);
+  void collectFunctionDirectiveAndName(pISA::FunctionDirectiveAndName &DN,
+                                       const pisa::pISAFunction &F);
   pISA::LinkageTy collectLinkage(const GlobalValue &V);
   std::optional<unsigned> collectDbgBinding(const Function &F);
   void collectGlobalVariable(pISA::GlobalVariable &PGV,
@@ -285,15 +288,6 @@ bool pISAAsmPrinter::doInitialization(Module &M) {
 
   // Emit Module level function decl
   for (auto &F : M.getFunctionList()) {
-    if (PTI) {
-      auto *Func = PTI->getpISAMod()->getFunction(F.getName().str());
-      if (Func->CC == pisa::Function) {
-        F.setCallingConv(CallingConv::PISA_FUNC);
-      } else {
-        assert(Func->CC == pisa::Kernel);
-        F.setCallingConv(CallingConv::PISA_KERNEL);
-      }
-    }
     if (!F.isDeclaration() ||
         F.getCallingConv() != CallingConv::PISA_FUNC) // avoid llvm builtins
       continue;
@@ -485,6 +479,7 @@ void pISAAsmPrinter::collectFunctionParameters(pISA::FunctionSignature &Sig) {
 
 void pISAAsmPrinter::collectFunctionDirectiveAndName(
     pISA::FunctionDirectiveAndName &DN, const Function &F) {
+
   DN.CC = F.getCallingConv();
   if (DN.CC == CallingConv::PISA_FUNC)
     DN.Linkage = collectLinkage(F);
@@ -497,10 +492,35 @@ void pISAAsmPrinter::collectFunctionDirectiveAndName(
   }
 }
 
+void pISAAsmPrinter::collectFunctionDirectiveAndName(
+    pISA::FunctionDirectiveAndName &DN, const pisa::pISAFunction &F) {
+  if (F.CC == pisa::Function) {
+    DN.CC = CallingConv::PISA_FUNC;
+  } else {
+    assert(F.CC == pisa::Kernel);
+    DN.CC = CallingConv::PISA_KERNEL;
+  }
+  DN.Name = F.Name;
+  if (DN.CC == CallingConv::PISA_FUNC) {
+    if (!F.Rets.empty()) {
+      assert(F.Rets.size() == 1);
+      assert(F.Rets[0].Ty.NumElts == 1);
+      DN.RetLLT = LLT::scalar(F.Rets[0].Ty.EltSize);
+    }
+  }
+}
+
 void pISAAsmPrinter::collectFunctionSignature(pISA::FunctionSignature &Sig) {
   Function &F = MF->getFunction();
-  collectFunctionDirectiveAndName(Sig.DN, F);
-  if (F.getCallingConv() != CallingConv::PISA_FUNC)
+  if (isFromTranslator()) {
+    collectFunctionDirectiveAndName(
+        Sig.DN, *PTI->getpISAMod()->getFunction(F.getName().str()));
+    if (Sig.DN.CC == CallingConv::PISA_FUNC)
+      Sig.DN.Linkage = collectLinkage(F);
+  } else {
+    collectFunctionDirectiveAndName(Sig.DN, F);
+  }
+  if (Sig.DN.CC != CallingConv::PISA_FUNC)
     collectKernelParameters(Sig);
   else
     collectFunctionParameters(Sig);