Use union in allocate_storage struct

Author: echuraev

include/tvm/runtime/vm/bytecode.h

@@ -197,16 +197,18 @@ struct Instruction {
       RegName* free_vars;
     };
     struct /* AllocStorage Operands */ {
-      /*! \brief The size of the allocation. */
-      RegName allocation_size;
       /*! \brief The alignment of the allocation. */
       Index alignment;
       /*! \brief The hint of the dtype. */
       DLDataType dtype_hint;
       /*! \brief The number of dimensions. */
       uint32_t ndim;
-      /*! \brief The shape of tensor. */
-      int64_t* shape;
+      union {
+        /*! \brief The shape of tensor. */
+        int64_t* shape;
+        /*! \brief The size of the allocation. */
+        RegName allocation_size;
+      };
       /*! \brief The index of the device on which the allocation will be made. */
       Index device_index;
     } alloc_storage;

src/runtime/vm/bytecode.cc

@@ -350,9 +350,11 @@ Instruction Instruction::AllocStorage(RegName size, Index alignment, DLDataType
   instr.alloc_storage.dtype_hint = dtype_hint;
   instr.alloc_storage.device_index = device_index;
   instr.alloc_storage.ndim = static_cast<uint32_t>(shape.size());
-  instr.alloc_storage.shape = new int64_t[shape.size()];
-  for (size_t i = 0; i < shape.size(); ++i) {
-    instr.alloc_storage.shape[i] = shape[i];
+  if (instr.alloc_storage.ndim > 0) {
+    instr.alloc_storage.shape = new int64_t[shape.size()];
+    for (size_t i = 0; i < shape.size(); ++i) {
+      instr.alloc_storage.shape[i] = shape[i];
+    }
   }
   return instr;
 }
@@ -626,10 +628,15 @@ void InstructionPrint(std::ostream& os, const Instruction& instr) {
       break;
     }
     case Opcode::AllocStorage: {
-      os << "alloc_storage $" << instr.dst << " $" << instr.alloc_storage.allocation_size << " "
-         << instr.alloc_storage.alignment << " "
-         << "[" << StrJoin<int64_t>(instr.alloc_storage.shape, 0, instr.alloc_storage.ndim) << "] "
-         << DLDataType2String(instr.alloc_storage.dtype_hint) << " "
+      os << "alloc_storage $" << instr.dst << " ";
+      if (instr.alloc_storage.ndim > 0) {
+        os << "[" << StrJoin<int64_t>(instr.alloc_storage.shape, 0, instr.alloc_storage.ndim)
+           << "] ";
+      } else {
+        os << "$" << instr.alloc_storage.allocation_size << " " << instr.alloc_storage.alignment
+           << " ";
+      }
+      os << DLDataType2String(instr.alloc_storage.dtype_hint) << " "
          << instr.alloc_storage.device_index;
       break;
     }

src/runtime/vm/profiler/vm.cc

@@ -129,9 +129,21 @@ void VirtualMachineDebug::OpStartHook(Instruction instr) {
           {{"Argument Shapes",
             profiling::ShapeString(shape_tensor, instr.alloc_tensor_reg.dtype)}});
     } else if (instr.op == Opcode::AllocStorage) {
-      auto size = LoadScalarInt(instr.alloc_storage.allocation_size);
       std::ostringstream shape;
-      shape << DLDataType2String(instr.alloc_storage.dtype_hint) << "[" << size << "]";
+      if (instr.alloc_storage.ndim > 0) {
+        std::string shape_str = "[";
+        for (uint32_t i = 0; i < instr.alloc_storage.ndim; ++i) {
+          if (i > 0) {
+            shape_str += ", ";
+          }
+          shape_str += std::to_string(instr.alloc_storage.shape[i]);
+        }
+        shape_str += "]";
+        shape << DLDataType2String(instr.alloc_storage.dtype_hint) << shape_str;
+      } else {
+        auto size = LoadScalarInt(instr.alloc_storage.allocation_size);
+        shape << DLDataType2String(instr.alloc_storage.dtype_hint) << "[" << size << "]";
+      }
       Device dev = GetDevice(instr.alloc_storage.device_index);
       prof_.operator*().StartCall("VM::AllocStorage", dev,
                                   {{"VM::Argument Shapes", String(shape.str())}});

src/runtime/vm/vm.cc

@@ -820,22 +820,35 @@ void VirtualMachine::RunLoop(const std::vector<Index>& output_tensor_reg_indices
       }
       case Opcode::AllocStorage: {
         OpStartHook(instr);
-        auto size = LoadScalarInt(instr.alloc_storage.allocation_size);
-        auto alignment = instr.alloc_storage.alignment;
 
         auto storage_obj = SimpleObjAllocator().make_object<StorageObj>();
         Allocator* allocator = GetAllocator(instr.alloc_storage.device_index);
         ICHECK(allocator) << "Did you forget to init the VirtualMachine with devices?";
-        VLOG(2) << "allocating with allocation_size=" << size << ", alignment=" << alignment
-                << ", dtype_hint=" << DLDataType2String(instr.alloc_storage.dtype_hint)
-                << ", device_index=" << instr.alloc_storage.device_index;
         std::string mem_scope = exec_->virtual_devices[instr.alloc_storage.device_index].second;
 
         if (instr.alloc_storage.ndim > 0) {
+          std::string shape = "[";
+          for (uint32_t i = 0; i < instr.alloc_storage.ndim; ++i) {
+            if (i > 0) {
+              shape += ", ";
+            }
+            shape += std::to_string(instr.alloc_storage.shape[i]);
+          }
+          shape += "]";
+          VLOG(2) << "allocating with ndims=" << instr.alloc_storage.ndim << ", shape=" << shape
+                  << ", dtype_hint=" << DLDataType2String(instr.alloc_storage.dtype_hint)
+                  << ", device_index=" << instr.alloc_storage.device_index
+                  << ", memory_scope=" << mem_scope;
           storage_obj->buffer =
               allocator->Alloc(instr.alloc_storage.ndim, instr.alloc_storage.shape,
                                instr.alloc_storage.dtype_hint, mem_scope);
         } else {
+          auto size = LoadScalarInt(instr.alloc_storage.allocation_size);
+          auto alignment = instr.alloc_storage.alignment;
+          VLOG(2) << "allocating with allocation_size=" << size << ", alignment=" << alignment
+                  << ", dtype_hint=" << DLDataType2String(instr.alloc_storage.dtype_hint)
+                  << ", device_index=" << instr.alloc_storage.device_index
+                  << ", memory_scope=" << mem_scope;
           storage_obj->buffer = allocator->Alloc(size, alignment, instr.alloc_storage.dtype_hint);
         }
         Storage storage(storage_obj);