Test on both pytorch-triton and triton-main (#63)

Summary:
Add check for both Triton-main and Triton-pytorch

Pull Request resolved: #63

Reviewed By: FindHao

Differential Revision: D66311149

Pulled By: xuzhao9

fbshipit-source-id: 06ae9cf052ff6a7e5ab3c5e9529cb41886df6ab4

Author: xuzhao9

.ci/tritonbench/install-triton-main.sh

@@ -28,5 +28,17 @@ conda activate "${CONDA_ENV}"
 cd /workspace
 git clone https://github.com/triton-lang/triton.git
 cd /workspace/triton
+# delete the original triton directory
+TRITON_PKG_DIR=$(python -c "import triton; import os; print(os.path.dirname(triton.__file__))")
+# make sure all pytorch triton has been uninstalled
+pip uninstall -y triton
+pip uninstall -y triton
+pip uninstall -y triton
+rm -rf "${TRITON_PKG_DIR}"
+
+# install main triton
 pip install ninja cmake wheel pybind11; # build-time dependencies
 pip install -e python
+
+# test main branch installation with importing experimental descriptor
+python -c "import triton.tools.experimental_descriptor"

.github/workflows/_linux-test-h100.yml

@@ -0,0 +1,32 @@
+name: linux-test-h100
+on:
+  workflow_call:
+    inputs:
+      conda_env:
+        required: True
+        type: string
+        description: |
+          Conda environment to activate when testing Triton
+
+jobs:
+  linux-test-h100:
+    if: github.repository_owner == 'pytorch-labs'
+    runs-on: [gcp-h100-runner]
+    timeout-minutes: 240
+    environment: docker-s3-upload
+    env:
+      SETUP_SCRIPT: "/workspace/setup_instance.sh"
+      CONDA_ENV: ${{ inputs.conda_env }}
+    steps:
+      - name: Checkout Tritonbench
+        uses: actions/checkout@v3
+        with:
+          submodules: recursive
+      - name: Tune Nvidia GPU
+        run: |
+          sudo nvidia-smi -pm 1
+          sudo ldconfig
+          nvidia-smi
+      - name: Test Tritonbench operators on H100 GPU
+        run: |
+          bash ./.ci/tritonbench/test-gpu.sh

.github/workflows/pr.yaml

@@ -7,28 +7,13 @@ on:
 
 jobs:
   h100-pytorch-test:
-    # Don't run on forked repos
-    if: github.repository_owner == 'pytorch-labs'
-    runs-on: [gcp-h100-runner]
-    timeout-minutes: 240
-    environment: docker-s3-upload
-    env:
-      CONDA_ENV: "pytorch"
-      SETUP_SCRIPT: "/workspace/setup_instance.sh"
-    steps:
-      - name: Checkout Tritonbench
-        uses: actions/checkout@v3
-        with:
-          # no need to checkout submodules recursively
-          submodules: true
-      - name: Tune Nvidia GPU
-        run: |
-          sudo nvidia-smi -pm 1
-          sudo ldconfig
-          nvidia-smi
-      - name: Test Tritonbench operators on H100 GPU
-        run: |
-          bash ./.ci/tritonbench/test-gpu.sh
+    uses: ./.github/workflows/_linux-test-h100.yml
+    with:
+      conda_env: "pytorch"
+  h100-triton-main-test:
+    uses: ./.github/workflows/_linux-test-h100.yml
+    with:
+      conda_env: "triton-main"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}