Adds infrastructure to automatically map operations to layers

Author: pranavm-nvidia

mlir-tensorrt/tensorrt/include/mlir-tensorrt-dialect/Target/TensorRTEncodingOpInterface/NetworkEncoder.h

@@ -89,13 +89,16 @@ class NvInferNetworkEncoder {
   /// Lookup the TRT ITensor* equivalents of a ValueRange.
   SmallVector<nvinfer1::ITensor *> lookupValues(ValueRange values);
 
-  /// Add a map from a Value to a TRT ITEnsor*.
+  /// Add a map from a Value to a TRT ITensor*.
   void map(Value from, nvinfer1::ITensor *to);
 
   /// Remap values in `from` to each layer in `to` using the output at index 0
   /// for each layer.
   void map(ValueRange from, ArrayRef<nvinfer1::ILayer *> to);
 
+  // Add a map from an Operation to a TRT ILayer*
+  void map(Operation *op, nvinfer1::ILayer *layer);
+
   /// Check whether the value map contains `v`.
   size_t contains(Value v) { return valueMap.count(v); }
 
@@ -135,6 +138,10 @@ class NvInferNetworkEncoder {
   /// and other temporary buffers.
   using WeightsMap = llvm::DenseMap<mlir::Attribute, std::vector<int8_t>>;
 
+  // Tracks the mapping of mlir::Operations to layers. Note that one operation
+  // may map to multiple layers.
+  using LayerMap = llvm::DenseMap<Operation *, std::vector<nvinfer1::ILayer *>>;
+
   using NamesSet = llvm::StringSet<>;
 
   TensorMap &getTensorMap() { return valueMap; }
@@ -210,6 +217,9 @@ class NvInferNetworkEncoder {
   // build ends.
   SmallVector<NvInferPluginPtr> pluginReferences;
 
+  // Tracks the mapping between mlir::Operations and TensorRT ILayers.
+  LayerMap layerMap;
+
   /// Holds the set of strings currently assigned as names to TensorRT ILayers.
   /// This is required because we must make new names unique. The TensorRT API
   /// does not have a set object to query names.

mlir-tensorrt/tensorrt/lib/Target/TensorRTEncodingOpInterface/NetworkEncoder.cpp

@@ -287,6 +287,13 @@ void NvInferNetworkEncoder::map(ValueRange from,
     valueMap.insert(v, l->getOutput(0));
 }
 
+void NvInferNetworkEncoder::map(Operation *op, nvinfer1::ILayer *layer) {
+  if (!layerMap.count(op))
+    layerMap[op] = {};
+
+  layerMap[op].push_back(layer);
+}
+
 bool NvInferNetworkEncoder::isStronglyTyped() const {
   if (!usesStronglyTyped)
     return false;

mlir-tensorrt/tools/MlirTensorRtTblgen.cpp

@@ -397,7 +397,11 @@ static bool emitLayerAddDefinitions(const llvm::RecordKeeper &recordKeeper,
       }
 
       // Emit the body with substitutions.
+      os << "auto const numLayersBefore = network->getNbLayers();\n";
       os << tblgen::tgfmt(expr, &ctx);
+      os << "auto const numLayersAfter = network->getNbLayers();\n";
+      os << "for (int64_t i = numLayersBefore; i < numLayersAfter; ++i) "
+            "encoder.map(tensorrtOp, network->getLayer(i));\n";
       os << "return success();\n";
       os.unindent();
       os << "}\n";
@@ -603,4 +607,4 @@ int main(int argc, char **argv) {
       });
 
   return mlir::MlirTblgenMain(argc, argv);
-}
+}