【2/N】add support wmma kernels for RDNA4(GFX1201)

.github/runner-config.yml

@@ -17,3 +17,7 @@ runners:
   linux-flydsl-mi355-8:
     gpu_arch: MI355
     gpu_count: 8
+
+  linux-flydsl-navi-2:
+    gpu_arch: gfx1201
+    gpu_count: 4

.github/workflows/flydsl.yaml

@@ -22,7 +22,7 @@ jobs:
   test:
     strategy:
       matrix:
-        runners: [ 'linux-flydsl-mi325-1', 'linux-flydsl-mi355-1' ]
+        runners: [ 'linux-flydsl-mi325-1', 'linux-flydsl-mi355-1', 'linux-flydsl-navi-2' ]
       fail-fast: false
     runs-on: ${{ matrix.runners }}
     steps:
@@ -74,7 +74,7 @@ jobs:
           path: mlir_install.tgz
           # Key includes LLVM commit and hashes of build scripts/workflow.
           # Repo is checked out under `flydsl-test/` (see actions/checkout path), so hash paths must include it.
-          key: mlir-install-${{ hashFiles('flydsl-test/thirdparty/llvm-hash.txt', 'flydsl-test/scripts/build_llvm.sh', 'flydsl-test/CMakeLists.txt', 'flydsl-test/.github/workflows/flydsl.yaml') }}
+          key: mlir-install-${{ matrix.runners }}-${{ hashFiles('flydsl-test/thirdparty/llvm-hash.txt', 'flydsl-test/scripts/build_llvm.sh', 'flydsl-test/CMakeLists.txt', 'flydsl-test/.github/workflows/flydsl.yaml') }}
 
       - name: Use cached MLIR install tarball (skip LLVM build)
         if: steps.mlir-cache.outputs.cache-hit == 'true'
@@ -127,6 +127,7 @@ jobs:
           retention-days: 7
 
       - name: Install aiter
+        if: ${{ !contains(matrix.runners, 'navi') }}
         run: |
           docker exec flydsl_test bash -c "git clone --depth 1 --recursive --shallow-submodules https://github.com/ROCm/aiter.git /tmp/aiter && cd /tmp/aiter && python3 setup.py develop"
 

kernels/layernorm_kernel.py

@@ -29,8 +29,11 @@
 
 EPS = 1e-5
 
+import math
+from kernels.kernels_common import get_warp_size
+
 BLOCK_THREADS = 256
-WARP_SIZE = 64
+WARP_SIZE = get_warp_size()
 VEC_WIDTH = 8
 USE_NONTEMPORAL = True
 VEC_ALIGN = 16
@@ -91,8 +94,8 @@ def layernorm_kernel(
         def wave_reduce_add(x):
             width_i32 = fx.Int32(WARP_SIZE)
             w = x
-            for sh in [32, 16, 8, 4, 2, 1]:
-                off = fx.Int32(sh)
+            for _sh_exp in range_constexpr(int(math.log2(WARP_SIZE))):
+                off = fx.Int32(WARP_SIZE // (2 << _sh_exp))
                 peer = w.shuffle_xor(off, width_i32)
                 w = w.addf(peer, fastmath=fm_fast)
             return w