From 4969495f31414ff24e92ba69e2d7343559d2fa60 Mon Sep 17 00:00:00 2001 From: Dheeraj Peri Date: Thu, 8 Jun 2023 14:21:02 -0700 Subject: [PATCH 1/3] feat: Implement dynamic shape support for floordiv, NumToTensor, layernorm Signed-off-by: Dheeraj Peri --- .../conversion/converters/impl/layer_norm.cpp | 6 +++--- core/conversion/evaluators/aten.cpp | 20 +++++++++++++++++++ core/conversion/evaluators/prim.cpp | 4 ++++ tests/core/conversion/converters/BUILD | 12 +++++------ 4 files changed, 33 insertions(+), 9 deletions(-) diff --git a/core/conversion/converters/impl/layer_norm.cpp b/core/conversion/converters/impl/layer_norm.cpp index 781e061a7f..0c00ee2c4d 100644 --- a/core/conversion/converters/impl/layer_norm.cpp +++ b/core/conversion/converters/impl/layer_norm.cpp @@ -20,8 +20,8 @@ auto layer_norm_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() /* Layer_Norm normalizes over last N dimensions. normalizaed_shape could be (C,H,W), (H,W), or (W). */ - auto normalized_shape = args[1].unwrapToIntList(); - auto normalized_shape_vec = util::toVec(util::toDims(normalized_shape)); + // This could be an IntList or ITensorList. We only need the size of this list. + auto normalized_shape = args[1].IValue()->toList(); // Unwrap eps. auto eps = args[4].unwrapToDouble(); @@ -30,7 +30,7 @@ auto layer_norm_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns() // Set up axis_ask for E[x]. uint32_t axis_mask = 0; - for (size_t i = 0; i < normalized_shape_vec.size(); i++) { + for (size_t i = 0; i < normalized_shape.size(); i++) { axis_mask |= 1 << (shape.size() - i - 1); } LOG_DEBUG("Axis Mask for E[x]" << std::bitset<32>(axis_mask)); diff --git a/core/conversion/evaluators/aten.cpp b/core/conversion/evaluators/aten.cpp index 838175461e..c797113905 100644 --- a/core/conversion/evaluators/aten.cpp +++ b/core/conversion/evaluators/aten.cpp @@ -9,6 +9,7 @@ #include "torch/csrc/jit/ir/ir.h" #include "torch/torch.h" +#include "core/conversion/converters/converter_util.h" #include "core/conversion/evaluators/eval_macros.h" #include "core/conversion/evaluators/eval_util.h" #include "core/conversion/evaluators/evaluators.h" @@ -677,6 +678,25 @@ auto aten_registrations TORCHTRT_UNUSED = .evaluator( {c10::Symbol::fromQualString("aten::floordiv"), [](ConversionCtx* ctx, const torch::jit::Node* n, kwargs& args) -> c10::optional { + // Dynamic version of aten::floordiv + if (args.at(n->input(0)).isITensor()) { + if (args.at(n->input(1)).IValue()->isInt()) { + auto int_tensor = scalar_to_tensor(args.at(n->input(1)).IValue()->toInt()); + auto int_itensor = converters::tensor_to_const(ctx, int_tensor, util::node_info(n) + "_constant"); + auto elementwise_layer = converters::add_elementwise( + ctx, + nvinfer1::ElementWiseOperation::kFLOOR_DIV, + args.at(n->input(0)).ITensor(), + int_itensor, + util::node_info(n)); + auto output_tensor = elementwise_layer->getOutput(0); + auto tensor_holder = TensorContainer(); + tensor_holder.hold_tensor(output_tensor); + auto output_ivalue = c10::IValue(std::move(c10::make_intrusive(tensor_holder))); + return output_ivalue; + } + } + // Static version if (args.at(n->input(0)).IValue()->isInt()) { auto a = args.at(n->input(0)).unwrapToInt(); auto b = args.at(n->input(1)).unwrapToInt(); diff --git a/core/conversion/evaluators/prim.cpp b/core/conversion/evaluators/prim.cpp index cbbc109982..456c47fa77 100644 --- a/core/conversion/evaluators/prim.cpp +++ b/core/conversion/evaluators/prim.cpp @@ -32,6 +32,10 @@ auto prim_registrations = .evaluator( {torch::jit::prim::NumToTensor, [](ConversionCtx* ctx, const torch::jit::Node* n, kwargs& args) -> c10::optional { + // Dynamic version receives an ITensor here so pass that as output directly. + if (args.at(n->input(0)).isITensor()) { + return args.at(n->input(0)).ITensor(); + } return evaluators::scalar_to_tensor(args.at(n->input(0)).IValue()->toScalar()); }}) .evaluator( diff --git a/tests/core/conversion/converters/BUILD b/tests/core/conversion/converters/BUILD index 1973c112fd..477774248d 100644 --- a/tests/core/conversion/converters/BUILD +++ b/tests/core/conversion/converters/BUILD @@ -224,33 +224,33 @@ test_suite( ":test_div", ":test_einsum", ":test_expand", + ":test_index", ":test_instance_norm", ":test_interpolate", - ":test_index", ":test_layer_norm", ":test_linear", ":test_lstm_cell", - ":test_matrix_multiply", ":test_masked_fill", + ":test_matrix_multiply", ":test_max", ":test_normalize", ":test_pooling", ":test_reduce", - ":test_roll", ":test_replication_pad", + ":test_roll", ":test_scatter", ":test_select", ":test_shuffle", + ":test_slice", ":test_softmax", + ":test_split", ":test_squeeze", ":test_stack", - ":test_split", - ":test_slice", ":test_topk", ":test_unary", - ":test_unsqueeze", ":test_unbind", ":test_unpack", + ":test_unsqueeze", ":test_where", ], ) From c62474a925a3f7dc4dfea9b39506de96be2676fd Mon Sep 17 00:00:00 2001 From: Dheeraj Peri Date: Wed, 21 Jun 2023 20:21:16 -0700 Subject: [PATCH 2/3] fix: Return static size for desired dimension if it's available Signed-off-by: Dheeraj Peri --- core/conversion/evaluators/aten.cpp | 16 +++++++++------- core/conversion/evaluators/eval_util.cpp | 4 ++++ 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/core/conversion/evaluators/aten.cpp b/core/conversion/evaluators/aten.cpp index c797113905..8cb6cf2dd6 100644 --- a/core/conversion/evaluators/aten.cpp +++ b/core/conversion/evaluators/aten.cpp @@ -299,7 +299,14 @@ auto aten_registrations TORCHTRT_UNUSED = } else { auto dim = args.at(n->input(1)).unwrapToInt(); if (tensor_var.isITensor()) { - if (ctx->input_is_dynamic) { + auto tensor = tensor_var.ITensor(); + auto dims = util::toVec(tensor->getDimensions()); + auto nbDims = tensor->getDimensions().nbDims; + if (dim < 0) { + dim += nbDims; + } + // Check if selected dimension size is -1 else return static size + if (ctx->input_is_dynamic && dims[dim] == -1) { if (ctx->settings.allow_shape_tensors) { return dynamic_size_layer(ctx, n, args); } else { @@ -307,12 +314,7 @@ auto aten_registrations TORCHTRT_UNUSED = "There may be undefined behavior using dynamic shape and aten::size without setting allow_shape_tensors"); } } - auto tensor = tensor_var.ITensor(); - auto dims = util::toVec(tensor->getDimensions()); - auto nbDims = tensor->getDimensions().nbDims; - if (dim < 0) { - dim += nbDims; - } + return dims[dim]; } else if (tensor_var.IValue()->isTensor()) { auto tensor = tensor_var.unwrapToTensor(); diff --git a/core/conversion/evaluators/eval_util.cpp b/core/conversion/evaluators/eval_util.cpp index 0a0b97cfe1..a40cd48523 100644 --- a/core/conversion/evaluators/eval_util.cpp +++ b/core/conversion/evaluators/eval_util.cpp @@ -45,6 +45,10 @@ c10::IValue dynamic_size_layer(ConversionCtx* ctx, const torch::jit::Node* n, kw // Handle negative axis by refering to nbDims of input Tensor dim = dim < 0 ? dim + maxDim : dim; LOG_DEBUG("Dimension to select: " << dim); + // Check if selected dimension size is -1 else return static size + if (input_dims.d[dim] != -1) { + return input_dims.d[dim]; + } shape_1d_tensor = index_layer(ctx, n, shape_1d_tensor, dim); LOG_DEBUG("Output tensor shape: " << shape_1d_tensor->getDimensions()); From c37eeecac7e1cc5e46017636cf7f3de57c061eda Mon Sep 17 00:00:00 2001 From: Dheeraj Peri Date: Wed, 21 Jun 2023 20:24:30 -0700 Subject: [PATCH 3/3] chore: deleting the previous fix Signed-off-by: Dheeraj Peri --- core/conversion/evaluators/eval_util.cpp | 4 ---- 1 file changed, 4 deletions(-) diff --git a/core/conversion/evaluators/eval_util.cpp b/core/conversion/evaluators/eval_util.cpp index a40cd48523..0a0b97cfe1 100644 --- a/core/conversion/evaluators/eval_util.cpp +++ b/core/conversion/evaluators/eval_util.cpp @@ -45,10 +45,6 @@ c10::IValue dynamic_size_layer(ConversionCtx* ctx, const torch::jit::Node* n, kw // Handle negative axis by refering to nbDims of input Tensor dim = dim < 0 ? dim + maxDim : dim; LOG_DEBUG("Dimension to select: " << dim); - // Check if selected dimension size is -1 else return static size - if (input_dims.d[dim] != -1) { - return input_dims.d[dim]; - } shape_1d_tensor = index_layer(ctx, n, shape_1d_tensor, dim); LOG_DEBUG("Output tensor shape: " << shape_1d_tensor->getDimensions());