Add ROCm runner

.github/workflows/fbgemmci.yml

@@ -308,6 +308,54 @@ jobs:
         python3 -c "import fbgemm_gpu"
         python3 -c "import fbgemm_gpu.split_embedding_codegen_lookup_invokers"
 
+  test_amd_gpu:
+    runs-on: rocm
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+
+    steps:
+    - name: pre-checkout
+      shell: bash
+      run: |
+        if [ -d ${{ github.workspace }} ]
+        then
+          sudo chown -R $USER:$USER ${{ github.workspace }}
+        fi
+        sudo add-apt-repository ppa:git-core/ppa
+        sudo apt update
+        sudo apt -y install --only-upgrade git
+
+    - uses: actions/checkout@v2
+      with:
+        ref: ${{ github.ref }}
+        submodules: 'true'
+
+    - name: build fbgemm_gpu and test
+      shell: bash
+      run: |
+        set -eux
+        env
+        ls -l
+        DOCKER_IMAGE=rocm/pytorch:rocm5.1.1_ubuntu20.04_py3.7_pytorch_staging_base
+        docker pull $DOCKER_IMAGE
+        JENKINS_REPO_DIR=fbgemm-private-jenkins
+        JENKINS_REPO_DIR_BAREMETAL=$PWD
+        JENKINS_REPO_DIR_DOCKER=/workspace/$JENKINS_REPO_DIR
+        DOCKER_OPTIONS="\
+        --user 0 \
+        --network=host \
+        --ipc=host \
+        --shm-size 16G \
+        --group-add video \
+        --cap-add=SYS_PTRACE \
+        --security-opt seccomp=unconfined \
+        --device=/dev/kfd \
+        --device=/dev/dri \
+        -v $JENKINS_REPO_DIR_BAREMETAL:$JENKINS_REPO_DIR_DOCKER
+        "
+        docker run $DOCKER_OPTIONS $DOCKER_IMAGE $JENKINS_REPO_DIR_DOCKER/.jenkins/rocm/build_and_test.sh $JENKINS_REPO_DIR_DOCKER
+
   build_cpu_only:
     runs-on: ${{ matrix.os }}
     strategy:

.jenkins/rocm/build_and_test.sh

@@ -0,0 +1,58 @@
+#!/bin/bash
+
+# exit immediately on failure, or if an undefined variable is used
+set -eux
+
+FBGEMM_REPO_DIR=${1:-/workspace/FBGEMM}
+
+git config --global --add safe.directory $FBGEMM_REPO_DIR
+git config --global --add safe.directory $FBGEMM_REPO_DIR/third_party/asmjit
+git config --global --add safe.directory $FBGEMM_REPO_DIR/third_party/cpuinfo
+git config --global --add safe.directory $FBGEMM_REPO_DIR/third_party/googletest
+git config --global --add safe.directory $FBGEMM_REPO_DIR/third_party/hipify_torch
+
+# Install dependencies
+apt-get update --allow-insecure-repositories && \
+  apt-get install -y --allow-unauthenticated \
+  git \
+  jq \
+  sshfs \
+  sshpass \
+  unzip
+
+apt-get install -y locales
+locale-gen en_US.UTF-8
+
+pip3 install click
+pip3 install jinja2
+pip3 install ninja
+pip3 install scikit-build
+pip3 install --upgrade hypothesis
+pip3 install --pre torch torchvision --extra-index-url https://download.pytorch.org/whl/nightly/rocm5.1.1/
+
+pip3 list
+
+# Build fbgemm_gpu
+cd $FBGEMM_REPO_DIR/fbgemm_gpu
+export MAX_JOBS=`nproc`
+export PYTORCH_ROCM_ARCH="gfx908"
+python setup.py build develop
+
+export FBGEMM_TEST_WITH_ROCM=1
+
+# Test fbgemm_gpu
+cd test
+
+python layout_transform_ops_test.py --verbose
+
+python permute_pooled_embedding_modules_test.py --verbose
+
+python sparse_ops_test.py --verbose
+
+python merge_pooled_embeddings_test.py --verbose
+
+python quantize_ops_test.py --verbose
+
+python split_embedding_inference_converter_test.py --verbose
+
+python split_table_batched_embeddings_test.py --verbose

fbgemm_gpu/CMakeLists.txt

@@ -74,7 +74,7 @@ if(USE_ROCM)
   include(Hipify)
 
   message("${message_line}")
-  message(STATUS "hip found ${ROCM_FOUND}")
+  message(STATUS "hip found ${HIP_FOUND}")
 endif()
 
 #