trt_ops.td 2.0 KB
Newer Older
1 2 3 4 5 6 7 8 9
#ifndef TRT_OPS
#define TRT_OPS

include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/LoopLikeInterface.td"
include "mlir/Interfaces/CallInterfaces.td"
include "mlir/IR/OpBase.td"
include "paddle/infrt/dialect/tensorrt/trt_op_base.td"

S
Shang Zhizhou 已提交
10 11 12

def TRT_CreateEngineOp : TRT_Op<"create_engine", [SingleBlockImplicitTerminator<"::infrt::ReturnOp">]> {
  let summary = "trt CreateEngine Op";
13 14 15 16
  let description = [{
    Describe a tensorrt subgraph.
  }];
  let regions = (region SizedRegion<1>:$body);
17
  let arguments = (ins Variadic<TRT_Tensor>:$inputs, DefaultValuedAttr<BoolAttr, "true">:$run_once);
S
Shang Zhizhou 已提交
18 19 20 21 22 23 24 25 26 27
  let results = (outs TRT_EngineType:$output);
}

def TRT_ExecuteOp : TRT_Op<"execute", [NoSideEffect]> {
  let summary = "trt execute Op";
  let description = [{
    Describe a tensorrt runtime.
  }];
  let arguments = (ins TRT_EngineType:$engine, Variadic<TRT_Tensor>:$inputs);
  let results = (outs Variadic<TRT_Tensor>:$output);
28
}
29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68

def TRT_ActivationOp : TRT_Op<"Activation", [NoSideEffect]> {
  let summary = "TensorRT IActivationLayer";
  let description = [{
    
    TensorRT IActivationLayer.
    
  }];
  let arguments = (ins  TRT_Tensor:$input, SI32Attr:$activation_type,
                        DefaultValuedAttr<F32Attr, "0.0">:$alpha,
                        DefaultValuedAttr<F32Attr, "0.0">:$beta);

  let results = (outs TRT_Tensor:$output);
}

def TRT_ElementWiseOp : TRT_Op<"ElementWise", [NoSideEffect]> {
  let summary = "TensorRT IElementWiseLayer";
  let description = [{
    
    TensorRT IElementWiseLayer.
    
  }];
  let arguments = (ins  TRT_Tensor:$input1, TRT_Tensor:$input2, SI32Attr:$elementwise_operation);

  let results = (outs TRT_Tensor:$output);
}

def TRT_MatrixMultiplyOp : TRT_Op<"MatrixMultiply", [NoSideEffect]> {
  let summary = "TensorRT IMatrixMultiplyLayer";
  let description = [{
    
    TensorRT IMatrixMultiplyLayer.
    
  }];
  let arguments = (ins  TRT_Tensor:$input1, BoolAttr:$transpose1,
                        TRT_Tensor:$input2, BoolAttr:$transpose2);

  let results = (outs TRT_Tensor:$output);
}

69
#endif  // TRT_OPS