Relay/TRT Integration (whole graph only)

[trevor-m]

This PR includes support for a version of Relay/TRT integration which only works when the entire model can be converted to TRT. It is enabled with the EnableTrt pass. If any op in the model cannot be converted to TRT, EnableTrt will return the original module unmodified.

How to use

1. Build TVM with cmake flag USE_TENSORRT=ON or USE_TENSORRT=/path/to/TensorRT. USE_CUDA should be enabled as well.

2. Convert the model into TensorRT. This step will determine if every node in the graph can be converted to TensorrRT and if so will mark the graph to use TensorRT and apply some specific optimization passes.

import tvm.relay.tensorrt
mod = relay.tensorrt.EnableTrt(mod, params)

3. Check if TRT was enabled. If not, it means some op in the graph is not supported by the TensorRT conversion. EnableTrt will output which particular ops are not supported and why.

assert mod['main'].attrs and mod['main'].attrs.Compiler == 'tensorrt'

4. Finish compilation.

with relay.build_config(opt_level=2, disabled_pass={"SimplifyInference"}):
  graph, lib, params = relay.build(mod, "cuda", params=params)

5. (Optional) Serialize/deserialize the compiled model. The model will be serialized to three files: compiled.json, compiled.params, and compiled.tensorrt.

# Serialize
with open('compiled.json', 'w') as f_graph_json:
  f_graph_json.write(graph)
with open('compiled.params', 'wb') as f_params:
  f_params.write(relay.save_param_dict(params))
lib.save('compiled.tensorrt')

# Deserialize
with open('compiled.json', 'r') as f_graph_json:
  graph = f_graph_json.read()
with open('compiled.params', 'rb') as f_params:
  params = tvm.relay.load_param_dict(f_params.read())
lib = tvm.module.load("compiled.tensorrt")

6. Run inference. The first invocation will trigger creation of the TensorRT engine. This could take up to a few minutes.

# Create graph runtime
mod = graph_runtime.create(graph, lib, ctx=tvm.gpu(0))
mod.set_input(**params)

i_data = np.random.uniform(0, 1, input_shape).astype(dtype)
# Build TensorRT engine
mod.run(data=i_data)

# Run inference
mod.run(data=i_data)
res = mod.get_output(0)

The tests tests/python/relay/test_tensorrt.py provide some deeper examples of how to use this feature.

The NNVM/TRT integration is still present.

[trevor-m]

@zhiics Could you please review?

[trevor-m]

@yongwww @reminisce @ziyu-guo Could you please review?

[zhiics]

Started reviewing. Will spend more time on it.

CC @comaniac who might be interested as well

[trevor-m]

@zhiics Could you take a quick look at my changes to use the external codegen/runtime?

[zhiics]

Need more time to review as this is a very big PR. Please ping a few more ppl to take a look.

[zhiics]

Some more comments

[trevor-m]

@comaniac @anijain2305 Could you please review? Thanks!

[comaniac]

Just took a quick pass. Will take another look later this week.
Except for the left comments, one issue I saw is the redundent functions such as GetShape and trt_version_ge defined in both Relay pass side and runtime side. This also brings anoher concern I have. Since TensorRT is a specific backend, it seems not a good practice to put its pass along with other Relay passes. It would be much better if all TRT related staffs can be put together.

In addition, please make up missing docstrings of C++ functions/classes with the formal format if time allows.

[apivovarov]

In all other places the log message ends with full stop

[apivovarov]

auto shape0_size = GetShape(call->type_args[0]).size();

[apivovarov]

I both of the messages above I'd output current pad_mode and pad_value.

LOG(INFO) << op_name << " not supported: pad mode is %s but must be constant.";

LOG(INFO) << op_name << " not supported: pad value is %f but must be 0.0."

[trevor-m]

Thanks, I went ahead and included current values for the rest of the ops messages as well.

[apivovarov]

I'd put i == 1 expr to the left of && because it can prevent unnecessary string comparisons

[apivovarov]

Should we say smth about trt_version_ here?

[trevor-m]

Added message to TrtVersionChecker, and calling TrtVersionChecker from other checkers when applicable.

[apivovarov]

Looks like compatible_ is not initialized after the object is constructed. It will be initialized only after user of the object calls Check.

[trevor-m]

Initializing it in constructor now

[apivovarov]

Can you use constructor(vector) or constructor(begin, end) to create and fill func_params here?

[apivovarov]

The call to std::move precludes the NRVO compiler optimization here.

[apivovarov]

Two checks above repeat 5 times in this module. Should we have utility method for this?

[apivovarov]

params=bytearray(f_params.read())
...
mod.load_params(params)

[apivovarov]

Should we use opt_level=3 here?

[apivovarov]

Looks like this is the only line which is different in TransposeRSCKtoCKRS in comparison to TransposeRSCKtoKCRS. Maybe define two functions to calculate output_index for CKRS and KCRS and use generic implementation + corresponding function to get one or another behavior?

[trevor-m]

Thanks, I created TransposeWeights4D which takes input and output strides as arguments.

[apivovarov]

Lets remove "TRT requires a constant input here." repetition here. e.g.

        CallNode* transpose = call->args[i].as<CallNode>();
        VarNode* weights;
        if (transpose != nullptr
            && transpose->op.as<OpNode>()->name == "transpose"
            && (weights = transpose->args[0].as<VarNode>()) != nullptr) {
          GetInputAsTransposedWeights(transpose, weights);
        } else {
          LOG(FATAL) << "TRT requires a constant input here.";
        }

[trevor-m]

Thanks, I was wondering how to clean that up.

[apivovarov]

We can make this "global" if to take just 5 lines and shift all code in the loop to the left

    if (converter->variable_input_count ||
        converter->input_types[i] != kWeight) {
      VisitExpr(call->args[i]);
      continue;
    }

[apivovarov]

Do we still need to use Legalize and relay.transpose? Can we use layout transform to convert source Relay graph to what TensorRT expects?

[trevor-m]

I think its better to leverage relay's pass ability to convert layout_transform op to the more standard transpose ops. This way we only need to write one TrtOpConverter for transpose. If we didn't perform this legalize, I would need to write an additional TrtOpConverter for layout_transform which would be nearly identical to the one for transpose.

This feature of relay is very useful. For example, TRT recently announced that they won't support INT8 for matmul/fully connected layer and they want everyone to just use 1x1 Conv instead (https://docs.nvidia.com/deeplearning/sdk/tensorrt-best-practices/index.html#optimize-layer). So in the future, I plan to have a similar pass to convert all matmul/dense layers into convolutions to take advantage of this. At that point I won't need a converter for dense anymore since everything would go to conv.