Add docker build

.github/workflows/build-tritonbench-docker.yml

@@ -8,6 +8,9 @@ on:
     - cron: '0 15 * * *'
   workflow_dispatch:
     inputs:
+      tritonbench_branch:
+        description: "Tritonbench branch name"
+        required: false
       nightly_date:
         description: "PyTorch nightly version"
         required: false
@@ -26,7 +29,8 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v3
         with:
-          path: benchmark
+          repository: pytorch-labs/tritonbench
+          path: tritonbench
       - name: Login to GitHub Container Registry
         if: ${{ env.WITH_PUSH == 'true' }}
         uses: docker/login-action@v2
@@ -38,10 +42,9 @@ jobs:
         run: |
           set -x
           export NIGHTLY_DATE="${{ github.event.inputs.nightly_date }}"
-          cd benchmark/docker
-          full_ref="${{ github.ref }}"
-          prefix="refs/heads/"
-          branch_name=${full_ref#$prefix}
+          cd tritonbench/docker
+          assigned_branch="${{ github.event.inputs.tritonbench_branch }}"
+          branch_name=${branch_name:-main}
           docker build . --build-arg TRITONBENCH_BRANCH="${branch_name}" --build-arg FORCE_DATE="${NIGHTLY_DATE}" \
               -f tritonbench-nightly.dockerfile -t ghcr.io/pytorch/tritonbench:latest
           # Extract pytorch version from the docker