[te] Fix casting of unsigned char, and abs(int)

ghstack-source-id: 9f939c4a78b832da338d851725a3b3b77c409b9e Pull Request resolved: #44157
pytorch · Sep 3, 2020 · a109d40 · a109d40
1 parent b457879
commit a109d40
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 10 deletions.
diff --git a/test/test_jit_fuser_te.py b/test/test_jit_fuser_te.py
@@ -1152,20 +1152,48 @@ def rand(dtype, device="cuda"):
             torch.int16,
             torch.int32,
             torch.int64,
-            torch.float16,
-            torch.bfloat16,
+            # torch.float16,
+            # torch.bfloat16,
             torch.float32,
             torch.float64,
-            torch.bool,
-            torch.complex32,
-            torch.complex64,
-            torch.complex128,
-            torch.qint8,
-            torch.quint8,
-            torch.qint32,
+            # torch.bool,
+            # torch.complex32,
+            # torch.complex64,
+            # torch.complex128,
+            # torch.qint8,
+            # torch.quint8,
+            # torch.qint32,
         ]
         unary_ops = [
             torch.sigmoid,
+            torch.reciprocal,
+            torch.neg,
+            torch.relu,
+            torch.log,
+            torch.log10,
+            torch.log2,
+            torch.exp,
+            torch.expm1,
+            torch.erf,
+            torch.erfc,
+            torch.cos,
+            torch.sin,
+            torch.tan,
+            torch.acos,
+            torch.asin,
+            torch.cosh,
+            torch.sinh,
+            torch.atan,
+            torch.tanh,
+            torch.sqrt,
+            torch.rsqrt,
+            torch.abs,
+            torch.ceil,
+            torch.floor,
+            torch.round,
+            torch.trunc,
+            torch.frac,
+            torch.lgamma,
         ]
         devices = [
             "cuda",

diff --git a/torch/csrc/jit/tensorexpr/cuda_codegen.cpp b/torch/csrc/jit/tensorexpr/cuda_codegen.cpp
@@ -186,7 +186,7 @@ void CudaPrinter::visit(const For* v) {
 }
 
 void CudaPrinter::visit(const Cast* v) {
-  os() << cudaDtypeCppString(v->dtype());
+  os() << "(" << cudaDtypeCppString(v->dtype()) << ")";
   os() << "(";
   v->src_value()->accept(this);
   os() << ")";
@@ -209,6 +209,9 @@ void CudaPrinter::visit(const Intrinsics* v) {
   if (returnType == ScalarType::Half || returnType == ScalarType::Float) {
     func_name = func_name + "f";
   }
+  if (v->op_type() == IntrinsicsOp::kFabs && is_integral(returnType)) {
+    func_name = "abs";
+  }
 
   os() << func_name << "(";
   for (int i = 0; i < v->nparams(); i++) {