1 files changed, 30 insertions, 13 deletions

diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index fbe6d8a84..e5037bf04 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -8168,24 +8168,34 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         // If the witness table is for a COM interface, always keep it alive.
         if (irWitnessTableBaseType->findDecoration<IRComInterfaceDecoration>())
         {
-            subBuilder->addPublicDecoration(irWitnessTable);
+            subBuilder->addHLSLExportDecoration(irWitnessTable);
         }
 
-        if (parentDecl->findModifier<HLSLExportModifier>())
+        for (auto mod : parentDecl->modifiers)
         {
-            subBuilder->addHLSLExportDecoration(irWitnessTable);
-            subBuilder->addKeepAliveDecoration(irWitnessTable);
+            if (as<HLSLExportModifier>(mod))
+            {
+                subBuilder->addHLSLExportDecoration(irWitnessTable);
+                subBuilder->addKeepAliveDecoration(irWitnessTable);
+            }
+            else if (as<AutoDiffBuiltinAttribute>(mod))
+            {
+                subBuilder->addAutoDiffBuiltinDecoration(irWitnessTable);
+            }
         }
 
         // Make sure that all the entries in the witness table have been filled in,
         // including any cases where there are sub-witness-tables for conformances
-        Dictionary<WitnessTable*, IRWitnessTable*> mapASTToIRWitnessTable;
-        lowerWitnessTable(
-            subContext,
-            inheritanceDecl->witnessTable,
-            irWitnessTable,
-            mapASTToIRWitnessTable);
-
+        bool isExplicitExtern = false;
+        if (!isImportedDecl(context, parentDecl, isExplicitExtern))
+        {
+            Dictionary<WitnessTable*, IRWitnessTable*> mapASTToIRWitnessTable;
+            lowerWitnessTable(
+                subContext,
+                inheritanceDecl->witnessTable,
+                irWitnessTable,
+                mapASTToIRWitnessTable);
+        }
         irWitnessTable->moveToEnd();
 
         return LoweredValInfo::simple(
@@ -11536,6 +11546,8 @@ RefPtr<IRModule> generateIRForTranslationUnit(
 
     RefPtr<IRModule> module = IRModule::create(session);
 
+    module->setName(translationUnit->getModuleDecl()->getName());
+
     IRBuilder builderStorage(module);
     IRBuilder* builder = &builderStorage;
 
@@ -11804,6 +11816,8 @@ RefPtr<IRModule> generateIRForTranslationUnit(
         stripOptions.stripSourceLocs = false;
         stripFrontEndOnlyInstructions(module, stripOptions);
 
+        stripImportedWitnessTable(module);
+
         // Stripping out decorations could leave some dead code behind
         // in the module, and in some cases that extra code is also
         // undesirable (e.g., the string literals referenced by name-hint
@@ -11847,6 +11861,8 @@ RefPtr<IRModule> generateIRForTranslationUnit(
             &writer);
     }
 
+    module->buildMangledNameToGlobalInstMap();
+
     return module;
 }
 
@@ -11884,7 +11900,7 @@ struct SpecializedComponentTypeIRGenContext : ComponentTypeVisitor
         context->irBuilder = builder;
 
         componentType->acceptVisitor(this, nullptr);
-
+        module->buildMangledNameToGlobalInstMap();
         return module;
     }
 
@@ -12040,6 +12056,7 @@ struct TypeConformanceIRGenContext
         {
             builder->addSequentialIDDecoration(witness, conformanceIdOverride);
         }
+        module->buildMangledNameToGlobalInstMap();
         return module;
     }
 };
@@ -12507,7 +12524,7 @@ RefPtr<IRModule> TargetProgram::createIRModuleForLayout(DiagnosticSink* sink)
         // Eliminate any dead code
         eliminateDeadCode(irModule, options);
     }
-
+    irModule->buildMangledNameToGlobalInstMap();
     m_irModuleForLayout = irModule;
     return irModule;
 }