Embedding executable source contents in binaries for tracing. (#16757)

This adds a new `CaptureExecutableSourcesPass` that allows for capture
of individual `hal.executable.variant` ops at any number of compilation
stages. Unlike the `DumpExecutableSourcesPass` this does not change the
original source locations in the IR and instead captures both the
textual IR and the remapped locations within it of each executable
export at the time the pass is run. The textual IR is stored as
resources and associated with the variants through linking and made
available to serialization for embedding in target-specific formats.

Because this increases compilation time (generating all of the sources
multiple times per executable is expensive) and bloats binaries the
capture is only enabled with the `--iree-hal-executable-debug-level=3`
or greater flag set (default is `=2`).

As part of this PR the CPU, Vulkan, and legacy ROCM formats have been
updated to store the new information and source it at runtime. This is a
breaking change to the executable library binary format. I'm not quite
happy with it, but it's probably good enough for the next 6mo-1yr.

To make this more usable a copy button has been added to the tracy
source view: wolfpld/tracy#750
Now clicking on a dispatch in the CPU or GPU timeline will show the
source and the copy button can be used to get it in the clipboard. The
source can then be run through `iree-compile
--compile-mode=hal-executable` to generate binaries.

![image](https://github.com/openxla/iree/assets/75337/e690e1f8-52a7-40db-a3aa-5fd7e781791b)

Fixes #15699. Closes #16223.

compiler/plugins/target/LLVMCPU/BUILD.bazel

@@ -39,6 +39,7 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/HAL/Target",
         "//compiler/src/iree/compiler/Dialect/HAL/Target:LLVMLinkerUtils",
         "//compiler/src/iree/compiler/Dialect/LinalgExt/IR",
+        "//compiler/src/iree/compiler/Dialect/Util/IR",
         "//compiler/src/iree/compiler/PluginAPI",
         "//compiler/src/iree/compiler/Utils",
         "//llvm-external-projects/iree-dialects:IREELinalgTransformDialect",
@@ -63,6 +64,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:ArmSMEDialect",
         "@llvm-project//mlir:ArmSMEToLLVMIRTranslation",
         "@llvm-project//mlir:BuiltinToLLVMIRTranslation",
+        "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LLVMDialect",
         "@llvm-project//mlir:LLVMToLLVMIRTranslation",
         "@llvm-project//mlir:PDLDialect",

compiler/plugins/target/LLVMCPU/CMakeLists.txt

@@ -43,6 +43,7 @@ iree_cc_library(
     MLIRArmSMEDialect
     MLIRArmSMEToLLVMIRTranslation
     MLIRBuiltinToLLVMIRTranslation
+    MLIRIR
     MLIRLLVMDialect
     MLIRLLVMToLLVMIRTranslation
     MLIRPDLDialect
@@ -57,6 +58,7 @@ iree_cc_library(
     iree::compiler::Dialect::HAL::Target
     iree::compiler::Dialect::HAL::Target::LLVMLinkerUtils
     iree::compiler::Dialect::LinalgExt::IR
+    iree::compiler::Dialect::Util::IR
     iree::compiler::PluginAPI
     iree::compiler::Utils
     iree::compiler::plugins::target::LLVMCPU::Builtins

compiler/plugins/target/LLVMCPU/LLVMCPUTarget.cpp

@@ -37,6 +37,8 @@
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
@@ -420,17 +422,45 @@ class LLVMCPUTargetBackend final : public TargetBackend {
                                     .value_or(APInt(64, 0))
                                     .getSExtValue();
 
-      std::string sourceFile = "";
-      int sourceLine = 0;
+      LibraryBuilder::SourceLocation sourceLocation;
       if (options.debugLevel >= 1) {
         if (auto loc = findFirstFileLoc(exportOp.getLoc())) {
-          sourceFile = loc->getFilename().str();
-          sourceLine = loc->getLine();
+          sourceLocation = {"", loc->getFilename().str(), loc->getLine()};
+        }
+      }
+      SmallVector<LibraryBuilder::SourceLocation> stageLocations;
+      if (options.debugLevel >= 3) {
+        if (auto locsAttr = exportOp.getSourceLocsAttr()) {
+          for (auto locAttr : locsAttr.getValue()) {
+            if (auto loc =
+                    findFirstFileLoc(cast<LocationAttr>(locAttr.getValue()))) {
+              stageLocations.push_back({
+                  locAttr.getName().str(),
+                  loc->getFilename().str(),
+                  loc->getLine(),
+              });
+            }
+          }
+        }
+      }
+      libraryBuilder.addExport(exportOp.getName(), std::move(sourceLocation),
+                               std::move(stageLocations), /*tag=*/"",
+                               LibraryBuilder::DispatchAttrs{localMemorySize},
+                               llvmFunc);
+    }
+
+    // Embed source files (if present).
+    if (auto sourcesAttr = variantOp.getSourcesAttr()) {
+      for (auto sourceAttr : sourcesAttr.getValue()) {
+        if (auto resourceAttr = dyn_cast_if_present<DenseResourceElementsAttr>(
+                sourceAttr.getValue())) {
+          auto handle = resourceAttr.getRawHandle();
+          SmallVector<char> rawData;
+          llvm::append_range(rawData, handle.getBlob()->getData());
+          libraryBuilder.addSourceFile(sourceAttr.getName(),
+                                       std::move(rawData));
         }
       }
-      libraryBuilder.addExport(
-          exportOp.getName(), sourceFile, sourceLine, /*tag=*/"",
-          LibraryBuilder::DispatchAttrs{localMemorySize}, llvmFunc);
     }
 
     auto queryFunctionName = std::string(kQueryFunctionName);