Skip to content

Cannot use tensorrt.broadcast with int8 inputs #580

@pranavm-nvidia

Description

@pranavm-nvidia

The following program (after disabling casting workarounds in full):

print(tp.ones(shape=(2, 2), dtype=tp.int8))

IRs:

==== Trace IR ====
def main() -> (
    %t3 : tensor<?x?xi8:gpu:0>
):
    %t1 = constant(shape=(), dtype=float32, device=gpu:0) : tensor<f32:gpu:0>
    %t2 = cast(%t1 : tensor<f32:gpu:0>, dtype=int8) : tensor<i8:gpu:0>
    %t0 = constant(shape=(2,), dtype=int32, device=gpu:0) : tensor<2xi32:gpu:0>
    %t3 = broadcast(%t2 : tensor<i8:gpu:0>, %t0 : tensor<2xi32:gpu:0>) : tensor<?x?xi8:gpu:0>
    return %t3

==== MLIR ====
module @"outs_%t3_0" {
  func.func @main() -> tensor<?x?xi8> {
    %cst_f32 = tensorrt.constant dense<1.000000e+00> : tensor<f32>
    %0 = tensorrt.cast %cst_f32 : tensor<f32> to tensor<i8>
    %cst_i32 = tensorrt.constant dense<2> : tensor<2xi32>
    %1 = tensorrt.broadcast %0 broadcast_dims<> shape(%cst_i32 : tensor<2xi32>) : tensor<i8> to tensor<?x?xi8>
    return %1 : tensor<?x?xi8>
  }
}

Fails with:

MTRTException: failed to run pass pipeline
    [optimizer.cpp::getFormatRequirements::3300] Error Code 2: Internal Error (Assertion !n->candidateRequirements.empty() failed. No supported formats for [tensorrt.slice] (%t3) , please check the TensorRT operator doc (https://docs.nvidia.com/deeplearning/tensorrt/operators/docs/index.html) for supported formats.)
    (%t3) error: failed to translate function 'tensorrt_cluster' to a TensorRT engine

Metadata

Metadata

Assignees

No one assigned

    Labels

    tripyPull request for the tripy project

    Type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions