Add tests for user-ops

as a part of test-on-install in pip.sh
Change: 120601648

3 files changed, 256 insertions, 10 deletions

diff --git a/tensorflow/g3doc/how_tos/adding_an_op/index.md b/tensorflow/g3doc/how_tos/adding_an_op/index.md
index 98e82037a5..ec5020bf11 100644
--- a/tensorflow/g3doc/how_tos/adding_an_op/index.md
+++ b/tensorflow/g3doc/how_tos/adding_an_op/index.md
@@ -142,6 +142,9 @@ TF_INC=$(python -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')
 g++ -std=c++11 -shared zero_out.cc -o zero_out.so -fPIC -I $TF_INC
 ```
 
+On Mac OS X, the additional flag "-undefined dynamic_lookup" is required when
+building the .so file.
+
 > Note on gcc version 5: gcc5 uses the new C++
 [ABI](https://gcc.gnu.org/gcc-5/changes.html#libstdcxx). The binary pip packages
 available on the TensorFlow website are built with gcc4 that uses the older ABI.
@@ -956,6 +959,14 @@ One thing to note, even when the GPU kernel version of `pad` is used, it still
 needs its `"paddings"` input in CPU memory.  To mark that inputs or outputs are
 kept on the CPU, add a `HostMemory()` call to the kernel registration, e.g.:
 
+```c++
+#define REGISTER_GPU_KERNEL(T)                         \
+  REGISTER_KERNEL_BUILDER(Name("Pad")                  \
+                              .Device(DEVICE_GPU)      \
+                              .TypeConstraint<T>("T")  \
+                              .HostMemory("paddings"), \
+                          PadOp<GPUDevice, T>)
+```
 
 ### Compiling the kernel for the GPU device
 
@@ -983,15 +994,6 @@ cuda_op_kernel.cu.o -I $TF_INC -fPIC -lcudart
 `cuda_op_kernel.so` produced above can be loaded as usual in Python, using the
 `tf.load_op_library` function.
 
-```c++
-#define REGISTER_GPU_KERNEL(T)                         \
-  REGISTER_KERNEL_BUILDER(Name("Pad")                  \
-                              .Device(DEVICE_GPU)      \
-                              .TypeConstraint<T>("T")  \
-                              .HostMemory("paddings"), \
-                          PadOp<GPUDevice, T>)
-```
-
 ## Implement the gradient in Python
 
 Given a graph of ops, TensorFlow uses automatic differentiation
diff --git a/tensorflow/tools/ci_build/builds/pip.sh b/tensorflow/tools/ci_build/builds/pip.sh
index 160fd1f701..5484abe59d 100755
--- a/tensorflow/tools/ci_build/builds/pip.sh
+++ b/tensorflow/tools/ci_build/builds/pip.sh
@@ -40,6 +40,9 @@
 # If NO_TEST_ON_INSTALL has any non-empty and non-0 value, the test-on-install
 # part will be skipped.
 #
+# If NO_TEST_USER_OPS has any non-empty and non-0 value, the testing of user-
+# defined ops against the installation will be skipped.
+#
 # Use --mavx or --mavx2 to let bazel use --copt=-mavx or --copt=-mavx2 options
 # while building the pip package, respectively.
 #
@@ -67,6 +70,13 @@ if [[ ! -z "${TF_BUILD_BAZEL_CLEAN}" ]] && \
   bazel clean
 fi
 
+DO_TEST_USER_OPS=1
+if [[ ! -z "${NO_TEST_USER_OPS}" ]] && \
+   [[ "${NO_TEST_USER_OPS}" != "0" ]]; then
+  echo "NO_TEST_USER_OPS=${NO_TEST_USER_OPS}: Will skip testing of user ops"
+  DO_TEST_USER_OPS=0
+fi
+
 DO_TEST_TUTORIALS=0
 DO_INTEGRATION_TESTS=0
 MAVX_FLAG=""
@@ -80,7 +90,7 @@ while true; do
   elif [[ "${1}" == "--mavx2" ]]; then
     MAVX_FLAG="--copt=-mavx2"
   fi
- 
+
   shift
   if [[ -z "${1}" ]]; then
     break
@@ -203,6 +213,12 @@ fi
 "${SCRIPT_DIR}/test_installation.sh" --virtualenv ${GPU_FLAG} ||
     die "PIP tests-on-install FAILED"
 
+# Test user ops
+if [[ "${DO_TEST_USER_OPS}" == "1" ]]; then
+  "${SCRIPT_DIR}/test_user_ops.sh" --virtualenv ${GPU_FLAG} || \
+      die "PIP user-op tests-on-install FAILED"
+fi
+
 # Optional: Run the tutorial tests
 if [[ "${DO_TEST_TUTORIALS}" == "1" ]]; then
   "${SCRIPT_DIR}/test_tutorials.sh" --virtualenv || \
diff --git a/tensorflow/tools/ci_build/builds/test_user_ops.sh b/tensorflow/tools/ci_build/builds/test_user_ops.sh
new file mode 100755
index 0000000000..57cbc0f1f2
--- /dev/null
+++ b/tensorflow/tools/ci_build/builds/test_user_ops.sh
@@ -0,0 +1,228 @@
+#!/usr/bin/env bash
+# Copyright 2016 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Test user-defined ops against installation of TensorFlow.
+#
+# Usage: test_user_ops.sh [--virtualenv] [--gpu]
+#
+# If the flag --virtualenv is set, the script will use "python" as the Python
+# binary path. Otherwise, it will use tools/python_bin_path.sh to determine
+# the Python binary path.
+#
+# The --gpu flag informs the script that this is a GPU build, so that the
+# appropriate test blacklists can be applied accordingly.
+#
+
+echo ""
+echo "=== Testing user ops ==="
+
+# Script directory
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+source "${SCRIPT_DIR}/builds_common.sh"
+
+
+# Process input arguments
+IS_VIRTUALENV=0
+IS_GPU=0
+while true; do
+  if [[ "$1" == "--virtualenv" ]]; then
+    IS_VIRTUALENV=1
+  elif [[ "$1" == "--gpu" ]]; then
+    IS_GPU=1
+  fi
+  shift
+
+  if [[ -z "$1" ]]; then
+    break
+  fi
+done
+
+TMP_DIR=$(mktemp -d)
+mkdir -p "${TMP_DIR}"
+
+cleanup() {
+  rm -rf "${TMP_DIR}"
+}
+
+die() {
+  echo $@
+  cleanup
+  exit 1
+}
+
+
+# Obtain the path to Python binary
+if [[ ${IS_VIRTUALENV} == "1" ]]; then
+  PYTHON_BIN_PATH="$(which python)"
+else
+  source tools/python_bin_path.sh
+  # Assume: PYTHON_BIN_PATH is exported by the script above
+fi
+echo "PYTHON_BIN_PATH: ${PYTHON_BIN_PATH}"
+
+
+pushd "${TMP_DIR}"
+
+# Obtain paths include and lib paths to the TensorFlow installation
+TF_INC=$("${PYTHON_BIN_PATH}" \
+         -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')
+
+if [[ -z "${TF_INC}" ]]; then
+  die "FAILED to determine TensorFlow include path"
+else
+  echo "TensorFlow include path: ${TF_INC}"
+fi
+
+# Check g++ availability
+GPP_BIN="g++"
+if [[ -z $(which "${GPP_BIN}") ]]; then
+  die "ERROR: ${GPP_BIN} not on path"
+fi
+
+echo ""
+echo "g++ version:"
+"${GPP_BIN}" -v
+echo ""
+
+
+IS_MAC=0
+if [[ $(uname) == "Darwin" ]]; then
+  echo "Detected Mac OS X environment"
+  IS_MAC=1
+fi
+
+EXTRA_GPP_FLAGS=""
+if [[ ${IS_MAC} == "1" ]]; then
+  # Extra flags required on Mac OS X, where dynamic_lookup is not the default
+  # behavior.
+  EXTRA_GPP_FLAGS="${EXTRA_GPP_FLAGS} -undefined dynamic_lookup"
+fi
+
+echo "Extra GPP flag: ${EXTRA_GPP_FLAGS}"
+
+# Input to the user op
+OP_INPUT="[42, 43, 44]"
+
+if [[ ${IS_GPU} == "0" ]]; then
+  echo "Testing user ops in CPU environment"
+
+  # Expected output from user op
+  EXPECTED_OUTPUT="[42, 0, 0]"
+
+  # Locate the op kernel C++ file
+  OP_KERNEL_CC="${SCRIPT_DIR}/../../../g3doc/how_tos/adding_an_op/zero_out_op_kernel_1.cc"
+  OP_KERNEL_CC=$(realpath "${OP_KERNEL_CC}")
+
+  if [[ ! -f "${OP_KERNEL_CC}" ]]; then
+    die "ERROR: Unable to find user-op kernel C++ file at: ${OP_KERNEL_CC}"
+  fi
+
+  # Copy the file to a non-TensorFlow source directory
+  cp "${OP_KERNEL_CC}" ./
+
+  # Compile the op kernel into an .so file
+  SRC_FILE=$(basename "${OP_KERNEL_CC}")
+
+  echo "Compiling user op C++ source file ${SRC_FILE}"
+
+  USER_OP_SO="zero_out.so"
+
+  "${GPP_BIN}" -std=c++11 ${EXTRA_GPP_FLAGS} \
+    -shared "${SRC_FILE}" -o "${USER_OP_SO}" \
+    -fPIC -I "${TF_INC}" || \
+    die "g++ compilation of ${SRC_FILE} FAILED"
+
+else
+  echo "Testing user ops in GPU environment"
+
+  # Expected output from user op
+  EXPECTED_OUTPUT="[43, 44, 45]"
+
+  # Check nvcc availability
+  NVCC_BIN="/usr/local/cuda/bin/nvcc"
+  if [[ -z $(which "${NVCC_BIN}") ]]; then
+    die "ERROR: ${NVCC_BIN} not on path"
+  fi
+
+  echo ""
+  echo "nvcc version:"
+  "${NVCC_BIN}" --version
+  echo ""
+
+  OP_KERNEL_CU="${SCRIPT_DIR}/../../../g3doc/how_tos/adding_an_op/cuda_op_kernel.cu.cc"
+  OP_KERNEL_CU=$(realpath "${OP_KERNEL_CU}")
+  if [[ ! -f "${OP_KERNEL_CU}" ]]; then
+    die "ERROR: Unable to find user-op kernel CUDA file at: ${OP_KERNEL_CU}"
+  fi
+
+  OP_KERNEL_CC="${SCRIPT_DIR}/../../../g3doc/how_tos/adding_an_op/cuda_op_kernel.cc"
+  OP_KERNEL_CC=$(realpath "${OP_KERNEL_CC}")
+  if [[ ! -f "${OP_KERNEL_CC}" ]]; then
+    die "ERROR: Unable to find user-op kernel C++ file at: ${OP_KERNEL_CC}"
+  fi
+
+  # Copy the file to a non-TensorFlow source directory
+  cp "${OP_KERNEL_CU}" ./
+  cp "${OP_KERNEL_CC}" ./
+
+  OP_KERNEL_O=$(echo "${OP_KERNEL_CC}" | sed -e 's/\.cc/\.o/')
+  "${NVCC_BIN}" -std=c++11 \
+      -c -o "${OP_KERNEL_O}" "${OP_KERNEL_CU}" \
+      -I "${TF_INC}" -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC || \
+      die "nvcc compilation of ${OP_KERNEL_CC} FAILED"
+
+  # USER_OP_SO=$(basename $(echo "${OP_KERNEL_CC}" | sed -e 's/\.cc/\.so/'))
+  USER_OP_SO="add_one.so"
+  "${GPP_BIN}" -std=c++11 ${EXTRA_GPP_FLAGS} \
+      -shared -o "${USER_OP_SO}" "${OP_KERNEL_CC}" \
+      "${OP_KERNEL_O}" -I "${TF_INC}" -L "/usr/local/cuda/lib64" \
+      -fPIC -lcudart || \
+      die "g++ compilation of ${OP_KERNEL_CC}" FAILED
+fi
+
+# Try running the op
+USER_OP=$(echo "${USER_OP_SO}" | sed -e 's/\.so//')
+echo "Invoking user op ${USER_OP} defined in file ${USER_OP_SO} "\
+"via pip installation"
+
+ORIG_OUTPUT=$("${PYTHON_BIN_PATH}" -c "import tensorflow as tf; print(tf.Session('').run(tf.load_op_library('./${USER_OP_SO}').${USER_OP}(${OP_INPUT})))")
+
+# Format OUTPUT for analysis
+if [[ -z $(echo "${ORIG_OUTPUT}" | grep -o ',') ]]; then
+  if [[ ${IS_MAC} == "1" ]]; then
+    OUTPUT=$(echo "${ORIG_OUTPUT}" | sed -E -e 's/[ \t]+/,/g')
+  else
+    OUTPUT=$(echo "${ORIG_OUTPUT}" | sed -r -e 's/[ \t]+/,/g')
+  fi
+else
+  OUTPUT="${ORIG_OUTPUT}"
+fi
+
+EQUALS_EXPECTED=$("${PYTHON_BIN_PATH}" -c "print(${OUTPUT} == ${EXPECTED_OUTPUT})")
+
+if [[ "${EQUALS_EXPECTED}" != "True" ]]; then
+  die "FAILED: Output from user op (${OUTPUT}) does not match expected "\
+"output ${EXPECTED_OUTPUT}"
+else
+  echo "Output from user op (${OUTPUT}) matches expected output"
+fi
+
+popd
+
+cleanup
+
+echo ""
+echo "SUCCESS: Testing of user ops PASSED"