Add Float16 support (#14)

pai-disc · Jul 20, 2022 · e0497e2 · e0497e2
1 parent ad6c81c
commit e0497e2
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/lib/Dialect/Torch/Utils/Utils.cpp b/lib/Dialect/Torch/Utils/Utils.cpp
@@ -55,6 +55,8 @@ torch_upstream::ScalarType Torch::getScalarTypeForType(Type type) {
     return torch_upstream::ScalarType::Bool;
   if (type.isBF16())
     return torch_upstream::ScalarType::BFloat16;
+  if (type.isF16())
+    return torch_upstream::ScalarType::Half;
   llvm::report_fatal_error("unhandled type for getScalarTypeForType");
 }
 
@@ -74,6 +76,8 @@ Type Torch::getTypeForScalarType(
     return IntegerType::get(context, 1);
   case torch_upstream::ScalarType::BFloat16:
     return mlir::FloatType::getBF16(context);
+  case torch_upstream::ScalarType::Half:
+    return mlir::FloatType::getF16(context);
   default:
     return Type();
   }

diff --git a/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/torch_to_mlir_utils.cpp b/python/torch_mlir/dialects/torch/importer/jit_ir/csrc/torch_to_mlir_utils.cpp
@@ -365,6 +365,10 @@ MlirAttribute torch_mlir::convertTensorToMlirElementsAttr(at::Tensor tensor,
   case ScalarType::BFloat16:
     return mlirDenseElementsAttrBFloat16Get(
         shapedType, numElements, static_cast<const uint16_t *>(tensorData));
+  case ScalarType::Half:
+    return mlirDenseElementsAttrFloat16Get(
+        shapedType, numElements, static_cast<const uint16_t *>(tensorData));
+
   default:
     throwUnsupportedTensorError();
   }