(1/2) Make TorchScript Preserve Fully Qualified Class Name for Python Exceptions: backend change

Summary: Reland for D33282878 (pytorch@911d527) . Land backend change first to maintain FC. Will wait for 2 weeks after this diff is in. And than land the front-end change in next diff.

Test Plan:
test in next diff

time buck test mode/dev-nosan fblearner/flow/projects/langtech/translation:tests -- test_e2e_base_training

Reviewed By: gmagogsfm

Differential Revision: D33342547

fbshipit-source-id: b3dee9a4bdfd78103848c12629e5fccafdd621e3
(cherry picked from commit ae1935f)

Author: pavithranrao

caffe2/serialize/versions.h

@@ -108,7 +108,7 @@ constexpr uint64_t kProducedBytecodeVersion = 0x7L;
 // we should support this model_version. For example, we provide a wrapper to
 // handle an updated operator.
 constexpr uint64_t kMinSupportedBytecodeVersion = 0x3L;
-constexpr uint64_t kMaxSupportedBytecodeVersion = 0x7L;
+constexpr uint64_t kMaxSupportedBytecodeVersion = 0x8L;
 
 } // namespace serialize
 } // namespace caffe2

test/cpp/jit/test_backend.cpp

@@ -187,6 +187,27 @@ TEST(BackendTest, TestComposite) {
   AT_ASSERT(res_jit.toTensor().equal(res_mobile.toTensor()));
 }
 
+TEST(BackendTest, TestPrimDtype) {
+  Module c("name");
+  c.define(R"(
+    def forward(self, x, y):
+      c = y.dtype
+      return c
+  )");
+
+  std::vector<IValue> inputs;
+  inputs.emplace_back(3.0 * torch::ones({}));
+  inputs.emplace_back(1.0 * torch::ones({}));
+  auto res_jit = c.forward(inputs);
+
+  std::stringstream ss;
+  c._save_for_mobile(ss);
+  auto mc = _load_for_mobile(ss);
+  auto res_mobile = mc.forward(inputs);
+
+  ASSERT_EQ(res_jit.toInt(), res_mobile.toInt());
+}
+
 Module getCompositeModuleWithSameNameSubModules() {
   // Two submodules with same module name but different forward and other
   // functions should be serialized and loaded correctly.

test/cpp/jit/test_lite_interpreter.cpp

@@ -186,6 +186,42 @@ TEST(LiteInterpreterTest, Tuple) {
   AT_ASSERT(output.toTupleRef().elements()[1].toInt() == 2);
 }
 
+TEST(LiteInterpreterTest, AtenFormat) {
+  Module m("m");
+  m.define(R"""(
+  def forward(self, fmt:str="first {} {}", num:str="abc"):
+    x = 2
+    x = x * x
+    return fmt.format(num, x)
+  )""");
+  std::stringstream ss;
+  m._save_for_mobile(ss);
+  mobile::Module bc = _load_for_mobile(ss);
+  std::vector<torch::jit::IValue> inputs;
+  auto output_bc = bc.get_method("forward")(inputs);
+  auto output_m = m.get_method("forward")(inputs);
+  // std::cout << output_m.toStringRef() << "\n"
+  //           << output_bc.toStringRef() << std::endl;
+  AT_ASSERT(output_m.toStringRef() == output_bc.toStringRef());
+}
+
+TEST(LiteInterpreterTest, PrimDevice) {
+  Module m("m");
+  m.define(R"""(
+  def forward(self, x:torch.Tensor):
+    return x.device
+  )""");
+  std::stringstream ss;
+  m._save_for_mobile(ss);
+  mobile::Module bc = _load_for_mobile(ss);
+  std::vector<torch::jit::IValue> inputs;
+  auto minput = 3.5 * torch::ones({});
+  inputs.emplace_back(minput);
+  auto output_bc = bc.get_method("forward")(inputs);
+  auto output_m = m.get_method("forward")(inputs);
+  AT_ASSERT(output_bc.toDevice().str() == output_m.toDevice().str());
+}
+
 TEST(LiteInterpreterTest, Dict) {
   Module m("m");
   m.define(R"JIT(

tools/build_variables.bzl

@@ -115,6 +115,7 @@ core_sources_common = [
     "torch/csrc/jit/runtime/instruction.cpp",
     "torch/csrc/jit/runtime/jit_exception.cpp",
     "torch/csrc/jit/runtime/operator.cpp",
+    "torch/csrc/jit/mobile/register_ops_common_utils.cpp",
     "torch/csrc/jit/runtime/print_handler.cpp",
     "torch/csrc/jit/runtime/slice_indices_adjust.cpp",
     "torch/csrc/jit/runtime/register_ops_utils.cpp",

torch/csrc/jit/mobile/interpreter.cpp

@@ -5,15 +5,15 @@
 #include <ATen/core/function.h>
 #include <ATen/core/jit_type.h>
 #include <ATen/core/operator_name.h>
-#include <torch/csrc/jit/mobile/function.h>
-#include <torch/csrc/jit/runtime/jit_exception.h>
-#include <torch/csrc/jit/runtime/vararg_functions.h>
-
 #include <ATen/record_function.h>
 #include <c10/util/Exception.h>
 #include <c10/util/irange.h>
 #include <torch/csrc/jit/backends/backend_exception.h>
+#include <torch/csrc/jit/mobile/function.h>
 #include <torch/csrc/jit/mobile/observer.h>
+#include <torch/csrc/jit/mobile/promoted_prim_ops.h>
+#include <torch/csrc/jit/runtime/jit_exception.h>
+#include <torch/csrc/jit/runtime/vararg_functions.h>
 
 namespace torch {
 namespace jit {
@@ -248,6 +248,62 @@ bool InterpreterState::run(Stack& stack) {
           tupleSlice(stack, inst.X, inst.X + inst.N);
           frame.step();
         } break;
+        case TUPLE_INDEX: {
+          tupleIndex(stack);
+          frame.step();
+        } break;
+        case RAISE_EXCEPTION: {
+          raiseExceptionWithMessage(stack);
+          frame.step();
+        } break;
+        case __IS__: {
+          is(stack);
+          frame.step();
+        } break;
+        case UN_INITIALIZED: {
+          unInitialized(stack);
+          frame.step();
+        } break;
+        case __ISNOT__: {
+          isNot(stack);
+          frame.step();
+        } break;
+        case FORMAT: {
+          format(stack, inst.X);
+          frame.step();
+        } break;
+        case DEVICE: {
+          device(stack);
+          frame.step();
+        } break;
+        case DTYPE: {
+          dtype(stack);
+          frame.step();
+        } break;
+        case DIM: {
+          dim(stack);
+          frame.step();
+        } break;
+        case __NOT__: {
+          _not(stack);
+          frame.step();
+        } break;
+        case DICT_INDEX: {
+          dictIndex(stack);
+          frame.step();
+        } break;
+        case TO_LIST: {
+          toList(stack);
+          frame.step();
+        } break;
+        case NUM_TO_TENSOR: {
+          numToTensorScalar(stack);
+          frame.step();
+        } break;
+        case IS_CUDA: {
+          isCuda(stack);
+          frame.step();
+        } break;
         case DICT_CONSTRUCT: {
           dictConstruct(stack, *code.types_[inst.X], inst.N);
           frame.step();

torch/csrc/jit/mobile/promoted_prim_ops.cpp

@@ -1,7 +1,8 @@
+#include <ATen/ScalarOps.h>
 #include <torch/csrc/jit/mobile/promoted_prim_ops.h>
-
 namespace torch {
 namespace jit {
+
 void tupleIndex(Stack& stack) {
   int64_t index = pop(stack).toInt();
   auto tuple = pop(stack).toTuple();
@@ -14,9 +15,23 @@ void tupleIndex(Stack& stack) {
 }
 
 void raiseException(Stack& stack) {
+  // this kernel supports RaiseException with only one argument: the error
+  // DEPRECATED from bytecode_version 8;
+  // Please do not make any changes to this to support BC
   throw JITException(pop(stack).toStringRef());
 }
 
+void raiseExceptionWithMessage(Stack& stack) {
+  // this kernel supports RaiseException with only two arguments: the error and
+  // the message Please make changes only to this kernel
+  c10::optional<std::string> qualified_class_name =
+      pop(stack).toOptional<std::string>();
+  std::string message;
+  pop(stack, message);
+
+  throw JITException(message, qualified_class_name);
+}
+
 void is(Stack& stack) {
   IValue self, obj;
   pop(stack, self, obj);
@@ -99,15 +114,15 @@ void toList(Stack& stack) {
 
   // Rebuild the output type using elem_ty_val and dim_val. Start
   // with the element type corresponding to elem_ty_val.
-  TypePtr out_ty;
+  at::TypePtr out_ty;
   if (elem_ty_val == 0) {
-    out_ty = IntType::get();
+    out_ty = at::IntType::get();
   } else if (elem_ty_val == 1) {
-    out_ty = FloatType::get();
+    out_ty = at::FloatType::get();
   } else if (elem_ty_val == 2) {
-    out_ty = BoolType::get();
+    out_ty = at::BoolType::get();
   } else if (elem_ty_val == 3) {
-    out_ty = ComplexType::get();
+    out_ty = at::ComplexType::get();
   } else {
     TORCH_CHECK(
         false,
@@ -120,8 +135,8 @@ void toList(Stack& stack) {
   // the elements will be casted to double/c10::complex<double>
   // later.
   TORCH_CHECK(
-      (out_ty == FloatType::get() && t.is_floating_point()) ||
-          (out_ty == ComplexType::get() && t.is_complex()) ||
+      (out_ty == at::FloatType::get() && t.is_floating_point()) ||
+          (out_ty == at::ComplexType::get() && t.is_complex()) ||
           tryScalarTypeFromJitType(*out_ty) == t.scalar_type(),
       "Output annotation element type and runtime tensor element type must match for tolist()");
 
@@ -134,7 +149,7 @@ void toList(Stack& stack) {
   // Wrap out_ty in a ListType dim times.
   for (const auto i : c10::irange(dim_val)) {
     (void)i; // Suppress unused variable warning
-    out_ty = ListType::create(out_ty);
+    out_ty = at::ListType::create(out_ty);
   }
 
   int64_t dim = t.dim();
@@ -150,7 +165,7 @@ void toList(Stack& stack) {
 void numToTensorScalar(Stack& stack) {
   at::Scalar s;
   pop(stack, s);
-  push(stack, at::scalar_to_tensor(s));
+  push(stack, c10::scalar_to_tensor(s));
 }
 
 void isCuda(Stack& stack) {
@@ -163,7 +178,7 @@ void numToTensorBool(Stack& stack) {
   // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
   bool b;
   pop(stack, b);
-  push(stack, at::scalar_to_tensor(b));
+  push(stack, c10::scalar_to_tensor(b));
 }
 
 void dictIndex(Stack& stack) {
@@ -181,7 +196,9 @@ static const C10_UNUSED std::array<mobile::prim_op_fn_register, 15> op_reg = {
     mobile::prim_op_fn_register("aten::Bool.Tensor", boolTensor),
     mobile::prim_op_fn_register("aten::format", aten_format),
     mobile::prim_op_fn_register("prim::NumToTensor.Scalar", numToTensorScalar),
-    mobile::prim_op_fn_register("prim::RaiseException", raiseException),
+    mobile::prim_op_fn_register(
+        "prim::RaiseException",
+        raiseExceptionWithMessage),
     mobile::prim_op_fn_register("prim::device", device),
     mobile::prim_op_fn_register("prim::dtype", dtype),
     mobile::prim_op_fn_register("aten::__not__", _not),

torch/csrc/jit/mobile/promoted_prim_ops.h

@@ -1,6 +1,6 @@
 #pragma once
 #include <torch/csrc/jit/mobile/prim_ops_registery.h>
-#include <torch/csrc/jit/runtime/register_ops_utils.h>
+#include <torch/csrc/jit/mobile/register_ops_common_utils.h>
 
 namespace torch {
 namespace jit {
@@ -41,5 +41,7 @@ void numToTensorBool(Stack& stack);
 
 void dictIndex(Stack& stack);
 
+void raiseExceptionWithMessage(Stack& stack);
+
 } // namespace jit
 } // namespace torch

torch/csrc/jit/mobile/register_ops_common_utils.cpp

@@ -0,0 +1,102 @@
+#include <ATen/core/dynamic_type.h>
+#include <ATen/core/type_factory.h>
+#include <torch/csrc/jit/mobile/register_ops_common_utils.h>
+
+namespace torch {
+namespace jit {
+
+int64_t normalizeIndex(int64_t idx, int64_t list_size) {
+  if (idx < 0) {
+    // Handle negative indexing
+    idx = list_size + idx;
+  }
+  return idx;
+}
+
+IValue tensorToListRecursive(
+    char* data,
+    int64_t cur_dim,
+    int64_t num_tensor_dims,
+    at::TypePtr ty,
+    at::ScalarType scalar_ty,
+    at::IntArrayRef sizes,
+    at::IntArrayRef strides,
+    size_t element_size) {
+  // If ty is a ListType, get the element type.
+  if (auto list_type = ty->cast<at::ListType>()) {
+    ty = list_type->getElementType();
+  } else {
+    // If the output type is a scalar, read and push one scalar of
+    // the right type onto the stack.
+    if (ty == at::IntType::get()) {
+      int64_t scalar = *(int64_t*)data;
+      return IValue(scalar);
+    } else if (ty == at::FloatType::get()) {
+      TORCH_INTERNAL_ASSERT(
+          scalar_ty == at::ScalarType::Float ||
+              scalar_ty == at::ScalarType::Double,
+          "Unexpected scalar type for Tensor");
+      double scalar =
+          scalar_ty == at::ScalarType::Float ? *(float*)data : *(double*)data;
+      return IValue(scalar);
+    } else if (ty == at::ComplexType::get()) {
+      TORCH_INTERNAL_ASSERT(
+          scalar_ty == at::ScalarType::ComplexFloat ||
+              scalar_ty == at::ScalarType::ComplexDouble,
+          "Unexpected scalar type for Tensor");
+      c10::complex<double> scalar = scalar_ty == at::ScalarType::ComplexFloat
+          ? *(c10::complex<float>*)data
+          : *(c10::complex<double>*)data;
+      return IValue(scalar);
+    } else if (ty == at::BoolType::get()) {
+      bool scalar = *(bool*)data;
+      return IValue(scalar);
+    } else {
+      TORCH_CHECK(
+          false,
+          ty->repr_str(),
+          " is not one of the supported types for tolist: int, float, bool");
+    }
+  }
+
+  // Make the result list consisting of elements of type ty. Since this
+  // invocation is processing dimension cur_dim, there will be sizes[cur_dim]
+  // output elements.
+  auto result = c10::impl::GenericList(ty);
+  result.reserve(sizes[cur_dim]);
+
+  // Since ty was a list type, tensorToListRecursive needs to be called
+  // recursively on each slice of the tensor in the current dimension.
+  for (int64_t i = 0, e = sizes[cur_dim]; i < e; ++i) {
+    auto inner_result = tensorToListRecursive(
+        data,
+        cur_dim + 1,
+        num_tensor_dims,
+        ty,
+        scalar_ty,
+        sizes,
+        strides,
+        element_size);
+
+    if (inner_result.isList()) {
+      result.emplace_back(inner_result.toList());
+    } else if (inner_result.isComplexDouble()) {
+      result.emplace_back(inner_result.toComplexDouble());
+    } else if (inner_result.isDouble()) {
+      result.emplace_back(inner_result.toDouble());
+    } else if (inner_result.isInt()) {
+      result.emplace_back(inner_result.toInt());
+    } else if (inner_result.isBool()) {
+      result.emplace_back(inner_result.toBool());
+    } else {
+      TORCH_INTERNAL_ASSERT("Unknown return type for tensorToListRecursive");
+    }
+
+    data += strides[cur_dim] * element_size;
+  }
+
+  return result;
+}
+
+} // namespace jit
+} // namespace torch