From 297d27e1c871ae93788bbfcd81229a3b9713b04c Mon Sep 17 00:00:00 2001
From: Shintaro Iwasaki <shintaro.iwasaki.work@gmail.com>
Date: Wed, 14 Sep 2022 23:09:56 -0700
Subject: [PATCH] [Triton-MLIR] add GitHub CI runners (#655)

This PR is to add GitHub Actions runners to the CI for better coverage.
---
 .github/workflows/integration-tests.yml       | 32 ++++++++++++++++---
 .../Conversion/TritonGPUToLLVM/PtxAsmFormat.h |  2 +-
 include/triton/driver/dispatch.h              |  9 +++++-
 .../TritonGPUToLLVM/PtxAsmFormat.cpp          |  2 +-
 4 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml
index 0727bf420..ea7d3f191 100644
--- a/.github/workflows/integration-tests.yml
+++ b/.github/workflows/integration-tests.yml
@@ -4,13 +4,33 @@ on:
   workflow_dispatch:
   pull_request:
     branches:
+      - main
       - triton-mlir
 
 jobs:
 
-  Integration-Tests:
+  Runner-Preparation:
+    runs-on: ubuntu-latest
+    outputs:
+      matrix: ${{ steps.set-matrix.outputs.matrix }}
+    steps:
+    - name: Prepare runner matrix
+      id: set-matrix
+      run: |
+        if [ x"${{ github.repository }}" == x"openai/triton" ]; then
+          echo '::set-output name=matrix::[["self-hosted", "A10"], "macos-latest"]'
+        else
+          echo '::set-output name=matrix::["ubuntu-latest", "macos-latest"]'
+        fi
 
-    runs-on: [self-hosted, A10]
+  Integration-Tests:
+    needs: Runner-Preparation
+
+    runs-on: ${{ matrix.runner }}
+
+    strategy:
+      matrix:
+        runner: ${{fromJson(needs.Runner-Preparation.outputs.matrix)}}
 
     steps:
 
@@ -19,26 +39,29 @@ jobs:
 
       - name: Clear cache
         run: |
-          rm -r ~/.triton/cache/
-        continue-on-error: true
+          rm -rf ~/.triton/cache/
 
       - name: Check imports
+        if: ${{ matrix.runner != 'macos-latest' }}
         run: |
           pip install isort
           isort -c ./python || ( echo '::error title=Imports not sorted::Please run \"isort ./python\"' ; exit 1 )
 
       - name: Check python style
+        if: ${{ matrix.runner != 'macos-latest' }}
         run: |
           pip install autopep8
           autopep8 -a -r -d --exit-code ./python || ( echo '::error title=Style issues::Please run \"autopep8 -a -r -i ./python\"' ; exit 1 )
 
       - name: Check cpp style
+        if: ${{ matrix.runner != 'macos-latest' }}
         run: |
           sudo apt-get install -y clang-format
           find . -regex '.*\.\(cpp\|hpp\|h\|cc\)' -not -path "./python/build/*" -not -path "./include/triton/external/*" -print0 | xargs -0 -n1 clang-format -style=file --dry-run -Werror -i ||
           (echo '::error title=Style issues:: Please run `find . -regex ".*\.\(cpp\|hpp\|h\|cc\)" -not -path "./python/build/*" -not -path "./include/triton/external/*" -print0 | xargs -0 -n1 clang-format -style=file -i`' ; exit 1)
 
       - name: Flake8
+        if: ${{ matrix.runner != 'macos-latest' }}
         run: |
           pip install flake8
           flake8 --config ./python/setup.cfg ./python || ( echo '::error::Flake8 failed; see logs for errors.' ; exit 1 )
@@ -59,6 +82,7 @@ jobs:
           lit -v "$LIT_TEST_DIR"
 
       - name: Run python tests
+        if: ${{ matrix.runner[0] == 'self-hosted' }}
         run: |
           cd python/tests
           pytest
diff --git a/include/triton/Conversion/TritonGPUToLLVM/PtxAsmFormat.h b/include/triton/Conversion/TritonGPUToLLVM/PtxAsmFormat.h
index 93b1011c7..7cecd59da 100644
--- a/include/triton/Conversion/TritonGPUToLLVM/PtxAsmFormat.h
+++ b/include/triton/Conversion/TritonGPUToLLVM/PtxAsmFormat.h
@@ -123,7 +123,7 @@ struct PTXBuilder {
 
   Operand *newAddrOperand(mlir::Value addr, StringRef constraint, int off = 0);
 
-  llvm::SmallVector<Operand *> getAllArgs() const;
+  llvm::SmallVector<Operand *, 4> getAllArgs() const;
 
   llvm::SmallVector<Value, 4> getAllMLIRArgs() const;
 
diff --git a/include/triton/driver/dispatch.h b/include/triton/driver/dispatch.h
index 85fc2cbc9..de0fa403c 100644
--- a/include/triton/driver/dispatch.h
+++ b/include/triton/driver/dispatch.h
@@ -48,8 +48,15 @@ protected:
     initializer();
     if (cache == nullptr) {
       cache = dlsym(lib_h, name);
-      if (cache == 0)
+      if (cache == 0) {
+#ifdef __EXCEPTIONS
         throw std::runtime_error("dlsym unable to load function");
+#else
+        std::cerr << "Triton: dlsym unable to load function `" << name << "`"
+                  << std::endl;
+        std::abort();
+#endif
+      }
     }
     FunPtrT fptr;
     *reinterpret_cast<void **>(&fptr) = cache;
diff --git a/lib/Conversion/TritonGPUToLLVM/PtxAsmFormat.cpp b/lib/Conversion/TritonGPUToLLVM/PtxAsmFormat.cpp
index 938cf8ce8..1299a083f 100644
--- a/lib/Conversion/TritonGPUToLLVM/PtxAsmFormat.cpp
+++ b/lib/Conversion/TritonGPUToLLVM/PtxAsmFormat.cpp
@@ -66,7 +66,7 @@ llvm::SmallVector<Value, 4> PTXBuilder::getAllMLIRArgs() const {
   return res;
 }
 
-SmallVector<PTXBuilder::Operand *> PTXBuilder::getAllArgs() const {
+SmallVector<PTXBuilder::Operand *, 4> PTXBuilder::getAllArgs() const {
   llvm::SmallVector<Operand *, 4> res;
   for (auto &x : argArchive)
     if (!x->isList())