From 45758022ee3f3711b74eff5bc1c5ee3507b33812 Mon Sep 17 00:00:00 2001
From: Illia Silin <98187287+illsilin@users.noreply.github.com>
Date: Fri, 8 Aug 2025 09:30:46 -0700
Subject: [PATCH] Add daily AITER tests on gfx942. (#2639)

* add option to select aiter branch, add tests on gfx942

[ROCm/composable_kernel commit: 7ac850ac72996ef462baf7c23efd9cfa3ba96bfe]
---
 Dockerfile.aiter | 10 +++++++---
 Jenkinsfile      | 20 ++++++++++++++------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/Dockerfile.aiter b/Dockerfile.aiter
index f6e66f460a..245e39fb75 100644
--- a/Dockerfile.aiter
+++ b/Dockerfile.aiter
@@ -1,16 +1,20 @@
 ARG BASE_DOCKER="rocm/pytorch:latest"
 FROM $BASE_DOCKER
-RUN groupadd -f render && \
+ARG AITER_BRANCH="main"
+ARG CK_AITER_BRANCH="develop"
+RUN groupadd -g 109 render && \
+    usermod -u 1001 jenkins && \
+    groupmod -g 1001 jenkins && \
     pip install pandas zmq einops && \
     pip install numpy==1.26.2 && \
     sudo mkdir /home/jenkins && \
     sudo mkdir /home/jenkins/workspace && \
     cd /home/jenkins/workspace && \
     rm -rf aiter && \
-    git clone --recursive https://github.com/ROCm/aiter.git && \
+    git clone -b "$AITER_BRANCH" --recursive https://github.com/ROCm/aiter.git && \
     cd aiter && \
     rm -rf 3rdparty/composable_kernel/ && \
-    git clone https://github.com/ROCm/composable_kernel.git 3rdparty/composable_kernel/ && \
+    git clone -b "$CK_AITER_BRANCH" https://github.com/ROCm/composable_kernel.git 3rdparty/composable_kernel/ && \
     python3 setup.py develop && \
     chown -R jenkins:jenkins /home/jenkins/workspace && \
     chmod -R a+rwx /home/jenkins/workspace && \
diff --git a/Jenkinsfile b/Jenkinsfile
index 0363b07d89..28e9f6dd00 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -190,7 +190,7 @@ def buildDocker(install_prefix){
     }
     else if(params.RUN_AITER_TESTS){
         image_name = "rocm/composable_kernel:ck_aiter"
-        dockerArgs = dockerArgs + " --no-cache -f Dockerfile.aiter . "
+        dockerArgs = dockerArgs + " --no-cache -f Dockerfile.aiter --build-arg AITER_BRANCH='${params.aiter_branch}' --build-arg CK_AITER_BRANCH='${params.ck_aiter_branch}' . "
     }
     else{
         dockerArgs = dockerArgs + " -f Dockerfile . "
@@ -843,10 +843,10 @@ def run_aiter_tests(Map conf=[:]){
     withDockerContainer(image: image, args: dockerOpts) {
         timeout(time: 45, unit: 'MINUTES'){
             try{
-                sh "python3 --version"
                 sh "rocminfo"
-                sh "python3 ../aiter/op_tests/test_gemm_a8w8_blockscale.py"
-                //sh "python3 ../aiter/op_tests/test_mha.py"
+                sh "python3 --version"
+                sh "python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8.py"
+                sh "python3 /home/jenkins/workspace/aiter/op_tests/test_gemm_a8w8_blockscale.py"
             }
             catch(e){
                 echo "Throwing error exception while running AITER tests"
@@ -1009,6 +1009,14 @@ pipeline {
             name: "RUN_AITER_TESTS",
             defaultValue: false,
             description: "Run AITER tests with latest CK develop branch (default: OFF)")
+        string(
+            name: 'aiter_branch',
+            defaultValue: 'main',
+            description: 'Specify which branch of AITER to use (default: main)')
+        string(
+            name: 'ck_aiter_branch',
+            defaultValue: 'develop',
+            description: 'Specify which branch of CK to test with AITER (default: develop)')
     }
     environment{
         dbuser = "${dbuser}"
@@ -1093,13 +1101,13 @@ pipeline {
         {
             parallel
             {
-                stage("Run AITER Tests on gfx90a")
+                stage("Run AITER Tests on gfx942")
                 {
                     when {
                         beforeAgent true
                         expression { params.RUN_AITER_TESTS.toBoolean() }
                     }
-                    agent{ label rocmnode("gfx90a")}
+                    agent{ label rocmnode("gfx942")}
                     steps{
                         run_aiter_tests()
                         cleanWs()