Update XLAPreAutograd keys. (#40265)

Ailing Zhang · facebook-github-bot · commit cfe1c6ef9e10 · 2020-06-19T21:12:50.000-07:00
Summary: Pull Request resolved: #40265 Differential Revision: D22137998 Pulled By: ailzhang fbshipit-source-id: 41edac06f8aafa5d4c1dcefd5da81be6c9ac4a9c
diff --git a/c10/core/DispatchKey.h b/c10/core/DispatchKey.h
@@ -230,13 +230,6 @@ static_assert(
 C10_API const char* toString(DispatchKey);
 C10_API std::ostream& operator<<(std::ostream&, DispatchKey);
 
-// For backwards compatibility with XLA repository
-// (I don't want to fix this in XLA right now because there might be
-// more renaming coming in the future.)
-static inline DispatchKey XLA() {
-  return DispatchKey::XLA;
-}
-
 // These are some convenience identifiers for dispatch keys which are
 // shorter to type than their long counterparts.  Note that some of these
 // dispatch keys directly correspond to DeviceType; and most APIs that
diff --git a/c10/core/DispatchKeySet.h b/c10/core/DispatchKeySet.h
@@ -132,4 +132,11 @@ static inline DispatchKey legacyExtractDispatchKey(DispatchKeySet s) {
   return s.highestPriorityTypeId();
 }
 
+// For backwards compatibility with XLA repository
+// (I don't want to fix this in XLA right now because there might be
+// more renaming coming in the future.)
+static inline DispatchKeySet XLA() {
+  return DispatchKeySet{DispatchKey::XLA, DispatchKey::XLAPreAutograd};
+}
+
 }
diff --git a/c10/core/TensorOptions.h b/c10/core/TensorOptions.h
@@ -643,6 +643,8 @@ inline DeviceType computeDeviceType(DispatchKey tid) {
     return DeviceType::MSNPU;
   } else if (tid == DispatchKey::XLA) {
     return DeviceType::XLA;
+  } else if (tid == DispatchKey::XLAPreAutograd) {
+    return DeviceType::XLA;
   } else if (tid == DispatchKey::SparseCPU) {
     return DeviceType::CPU;
   } else if (tid == DispatchKey::SparseCUDA) {
diff --git a/torch/csrc/utils/tensor_new.cpp b/torch/csrc/utils/tensor_new.cpp
@@ -336,7 +336,7 @@ void check_base_legacy_new(c10::DispatchKey dispatch_key, at::Layout expected_la
     TORCH_CHECK(dispatch_key == c10::DispatchKey::CPU
                 || dispatch_key == c10::DispatchKey::CUDA
                 || dispatch_key == c10::DispatchKey::HIP
-                || dispatch_key == c10::XLA(),
+                || c10::XLA().has(dispatch_key),
                 "new(): expected DispatchKey: ", c10::DispatchKey::CPU,
                 " or ", c10::DispatchKey::CUDA,
                 " or ", c10::DispatchKey::HIP,

Original file line number	Diff line number	Diff line change
`@@ -132,4 +132,11 @@ static inline DispatchKey legacyExtractDispatchKey(DispatchKeySet s) {`
`132`	`132`	`return s.highestPriorityTypeId();`
`133`	`133`	`}`
`134`	`134`
	`135`	`+// For backwards compatibility with XLA repository`
	`136`	`+// (I don't want to fix this in XLA right now because there might be`
	`137`	`+// more renaming coming in the future.)`
	`138`	`+static inline DispatchKeySet XLA() {`
	`139`	`+ return DispatchKeySet{DispatchKey::XLA, DispatchKey::XLAPreAutograd};`
	`140`	`+}`
	`141`	`+`
`135`	`142`	`}`