From 1c8829d48071d378acabee0656127b951bf00965 Mon Sep 17 00:00:00 2001 From: Matthew Francis-Landau Date: Tue, 21 Oct 2025 07:29:44 -0700 Subject: [PATCH] [mlir-tensorrt] fix the way that layer names are set for TensorRT Previously the layer name was set to the root caller and ignored callee. This change updates the layer name to capture the last 3 callees to enable better tracking of locations in the generated TensorRT engine. --- .../NetworkEncoder.cpp | 24 +++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/mlir-tensorrt/tensorrt/lib/Target/TensorRTEncodingOpInterface/NetworkEncoder.cpp b/mlir-tensorrt/tensorrt/lib/Target/TensorRTEncodingOpInterface/NetworkEncoder.cpp index 10f062533..39a82ff5a 100644 --- a/mlir-tensorrt/tensorrt/lib/Target/TensorRTEncodingOpInterface/NetworkEncoder.cpp +++ b/mlir-tensorrt/tensorrt/lib/Target/TensorRTEncodingOpInterface/NetworkEncoder.cpp @@ -233,9 +233,28 @@ static std::string getUniqueName(NvInferNetworkEncoder::NamesSet &names, /// an open system of location attributes, there may be some location types that /// we cannot handle. We do not use the location's builtin printer because it /// could be extremely verbose. +static void getCallSiteLocs(Location loc, SmallVector &locs) { + if (auto callLoc = dyn_cast(loc)) { + getCallSiteLocs(callLoc.getCaller(), locs); + getCallSiteLocs(callLoc.getCallee(), locs); + } else { + locs.push_back(loc); + } +} + static void translateLocation(Location loc, llvm::raw_ostream &os) { if (auto callLoc = dyn_cast(loc)) { - translateLocation(callLoc.getCaller(), os); + SmallVector locs; + getCallSiteLocs(callLoc, locs); + // only include the last 3 locations in the names as this should be + // sufficient to identify the call site for an op + for (size_t i = locs.size() > 3 ? locs.size() - 3 : 0; i < locs.size(); + i++) { + translateLocation(locs[i], os); + if (i < locs.size() - 1) { + os << " -> "; + } + } return; } if (auto fileLoc = dyn_cast(loc)) { @@ -274,7 +293,8 @@ static std::string createName(NvInferNetworkEncoder::NamesSet &names, ss.flush(); } // Truncate to TRT limit. - static constexpr size_t kLayerNameSizeLimit = 2048; + static constexpr size_t kLayerNameSizeLimit = + 2048 - 6; // -6 to give some space for UniqueName if (name.size() > kLayerNameSizeLimit) name = name.substr(0, kLayerNameSizeLimit);