refactor(//core/conversion/var): Modify var contract

This modifies the variable contract to allow requesting of
mutable handles to `IValue`s which previously were provided
only as `const`. Now every IValue pointer is not `const` within
`Var`, but in order to get a `const` handle you use IValueMut.

This supports the evaluation of inplace operations and makes
state changes explicit rather than implicit.

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

core/conversion/conversion.cpp

@@ -416,7 +416,7 @@ void ConvertBlockToNetDef(
       EvaluateConditionalBlock(ctx, n);
     } else if (to_eval) {
       auto eval = EvaluateNode(ctx, n);
-      if (eval && n->outputs().size() > 0) {
+      if (eval) {
         if (n->outputs().size() > 1) { // For ListUnpack scenario
           if (eval.value().isTuple()) {
             auto eval_list = eval.value().toTuple();

core/conversion/evaluators/aten.cpp

@@ -288,14 +288,18 @@ auto aten_registrations TORCHTRT_UNUSED =
         .evaluator({c10::Symbol::fromQualString("aten::extend"),
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
                       if (args.at(n->input(0)).IValue()->isList() && args.at(n->input(1)).IValue()->isList()) {
-                        auto self = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
+                        c10::IValue* self_ptr = args.at(n->input(0)).IValueMut();
+                        auto self = self_ptr->to<c10::List<c10::IValue>>();
                         auto other = args.at(n->input(1)).IValue()->to<c10::List<c10::IValue>>();
                         const int64_t other_size = other.size();
 
+                        // Modify value in place
                         for (int64_t i = 0; i < other_size; i++) {
                           self.push_back(other.get(i));
                         }
-                        return self;
+
+                        *self_ptr = c10::IValue(self);
+                        return {};
                       } else {
                         TORCHTRT_THROW_ERROR(
                             "Unimplemented data type for aten::extend.t evaluator: "

core/conversion/var/Var.cpp

@@ -13,7 +13,7 @@ Var::Var() {
   type_ = Type::kNone;
 }
 
-Var::Var(const torch::jit::IValue* p) : type_(Type::kIValue) {
+Var::Var(torch::jit::IValue* p) : type_(Type::kIValue) {
   ptr_.ivalue = p;
 }
 
@@ -56,7 +56,7 @@ Var& Var::operator=(const Var& a) {
   return (*this);
 }
 
-Var& Var::operator=(const torch::jit::IValue* in) {
+Var& Var::operator=(torch::jit::IValue* in) {
   ptr_.ivalue = in;
   type_ = Type::kIValue;
   return (*this);
@@ -116,6 +116,10 @@ nvinfer1::ITensor* Var::ITensorOrFreeze(ConversionCtx* ctx) {
 }
 
 const torch::jit::IValue* Var::IValue() const {
+  return IValueMut();
+}
+
+torch::jit::IValue* Var::IValueMut() const {
   TORCHTRT_CHECK(isIValue(), "Requested IValue from Var, however Var type is " << type_name());
   if (type_ == Type::kIValue) {
     return ptr_.ivalue;

core/conversion/var/Var.h

@@ -17,13 +17,14 @@ class Var : torch::CustomClassHolder {
   enum Type { kITensor, kIValue, kNone };
 
   Var();
-  Var(const torch::jit::IValue* p);
+  Var(torch::jit::IValue* p);
   Var(nvinfer1::ITensor* p);
   Var(const Var& a);
   Var& operator=(const Var& a);
-  Var& operator=(const torch::jit::IValue* in);
+  Var& operator=(torch::jit::IValue* in);
   Var& operator=(nvinfer1::ITensor* in);
   const torch::jit::IValue* IValue() const;
+  torch::jit::IValue* IValueMut() const;
   nvinfer1::ITensor* ITensor() const;
 
   // TODO: Can we consolidate this in a way that prevents requesting invalid
@@ -63,7 +64,7 @@ class Var : torch::CustomClassHolder {
 
  private:
   union VarContainer {
-    const torch::jit::IValue* ivalue;
+    torch::jit::IValue* ivalue;
     nvinfer1::ITensor* tensor;
     void* none;
   };