From ebfcd28cd1349cadbccee61218da234a6f53ed7c Mon Sep 17 00:00:00 2001
From: Tianqi Chen <tqchen@users.noreply.github.com>
Date: Thu, 31 Oct 2019 11:13:32 -0700
Subject: [PATCH] [CUDA] Fix fp16 intrin, disable bad fp16 vecadd test for now
 (#4239)

---
 docker/bash.sh                             |  4 ++--
 src/codegen/codegen_cuda.cc                | 28 ++++++++++++++--------------
 tests/python/unittest/test_codegen_cuda.py |  4 ++++
 3 files changed, 20 insertions(+), 16 deletions(-)

diff --git a/docker/bash.sh b/docker/bash.sh
index 0fdc476..53bf03d 100755
--- a/docker/bash.sh
+++ b/docker/bash.sh
@@ -58,7 +58,7 @@ else
 fi
 
 if [[ "${DOCKER_IMAGE_NAME}" == *"gpu"* ]]; then
-    if ! type "nvidia-docker" > /dev/null
+    if ! type "nvidia-docker" 1> /dev/null 2> /dev/null
     then
         DOCKER_BINARY="docker"
         CUDA_ENV=" --gpus all "${CUDA_ENV}
@@ -79,7 +79,6 @@ echo "Running '${COMMAND[@]}' inside ${DOCKER_IMAGE_NAME}..."
 # By default we cleanup - remove the container once it finish running (--rm)
 # and share the PID namespace (--pid=host) so the process inside does not have
 # pid 1 and SIGKILL is propagated to the process inside (jenkins can kill it).
-echo ${DOCKER_BINARY}
 ${DOCKER_BINARY} run --rm --pid=host\
     -v ${WORKSPACE}:/workspace \
     -v ${SCRIPT_DIR}:/docker \
@@ -95,3 +94,4 @@ ${DOCKER_BINARY} run --rm --pid=host\
     ${DOCKER_IMAGE_NAME}\
     bash --login /docker/with_the_same_user \
     ${COMMAND[@]}
+
diff --git a/src/codegen/codegen_cuda.cc b/src/codegen/codegen_cuda.cc
index 39a3ab7..5f04dd0 100644
--- a/src/codegen/codegen_cuda.cc
+++ b/src/codegen/codegen_cuda.cc
@@ -51,20 +51,20 @@ void CodeGenCUDA::AddFunction(LoweredFunc f) {
 std::string CodeGenCUDA::Finish() {
   if (enable_fp16_) {
     decl_stream << "#include <cuda_fp16.h>\n";
-    decl_stream << "__device__ half max" \
-                    "(const half a, const half b)\n"
-                    "{\n  return __hgt(__half(a), __half(b)) ? a : b;\n}\n";
-    decl_stream << "__device__ half min(const half a, const half b)\n"
-                    "{\n  return __hlt(__half(a), __half(b)) ? a : b;\n}\n";
-    decl_stream << "__device__ half operator+" \
-                    "(const volatile __half &a,  const volatile __half &b)\n"
-                    "{\n  return __hadd(a, b);\n}\n";
-    decl_stream << "__device__ half operator<=" \
-                   "(const volatile __half &a,  const volatile __half &b)\n"
-                    "{\n  return __hlt(a, b);\n}\n";
-    decl_stream << "__device__ half operator*" \
-                    "(const volatile __half &a,  const volatile __half &b)\n"
-                    "{\n  return __hmul(a, b);\n}\n";
+    decl_stream << "__device__ half max"
+                << "(half a, half b)\n"
+                << "{\n  return __hgt(__half(a), __half(b)) ? a : b;\n}\n";
+    decl_stream << "__device__ half min(half a, half b)\n"
+                << "{\n  return __hlt(__half(a), __half(b)) ? a : b;\n}\n";
+    decl_stream << "__device__ half operator<="
+                << "(__half a,  __half b)\n"
+                << "{\n  return __hlt(a, b);\n}\n";
+    decl_stream << "__device__ half operator+"
+                << "(__half a,  __half &b)\n"
+                <<"{\n  return __hadd(a, b);\n}\n";
+    decl_stream << "__device__ half operator*"
+                << "(__half a, __half b)\n"
+                <<   "{\n  return __hmul(a, b);\n}\n";
   }
 
   if (enable_int8_) {
diff --git a/tests/python/unittest/test_codegen_cuda.py b/tests/python/unittest/test_codegen_cuda.py
index a21a58f..7991c60 100644
--- a/tests/python/unittest/test_codegen_cuda.py
+++ b/tests/python/unittest/test_codegen_cuda.py
@@ -54,6 +54,10 @@ def test_cuda_vectorize_add():
     check_cuda("int8", 64, 4)
     # check_cuda("float16", 64, 2)
 
+    # TODO(tvm-team) fix fp16 codegen here
+    # or hit an error if it is less frequently used.
+    # check_cuda("float16", 64, 2)
+
 
 def test_cuda_multiply_add():
     num_thread = 8
-- 
2.7.4