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..8cb6cf2dd6 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" @@ -298,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 { @@ -306,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(); @@ -677,6 +680,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", ], )