WIP: Adding TRT options/task

[pranavm-nvidia]

No description provided.

[pranavm-nvidia]

I can move TensorRTTranslationOptions to make them an options provider if that makes sense to do.

[pranavm-nvidia]

Not sure if we want to require extensions for all options types or if we need to handle both cases in the options registry. If it's the former, then I can just assert that the extensions are empty here (or maybe even just add support?). If it's the latter, we could have a setExtensions method so it becomes optional instead of having it part of the constructor.

[christopherbate]

This catch-all library has been removed in a commit which hasn't been sync'd up yet. Just remove this and add in dependent libraries explicitly.

[christopherbate]

You will want to add a host-target flag which can take options executor, llvm or emitc. This executor lowering pass is just for the executor option. The pipeliens for the other branches will be sync'd up Friday

[yizhuoz004]

@christopherbate I had to move getOrCreateTensorRTModuleOp here to make this pass work, otherwise in createOutlinedFunc, at line SymbolTable(module) will hit seg fault. Could you check if this is correct?

An example run, input:

func.func @trt_gather_default1(%arg0: tensor<10x20x30xf32>, %arg1: tensor<2x5xi32>,
                %arg2: tensor<10x2x5x30xf32>) -> tensor<10x2x5x30xf32> {
  %0 = tensorrt.gather {
    axis = 1 : i64
  } ins(%arg0, %arg1 : tensor<10x20x30xf32>, tensor<2x5xi32>) -> tensor<10x2x5x30xf32>
  %1 = tensorrt.element_wise <kSUM>(%0, %arg2 : tensor<10x2x5x30xf32>, tensor<10x2x5x30xf32>) -> tensor<10x2x5x30xf32>
  return %1 : tensor<10x2x5x30xf32>
}

Output:

module {
  func.func @trt_gather_default1(%arg0: tensor<10x20x30xf32>, %arg1: tensor<2x5xi32>, %arg2: tensor<10x2x5x30xf32>) -> tensor<10x2x5x30xf32> {
    %0 = tensorrt.call_alloc @trt_engines::@tensorrt_cluster(%arg0, %arg1, %arg2 : tensor<10x20x30xf32>, tensor<2x5xi32>, tensor<10x2x5x30xf32>) -> tensor<10x2x5x30xf32>
    return %0 : tensor<10x2x5x30xf32>
  }
  tensorrt.module @trt_engines {
    func.func @tensorrt_cluster(%arg0: tensor<10x20x30xf32>, %arg1: tensor<2x5xi32>, %arg2: tensor<10x2x5x30xf32>) -> tensor<10x2x5x30xf32> {
      %0 = tensorrt.gather {axis = 1 : i64} ins(%arg0, %arg1 : tensor<10x20x30xf32>, tensor<2x5xi32>) -> tensor<10x2x5x30xf32>
      %1 = tensorrt.element_wise <kSUM>(%0, %arg2 : tensor<10x2x5x30xf32>, tensor<10x2x5x30xf32>) -> tensor<10x2x5x30xf32>
      return %1 : tensor<10x2x5x30xf32>
    }
  }
}

[christopherbate]

Commented above on how to fix. You should move it out of outlineOp back to its original position in order to avoid performing multiple linear scans unnecessarily.