2 files changed, 36 insertions, 3 deletions
diff --git a/src/gallium/drivers/swr/rasterizer/jitter/JitManager.cpp b/src/gallium/drivers/swr/rasterizer/jitter/JitManager.cpp
index 0312fc47fb6..58d30d4e119 100644
--- a/src/gallium/drivers/swr/rasterizer/jitter/JitManager.cpp
+++ b/src/gallium/drivers/swr/rasterizer/jitter/JitManager.cpp
@@ -443,7 +443,7 @@ std::string JitManager::GetOutputDir()
 
 //////////////////////////////////////////////////////////////////////////
 /// @brief Dump function to file.
-void JitManager::DumpToFile(Module* M, const char* fileName)
+void JitManager::DumpToFile(Module* M, const char* fileName, llvm::AssemblyAnnotationWriter* annotater)
 {
     if (KNOB_DUMP_SHADER_IR)
     {
@@ -458,7 +458,7 @@ void JitManager::DumpToFile(Module* M, const char* fileName)
         sprintf(fName, "%s.%s.ll", funcName, fileName);
 #endif
         raw_fd_ostream fd(fName, EC, llvm::sys::fs::F_None);
-        M->print(fd, nullptr);
+        M->print(fd, annotater);
         fd.flush();
     }
 }
@@ -758,3 +758,26 @@ std::unique_ptr<llvm::MemoryBuffer> JitCache::getObject(const llvm::Module* M)
 
     return pBuf;
 }
+
+void InterleaveAssemblyAnnotater::emitInstructionAnnot(const llvm::Instruction *pInst, llvm::formatted_raw_ostream &OS)
+{
+    auto dbgLoc = pInst->getDebugLoc();
+    if(dbgLoc)
+    {
+        unsigned int line = dbgLoc.getLine();
+        if(line != mCurrentLineNo)
+        {
+            if(line > 0 && line <= mAssembly.size())
+            {
+                // HACK: here we assume that OS is a formatted_raw_ostream(ods())
+                // and modify the color accordingly. We can't do the color
+                // modification on OS because formatted_raw_ostream strips
+                // the color information. The only way to fix this behavior
+                // is to patch LLVM.
+                OS << "\n; " << line << ": " << mAssembly[line-1] << "\n";
+            }
+            mCurrentLineNo = line;
+        }
+    }
+}
+
diff --git a/src/gallium/drivers/swr/rasterizer/jitter/JitManager.h b/src/gallium/drivers/swr/rasterizer/jitter/JitManager.h
index a5b6af91f06..2f479314c76 100644
--- a/src/gallium/drivers/swr/rasterizer/jitter/JitManager.h
+++ b/src/gallium/drivers/swr/rasterizer/jitter/JitManager.h
@@ -31,6 +31,7 @@
 
 #include "jit_pch.hpp"
 #include "common/isa.hpp"
+#include <llvm/IR/AssemblyAnnotationWriter.h>
 
 
 //////////////////////////////////////////////////////////////////////////
@@ -151,7 +152,7 @@ struct JitManager
 
     void               DumpAsm(llvm::Function* pFunction, const char* fileName);
     static void        DumpToFile(llvm::Function* f, const char* fileName);
-    static void        DumpToFile(llvm::Module* M, const char* fileName);
+    static void        DumpToFile(llvm::Module* M, const char* fileName, llvm::AssemblyAnnotationWriter* annotater = nullptr);
     static std::string GetOutputDir();
 
     // Debugging support methods
@@ -178,3 +179,12 @@ struct JitManager
                           uint32_t                                             lineNum,
                           const std::vector<std::pair<std::string, uint32_t>>& members);
 };
+
+class InterleaveAssemblyAnnotater : public llvm::AssemblyAnnotationWriter
+{
+public:
+    void emitInstructionAnnot(const llvm::Instruction *pInst, llvm::formatted_raw_ostream &OS) override;
+    std::vector<std::string> mAssembly;
+private:
+    uint32_t mCurrentLineNo = 0;
+};