From 10e036b8071eeb1c24917033533a22bc85b96537 Mon Sep 17 00:00:00 2001 From: Michael Feliz Date: Wed, 3 Aug 2022 10:44:48 -0700 Subject: [PATCH 1/5] Add automatic type promotion to element-wise ops Adds automatic type promotion to match the default torch-script behavior for element-wise ops. Debug messages added for the type mismatch and cast. Messages written to log. ``` DEBUG: [Torch-TensorRT] - Type mismatch for inputs in element-wise operation %3 : Tensor = aten::add(%0, %1, %2): Int32, Float32 DEBUG: [Torch-TensorRT] - Element-wise op type promotion adding cast from Int32 to Float32 for layer %3 : Tensor = aten::add(%0, %1, %2) ``` Fixes # (issue) Please delete options that are not relevant and/or add your own. - Bug fix (non-breaking change which fixes an issue) - New feature (non-breaking change which adds functionality) - Breaking change (fix or feature that would cause existing functionality to not work as expected) - This change requires a documentation update - [ ] My code follows the style guidelines of this project (You can use the linters) - [ ] I have performed a self-review of my own code - [ ] I have commented my code, particularly in hard-to-understand areas and hacks - [ ] I have made corresponding changes to the documentation - [ ] I have added tests to verify my fix or my feature - [ ] New and existing unit tests pass locally with my changes - [ ] I have added the relevant labels to my PR in so that relevant reviewers are notified Signed-off-by: Michael Feliz --- core/conversion/converters/converter_util.cpp | 22 + .../converters/impl/element_wise.cpp | 1518 ++++++++--------- .../converters/test_element_wise.cpp | 128 +- 3 files changed, 854 insertions(+), 814 deletions(-) diff --git a/core/conversion/converters/converter_util.cpp b/core/conversion/converters/converter_util.cpp index 94ac827ef4..9f1d2a7124 100644 --- a/core/conversion/converters/converter_util.cpp +++ b/core/conversion/converters/converter_util.cpp @@ -59,6 +59,14 @@ nvinfer1::ITensor* addUnpadding( } } +nvinfer1::DataType promote_types(nvinfer1::DataType type_a, nvinfer1::DataType type_b){ + auto torch_type_a = util::TRTDataTypeToScalarType(type_a); + auto torch_type_b = util::TRTDataTypeToScalarType(type_b); + auto promo_type = at::promote_types(torch_type_a, torch_type_b); + auto trt_promo_type = util::ScalarTypeToTRTDataType(promo_type); + return trt_promo_type; +} + nvinfer1::ILayer* add_elementwise( ConversionCtx* ctx, nvinfer1::ElementWiseOperation op, @@ -78,6 +86,20 @@ nvinfer1::ILayer* add_elementwise( std::swap(self, other); swapSelfOther = true; } + + if(self->getType() != other->getType()){ + LOG_DEBUG("Type mismatch for inputs in element-wise operation " << name << ": " << self->getType() << ", " << other->getType()); + auto promo_type = promote_types(self->getType(), other->getType()); + if(self->getType() != promo_type){ + LOG_DEBUG("Element-wise op type promotion adding cast from " << self->getType() << " to " << promo_type << " for layer " << name); + self = castITensor(ctx, self, promo_type); + } + if(other->getType() != promo_type){ + LOG_DEBUG("Element-wise op type promotion adding cast from " << other->getType() << " to " << promo_type << " for layer " << name); + other = castITensor(ctx, other, promo_type); + } + } + auto selfDim = util::toVec(self->getDimensions()); auto otherDim = util::toVec(other->getDimensions()); if (selfDim.size() != otherDim.size()) { diff --git a/core/conversion/converters/impl/element_wise.cpp b/core/conversion/converters/impl/element_wise.cpp index f2770508ca..d58676a00b 100644 --- a/core/conversion/converters/impl/element_wise.cpp +++ b/core/conversion/converters/impl/element_wise.cpp @@ -48,775 +48,755 @@ nvinfer1::ITensor* scalar_to_tensor(ConversionCtx* ctx, at::Scalar s) { auto element_wise_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() - .pattern( - {"aten::add.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> " - "Tensor", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self + alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto scalar = args[2].unwrapToScalar().to(); - - if (1 != scalar) { - auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - other, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - other = scaleLayer->getOutput(0); - } - - auto add = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); - TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); - - add->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar " - "alpha=1) -> (Tensor(a!))", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self + alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto scalar = args[2].unwrapToScalar().to(); - - if (1 != scalar) { - auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - other, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - other = scaleLayer->getOutput(0); - } - - auto add = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); - TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); - - add->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self + alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[2].unwrapToScalar().to() * args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - - auto add = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); - TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); - - add->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Compute min(max(min_threshold, input), max_threshold) - auto self = args[0].ITensorOrFreeze(ctx); - auto clamp_layer_out = self; - - if (args[1].isIValue() && args[1].IValue()->isScalar() && args[2].isIValue() && - args[2].IValue()->isScalar()) { - auto alpha = args[1].unwrapToScalar().to(); - auto beta = args[2].unwrapToScalar().to(); - auto clip_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kCLIP); - TORCHTRT_CHECK(clip_layer, "Unable to create clip layer for node: " << *n); - clip_layer->setAlpha(alpha); - clip_layer->setBeta(beta); - clamp_layer_out = clip_layer->getOutput(0); - } else if (args[1].isIValue() && args[1].IValue()->isScalar()) { - auto limit = args[1].unwrapToScalar().to(); - clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max"); - } else if (args[2].isIValue() && args[2].IValue()->isScalar()) { - auto limit = args[2].unwrapToScalar().to(); - clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min"); - } - - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); - LOG_DEBUG("Clamp layer output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::clamp_min(Tensor self, Scalar min) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Compute min(max(min_threshold, input), max_threshold) - auto self = args[0].ITensorOrFreeze(ctx); - auto clamp_layer_out = self; - if (args[1].isIValue() && args[1].IValue()->isScalar()) { - auto limit = args[1].unwrapToScalar().to(); - clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max"); - } - - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); - LOG_DEBUG("clamp_min layer output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::clamp_max(Tensor self, Scalar max) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Compute min(max(min_threshold, input), max_threshold) - auto self = args[0].ITensorOrFreeze(ctx); - auto clamp_layer_out = self; - if (args[1].isIValue() && args[1].IValue()->isScalar()) { - auto limit = args[1].unwrapToScalar().to(); - clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min"); - } - - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); - LOG_DEBUG("clamp_max layer output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::sub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> " - "Tensor", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self - alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto scalar = args[2].unwrapToScalar().to(); - auto other = args[1].ITensorOrFreeze(ctx); - - if (1 != scalar) { - auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - other, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - other = scaleLayer->getOutput(0); - } - - auto sub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n)); - TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); - - sub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self - alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].unwrapToScalar().to(); - auto alpha = args[2].unwrapToScalar().to(); - auto scaled_val = other * alpha; - - auto scaled_other_tensor = tensor_to_const(ctx, torch::tensor({scaled_val})); - auto sub = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kSUB, self, scaled_other_tensor, util::node_info(n)); - TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); - sub->setName(util::node_info(n).c_str()); - LOG_DEBUG("Output tensor shape: " << sub->getOutput(0)->getDimensions()); - ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); - - return true; - }}) - .pattern( - {"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar " - "alpha=1) -> (Tensor(a!))", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self - alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto scalar = args[2].unwrapToScalar().to(); - auto other = args[1].ITensorOrFreeze(ctx); - - if (1 != scalar) { - auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - other, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - other = scaleLayer->getOutput(0); - } - - auto sub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n)); - TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); - - sub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::rsub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement other - alpha * self - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - auto scalar = args[2].unwrapToScalar().to(); - - if (1 != scalar) { - auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - self, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - self = scaleLayer->getOutput(0); - } - - auto rsub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n)); - TORCHTRT_CHECK(rsub, "Unable to create rsub layer from node: " << *n); - - rsub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::rsub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement other - alpha * self - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto scalar = args[2].unwrapToScalar().to(); - - if (1 != scalar) { - auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - self, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - self = scaleLayer->getOutput(0); - } - - auto rsub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n)); - TORCHTRT_CHECK(rsub, "Unable to create rsub layer from node: " << *n); - - rsub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::div.Tensor(Tensor self, Tensor other) -> Tensor", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self / other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); - - TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); - - div->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::div.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self / other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - std::string rounding_mode = "default"; - if (args[2].isIValue() && args[2].IValue()->isString()) { - rounding_mode = args[2].unwrapToString(); - } - nvinfer1::ILayer* div = nullptr; - if (rounding_mode == "floor") { - div = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); - } else if (rounding_mode == "trunc") { - // trunc = floor(abs(div)) * sign(div) - auto tmp_div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, "tmp_div"); - auto abs = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kABS); - auto floor = ctx->net->addUnary(*abs->getOutput(0), nvinfer1::UnaryOperation::kFLOOR); - auto sign = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kSIGN); - div = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - floor->getOutput(0), - sign->getOutput(0), - util::node_info(n)); - } else { - div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); - } - - TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); - - div->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::div.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); - TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); - - div->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::div_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); - - TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); - - div->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); - TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); - - div->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::mul.Tensor(Tensor self, Tensor other) -> Tensor", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self * other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - - auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); - TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); - - mul->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::mul.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); - - auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); - TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); - - mul->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::mul_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); - TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); - - mul->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::ne.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto equal = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kEQUAL, - self, - other, - util::node_info(n) + std::string("is_equal")); - TORCHTRT_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n); - // XOR with ones negates and produces not_equal result - auto options = torch::TensorOptions().dtype(torch::kFloat32); - auto ones = at::full({1}, 1, {options}); - auto ones_tensor = tensor_to_const(ctx, ones); - nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor); - cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL); - - auto sub = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kXOR, - cast_layer->getOutput(0), - equal->getOutput(0), - util::node_info(n)); - TORCHTRT_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n); - - sub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); - LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::ne.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto scalar = args[1].unwrapToScalar(); - nvinfer1::ITensor* scalar_tensor; - if (self->getType() == nvinfer1::DataType::kFLOAT || self->getType() == nvinfer1::DataType::kHALF) { - scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to()})); - } else { - scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to()})); - } - auto equal = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kEQUAL, - self, - scalar_tensor, - util::node_info(n) + std::string("is_equal")); - TORCHTRT_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n); - // XOR with ones negates and produces not_equal result - auto options = torch::TensorOptions().dtype(torch::kFloat32); - auto ones = at::full({1}, 1, {options}); - auto ones_tensor = tensor_to_const(ctx, ones); - nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor); - cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL); - - auto sub = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kXOR, - cast_layer->getOutput(0), - equal->getOutput(0), - util::node_info(n)); - TORCHTRT_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n); - - sub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); - LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::pow.Tensor_Tensor(Tensor self, Tensor exponent) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto exponent = args[1].ITensorOrFreeze(ctx); - auto pow = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n)); - TORCHTRT_CHECK(pow, "Unable to create Power layer from node: " << *n); - - pow->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::pow.Tensor_Scalar(Tensor self, Scalar exponent) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto exponentScalar = args[1].unwrapToScalar().to(); - auto exponent = tensor_to_const(ctx, torch::tensor({exponentScalar})); - auto pow = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n)); - TORCHTRT_CHECK(pow, "Unable to create Power layer from node: " << *n); - - pow->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::floor_divide(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto floor_divide = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); - TORCHTRT_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n); - - floor_divide->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::floor_divide.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - auto floor_divide = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); - TORCHTRT_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n); - - floor_divide->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::max.other(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto max = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMAX, self, other, util::node_info(n)); - TORCHTRT_CHECK(max, "Unable to create max layer from node: " << *n); - - max->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], max->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::min.other(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto min = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMIN, self, other, util::node_info(n)); - TORCHTRT_CHECK(min, "Unable to create min layer from node: " << *n); - - min->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], min->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::gt.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto gt = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n)); - TORCHTRT_CHECK(gt, "Unable to create greater layer from node: " << *n); - - gt->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], gt->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::gt.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); - if (self->getType() != other->getType()) { - other = castITensor(ctx, other, self->getType()); - } - auto gt = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n)); - TORCHTRT_CHECK(gt, "Unable to create greater layer from node: " << *n); - - gt->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], gt->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::lt.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto lt = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n)); - TORCHTRT_CHECK(lt, "Unable to create less layer from node: " << *n); - - lt->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], lt->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::lt.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); - if (self->getType() != other->getType()) { - other = castITensor(ctx, other, self->getType()); - } - auto lt = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n)); - TORCHTRT_CHECK(lt, "Unable to create less layer from node: " << *n); - - lt->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], lt->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::eq.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto eq = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n)); - TORCHTRT_CHECK(eq, "Unable to create equal layer from node: " << *n); - - eq->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], eq->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::eq.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - if (self->getType() == nvinfer1::DataType::kBOOL) { - if (otherScalar == 0 || otherScalar == 1) { - LOG_DEBUG("Since input tensor is type bool, casting input tensor and scalar to int32"); - other = castITensor(ctx, other, nvinfer1::DataType::kINT32); - self = castITensor(ctx, self, nvinfer1::DataType::kINT32); - } else { - LOG_WARNING("Input Tensor has type bool, but scalar is not 0 or 1. Found: " << otherScalar); - } - } - if (self->getType() != other->getType()) { - other = castITensor(ctx, other, self->getType()); - } - auto eq = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n)); - TORCHTRT_CHECK(eq, "Unable to create equal layer from node: " << *n); - - eq->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], eq->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::ge.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - - auto greater = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater"); - TORCHTRT_CHECK(greater, "Unable to create Greater layer from node: " << *n); - - auto equal = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); - TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); - - auto or_op = ctx->net->addElementWise( - *greater->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); - - TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); - or_op->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::ge.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); - if (self->getType() != other->getType()) { - other = castITensor(ctx, other, self->getType()); - } - - auto greater = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater"); - TORCHTRT_CHECK(greater, "Unable to create Greater layer from node: " << *n); - - auto equal = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); - TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); - - auto or_op = ctx->net->addElementWise( - *greater->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); - - TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); - or_op->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::le.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - - auto less = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less"); - TORCHTRT_CHECK(less, "Unable to create Less layer from node: " << *n); - - auto equal = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); - TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); - - auto or_op = ctx->net->addElementWise( - *less->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); - - TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); - or_op->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern( - {"aten::le.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); - if (self->getType() != other->getType()) { - other = castITensor(ctx, other, self->getType()); - } - - auto less = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less"); - TORCHTRT_CHECK(less, "Unable to create Less layer from node: " << *n); - - auto equal = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); - TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); - - auto or_op = ctx->net->addElementWise( - *less->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); - - TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); - or_op->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}); + .pattern({"aten::add.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> " + "Tensor", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self + alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto scalar = args[2].unwrapToScalar(); + + if (1 != scalar.to()) { + auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + other, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + other = scaleLayer->getOutput(0); + } + + auto add = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); + TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); + + add->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar " + "alpha=1) -> (Tensor(a!))", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self + alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto scalar = args[2].unwrapToScalar(); + + if (1 != scalar.to()) { + auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + other, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + other = scaleLayer->getOutput(0); + } + + auto add = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); + TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); + + add->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self + alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto otherScalar = args[2].unwrapToScalar().to() * args[1].unwrapToScalar().to(); + auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + + auto add = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); + TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); + + add->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Compute min(max(min_threshold, input), max_threshold) + auto self = args[0].ITensorOrFreeze(ctx); + auto clamp_layer_out = self; + + if (args[1].isIValue() && args[1].IValue()->isScalar() && args[2].isIValue() && + args[2].IValue()->isScalar()) { + auto alpha = args[1].unwrapToScalar().to(); + auto beta = args[2].unwrapToScalar().to(); + auto clip_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kCLIP); + TORCHTRT_CHECK(clip_layer, "Unable to create clip layer for node: " << *n); + clip_layer->setAlpha(alpha); + clip_layer->setBeta(beta); + clamp_layer_out = clip_layer->getOutput(0); + } else if (args[1].isIValue() && args[1].IValue()->isScalar()) { + auto limit = args[1].unwrapToScalar().to(); + clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max"); + } else if (args[2].isIValue() && args[2].IValue()->isScalar()) { + auto limit = args[2].unwrapToScalar().to(); + clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min"); + } + + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); + LOG_DEBUG("Clamp layer output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::clamp_min(Tensor self, Scalar min) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Compute min(max(min_threshold, input), max_threshold) + auto self = args[0].ITensorOrFreeze(ctx); + auto clamp_layer_out = self; + if (args[1].isIValue() && args[1].IValue()->isScalar()) { + auto limit = args[1].unwrapToScalar().to(); + clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max"); + } + + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); + LOG_DEBUG("clamp_min layer output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::clamp_max(Tensor self, Scalar max) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Compute min(max(min_threshold, input), max_threshold) + auto self = args[0].ITensorOrFreeze(ctx); + auto clamp_layer_out = self; + if (args[1].isIValue() && args[1].IValue()->isScalar()) { + auto limit = args[1].unwrapToScalar().to(); + clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min"); + } + + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); + LOG_DEBUG("clamp_max layer output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::sub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> " + "Tensor", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self - alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto scalar = args[2].unwrapToScalar().to(); + auto other = args[1].ITensorOrFreeze(ctx); + + if (1 != scalar) { + auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + other, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + other = scaleLayer->getOutput(0); + } + + auto sub = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n)); + TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); + + sub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self - alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].unwrapToScalar().to(); + auto alpha = args[2].unwrapToScalar().to(); + auto scaled_val = other * alpha; + + auto scaled_other_tensor = tensor_to_const(ctx, torch::tensor({scaled_val})); + auto sub = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kSUB, self, scaled_other_tensor, util::node_info(n)); + TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); + sub->setName(util::node_info(n).c_str()); + LOG_DEBUG("Output tensor shape: " << sub->getOutput(0)->getDimensions()); + ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); + + return true; + }}) + .pattern({"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar " + "alpha=1) -> (Tensor(a!))", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self - alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto scalar = args[2].unwrapToScalar().to(); + auto other = args[1].ITensorOrFreeze(ctx); + + if (1 != scalar) { + auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + other, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + other = scaleLayer->getOutput(0); + } + + auto sub = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n)); + TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); + + sub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::rsub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement other - alpha * self + auto self = args[0].ITensorOrFreeze(ctx); + auto other = impl::scalar_to_tensor(ctx, args[1].unwrapToScalar()); + auto scalar = args[2].unwrapToScalar(); + + if (1 != scalar.to()) { + auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + self, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + self = scaleLayer->getOutput(0); + } + + auto rsub = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n)); + TORCHTRT_CHECK(rsub, "Unable to create rsub layer from node: " << *n); + + rsub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::rsub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement other - alpha * self + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto scalar = args[2].unwrapToScalar(); + + if (1 != scalar.to()) { + auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + self, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + self = scaleLayer->getOutput(0); + } + + auto rsub = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n)); + TORCHTRT_CHECK(rsub, "Unable to create rsub layer from node: " << *n); + + rsub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::div.Tensor(Tensor self, Tensor other) -> Tensor", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self / other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto div = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); + + TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); + + div->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::div.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self / other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + std::string rounding_mode = "default"; + if (args[2].isIValue() && args[2].IValue()->isString()) { + rounding_mode = args[2].unwrapToString(); + } + nvinfer1::ILayer* div = nullptr; + if (rounding_mode == "floor") { + div = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); + } else if (rounding_mode == "trunc") { + // trunc = floor(abs(div)) * sign(div) + auto tmp_div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, "tmp_div"); + auto abs = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kABS); + auto floor = ctx->net->addUnary(*abs->getOutput(0), nvinfer1::UnaryOperation::kFLOOR); + auto sign = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kSIGN); + div = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + floor->getOutput(0), + sign->getOutput(0), + util::node_info(n)); + } else { + div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); + } + + TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); + + div->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::div.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto otherScalar = args[1].unwrapToScalar().to(); + auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + auto div = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); + TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); + + div->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::div_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto div = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); + + TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); + + div->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto otherScalar = args[1].unwrapToScalar().to(); + auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + auto div = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); + TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); + + div->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::mul.Tensor(Tensor self, Tensor other) -> Tensor", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self * other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto mul = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); + TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); + + mul->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::mul.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); + auto mul = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); + TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); + + mul->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::mul_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto mul = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); + TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); + + mul->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::ne.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto equal = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kEQUAL, + self, + other, + util::node_info(n) + std::string("is_equal")); + TORCHTRT_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n); + // XOR with ones negates and produces not_equal result + auto options = torch::TensorOptions().dtype(torch::kFloat32); + auto ones = at::full({1}, 1, {options}); + auto ones_tensor = tensor_to_const(ctx, ones); + nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor); + cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL); + + auto sub = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kXOR, + cast_layer->getOutput(0), + equal->getOutput(0), + util::node_info(n)); + TORCHTRT_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n); + + sub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); + LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::ne.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto scalar = args[1].unwrapToScalar(); + nvinfer1::ITensor* scalar_tensor; + if (self->getType() == nvinfer1::DataType::kFLOAT || self->getType() == nvinfer1::DataType::kHALF) { + scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to()})); + } else { + scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to()})); + } + auto equal = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kEQUAL, + self, + scalar_tensor, + util::node_info(n) + std::string("is_equal")); + TORCHTRT_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n); + // XOR with ones negates and produces not_equal result + auto options = torch::TensorOptions().dtype(torch::kFloat32); + auto ones = at::full({1}, 1, {options}); + auto ones_tensor = tensor_to_const(ctx, ones); + nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor); + cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL); + + auto sub = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kXOR, + cast_layer->getOutput(0), + equal->getOutput(0), + util::node_info(n)); + TORCHTRT_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n); + + sub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); + LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::pow.Tensor_Tensor(Tensor self, Tensor exponent) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto exponent = args[1].ITensorOrFreeze(ctx); + auto pow = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n)); + TORCHTRT_CHECK(pow, "Unable to create Power layer from node: " << *n); + + pow->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::pow.Tensor_Scalar(Tensor self, Scalar exponent) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto exponentScalar = args[1].unwrapToScalar().to(); + auto exponent = tensor_to_const(ctx, torch::tensor({exponentScalar})); + auto pow = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n)); + TORCHTRT_CHECK(pow, "Unable to create Power layer from node: " << *n); + + pow->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::floor_divide(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto floor_divide = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); + TORCHTRT_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n); + + floor_divide->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::floor_divide.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto otherScalar = args[1].unwrapToScalar().to(); + auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + auto floor_divide = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); + TORCHTRT_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n); + + floor_divide->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::max.other(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto max = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMAX, self, other, util::node_info(n)); + TORCHTRT_CHECK(max, "Unable to create max layer from node: " << *n); + + max->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], max->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::min.other(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto min = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMIN, self, other, util::node_info(n)); + TORCHTRT_CHECK(min, "Unable to create min layer from node: " << *n); + + min->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], min->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::gt.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto gt = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n)); + TORCHTRT_CHECK(gt, "Unable to create greater layer from node: " << *n); + + gt->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], gt->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::gt.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); + if (self->getType() != other->getType()) { + other = castITensor(ctx, other, self->getType()); + } + auto gt = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n)); + TORCHTRT_CHECK(gt, "Unable to create greater layer from node: " << *n); + + gt->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], gt->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::lt.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto lt = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n)); + TORCHTRT_CHECK(lt, "Unable to create less layer from node: " << *n); + + lt->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], lt->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::lt.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); + if (self->getType() != other->getType()) { + other = castITensor(ctx, other, self->getType()); + } + auto lt = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n)); + TORCHTRT_CHECK(lt, "Unable to create less layer from node: " << *n); + + lt->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], lt->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::eq.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto eq = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n)); + TORCHTRT_CHECK(eq, "Unable to create equal layer from node: " << *n); + + eq->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], eq->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::eq.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto otherScalar = args[1].unwrapToScalar().to(); + auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + if (self->getType() == nvinfer1::DataType::kBOOL) { + if (otherScalar == 0 || otherScalar == 1) { + LOG_DEBUG("Since input tensor is type bool, casting input tensor and scalar to int32"); + other = castITensor(ctx, other, nvinfer1::DataType::kINT32); + self = castITensor(ctx, self, nvinfer1::DataType::kINT32); + } else { + LOG_WARNING("Input Tensor has type bool, but scalar is not 0 or 1. Found: " << otherScalar); + } + } + if (self->getType() != other->getType()) { + other = castITensor(ctx, other, self->getType()); + } + auto eq = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n)); + TORCHTRT_CHECK(eq, "Unable to create equal layer from node: " << *n); + + eq->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], eq->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::ge.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + + auto greater = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater"); + TORCHTRT_CHECK(greater, "Unable to create Greater layer from node: " << *n); + + auto equal = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); + TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); + + auto or_op = ctx->net->addElementWise( + *greater->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); + + TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); + or_op->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::ge.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); + if (self->getType() != other->getType()) { + other = castITensor(ctx, other, self->getType()); + } + + auto greater = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater"); + TORCHTRT_CHECK(greater, "Unable to create Greater layer from node: " << *n); + + auto equal = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); + TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); + + auto or_op = ctx->net->addElementWise( + *greater->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); + + TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); + or_op->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::le.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + + auto less = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less"); + TORCHTRT_CHECK(less, "Unable to create Less layer from node: " << *n); + + auto equal = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); + TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); + + auto or_op = ctx->net->addElementWise( + *less->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); + + TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); + or_op->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern({"aten::le.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); + if (self->getType() != other->getType()) { + other = castITensor(ctx, other, self->getType()); + } + + auto less = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less"); + TORCHTRT_CHECK(less, "Unable to create Less layer from node: " << *n); + + auto equal = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); + TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); + + auto or_op = ctx->net->addElementWise( + *less->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); + + TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); + or_op->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}); } // namespace } // namespace impl diff --git a/tests/core/conversion/converters/test_element_wise.cpp b/tests/core/conversion/converters/test_element_wise.cpp index 061eaa500e..6a1bf400e2 100644 --- a/tests/core/conversion/converters/test_element_wise.cpp +++ b/tests/core/conversion/converters/test_element_wise.cpp @@ -12,40 +12,29 @@ void pointwise_test_helper( std::vector shape1 = {5}, std::vector shape2 = {5}, bool negative_input = false, - bool int_tensors = false, - bool float_int_tensors = false, - bool int_float_tensors = false) { + at::ScalarType type1 = at::kFloat, + at::ScalarType type2 = at::kFloat) { auto g = std::make_shared(); torch::jit::parseIR(graph_ir, g.get()); // singleInput case is enabled when elementwise operation is performed // with an input and a constant embedded in graph std::vector torch_inputs; - if (negative_input) { - torch_inputs.push_back(at::randint(-5, 5, shape1, {at::kCUDA})); - } else { - torch_inputs.push_back(at::randint(1, 5, shape1, {at::kCUDA})); + int first_min = negative_input ? -5 : 1; + int first_max = 5; + int second_min = 1; + int second_max = 5; + if(type1 == at::kBool){ + first_min = 0; + first_max = 1; } - if (!singleInput) { - torch_inputs.push_back(at::randint(1, 5, shape2, {at::kCUDA})); + if(type2 == at::kBool){ + second_min = 0; + second_max = 1; } - - TORCHTRT_CHECK( - !((int_tensors && (float_int_tensors || int_float_tensors)) || (float_int_tensors && int_float_tensors)), - "Invalid test configuration, only one of int_tensors, float_int_tensors, int_float_tensors can be true"); - - if (int_tensors) { - for (size_t i = 0UL; i < torch_inputs.size(); ++i) { - torch_inputs[i] = torch_inputs[i].to(at::kInt); - } - } else if (float_int_tensors) { - TORCHTRT_CHECK(!singleInput, "float_int_tensors tests require two inputs"); - torch_inputs[0] = torch_inputs[0].to(at::kFloat); - torch_inputs[1] = torch_inputs[1].to(at::kInt); - } else if (int_float_tensors) { - TORCHTRT_CHECK(!singleInput, "int_float_tensors tests require two inputs"); - torch_inputs[0] = torch_inputs[0].to(at::kInt); - torch_inputs[1] = torch_inputs[1].to(at::kFloat); + torch_inputs.push_back(at::randint(first_min, first_max, shape1, at::TensorOptions(at::kCUDA).dtype(type1))); + if (!singleInput) { + torch_inputs.push_back(at::randint(second_min, second_max, shape2, at::TensorOptions(at::kCUDA).dtype(type2))); } auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {}); @@ -78,8 +67,6 @@ TEST(Converters, ATenAddConvertsCorrectly) { pointwise_test_helper(graph, false, false, {4}, {3, 4}); pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}); pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, false, true); } TEST(Converters, ATenAddWithAlphaConvertsCorrectly) { @@ -93,8 +80,8 @@ TEST(Converters, ATenAddWithAlphaConvertsCorrectly) { pointwise_test_helper(graph, false, false, {4}, {3, 4}); pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}); pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, false, true); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat); } TEST(Converters, ATenAddInplaceWithAlphaConvertsCorrectly) { @@ -106,6 +93,17 @@ TEST(Converters, ATenAddInplaceWithAlphaConvertsCorrectly) { pointwise_test_helper(graph, false); pointwise_test_helper(graph, false, false, {3, 4}, {4}); pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}); + pointwise_test_helper(graph, false, false, {3, 4, 3}, {4, 3}, false, at::kFloat, at::kInt); +} + +TEST(Converters, ATenAddImplicitWithIntAlphaConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor, %1 : Tensor): + %2 : int = prim::Constant[value=42]() + %3 : Tensor = aten::add_(%0, %1, %2) + return (%3))IR"; + pointwise_test_helper(graph, false, false, {2, 2}, {2, 2}, false, at::kInt, at::kInt); + pointwise_test_helper(graph, false, false, {3, 4, 3}, {4, 3}, false, at::kInt, at::kInt); } TEST(Converters, ATenAddWithScalarConvertsCorrectly) { @@ -129,8 +127,8 @@ TEST(Converters, ATenSubConvertsCorrectly) { pointwise_test_helper(graph, false, false, {4}, {3, 4}); pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}); pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, false, true); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat); } TEST(Converters, ATenMulConvertsCorrectly) { @@ -143,8 +141,8 @@ TEST(Converters, ATenMulConvertsCorrectly) { pointwise_test_helper(graph, false, false, {4}, {3, 4}); pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}); pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, false, true); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat); } TEST(Converters, ATenMulWithScalarConvertsCorrectly) { @@ -162,7 +160,7 @@ TEST(Converters, ATenMulWithIntScalarConvertsCorrectly) { %scalar : int = prim::Constant[value=2]() %1 : Tensor = aten::mul(%0, %scalar) return (%1))IR"; - pointwise_test_helper(graph, true, false, {5}, {5}, false, true); + pointwise_test_helper(graph, true, false, {5}, {5}, false, at::kInt); } TEST(Converters, ATenDivConvertsCorrectly) { @@ -175,8 +173,6 @@ TEST(Converters, ATenDivConvertsCorrectly) { pointwise_test_helper(graph, false, false, {4}, {3, 4}); pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}); pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, false, true); } TEST(Converters, ATenDivWithScalarConvertsCorrectly) { @@ -199,8 +195,8 @@ TEST(Converters, ATenDivRoundingFloorConvertsCorrectly) { pointwise_test_helper(graph, false, false, {4}, {3, 4}, true); pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}, true); pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, false, true); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat); } TEST(Converters, ATenDivRoundingTruncConvertsCorrectly) { @@ -214,8 +210,8 @@ TEST(Converters, ATenDivRoundingTruncConvertsCorrectly) { pointwise_test_helper(graph, false, false, {4}, {3, 4}, true); pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}, true); pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, false, true); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat); } TEST(Converters, ATenDivRoundingNoneConvertsCorrectly) { @@ -241,8 +237,8 @@ TEST(Converters, ATenPowTensorConvertsCorrectly) { pointwise_test_helper(graph, false, false, {4}, {3, 4}); pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}); pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, false, true); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat); } TEST(Converters, ATenPowScalarConvertsCorrectly) { @@ -283,8 +279,8 @@ TEST(Converters, ATenFloorDivideConvertsCorrectly) { pointwise_test_helper(graph, false, false, {4}, {3, 4}); pointwise_test_helper(graph, false, true, {3, 4, 3}, {4, 3}); pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, true); - pointwise_test_helper(graph, false, true, {5}, {5}, false, false, false, true); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kFloat, at::kInt); + pointwise_test_helper(graph, false, true, {5}, {5}, false, at::kInt, at::kFloat); } TEST(Converters, ATenFloorDivideWithScalarConvertsCorrectly) { @@ -329,6 +325,8 @@ TEST(Converters, ATenRsubWithTensorConvertsCorrectly) { pointwise_test_helper(graph, false, false, {3, 4}, {4}); pointwise_test_helper(graph, false, false, {4}, {3, 4}); pointwise_test_helper(graph, false, true, {4, 3, 3, 3}, {4, 3, 3, 3}); + pointwise_test_helper(graph, false, false, {4, 3, 3, 3}, {4, 3, 3, 3}, false, at::kInt, at::kFloat); + pointwise_test_helper(graph, false, false, {4, 3, 3, 3}, {4, 3, 3, 3}, false, at::kInt, at::kInt); } TEST(Converters, ATenRsubWithScalarConvertsCorrectly) { @@ -341,6 +339,46 @@ TEST(Converters, ATenRsubWithScalarConvertsCorrectly) { pointwise_test_helper(graph, true, false, {4, 3, 3, 3}); } +TEST(Converters, ATenRsubWithIntScalarConvertsCorrectly) { + const auto graph = R"IR( + graph(%0 : Tensor): + %2 : int = prim::Constant[value=2]() + %scalar : int = prim::Constant[value=8]() + %3 : Tensor = aten::rsub(%0, %scalar, %2) + return (%3))IR"; + pointwise_test_helper(graph, true, false, {4, 3, 3, 3}, {}, false, at::kInt); +} + +TEST(Converters, ATenClipMinConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : float = prim::Constant[value=1.5]() + %3 : None = prim::Constant() + %4 : Tensor = aten::clip(%x.1, %2, %3) + return (%4))IR"; + pointwise_test_helper(graph, true); +} + +TEST(Converters, ATenClipMaxConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : float = prim::Constant[value=3.5]() + %3 : None = prim::Constant() + %4 : Tensor = aten::clip(%x.1, %3, %2) + return (%4))IR"; + pointwise_test_helper(graph, true); +} + +TEST(Converters, ATenClipMinMaxConvertsCorrectly) { + const auto graph = R"IR( + graph(%x.1 : Tensor): + %2 : float = prim::Constant[value=3.5]() + %3 : float = prim::Constant[value=1.5]() + %4 : Tensor = aten::clip(%x.1, %3, %2) + return (%4))IR"; + pointwise_test_helper(graph, true); +} + TEST(Converters, ATenClampMinConvertsCorrectly) { const auto graph = R"IR( graph(%x.1 : Tensor): From 5d76b024dd9baa3e31f38fec7534d444e46af3a3 Mon Sep 17 00:00:00 2001 From: Michael Feliz Date: Tue, 9 Aug 2022 15:51:17 -0700 Subject: [PATCH 2/5] remove clip tests --- .../converters/test_element_wise.cpp | 30 ------------------- 1 file changed, 30 deletions(-) diff --git a/tests/core/conversion/converters/test_element_wise.cpp b/tests/core/conversion/converters/test_element_wise.cpp index 6a1bf400e2..99be3cbffa 100644 --- a/tests/core/conversion/converters/test_element_wise.cpp +++ b/tests/core/conversion/converters/test_element_wise.cpp @@ -349,36 +349,6 @@ TEST(Converters, ATenRsubWithIntScalarConvertsCorrectly) { pointwise_test_helper(graph, true, false, {4, 3, 3, 3}, {}, false, at::kInt); } -TEST(Converters, ATenClipMinConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : float = prim::Constant[value=1.5]() - %3 : None = prim::Constant() - %4 : Tensor = aten::clip(%x.1, %2, %3) - return (%4))IR"; - pointwise_test_helper(graph, true); -} - -TEST(Converters, ATenClipMaxConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : float = prim::Constant[value=3.5]() - %3 : None = prim::Constant() - %4 : Tensor = aten::clip(%x.1, %3, %2) - return (%4))IR"; - pointwise_test_helper(graph, true); -} - -TEST(Converters, ATenClipMinMaxConvertsCorrectly) { - const auto graph = R"IR( - graph(%x.1 : Tensor): - %2 : float = prim::Constant[value=3.5]() - %3 : float = prim::Constant[value=1.5]() - %4 : Tensor = aten::clip(%x.1, %3, %2) - return (%4))IR"; - pointwise_test_helper(graph, true); -} - TEST(Converters, ATenClampMinConvertsCorrectly) { const auto graph = R"IR( graph(%x.1 : Tensor): From 3d754821e16acc814b3cda338f1f3588864032ea Mon Sep 17 00:00:00 2001 From: Michael Feliz Date: Wed, 10 Aug 2022 14:03:07 -0700 Subject: [PATCH 3/5] lint element_wise.cpp --- .../converters/impl/element_wise.cpp | 1515 +++++++++-------- 1 file changed, 766 insertions(+), 749 deletions(-) diff --git a/core/conversion/converters/impl/element_wise.cpp b/core/conversion/converters/impl/element_wise.cpp index d58676a00b..4d32ce8bb6 100644 --- a/core/conversion/converters/impl/element_wise.cpp +++ b/core/conversion/converters/impl/element_wise.cpp @@ -48,755 +48,772 @@ nvinfer1::ITensor* scalar_to_tensor(ConversionCtx* ctx, at::Scalar s) { auto element_wise_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() - .pattern({"aten::add.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> " - "Tensor", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self + alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto scalar = args[2].unwrapToScalar(); - - if (1 != scalar.to()) { - auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - other, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - other = scaleLayer->getOutput(0); - } - - auto add = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); - TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); - - add->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar " - "alpha=1) -> (Tensor(a!))", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self + alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto scalar = args[2].unwrapToScalar(); - - if (1 != scalar.to()) { - auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - other, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - other = scaleLayer->getOutput(0); - } - - auto add = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); - TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); - - add->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self + alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[2].unwrapToScalar().to() * args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - - auto add = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); - TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); - - add->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Compute min(max(min_threshold, input), max_threshold) - auto self = args[0].ITensorOrFreeze(ctx); - auto clamp_layer_out = self; - - if (args[1].isIValue() && args[1].IValue()->isScalar() && args[2].isIValue() && - args[2].IValue()->isScalar()) { - auto alpha = args[1].unwrapToScalar().to(); - auto beta = args[2].unwrapToScalar().to(); - auto clip_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kCLIP); - TORCHTRT_CHECK(clip_layer, "Unable to create clip layer for node: " << *n); - clip_layer->setAlpha(alpha); - clip_layer->setBeta(beta); - clamp_layer_out = clip_layer->getOutput(0); - } else if (args[1].isIValue() && args[1].IValue()->isScalar()) { - auto limit = args[1].unwrapToScalar().to(); - clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max"); - } else if (args[2].isIValue() && args[2].IValue()->isScalar()) { - auto limit = args[2].unwrapToScalar().to(); - clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min"); - } - - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); - LOG_DEBUG("Clamp layer output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::clamp_min(Tensor self, Scalar min) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Compute min(max(min_threshold, input), max_threshold) - auto self = args[0].ITensorOrFreeze(ctx); - auto clamp_layer_out = self; - if (args[1].isIValue() && args[1].IValue()->isScalar()) { - auto limit = args[1].unwrapToScalar().to(); - clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max"); - } - - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); - LOG_DEBUG("clamp_min layer output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::clamp_max(Tensor self, Scalar max) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Compute min(max(min_threshold, input), max_threshold) - auto self = args[0].ITensorOrFreeze(ctx); - auto clamp_layer_out = self; - if (args[1].isIValue() && args[1].IValue()->isScalar()) { - auto limit = args[1].unwrapToScalar().to(); - clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min"); - } - - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); - LOG_DEBUG("clamp_max layer output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::sub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> " - "Tensor", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self - alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto scalar = args[2].unwrapToScalar().to(); - auto other = args[1].ITensorOrFreeze(ctx); - - if (1 != scalar) { - auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - other, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - other = scaleLayer->getOutput(0); - } - - auto sub = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n)); - TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); - - sub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self - alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].unwrapToScalar().to(); - auto alpha = args[2].unwrapToScalar().to(); - auto scaled_val = other * alpha; - - auto scaled_other_tensor = tensor_to_const(ctx, torch::tensor({scaled_val})); - auto sub = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kSUB, self, scaled_other_tensor, util::node_info(n)); - TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); - sub->setName(util::node_info(n).c_str()); - LOG_DEBUG("Output tensor shape: " << sub->getOutput(0)->getDimensions()); - ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); - - return true; - }}) - .pattern({"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar " - "alpha=1) -> (Tensor(a!))", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self - alpha * other - auto self = args[0].ITensorOrFreeze(ctx); - auto scalar = args[2].unwrapToScalar().to(); - auto other = args[1].ITensorOrFreeze(ctx); - - if (1 != scalar) { - auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - other, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - other = scaleLayer->getOutput(0); - } - - auto sub = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n)); - TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); - - sub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::rsub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement other - alpha * self - auto self = args[0].ITensorOrFreeze(ctx); - auto other = impl::scalar_to_tensor(ctx, args[1].unwrapToScalar()); - auto scalar = args[2].unwrapToScalar(); - - if (1 != scalar.to()) { - auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - self, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - self = scaleLayer->getOutput(0); - } - - auto rsub = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n)); - TORCHTRT_CHECK(rsub, "Unable to create rsub layer from node: " << *n); - - rsub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::rsub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement other - alpha * self - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto scalar = args[2].unwrapToScalar(); - - if (1 != scalar.to()) { - auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); - auto scaleLayer = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - self, - alphaTensor, - util::node_info(n) + std::string("_AlphaMultiplier")); - TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); - self = scaleLayer->getOutput(0); - } - - auto rsub = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n)); - TORCHTRT_CHECK(rsub, "Unable to create rsub layer from node: " << *n); - - rsub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::div.Tensor(Tensor self, Tensor other) -> Tensor", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self / other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto div = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); - - TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); - - div->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::div.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self / other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - std::string rounding_mode = "default"; - if (args[2].isIValue() && args[2].IValue()->isString()) { - rounding_mode = args[2].unwrapToString(); - } - nvinfer1::ILayer* div = nullptr; - if (rounding_mode == "floor") { - div = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); - } else if (rounding_mode == "trunc") { - // trunc = floor(abs(div)) * sign(div) - auto tmp_div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, "tmp_div"); - auto abs = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kABS); - auto floor = ctx->net->addUnary(*abs->getOutput(0), nvinfer1::UnaryOperation::kFLOOR); - auto sign = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kSIGN); - div = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kPROD, - floor->getOutput(0), - sign->getOutput(0), - util::node_info(n)); - } else { - div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); - } - - TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); - - div->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::div.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - auto div = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); - TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); - - div->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::div_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto div = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); - - TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); - - div->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - auto div = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); - TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); - - div->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::mul.Tensor(Tensor self, Tensor other) -> Tensor", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // Should implement self * other - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto mul = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); - TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); - - mul->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::mul.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); - auto mul = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); - TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); - - mul->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::mul_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto mul = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); - TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); - - mul->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::ne.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto equal = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kEQUAL, - self, - other, - util::node_info(n) + std::string("is_equal")); - TORCHTRT_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n); - // XOR with ones negates and produces not_equal result - auto options = torch::TensorOptions().dtype(torch::kFloat32); - auto ones = at::full({1}, 1, {options}); - auto ones_tensor = tensor_to_const(ctx, ones); - nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor); - cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL); - - auto sub = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kXOR, - cast_layer->getOutput(0), - equal->getOutput(0), - util::node_info(n)); - TORCHTRT_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n); - - sub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); - LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::ne.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto scalar = args[1].unwrapToScalar(); - nvinfer1::ITensor* scalar_tensor; - if (self->getType() == nvinfer1::DataType::kFLOAT || self->getType() == nvinfer1::DataType::kHALF) { - scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to()})); - } else { - scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to()})); - } - auto equal = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kEQUAL, - self, - scalar_tensor, - util::node_info(n) + std::string("is_equal")); - TORCHTRT_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n); - // XOR with ones negates and produces not_equal result - auto options = torch::TensorOptions().dtype(torch::kFloat32); - auto ones = at::full({1}, 1, {options}); - auto ones_tensor = tensor_to_const(ctx, ones); - nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor); - cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL); - - auto sub = add_elementwise( - ctx, - nvinfer1::ElementWiseOperation::kXOR, - cast_layer->getOutput(0), - equal->getOutput(0), - util::node_info(n)); - TORCHTRT_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n); - - sub->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); - LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::pow.Tensor_Tensor(Tensor self, Tensor exponent) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto exponent = args[1].ITensorOrFreeze(ctx); - auto pow = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n)); - TORCHTRT_CHECK(pow, "Unable to create Power layer from node: " << *n); - - pow->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::pow.Tensor_Scalar(Tensor self, Scalar exponent) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto exponentScalar = args[1].unwrapToScalar().to(); - auto exponent = tensor_to_const(ctx, torch::tensor({exponentScalar})); - auto pow = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n)); - TORCHTRT_CHECK(pow, "Unable to create Power layer from node: " << *n); - - pow->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::floor_divide(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto floor_divide = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); - TORCHTRT_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n); - - floor_divide->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::floor_divide.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - auto floor_divide = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); - TORCHTRT_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n); - - floor_divide->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::max.other(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto max = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMAX, self, other, util::node_info(n)); - TORCHTRT_CHECK(max, "Unable to create max layer from node: " << *n); - - max->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], max->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::min.other(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - // TODO: Remove with functionalization - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto min = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMIN, self, other, util::node_info(n)); - TORCHTRT_CHECK(min, "Unable to create min layer from node: " << *n); - - min->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], min->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::gt.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto gt = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n)); - TORCHTRT_CHECK(gt, "Unable to create greater layer from node: " << *n); - - gt->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], gt->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::gt.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); - if (self->getType() != other->getType()) { - other = castITensor(ctx, other, self->getType()); - } - auto gt = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n)); - TORCHTRT_CHECK(gt, "Unable to create greater layer from node: " << *n); - - gt->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], gt->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::lt.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto lt = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n)); - TORCHTRT_CHECK(lt, "Unable to create less layer from node: " << *n); - - lt->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], lt->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::lt.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); - if (self->getType() != other->getType()) { - other = castITensor(ctx, other, self->getType()); - } - auto lt = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n)); - TORCHTRT_CHECK(lt, "Unable to create less layer from node: " << *n); - - lt->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], lt->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::eq.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - auto eq = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n)); - TORCHTRT_CHECK(eq, "Unable to create equal layer from node: " << *n); - - eq->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], eq->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::eq.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto otherScalar = args[1].unwrapToScalar().to(); - auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); - if (self->getType() == nvinfer1::DataType::kBOOL) { - if (otherScalar == 0 || otherScalar == 1) { - LOG_DEBUG("Since input tensor is type bool, casting input tensor and scalar to int32"); - other = castITensor(ctx, other, nvinfer1::DataType::kINT32); - self = castITensor(ctx, self, nvinfer1::DataType::kINT32); - } else { - LOG_WARNING("Input Tensor has type bool, but scalar is not 0 or 1. Found: " << otherScalar); - } - } - if (self->getType() != other->getType()) { - other = castITensor(ctx, other, self->getType()); - } - auto eq = - add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n)); - TORCHTRT_CHECK(eq, "Unable to create equal layer from node: " << *n); - - eq->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], eq->getOutput(0)); - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::ge.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - - auto greater = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater"); - TORCHTRT_CHECK(greater, "Unable to create Greater layer from node: " << *n); - - auto equal = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); - TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); - - auto or_op = ctx->net->addElementWise( - *greater->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); - - TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); - or_op->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::ge.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); - if (self->getType() != other->getType()) { - other = castITensor(ctx, other, self->getType()); - } - - auto greater = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater"); - TORCHTRT_CHECK(greater, "Unable to create Greater layer from node: " << *n); - - auto equal = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); - TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); - - auto or_op = ctx->net->addElementWise( - *greater->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); - - TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); - or_op->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::le.Tensor(Tensor self, Tensor other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = args[1].ITensorOrFreeze(ctx); - - auto less = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less"); - TORCHTRT_CHECK(less, "Unable to create Less layer from node: " << *n); - - auto equal = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); - TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); - - auto or_op = ctx->net->addElementWise( - *less->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); - - TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); - or_op->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}) - .pattern({"aten::le.Scalar(Tensor self, Scalar other) -> (Tensor)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); - if (self->getType() != other->getType()) { - other = castITensor(ctx, other, self->getType()); - } - - auto less = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less"); - TORCHTRT_CHECK(less, "Unable to create Less layer from node: " << *n); - - auto equal = add_elementwise( - ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); - TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); - - auto or_op = ctx->net->addElementWise( - *less->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); - - TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); - or_op->setName(util::node_info(n).c_str()); - auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); - - LOG_DEBUG("Output tensor shape: " << out->getDimensions()); - return true; - }}); + .pattern( + {"aten::add.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> " + "Tensor", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self + alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto scalar = args[2].unwrapToScalar(); + + if (1 != scalar.to()) { + auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + other, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + other = scaleLayer->getOutput(0); + } + + auto add = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); + TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); + + add->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar " + "alpha=1) -> (Tensor(a!))", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self + alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto scalar = args[2].unwrapToScalar(); + + if (1 != scalar.to()) { + auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + other, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + other = scaleLayer->getOutput(0); + } + + auto add = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); + TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); + + add->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self + alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto otherScalar = args[2].unwrapToScalar().to() * args[1].unwrapToScalar().to(); + auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + + auto add = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n)); + TORCHTRT_CHECK(add, "Unable to create add layer from node: " << *n); + + add->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Compute min(max(min_threshold, input), max_threshold) + auto self = args[0].ITensorOrFreeze(ctx); + auto clamp_layer_out = self; + + if (args[1].isIValue() && args[1].IValue()->isScalar() && args[2].isIValue() && + args[2].IValue()->isScalar()) { + auto alpha = args[1].unwrapToScalar().to(); + auto beta = args[2].unwrapToScalar().to(); + auto clip_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kCLIP); + TORCHTRT_CHECK(clip_layer, "Unable to create clip layer for node: " << *n); + clip_layer->setAlpha(alpha); + clip_layer->setBeta(beta); + clamp_layer_out = clip_layer->getOutput(0); + } else if (args[1].isIValue() && args[1].IValue()->isScalar()) { + auto limit = args[1].unwrapToScalar().to(); + clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max"); + } else if (args[2].isIValue() && args[2].IValue()->isScalar()) { + auto limit = args[2].unwrapToScalar().to(); + clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min"); + } + + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); + LOG_DEBUG("Clamp layer output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::clamp_min(Tensor self, Scalar min) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Compute min(max(min_threshold, input), max_threshold) + auto self = args[0].ITensorOrFreeze(ctx); + auto clamp_layer_out = self; + if (args[1].isIValue() && args[1].IValue()->isScalar()) { + auto limit = args[1].unwrapToScalar().to(); + clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max"); + } + + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); + LOG_DEBUG("clamp_min layer output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::clamp_max(Tensor self, Scalar max) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Compute min(max(min_threshold, input), max_threshold) + auto self = args[0].ITensorOrFreeze(ctx); + auto clamp_layer_out = self; + if (args[1].isIValue() && args[1].IValue()->isScalar()) { + auto limit = args[1].unwrapToScalar().to(); + clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min"); + } + + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out); + LOG_DEBUG("clamp_max layer output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::sub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> " + "Tensor", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self - alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto scalar = args[2].unwrapToScalar().to(); + auto other = args[1].ITensorOrFreeze(ctx); + + if (1 != scalar) { + auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + other, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + other = scaleLayer->getOutput(0); + } + + auto sub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n)); + TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); + + sub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self - alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].unwrapToScalar().to(); + auto alpha = args[2].unwrapToScalar().to(); + auto scaled_val = other * alpha; + + auto scaled_other_tensor = tensor_to_const(ctx, torch::tensor({scaled_val})); + auto sub = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kSUB, self, scaled_other_tensor, util::node_info(n)); + TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); + sub->setName(util::node_info(n).c_str()); + LOG_DEBUG("Output tensor shape: " << sub->getOutput(0)->getDimensions()); + ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); + + return true; + }}) + .pattern( + {"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar " + "alpha=1) -> (Tensor(a!))", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self - alpha * other + auto self = args[0].ITensorOrFreeze(ctx); + auto scalar = args[2].unwrapToScalar().to(); + auto other = args[1].ITensorOrFreeze(ctx); + + if (1 != scalar) { + auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar})); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + other, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + other = scaleLayer->getOutput(0); + } + + auto sub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n)); + TORCHTRT_CHECK(sub, "Unable to create sub layer from node: " << *n); + + sub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::rsub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement other - alpha * self + auto self = args[0].ITensorOrFreeze(ctx); + auto other = impl::scalar_to_tensor(ctx, args[1].unwrapToScalar()); + auto scalar = args[2].unwrapToScalar(); + + if (1 != scalar.to()) { + auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + self, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + self = scaleLayer->getOutput(0); + } + + auto rsub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n)); + TORCHTRT_CHECK(rsub, "Unable to create rsub layer from node: " << *n); + + rsub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::rsub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement other - alpha * self + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto scalar = args[2].unwrapToScalar(); + + if (1 != scalar.to()) { + auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto scaleLayer = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + self, + alphaTensor, + util::node_info(n) + std::string("_AlphaMultiplier")); + TORCHTRT_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n); + self = scaleLayer->getOutput(0); + } + + auto rsub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n)); + TORCHTRT_CHECK(rsub, "Unable to create rsub layer from node: " << *n); + + rsub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::div.Tensor(Tensor self, Tensor other) -> Tensor", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self / other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); + + TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); + + div->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::div.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self / other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + std::string rounding_mode = "default"; + if (args[2].isIValue() && args[2].IValue()->isString()) { + rounding_mode = args[2].unwrapToString(); + } + nvinfer1::ILayer* div = nullptr; + if (rounding_mode == "floor") { + div = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); + } else if (rounding_mode == "trunc") { + // trunc = floor(abs(div)) * sign(div) + auto tmp_div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, "tmp_div"); + auto abs = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kABS); + auto floor = ctx->net->addUnary(*abs->getOutput(0), nvinfer1::UnaryOperation::kFLOOR); + auto sign = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kSIGN); + div = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kPROD, + floor->getOutput(0), + sign->getOutput(0), + util::node_info(n)); + } else { + div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); + } + + TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); + + div->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::div.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto otherScalar = args[1].unwrapToScalar().to(); + auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); + TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); + + div->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::div_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); + + TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); + + div->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto otherScalar = args[1].unwrapToScalar().to(); + auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n)); + TORCHTRT_CHECK(div, "Unable to create div layer from node: " << *n); + + div->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::mul.Tensor(Tensor self, Tensor other) -> Tensor", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // Should implement self * other + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); + TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); + + mul->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::mul.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); + auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); + TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); + + mul->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::mul_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n)); + TORCHTRT_CHECK(mul, "Unable to create mul layer from node: " << *n); + + mul->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::ne.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto equal = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kEQUAL, + self, + other, + util::node_info(n) + std::string("is_equal")); + TORCHTRT_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n); + // XOR with ones negates and produces not_equal result + auto options = torch::TensorOptions().dtype(torch::kFloat32); + auto ones = at::full({1}, 1, {options}); + auto ones_tensor = tensor_to_const(ctx, ones); + nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor); + cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL); + + auto sub = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kXOR, + cast_layer->getOutput(0), + equal->getOutput(0), + util::node_info(n)); + TORCHTRT_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n); + + sub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); + LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::ne.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto scalar = args[1].unwrapToScalar(); + nvinfer1::ITensor* scalar_tensor; + if (self->getType() == nvinfer1::DataType::kFLOAT || self->getType() == nvinfer1::DataType::kHALF) { + scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to()})); + } else { + scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar.to()})); + } + auto equal = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kEQUAL, + self, + scalar_tensor, + util::node_info(n) + std::string("is_equal")); + TORCHTRT_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n); + // XOR with ones negates and produces not_equal result + auto options = torch::TensorOptions().dtype(torch::kFloat32); + auto ones = at::full({1}, 1, {options}); + auto ones_tensor = tensor_to_const(ctx, ones); + nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor); + cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL); + + auto sub = add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kXOR, + cast_layer->getOutput(0), + equal->getOutput(0), + util::node_info(n)); + TORCHTRT_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n); + + sub->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0)); + LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::pow.Tensor_Tensor(Tensor self, Tensor exponent) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto exponent = args[1].ITensorOrFreeze(ctx); + auto pow = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n)); + TORCHTRT_CHECK(pow, "Unable to create Power layer from node: " << *n); + + pow->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::pow.Tensor_Scalar(Tensor self, Scalar exponent) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto exponentScalar = args[1].unwrapToScalar().to(); + auto exponent = tensor_to_const(ctx, torch::tensor({exponentScalar})); + auto pow = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n)); + TORCHTRT_CHECK(pow, "Unable to create Power layer from node: " << *n); + + pow->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::floor_divide(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto floor_divide = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); + TORCHTRT_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n); + + floor_divide->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::floor_divide.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto otherScalar = args[1].unwrapToScalar().to(); + auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + auto floor_divide = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n)); + TORCHTRT_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n); + + floor_divide->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::max.other(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto max = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMAX, self, other, util::node_info(n)); + TORCHTRT_CHECK(max, "Unable to create max layer from node: " << *n); + + max->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], max->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::min.other(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + // TODO: Remove with functionalization + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto min = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMIN, self, other, util::node_info(n)); + TORCHTRT_CHECK(min, "Unable to create min layer from node: " << *n); + + min->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], min->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::gt.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto gt = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n)); + TORCHTRT_CHECK(gt, "Unable to create greater layer from node: " << *n); + + gt->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], gt->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::gt.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); + if (self->getType() != other->getType()) { + other = castITensor(ctx, other, self->getType()); + } + auto gt = + add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n)); + TORCHTRT_CHECK(gt, "Unable to create greater layer from node: " << *n); + + gt->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], gt->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::lt.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto lt = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n)); + TORCHTRT_CHECK(lt, "Unable to create less layer from node: " << *n); + + lt->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], lt->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::lt.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); + if (self->getType() != other->getType()) { + other = castITensor(ctx, other, self->getType()); + } + auto lt = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n)); + TORCHTRT_CHECK(lt, "Unable to create less layer from node: " << *n); + + lt->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], lt->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::eq.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + auto eq = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n)); + TORCHTRT_CHECK(eq, "Unable to create equal layer from node: " << *n); + + eq->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], eq->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::eq.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto otherScalar = args[1].unwrapToScalar().to(); + auto other = tensor_to_const(ctx, torch::tensor({otherScalar})); + if (self->getType() == nvinfer1::DataType::kBOOL) { + if (otherScalar == 0 || otherScalar == 1) { + LOG_DEBUG("Since input tensor is type bool, casting input tensor and scalar to int32"); + other = castITensor(ctx, other, nvinfer1::DataType::kINT32); + self = castITensor(ctx, self, nvinfer1::DataType::kINT32); + } else { + LOG_WARNING("Input Tensor has type bool, but scalar is not 0 or 1. Found: " << otherScalar); + } + } + if (self->getType() != other->getType()) { + other = castITensor(ctx, other, self->getType()); + } + auto eq = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n)); + TORCHTRT_CHECK(eq, "Unable to create equal layer from node: " << *n); + + eq->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], eq->getOutput(0)); + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::ge.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + + auto greater = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater"); + TORCHTRT_CHECK(greater, "Unable to create Greater layer from node: " << *n); + + auto equal = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); + TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); + + auto or_op = ctx->net->addElementWise( + *greater->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); + + TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); + or_op->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::ge.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); + if (self->getType() != other->getType()) { + other = castITensor(ctx, other, self->getType()); + } + + auto greater = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater"); + TORCHTRT_CHECK(greater, "Unable to create Greater layer from node: " << *n); + + auto equal = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); + TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); + + auto or_op = ctx->net->addElementWise( + *greater->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); + + TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); + or_op->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::le.Tensor(Tensor self, Tensor other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = args[1].ITensorOrFreeze(ctx); + + auto less = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less"); + TORCHTRT_CHECK(less, "Unable to create Less layer from node: " << *n); + + auto equal = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); + TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); + + auto or_op = ctx->net->addElementWise( + *less->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); + + TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); + or_op->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}) + .pattern( + {"aten::le.Scalar(Tensor self, Scalar other) -> (Tensor)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); + if (self->getType() != other->getType()) { + other = castITensor(ctx, other, self->getType()); + } + + auto less = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less"); + TORCHTRT_CHECK(less, "Unable to create Less layer from node: " << *n); + + auto equal = add_elementwise( + ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal"); + TORCHTRT_CHECK(equal, "Unable to create Equal layer from node: " << *n); + + auto or_op = ctx->net->addElementWise( + *less->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR); + + TORCHTRT_CHECK(or_op, "Unable to create Or layer from node: " << *n); + or_op->setName(util::node_info(n).c_str()); + auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0)); + + LOG_DEBUG("Output tensor shape: " << out->getDimensions()); + return true; + }}); } // namespace } // namespace impl From cb1f504ba875d0e87f7dcb1ee8bb2552fc8d82d6 Mon Sep 17 00:00:00 2001 From: Michael Feliz Date: Wed, 10 Aug 2022 14:09:23 -0700 Subject: [PATCH 4/5] more linting --- core/conversion/converters/converter_util.cpp | 20 ++++++++++++------- .../converters/test_element_wise.cpp | 4 ++-- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/core/conversion/converters/converter_util.cpp b/core/conversion/converters/converter_util.cpp index 9f1d2a7124..ded4671e56 100644 --- a/core/conversion/converters/converter_util.cpp +++ b/core/conversion/converters/converter_util.cpp @@ -59,7 +59,7 @@ nvinfer1::ITensor* addUnpadding( } } -nvinfer1::DataType promote_types(nvinfer1::DataType type_a, nvinfer1::DataType type_b){ +nvinfer1::DataType promote_types(nvinfer1::DataType type_a, nvinfer1::DataType type_b) { auto torch_type_a = util::TRTDataTypeToScalarType(type_a); auto torch_type_b = util::TRTDataTypeToScalarType(type_b); auto promo_type = at::promote_types(torch_type_a, torch_type_b); @@ -87,15 +87,21 @@ nvinfer1::ILayer* add_elementwise( swapSelfOther = true; } - if(self->getType() != other->getType()){ - LOG_DEBUG("Type mismatch for inputs in element-wise operation " << name << ": " << self->getType() << ", " << other->getType()); + if (self->getType() != other->getType()) { + LOG_DEBUG( + "Type mismatch for inputs in element-wise operation " << name << ": " << self->getType() << ", " + << other->getType()); auto promo_type = promote_types(self->getType(), other->getType()); - if(self->getType() != promo_type){ - LOG_DEBUG("Element-wise op type promotion adding cast from " << self->getType() << " to " << promo_type << " for layer " << name); + if (self->getType() != promo_type) { + LOG_DEBUG( + "Element-wise op type promotion adding cast from " << self->getType() << " to " << promo_type << " for layer " + << name); self = castITensor(ctx, self, promo_type); } - if(other->getType() != promo_type){ - LOG_DEBUG("Element-wise op type promotion adding cast from " << other->getType() << " to " << promo_type << " for layer " << name); + if (other->getType() != promo_type) { + LOG_DEBUG( + "Element-wise op type promotion adding cast from " << other->getType() << " to " << promo_type + << " for layer " << name); other = castITensor(ctx, other, promo_type); } } diff --git a/tests/core/conversion/converters/test_element_wise.cpp b/tests/core/conversion/converters/test_element_wise.cpp index 99be3cbffa..3ecfdb2019 100644 --- a/tests/core/conversion/converters/test_element_wise.cpp +++ b/tests/core/conversion/converters/test_element_wise.cpp @@ -24,11 +24,11 @@ void pointwise_test_helper( int first_max = 5; int second_min = 1; int second_max = 5; - if(type1 == at::kBool){ + if (type1 == at::kBool) { first_min = 0; first_max = 1; } - if(type2 == at::kBool){ + if (type2 == at::kBool) { second_min = 0; second_max = 1; } From 5a4907799000b342303e443443826f3bbbfa5a75 Mon Sep 17 00:00:00 2001 From: Michael Feliz Date: Thu, 11 Aug 2022 10:30:17 -0700 Subject: [PATCH 5/5] Remove impl:: namespace from scalar_to_tensor calls --- core/conversion/converters/impl/element_wise.cpp | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) mode change 100644 => 100755 core/conversion/converters/impl/element_wise.cpp diff --git a/core/conversion/converters/impl/element_wise.cpp b/core/conversion/converters/impl/element_wise.cpp old mode 100644 new mode 100755 index 4d32ce8bb6..5c5966841a --- a/core/conversion/converters/impl/element_wise.cpp +++ b/core/conversion/converters/impl/element_wise.cpp @@ -58,7 +58,7 @@ auto element_wise_registrations TORCHTRT_UNUSED = auto scalar = args[2].unwrapToScalar(); if (1 != scalar.to()) { - auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto alphaTensor = scalar_to_tensor(ctx, scalar); auto scaleLayer = add_elementwise( ctx, nvinfer1::ElementWiseOperation::kPROD, @@ -87,7 +87,7 @@ auto element_wise_registrations TORCHTRT_UNUSED = auto scalar = args[2].unwrapToScalar(); if (1 != scalar.to()) { - auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto alphaTensor = scalar_to_tensor(ctx, scalar); auto scaleLayer = add_elementwise( ctx, nvinfer1::ElementWiseOperation::kPROD, @@ -262,11 +262,11 @@ auto element_wise_registrations TORCHTRT_UNUSED = [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { // Should implement other - alpha * self auto self = args[0].ITensorOrFreeze(ctx); - auto other = impl::scalar_to_tensor(ctx, args[1].unwrapToScalar()); + auto other = scalar_to_tensor(ctx, args[1].unwrapToScalar()); auto scalar = args[2].unwrapToScalar(); if (1 != scalar.to()) { - auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto alphaTensor = scalar_to_tensor(ctx, scalar); auto scaleLayer = add_elementwise( ctx, nvinfer1::ElementWiseOperation::kPROD, @@ -294,7 +294,7 @@ auto element_wise_registrations TORCHTRT_UNUSED = auto scalar = args[2].unwrapToScalar(); if (1 != scalar.to()) { - auto alphaTensor = impl::scalar_to_tensor(ctx, scalar); + auto alphaTensor = scalar_to_tensor(ctx, scalar); auto scaleLayer = add_elementwise( ctx, nvinfer1::ElementWiseOperation::kPROD,