[TensorExpr] Add a class for representing data type. (pytorch#33217)

Mikhail Zolotukhin · facebook-github-bot · commit 1a4f99717857 · 2020-02-21T13:10:12.000-08:00
Summary: Pull Request resolved: pytorch#33217 Test Plan: Imported from OSS Differential Revision: D19848380 Pulled By: ZolotukhinM fbshipit-source-id: d8683f8fc4555d2456cd2a7c827d8e8231915b49
diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
@@ -456,6 +456,7 @@ if (NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
     ${TORCH_SRC_DIR}/csrc/jit/vararg_functions.cpp
 
     ${TORCH_SRC_DIR}/csrc/jit/tensorexpr/mem_arena.cpp
+    ${TORCH_SRC_DIR}/csrc/jit/tensorexpr/types.cpp
     )
 
   if (NOT INTERN_DISABLE_MOBILE_INTERP)
@@ -757,6 +758,7 @@ ENDIF()
 
   if (BUILD_TEST AND NOT MSVC AND NOT USE_ROCM)
     add_subdirectory(${TORCH_ROOT}/test/cpp/jit ${CMAKE_BINARY_DIR}/test_jit)
+    add_subdirectory(${TORCH_ROOT}/test/cpp/tensorexpr ${CMAKE_BINARY_DIR}/test_tensorexpr)
     if (USE_DISTRIBUTED)
       add_subdirectory(${TORCH_ROOT}/test/cpp/rpc ${CMAKE_BINARY_DIR}/test_cpp_rpc)
     endif()
diff --git a/test/cpp/tensorexpr/CMakeLists.txt b/test/cpp/tensorexpr/CMakeLists.txt
@@ -0,0 +1,39 @@
+set(TENSOREXPR_TEST_ROOT ${TORCH_ROOT}/test/cpp/tensorexpr)
+
+file(GLOB TENSOREXPR_TEST_SRCS ${TENSOREXPR_TEST_ROOT}/test_*.cpp)
+set(TENSOREXPR_TEST_SRCS ${TENSOREXPR_TEST_SRCS} PARENT_SCOPE)
+
+add_executable(test_tensorexpr
+  ${TORCH_ROOT}/test/cpp/common/main.cpp
+  ${TENSOREXPR_TEST_ROOT}/gtest.cpp
+  ${TENSOREXPR_TEST_SRCS})
+
+target_link_libraries(test_tensorexpr PRIVATE torch gtest)
+target_include_directories(test_tensorexpr PRIVATE ${ATen_CPU_INCLUDE})
+
+if (USE_CUDA)
+  target_link_libraries(test_tensorexpr PRIVATE
+    ${CUDA_LIBRARIES}
+    ${CUDA_NVRTC_LIB}
+    ${CUDA_CUDA_LIB}
+    ${TORCH_CUDA_LIBRARIES})
+
+  target_compile_definitions(test_tensorexpr PRIVATE USE_CUDA)
+elseif (USE_ROCM)
+  target_link_libraries(test_tensorexpr PRIVATE
+    ${ROCM_HIPRTC_LIB}
+    ${PYTORCH_HIP_HCC_LIBRARIES}
+    ${TORCH_CUDA_LIBRARIES})
+
+  target_link_libraries(test_tensorexpr PRIVATE caffe2_gpu)
+
+  target_compile_definitions(test_tensorexpr PRIVATE USE_ROCM)
+endif()
+
+if (INSTALL_TEST)
+  install(TARGETS test_tensorexpr DESTINATION bin)
+  # Install PDB files for MSVC builds
+  if (MSVC AND BUILD_SHARED_LIBS)
+    install(FILES $<TARGET_PDB_FILE:test_tensorexpr> DESTINATION bin OPTIONAL)
+  endif()
+endif()
diff --git a/test/cpp/tensorexpr/README.md b/test/cpp/tensorexpr/README.md
@@ -0,0 +1,55 @@
+# TensorExpr C++ Tests
+
+## How to add a new test
+First, create a new test file. Test files should have be placed in this
+directory, with a name that starts with `test_`, like `test_foo.cpp`.
+
+Here is an example test file you can copy-paste.
+```cpp
+#include <test/cpp/tensorexpr/test_base.h>
+
+// Tests go in torch::jit
+namespace torch {
+namespace jit {
+
+// 1. Test cases are void() functions.
+// 2. They start with the prefix `test`
+void testCaseOne() {
+    // ...
+}
+
+void testCaseTwo() {
+    // ...
+}
+}
+}
+```
+
+Then, register your test in `tests.h`:
+```cpp
+// Add to TH_FORALL_TESTS_CUDA instead for CUDA-requiring tests
+#define TH_FORALL_TESTS(_)             \
+  _(ADFormulas)                        \
+  _(Attributes)                        \
+  ...
+  _(CaseOne)  // note that the `test` prefix is omitted.
+  _(CaseTwo)
+```
+
+We glob all the test files together in `CMakeLists.txt` so that you don't
+have to edit it every time you add a test. Unfortunately, this means that in
+order to get the build to pick up your new test file, you need to re-run
+cmake:
+```
+python setup.py build --cmake
+```
+
+## How do I run the tests?
+The following commands assume you are in PyTorch root.
+
+ ```bash
+ # (re)build the test binary
+ ninja build/bin/test_tensorexpr
+ # run
+ build/bin/test_tensorexpr --gtest_filter='glob_style_filter*'
+ ```
diff --git a/test/cpp/tensorexpr/__init__.py b/test/cpp/tensorexpr/__init__.py
diff --git a/test/cpp/tensorexpr/gtest.cpp b/test/cpp/tensorexpr/gtest.cpp
@@ -0,0 +1,16 @@
+#include <test/cpp/tensorexpr/tests.h>
+
+#include <gtest/gtest.h>
+
+namespace torch {
+namespace jit {
+
+#define TENSOREXPR_GTEST(name) \
+  TEST(TensorExprTest, name) { \
+    test##name();       \
+  }
+TH_FORALL_TESTS(TENSOREXPR_GTEST)
+#undef TENSOREXPR_GTEST
+
+} // namespace jit
+} // namespace torch
diff --git a/test/cpp/tensorexpr/test_base.h b/test/cpp/tensorexpr/test_base.h
@@ -0,0 +1,31 @@
+#pragma once
+
+#include <gtest/gtest.h>
+#include <test/cpp/common/support.h>
+
+namespace torch {
+namespace jit {
+namespace tensorexpr {
+
+template <typename U, typename V>
+void ExpectAllNear(
+    const std::vector<U>& v1,
+    const std::vector<U>& v2,
+    V threshold,
+    const std::string& name = "") {
+  ASSERT_EQ(v1.size(), v2.size());
+  for (int i = 0; i < v1.size(); i++) {
+    EXPECT_NEAR(v1[i], v2[i], threshold)
+        << "element index: " << i << ", name: " << name;
+  }
+}
+
+template <typename T>
+static void assertAllEqual(const std::vector<T>& vec, const T& val) {
+  for (auto const& elt : vec) {
+    ASSERT_EQ(elt, val);
+  }
+}
+} // namespace tensorexpr
+} // namespace jit
+} // namespace torch
diff --git a/test/cpp/tensorexpr/test_type.cpp b/test/cpp/tensorexpr/test_type.cpp
@@ -0,0 +1,37 @@
+#include "test/cpp/tensorexpr/test_base.h"
+#include "torch/csrc/jit/tensorexpr/mem_arena.h"
+#include "torch/csrc/jit/tensorexpr/types.h"
+
+namespace torch {
+namespace jit {
+using namespace torch::jit::tensorexpr;
+
+void testTypeTest01() {
+  KernelScope kernel_scope;
+  {
+    Dtype dt1 = kInt32;
+    EXPECT_EQ(dt1, kInt32);
+  }
+  {
+    Dtype dt2_a(kInt32, 8);
+    Dtype dt2_b(kInt32, 4);
+    Dtype dt2_c(kInt32, 8);
+    EXPECT_EQ(dt2_a, dt2_c);
+    EXPECT_NE(dt2_a, dt2_b);
+  }
+  {
+    EXPECT_EQ(kInt32, ToDtype<int>());
+    EXPECT_EQ(kFloat32, ToDtype<float>());
+  }
+  {
+    Dtype int32x8(kInt32, 8);
+    Dtype float32x8(kFloat32, 8);
+    EXPECT_NE(int32x8, float32x8);
+    EXPECT_EQ(float32x8, BinaryOpDtype(int32x8, float32x8));
+    EXPECT_EQ(float32x8, BinaryOpDtype(float32x8, int32x8));
+    EXPECT_EQ(int32x8, BinaryOpDtype(int32x8, int32x8));
+    EXPECT_EQ(float32x8, BinaryOpDtype(float32x8, float32x8));
+  }
+}
+} // namespace jit
+} // namespace torch
diff --git a/test/cpp/tensorexpr/test_utils.h b/test/cpp/tensorexpr/test_utils.h
@@ -0,0 +1,14 @@
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "test/cpp/tensorexpr/test_base.h"
+#include "torch/csrc/jit/testing/file_check.h"
+
+namespace torch {
+namespace jit {
+using namespace torch::jit::tensorexpr;
+
+} // namespace jit
+} // namespace torch
diff --git a/test/cpp/tensorexpr/tests.h b/test/cpp/tensorexpr/tests.h
@@ -0,0 +1,24 @@
+#pragma once
+
+/**
+ * See README.md for instructions on how to add a new test.
+ */
+#include <c10/macros/Export.h>
+#include <torch/csrc/WindowsTorchApiMacro.h>
+
+namespace torch {
+namespace jit {
+#define TH_FORALL_TESTS(_)      \
+  _(TypeTest01)                 \
+
+#define TH_FORALL_TESTS_CUDA(_) \
+
+#define DECLARE_TENSOREXPR_TEST(name) void test##name();
+TH_FORALL_TESTS(DECLARE_TENSOREXPR_TEST)
+#ifdef USE_CUDA
+TH_FORALL_TESTS_CUDA(DECLARE_TENSOREXPR_TEST)
+#endif
+#undef DECLARE_TENSOREXPR_TEST
+
+} // namespace jit
+} // namespace torch
diff --git a/tools/build_variables.bzl b/tools/build_variables.bzl
@@ -191,10 +191,12 @@ libtorch_sources = [
     "torch/csrc/jit/mobile/interpreter.cpp",
     "torch/csrc/jit/mobile/type_parser.cpp",
     "torch/csrc/jit/tensorexpr/mem_arena.cpp",
+    "torch/csrc/jit/tensorexpr/types.cpp",
     "torch/csrc/utils/byte_order.cpp",
     "torch/csrc/utils/tensor_flatten.cpp",
     "torch/csrc/utils/variadic.cpp",
     "torch/csrc/jit/tensorexpr/mem_arena.cpp",
+    "torch/csrc/jit/tensorexpr/types.cpp",
 ]
 
 libtorch_cuda_sources = [
diff --git a/torch/csrc/jit/tensorexpr/types.cpp b/torch/csrc/jit/tensorexpr/types.cpp
@@ -0,0 +1,100 @@
+#include "torch/csrc/jit/tensorexpr/types.h"
+#include <torch/csrc/WindowsTorchApiMacro.h>
+
+#include <c10/util/Logging.h>
+
+namespace torch {
+namespace jit {
+namespace tensorexpr {
+
+enum ScalarType {
+  kScalarUninitialized,
+  kScalarHandle,
+  kScalarInt32,
+  kScalarFloat32,
+};
+
+Dtype Dtype::scalar_type() const {
+  switch (static_cast<ScalarType>(scalar_type_)) {
+    case kScalarUninitialized:
+      return kUninitialized;
+    case kScalarHandle:
+      return kHandle;
+    case kScalarInt32:
+      return kInt32;
+    case kScalarFloat32:
+      return kFloat32;
+    default:
+      LOG(FATAL) << "invalid scalar type: " << scalar_type_;
+      return kUninitialized;
+  }
+}
+
+TORCH_API Dtype kInt32(kScalarInt32, 1);
+TORCH_API Dtype kFloat32(kScalarFloat32, 1);
+TORCH_API Dtype kHandle(kScalarHandle, 1);
+TORCH_API Dtype kUninitialized(kScalarUninitialized, 1);
+
+TORCH_API std::ostream& operator<<(std::ostream& stream, const Dtype& dtype) {
+  switch (static_cast<ScalarType>(dtype.scalar_type_)) {
+    case kScalarUninitialized:
+      stream << "uninitialized";
+      break;
+    case kScalarHandle:
+      stream << "handle";
+      break;
+    case kScalarInt32:
+      stream << "int32";
+      break;
+    case kScalarFloat32:
+      stream << "float32";
+      break;
+    default:
+      LOG(FATAL) << "invalid scalar type: " << dtype.scalar_type_;
+  }
+  if (dtype.lanes() > 1) {
+    stream << "x" << dtype.lanes();
+    ;
+  }
+  return stream;
+}
+
+int Dtype::byte_size() const {
+  int scalar_size = -1;
+  switch (scalar_type_) {
+    case kScalarInt32:
+      scalar_size = sizeof(int32);
+      break;
+    case kScalarFloat32:
+      scalar_size = sizeof(float);
+      break;
+    default:
+      throw std::runtime_error(
+          "invalid scalar type; " + std::to_string(scalar_type_));
+  }
+  return scalar_size * lanes();
+}
+
+std::string Dtype::ToCppString() const {
+  if (scalar_type_ == kScalarInt32) {
+    return "int";
+  } else if (scalar_type_ == kScalarFloat32) {
+    return "float";
+  } else {
+    throw std::runtime_error("Invalid dtype: " + std::to_string(scalar_type_));
+  }
+}
+
+} // namespace tensorexpr
+} // namespace jit
+} // namespace torch
+
+namespace std {
+
+std::string to_string(const Dtype& dtype) {
+  std::ostringstream oss;
+  oss << dtype;
+  return oss.str();
+}
+
+} // namespace std
diff --git a/torch/csrc/jit/tensorexpr/types.h b/torch/csrc/jit/tensorexpr/types.h
