WriteSemanticDefines before clang codeGen to update codeGen option (#4261)

* WriteSemanticDefines before clang codeGen to update codeGen option based on HLSLOptimizationToggles.

include/dxc/HLSL/HLSLExtensionsCodegenHelper.h

@@ -14,6 +14,10 @@
 #include <vector>
 #include <string>
 
+namespace clang {
+class CodeGenOptions;
+}
+
 namespace llvm {
 class CallInst;
 class Value;
@@ -61,8 +65,8 @@ public:
   typedef std::vector<SemanticDefineError> SemanticDefineErrorList;
 
   // Write semantic defines as metadata in the module.
-  virtual SemanticDefineErrorList WriteSemanticDefines(llvm::Module *M) = 0;
-
+  virtual void WriteSemanticDefines(llvm::Module *M) = 0;
+  virtual void UpdateCodeGenOptions(clang::CodeGenOptions &CGO) = 0;
   // Query the named option enable
   // Needed because semantic defines may have set it since options were copied
   virtual bool IsOptionEnabled(std::string option) = 0;

tools/clang/lib/CodeGen/CGHLSLMSFinishCodeGen.cpp

@@ -2937,22 +2937,7 @@ void AddRegBindingsForResourceInConstantBuffer(
 
 // extension codegen.
 void ExtensionCodeGen(HLModule &HLM, clang::CodeGen::CodeGenModule &CGM) {
-  // Add semantic defines for extensions if any are available.
-  HLSLExtensionsCodegenHelper::SemanticDefineErrorList errors =
-      CGM.getCodeGenOpts().HLSLExtensionsCodegen->WriteSemanticDefines(
-          HLM.GetModule());
-
-  clang::DiagnosticsEngine &Diags = CGM.getDiags();
-  for (const HLSLExtensionsCodegenHelper::SemanticDefineError &error : errors) {
-    clang::DiagnosticsEngine::Level level = clang::DiagnosticsEngine::Error;
-    if (error.IsWarning())
-      level = clang::DiagnosticsEngine::Warning;
-    unsigned DiagID = Diags.getCustomDiagID(level, "%0");
-    Diags.Report(clang::SourceLocation::getFromRawEncoding(error.Location()),
-                 DiagID)
-        << error.Message();
-  }
-
+  auto &Diags = CGM.getDiags();
   // Add root signature from a #define. Overrides root signature in function
   // attribute.
   {

tools/clang/lib/CodeGen/ModuleBuilder.cpp

@@ -211,6 +211,22 @@ namespace {
         M.reset();
         return;
       }
+      // HLSL Change Begins.
+      if (&Ctx == this->Ctx) {
+        if (Builder) {
+          // Add semantic defines for extensions if any are available.
+          auto &CodeGenOpts =
+              const_cast<CodeGenOptions &>(Builder->getCodeGenOpts());
+          if (CodeGenOpts.HLSLExtensionsCodegen) {
+            CodeGenOpts.HLSLExtensionsCodegen->WriteSemanticDefines(
+                    M.get());
+            // Builder->CodeGenOpts is a copy. So update it for every Builder.
+            CodeGenOpts.HLSLExtensionsCodegen->UpdateCodeGenOptions(
+                CodeGenOpts);
+          }
+        }
+      }
+      // HLSL Change Ends.
       if (Builder)
         Builder->Release();
 

tools/clang/tools/dxcompiler/dxcompilerobj.cpp

@@ -479,7 +479,7 @@ public:
   {}
 
   // Write semantic defines as metadata in the module.
-  virtual std::vector<SemanticDefineError> WriteSemanticDefines(llvm::Module *M) override {
+  virtual void WriteSemanticDefines(llvm::Module *M) override {
     // Grab the semantic defines seen by the parser.
     ParsedSemanticDefineList defines =
       CollectSemanticDefinesParsedByCompiler(m_CI, &m_langExtensionsHelper);
@@ -487,14 +487,31 @@ public:
     // Nothing to do if we have no defines.
     SemanticDefineErrorList errors;
     if (!defines.size())
-      return errors;
+      return;
 
     ParsedSemanticDefineList validated;
     GetValidatedSemanticDefines(defines, validated, errors);
     WriteSemanticDefines(M, validated);
-    return errors;
-  }
 
+    auto &Diags = m_CI.getDiagnostics();
+    for (const auto &error : errors) {
+      clang::DiagnosticsEngine::Level level = clang::DiagnosticsEngine::Error;
+      if (error.IsWarning())
+        level = clang::DiagnosticsEngine::Warning;
+      unsigned DiagID = Diags.getCustomDiagID(level, "%0");
+      Diags.Report(clang::SourceLocation::getFromRawEncoding(error.Location()),
+                   DiagID)
+          << error.Message();
+    }
+
+  }
+  // Update CodeGenOption based on HLSLOptimizationToggles.
+  void UpdateCodeGenOptions(clang::CodeGenOptions &CGO) override {
+    auto &CodeGenOpts = m_CI.getCodeGenOpts();
+    CGO.HLSLEnableLifetimeMarkers &=
+        (!CodeGenOpts.HLSLOptimizationToggles.count("lifetime-markers") ||
+         CodeGenOpts.HLSLOptimizationToggles.find("lifetime-markers")->second);
+  }
   virtual bool IsOptionEnabled(std::string option) override {
     return m_CI.getCodeGenOpts().HLSLOptimizationToggles.count(option) &&
       m_CI.getCodeGenOpts().HLSLOptimizationToggles.find(option)->second;

tools/clang/unittests/HLSL/ExtensionTest.cpp

@@ -478,10 +478,13 @@ public:
     return Compile(program, {}, {});
   }
 
-  IDxcOperationResult *Compile(const char *program, const std::vector<LPCWSTR> &arguments, const std::vector<DxcDefine> defs ) {
+  IDxcOperationResult *Compile(const char *program,
+                               const std::vector<LPCWSTR> &arguments,
+                               const std::vector<DxcDefine> defs,
+                               LPCWSTR target = L"ps_6_0") {
     Utf8ToBlob(m_dllSupport, program, &pCodeBlob);
     VERIFY_SUCCEEDED(pCompiler->Compile(pCodeBlob, L"hlsl.hlsl", L"main",
-      L"ps_6_0",
+      target,
       const_cast<LPCWSTR *>(arguments.data()), arguments.size(),
       defs.data(), defs.size(),
       nullptr, &pCompileResult));
@@ -536,6 +539,7 @@ public:
   TEST_METHOD(DefineOverrideDeterministicOutput)
   TEST_METHOD(OptionFromDefineGVN)
   TEST_METHOD(OptionFromDefineStructurizeReturns)
+  TEST_METHOD(OptionFromDefineLifetimeMarkers)
   TEST_METHOD(TargetTriple)
   TEST_METHOD(IntrinsicWhenAvailableThenUsed)
   TEST_METHOD(CustomIntrinsicName)
@@ -1052,6 +1056,32 @@ TEST_F(ExtensionTest, OptionFromDefineStructurizeReturns) {
   VERIFY_IS_TRUE(regex.match(disassembly));
 }
 
+// Test setting of codegen options from semantic defines
+TEST_F(ExtensionTest, OptionFromDefineLifetimeMarkers) {
+  std::string shader = "\n"
+      "float foo(float a) {\n"
+    "float res[2] = {a, 2 * 2};\n"
+    "return res[a];\n"
+    "}\n"
+  "float4 main(float a : A) : SV_Target { return foo(a); }\n";
+
+  Compiler c(m_dllSupport);
+  c.RegisterSemanticDefine(L"FOO*");
+  c.Compile(shader.data(), {L"/Vd", L"-DFOO_DISABLE_LIFETIME_MARKERS"}, {},
+            L"ps_6_6");
+
+  std::string disassembly = c.Disassemble();
+  Compiler c2(m_dllSupport);
+  c2.Compile(shader.data(), {L"/Vd", L""}, {}, L"ps_6_6");
+  std::string disassembly2 = c2.Disassemble();
+  // Make sure lifetime marker not exist with FOO_DISABLE_LIFETIME_MARKERS.
+  VERIFY_IS_TRUE(disassembly.find("lifetime") == std::string::npos);
+  VERIFY_IS_TRUE(disassembly.find("FOO_DISABLE_LIFETIME_MARKERS\", !\"1\"") !=
+                 std::string::npos);
+  // Make sure lifetime marker exist by default.
+  VERIFY_IS_TRUE(disassembly2.find("lifetime") != std::string::npos);
+}
+
 
 TEST_F(ExtensionTest, TargetTriple) {
   Compiler c(m_dllSupport);