From d64ee86a3130f8eeb75d09193c38c621d7565eba Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Sun, 26 Mar 2023 13:59:11 -0700
Subject: Add PyTorch C++ binding generation. (#2734)

* Add PyTorch C++ binding generation.

* fix

---------

Co-authored-by: Yong He <yhe@nvidia.com>
---
 source/slang/slang-lower-to-ir.cpp | 94 ++++++++++++++++++++++----------------
 1 file changed, 55 insertions(+), 39 deletions(-)

(limited to 'source/slang/slang-lower-to-ir.cpp')

diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 7144b3450..9d424d1e8 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -1150,51 +1150,67 @@ static void addLinkageDecoration(
     {
         builder->addExportDecoration(inst, mangledName);
     }
-    if (decl->findModifier<PublicModifier>())
+    for (auto modifier : decl->modifiers)
     {
-        builder->addPublicDecoration(inst);
-        builder->addKeepAliveDecoration(inst);
-    }
-    if (decl->findModifier<HLSLExportModifier>())
-    {
-        builder->addHLSLExportDecoration(inst);
-        builder->addKeepAliveDecoration(inst);
-    }
-    if (decl->findModifier<ExternCppModifier>())
-    {
-        builder->addExternCppDecoration(inst, mangledName);
+        if (as<PublicModifier>(modifier))
+        {
+            builder->addPublicDecoration(inst);
+            builder->addKeepAliveDecoration(inst);
+        }
+        else if (as<HLSLExportModifier>(modifier))
+        {
+            builder->addHLSLExportDecoration(inst);
+            builder->addKeepAliveDecoration(inst);
+        }
+        else if (as<ExternCppModifier>(modifier))
+        {
+            builder->addExternCppDecoration(inst, mangledName);
+        }
+        else if (auto dllImportModifier = as<DllImportAttribute>(modifier))
+        {
+            auto libraryName = dllImportModifier->modulePath;
+            auto functionName = dllImportModifier->functionName.getLength()
+                ? dllImportModifier->functionName.getUnownedSlice()
+                : decl->getName()->text.getUnownedSlice();
+            builder->addDllImportDecoration(inst, libraryName.getUnownedSlice(), functionName);
+        }
+        else if (as<DllExportAttribute>(modifier))
+        {
+            builder->addDllExportDecoration(inst, decl->getName()->text.getUnownedSlice());
+            builder->addPublicDecoration(inst);
+        }
+        else if (as<CudaDeviceExportAttribute>(modifier))
+        {
+            builder->addCudaDeviceExportDecoration(inst, decl->getName()->text.getUnownedSlice());
+            builder->addPublicDecoration(inst);
+            builder->addExternCppDecoration(inst, decl->getName()->text.getUnownedSlice());
+        }
+        else if (as<CudaHostAttribute>(modifier))
+        {
+            builder->addCudaHostDecoration(inst);
+            builder->addExternCppDecoration(inst, decl->getName()->text.getUnownedSlice());
+        }
+        else if (as<CudaKernelAttribute>(modifier))
+        {
+            builder->addCudaKernelDecoration(inst);
+            builder->addExternCppDecoration(inst, decl->getName()->text.getUnownedSlice());
+            builder->addPublicDecoration(inst);
+            builder->addKeepAliveDecoration(inst);
+        }
+        else if (as<TorchEntryPointAttribute>(modifier))
+        {
+            builder->addTorchEntryPointDecoration(inst, decl->getName()->text.getUnownedSlice());
+            builder->addCudaHostDecoration(inst);
+            builder->addPublicDecoration(inst);
+            builder->addExternCppDecoration(inst, decl->getName()->text.getUnownedSlice());
+        }
     }
     if (as<InterfaceDecl>(decl->parentDecl) &&
-        decl->parentDecl->hasModifier<ComInterfaceAttribute>())
+        decl->parentDecl->hasModifier<ComInterfaceAttribute>() &&
+        !inst->findDecoration<IRExternCppDecoration>())
     {
         builder->addExternCppDecoration(inst, decl->getName()->text.getUnownedSlice());
     }
-    if (auto dllImportModifier = decl->findModifier<DllImportAttribute>())
-    {
-        auto libraryName = dllImportModifier->modulePath;
-        auto functionName = dllImportModifier->functionName.getLength()
-            ? dllImportModifier->functionName.getUnownedSlice()
-            : decl->getName()->text.getUnownedSlice();
-        builder->addDllImportDecoration(inst, libraryName.getUnownedSlice(), functionName);
-    }
-    if (decl->findModifier<DllExportAttribute>())
-    {
-        builder->addDllExportDecoration(inst, decl->getName()->text.getUnownedSlice());
-        builder->addPublicDecoration(inst);
-    }
-    if (decl->findModifier<CudaDeviceExportAttribute>())
-    {
-        builder->addCudaDeviceExportDecoration(inst, decl->getName()->text.getUnownedSlice());
-        builder->addPublicDecoration(inst);
-    }
-    if (decl->findModifier<CudaHostAttribute>())
-    {
-        builder->addCudaHostDecoration(inst);
-    }
-    if (decl->findModifier<CudaKernelAttribute>())
-    {
-        builder->addCudaKernelDecoration(inst);
-    }
 }
 
 static void addLinkageDecoration(
-- 
cgit v1.2.3