Added hash for device (#9246)

vishwakftw · facebook-github-bot · commit 1c3580b6fef0 · 2018-07-17T17:10:17.000-07:00
Summary: If this is good, I could write some tests to ensure collision doesn't occur within a given range. Closes #7228 Pull Request resolved: #9246 Differential Revision: D8872608 Pulled By: ezyang fbshipit-source-id: 0ed29a73188f4167b42756f59a5c9a3d5cb37326
diff --git a/aten/src/ATen/Device.h b/aten/src/ATen/Device.h
@@ -6,6 +6,7 @@
 #include <cstddef>
 #include <iosfwd>
 #include <string>
+#include <functional>
 
 namespace at {
 /// Represents a a compute device on which a tensor is located. A device is
@@ -112,3 +113,16 @@ struct Device {
 
 std::ostream& operator<<(std::ostream& stream, at::Device::Type type);
 std::ostream& operator<<(std::ostream& stream, const at::Device& device);
+
+namespace std {
+  template<> struct hash<at::Device>
+  {
+    size_t operator()(const at::Device& device) const noexcept {
+      size_t hash_val = static_cast<size_t>(device.index() + 1);
+      if (device.is_cuda()) {
+        hash_val += 2;
+      }
+      return hash_val;
+    }
+  };
+} // namespace std
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -1909,6 +1909,12 @@ def test_device(self):
         self.assertRaises(TypeError, lambda: torch.device('other'))
         self.assertRaises(TypeError, lambda: torch.device('other:0'))
 
+        device_set = {'cpu', 'cpu:0', 'cuda', 'cuda:0', 'cuda:1', 'cuda:10', 'cuda:100'}
+        device_hash_set = set()
+        for device in list(device_set):
+            device_hash_set.add(hash(torch.device(device)))
+        self.assertEqual(len(device_set), len(device_hash_set))
+
     def test_tensor_device(self):
         def assertEqual(device_str, fn):
             self.assertEqual(torch.device(device_str), fn().device)
diff --git a/torch/csrc/Device.cpp b/torch/csrc/Device.cpp
@@ -4,12 +4,14 @@
 #include "torch/csrc/utils/object_ptr.h"
 #include "torch/csrc/utils/python_arg_parser.h"
 #include "torch/csrc/utils/python_strings.h"
+#include "torch/csrc/utils/python_numbers.h"
 #include "torch/csrc/utils/pybind.h"
 
 #include <ATen/Device.h>
 #include <ATen/Error.h>
 
 #include <cstring>
+#include <limits>
 #include <structmember.h>
 #include <sstream>
 
@@ -95,6 +97,13 @@ PyObject *THPDevice_index(THPDevice *self)
   END_HANDLE_TH_ERRORS
 }
 
+static Py_ssize_t THPDevice_hash(THPDevice *self)
+{
+  HANDLE_TH_ERRORS
+  return static_cast<Py_ssize_t>(std::hash<at::Device>{}(self->device) % std::numeric_limits<Py_ssize_t>::max());
+  END_HANDLE_TH_ERRORS_RET(-1)
+}
+
 PyObject *THPDevice_rc(PyObject *a, PyObject *b, int op) {
   HANDLE_TH_ERRORS
   if (!THPDevice_Check(a) || !THPDevice_Check(b)) {
@@ -181,7 +190,7 @@ PyTypeObject THPDeviceType = {
   0,                                     /* tp_as_number */
   0,                                     /* tp_as_sequence */
   0,                                     /* tp_as_mapping */
-  0,                                     /* tp_hash  */
+  (hashfunc)THPDevice_hash,              /* tp_hash  */
   0,                                     /* tp_call */
   (reprfunc)THPDevice_str,               /* tp_str */
   0,                                     /* tp_getattro */
