Merge pull request #4 from o3de/specialization_constants

Add DXSC tool to patch specialization constant on DXIL
o3de · Jun 28, 2024 · 87f4e45 · 87f4e45
2 parents 8c9f508 + 1d814a9
commit 87f4e45
Show file tree

Hide file tree

Showing 7 changed files with 498 additions and 3 deletions.
diff --git a/include/llvm/Bitcode/BitstreamWriter.h b/include/llvm/Bitcode/BitstreamWriter.h
@@ -21,6 +21,11 @@
 #include "llvm/Support/Endian.h"
 #include <vector>
 
+// O3DE change start
+#include "llvm/Bitcode/LLVMBitCodes.h"
+#include <functional>
+// O3DE change end
+
 namespace llvm {
 
 class BitstreamWriter {
@@ -88,6 +93,11 @@ class BitstreamWriter {
   }
 
 public:
+  // O3DE change start
+  typedef std::function<void(uint64_t, uint64_t)> ConstantHandlerFn;
+  ConstantHandlerFn WriteConstantCallback = nullptr;
+  // O3DE change end
+
   explicit BitstreamWriter(SmallVectorImpl<char> &O)
     : Out(O), CurBit(0), CurValue(0), CurCodeSize(2) {}
 
@@ -370,6 +380,22 @@ class BitstreamWriter {
           WriteByte(0);
       } else {  // Single scalar field.
         assert(RecordIdx < Vals.size() && "Invalid abbrev/record");
+
+        // O3DE Change Start
+        if (WriteConstantCallback &&
+            Vals[0] == bitc::CST_CODE_INTEGER &&
+            Op.getEncoding() == BitCodeAbbrevOp::VBR) {
+          uint64_t SCBitOffset = (uint64_t)Out.size_in_bytes() * 8 + CurBit;
+          uint64_t SCVal = Vals[RecordIdx];
+          if (Vals[RecordIdx] & 1) {
+            SCVal = -(SCVal >> 1);
+          } else {
+            SCVal = SCVal >> 1;
+          }
+
+          WriteConstantCallback(SCVal, SCBitOffset);
+        }
+        // O3DE Change End
         EmitAbbreviatedField(Op, Vals[RecordIdx]);
         ++RecordIdx;
       }

diff --git a/include/llvm/Bitcode/ReaderWriter.h b/include/llvm/Bitcode/ReaderWriter.h
@@ -18,11 +18,13 @@
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/MemoryBuffer.h"
+// O3DE change start
+#include "llvm/Bitcode/BitstreamWriter.h"
+// O3DE change end
 #include <memory>
 #include <string>
 
 namespace llvm {
-  class BitstreamWriter;
   class DataStreamer;
   class LLVMContext;
   class Module;
@@ -69,7 +71,8 @@ namespace llvm {
   /// Value in \c M.  These will be reconstructed exactly when \a M is
   /// deserialized.
   void WriteBitcodeToFile(const Module *M, raw_ostream &Out,
-                          bool ShouldPreserveUseListOrder = false);
+                          bool ShouldPreserveUseListOrder = false,
+                          BitstreamWriter::ConstantHandlerFn WriteCallback = nullptr); // O3DE change
 
   /// isBitcodeWrapper - Return true if the given bytes are the magic bytes
   /// for an LLVM IR bitcode wrapper.

diff --git a/lib/Bitcode/Writer/BitcodeWriter.cpp b/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -2494,7 +2494,8 @@ static void EmitDarwinBCHeaderAndTrailer(SmallVectorImpl<char> &Buffer,
 /// WriteBitcodeToFile - Write the specified module to the specified output
 /// stream.
 void llvm::WriteBitcodeToFile(const Module *M, raw_ostream &Out,
-                              bool ShouldPreserveUseListOrder) {
+                              bool ShouldPreserveUseListOrder,
+                              BitstreamWriter::ConstantHandlerFn WriteCallback) { // O3DE change
   SmallVector<char, 0> Buffer;
   Buffer.reserve(256*1024);
 
@@ -2507,6 +2508,9 @@ void llvm::WriteBitcodeToFile(const Module *M, raw_ostream &Out,
   // Emit the module into the buffer.
   {
     BitstreamWriter Stream(Buffer);
+    // O3DE change start
+    Stream.WriteConstantCallback = WriteCallback;
+    // O3DE change end
 
     // Emit the file header.
     Stream.Emit((unsigned)'B', 8);

diff --git a/tools/clang/tools/CMakeLists.txt b/tools/clang/tools/CMakeLists.txt
@@ -41,3 +41,7 @@ add_subdirectory(dxlib-sample)
 add_subdirectory(dotnetc)
 endif (WIN32)
 # HLSL Change Ends
+
+# O3DE Change Starts
+add_subdirectory(dxsc)
+# O3DE Change Ends
diff --git a/tools/clang/tools/dxsc/CMakeLists.txt b/tools/clang/tools/dxsc/CMakeLists.txt
@@ -0,0 +1,39 @@
+# Copyright (C) Microsoft Corporation. All rights reserved.
+# This file is distributed under the University of Illinois Open Source License. See LICENSE.TXT for details.
+# Builds dxsc.exe
+
+set( LLVM_LINK_COMPONENTS
+  ${LLVM_TARGETS_TO_BUILD}
+  DXIL
+  DxilContainer
+  DxilRootSignature
+  HLSL
+  dxcsupport
+  Option     # option library
+  MSSupport  # for CreateMSFileSystemForDisk
+  )
+
+add_clang_executable(dxsc
+  dxsc.cpp
+  )
+
+target_link_libraries(dxsc
+  dxcompiler
+  )
+
+set_target_properties(dxsc PROPERTIES VERSION ${CLANG_EXECUTABLE_VERSION})
+
+add_dependencies(dxsc dxcompiler)
+
+if(UNIX)
+  set(CLANGXX_LINK_OR_COPY create_symlink)
+# Create a relative symlink
+  set(dxsc_binary "dxsc${CMAKE_EXECUTABLE_SUFFIX}")
+else()
+  set(CLANGXX_LINK_OR_COPY copy)
+  set(dxsc_binary "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${CMAKE_CFG_INTDIR}/dxsc${CMAKE_EXECUTABLE_SUFFIX}")
+endif()
+
+install(TARGETS dxsc
+  RUNTIME DESTINATION bin)
+