Make ValType hashable. (#35917)

csarofeen · facebook-github-bot · commit aeb13f212bec · 2020-04-03T00:16:56.000-07:00
Summary: Build fix stemming from #34785 Pull Request resolved: #35917 Differential Revision: D20829353 Pulled By: soumith fbshipit-source-id: 4ba84ecedd354efbc9ac47c9b0f0e3871b404f13
diff --git a/torch/csrc/jit/codegen/cuda/fusion.h b/torch/csrc/jit/codegen/cuda/fusion.h
@@ -14,6 +14,14 @@ namespace torch {
 namespace jit {
 namespace fuser {
 
+// https://stackoverflow.com/questions/18837857/cant-use-enum-class-as-unordered-map-key
+struct TypeHash {
+  template <typename T>
+  std::size_t operator()(T t) const {
+    return static_cast<std::size_t>(t);
+  }
+};
+
 /*
  * Usage: FusionGuard and Fusion are required user interfaces for any operation
  * underlying the code generator. In order to create values, expressions, and
@@ -175,7 +183,7 @@ struct TORCH_CUDA_API Fusion : public IRInputOutput {
   StmtNameType getExprName();
 
   // map from valtype to individual name counters
-  std::unordered_map<ValType, StmtNameType> val_type_name_map = {
+  std::unordered_map<ValType, StmtNameType, TypeHash> val_type_name_map = {
       {ValType::TensorView, 0},
       {ValType::TensorDomain, 0},
       {ValType::IterDomain, 0},
