Tensor count insertion pass

lit_tests/kernel/wave/codegen.py

@@ -1176,7 +1176,7 @@ def schedule_ops(a: tkl.Memory[M, N, ADDRESS_SPACE, tkl.f16]):
     print(schedule_ops.asm)
 
     # CHECK-LABEL:    func.func @schedule_ops
-    # CHECK:            rocdl.s.wait.tensorcnt 0
+    # CHECK:            amdgpu.memory_counter_wait tensor(0)
     # CHECK:            rocdl.s.wait.dscnt 0
     # CHECK:            rocdl.s.barrier.signal id = -1
     # CHECK:            rocdl.s.barrier.wait id = -1

lit_tests/kernel/wave/gemm.py

@@ -1421,7 +1421,7 @@ def test_gemm_four_stage_global_to_lds():
     # Verify prologue stores to shared memory
     # CHECK: amdgpu.tensor_load_to_lds
 
-    # CHECK: rocdl.s.wait.tensorcnt 0
+    # CHECK: amdgpu.memory_counter_wait tensor(0)
     # CHECK: rocdl.s.wait.dscnt 0
     # CHECK: rocdl.s.barrier.signal id = -1
     # CHECK: rocdl.s.barrier.wait id = -1
@@ -1440,7 +1440,7 @@ def test_gemm_four_stage_global_to_lds():
     # Verify WMMA exists
     # CHECK: rocdl.wmma.f32.16x16x32.f16 %{{.*}}, %{{.*}}, %{{.*}}
 
-    # CHECK: rocdl.s.wait.tensorcnt 0
+    # CHECK: amdgpu.memory_counter_wait tensor(0)
     # CHECK: rocdl.s.wait.dscnt 0
     # CHECK: rocdl.s.barrier.signal id = -1
     # CHECK: rocdl.s.barrier.wait id = -1
@@ -1459,7 +1459,7 @@ def test_gemm_four_stage_global_to_lds():
     # Epilogue:
     # CHECK: rocdl.wmma.f32.16x16x32.f16 %{{.*}}, %{{.*}}, %{{.*}}
 
-    # CHECK: rocdl.s.wait.tensorcnt 0
+    # CHECK: amdgpu.memory_counter_wait tensor(0)
     # CHECK: rocdl.s.wait.dscnt 0
     # CHECK: rocdl.s.barrier.signal id = -1
     # CHECK: rocdl.s.barrier.wait id = -1

lit_tests/kernel/wave/mma.py

@@ -684,7 +684,7 @@ def mma(
 
     ### resource provider
     # CHECK:        amdgpu.tensor_load_to_lds %[[DESC_FUSED:.*]]
-    # CHECK:        rocdl.s.wait.tensorcnt 0
+    # CHECK:        amdgpu.memory_counter_wait tensor(0)
     # CHECK:        rocdl.s.wait.dscnt 0
     # CHECK:        rocdl.s.barrier.signal id = -1
 

tests/kernel/wave_gemm_test.py

@@ -2874,7 +2874,7 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
     asm = gemm.asm
 
     assert (
-        "wait.tensorcnt" in asm
+        "memory_counter_wait tensor" in asm
     ), "tensor waitcnts are not found in asm: required for tensor load instructions."
 
     validate_gemm_result(a, b, c, options)
@@ -3385,7 +3385,10 @@ def testSpecializeGemm(
 @require_gfx1250
 @pytest.mark.parametrize("shape", [(1024, 1024, 1024)])
 @pytest.mark.parametrize("mfma_variant", [MMAType.GFX1250_F32_16x16x32_F16])
-def test_gfx1250_tbuf_gemm(shape: tuple[int], mfma_variant: MMAType):
+@use_water_backend_bool("use_water_backend")
+def test_gfx1250_tbuf_gemm(
+    shape: tuple[int, int, int], mfma_variant: MMAType, use_water_backend: bool
+):
     gemm, options = get_tagged_BxA_T_gemm(
         shape=shape,
         block_shape=(256, 256, 64),
@@ -3395,8 +3398,11 @@ def test_gfx1250_tbuf_gemm(shape: tuple[int], mfma_variant: MMAType):
         compile_to_mlir=False,
     )
 
-    schedule = get_gfx1250_tbuf_gemm_schedule()
+    schedule = get_gfx1250_tbuf_gemm_schedule(
+        insert_tensor_waitcount=not use_water_backend
+    )
     options = set_default_run_config(options)
+    options.use_water_backend = use_water_backend
     gemm = wave_compile(options, gemm, schedule)
 
     a = device_randn(shape[0], shape[2], dtype=torch.float16)
@@ -3420,7 +3426,9 @@ def test_gfx1250_tbuf_gemm_codegen(use_water_backend: bool, tmp_path: Path):
         compile_to_mlir=False,
     )
 
-    schedule = get_gfx1250_tbuf_gemm_schedule()
+    schedule = get_gfx1250_tbuf_gemm_schedule(
+        insert_tensor_waitcount=not use_water_backend
+    )
     options.target = "gfx1250"
     options.dump_intermediates = tmp_path
     options.use_water_backend = use_water_backend
@@ -3441,7 +3449,6 @@ def test_gfx1250_tbuf_gemm_codegen(use_water_backend: bool, tmp_path: Path):
             "s_wait_xcnt 0x0",
             "s_wait_kmcnt 0x0",
             "s_wait_tensorcnt 0x1",
-            "s_wait_tensorcnt 0x1",
             "s_wait_dscnt 0x0",
             "s_wait_tensorcnt 0x1",
             "s_wait_dscnt 0xe",

water/include/water/Transforms/Passes.td

@@ -169,4 +169,21 @@ def WaterMemrefDecompositionPass : Pass<"water-memref-decomposition"> {
   ];
 }
 
+def WaterInsertWaitcnt : Pass<"water-insert-waitcnt"> {
+  let summary = "Insert wait instructions for asynchronous memory operations";
+  let description = [{
+    This pass analyzes asynchronous memory operations and inserts appropriate
+    wait/synchronization instructions to ensure memory operations complete
+    before their results are used.
+
+    The pass tracks dependencies between async memory operations,
+    maintaining scoreboards to determine when waits are necessary. It handles:
+    - Read-after-write (RAW) dependencies
+    - Write-after-read (WAR) dependencies
+  }];
+  let dependentDialects = [
+    "::mlir::amdgpu::AMDGPUDialect",
+  ];
+}
+
 #endif // WATER_PASSES

water/lib/Transforms/CMakeLists.txt

@@ -8,6 +8,7 @@ add_mlir_dialect_library(MLIRWaterTransforms
   GPUToGPURuntime.cpp
   MemrefDecomposition.cpp
   SLPVectorizer.cpp
+  WaterInsertWaitcnt.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/include/water