update tests

Signed-off-by: Tim Gymnich <[email protected]>

Author: tgymnich

lit_tests/kernel/wave/mma.py

@@ -644,6 +644,8 @@ def mma(
     # CHECK:          func.func @mma
 
     ### global buffer is bound to %0, %1 and %2 : MK, NK, MN
+    # CHECK:        %[[C0:.*]] = arith.constant 0 : index
+
     # CHECK:        %[[SUBSPAN0:.*]] = stream.binding.subspan
     # CHECK:        %[[SUBSPAN1:.*]] = stream.binding.subspan
     # CHECK:        %[[SUBSPAN2:.*]] = stream.binding.subspan
@@ -671,30 +673,17 @@ def mma(
     # CHECK:        %[[MASK3:.*]] = arith.select %[[COND1]], %{{.*}}, %[[MASK2]] : index
     # CHECK:        %[[MASK4:.*]] = arith.select %[[COND0]], %{{.*}}, %[[MASK3]] : index
 
-    ### pack descriptors and invoke tensor load
-
-    # CHECK:        %[[TENSOR_DESC_0:.*]] = arith.select
-    # CHECK-NOT:    amdgpu.tensor_load_to_lds
-    # CHECK-NOT:    rocdl.s.wait.tensorcnt
-    # CHECK-NOT:    amdgpu.lds_barrier
-
-    ### get shared buffer pointer
-    # CHECK:        %[[CAST_4:.*]] = memref.reinterpret_cast %[[VIEW0]]
-    # CHECK:        %[[INT_PTR_2:.+]] = memref.extract_aligned_pointer_as_index %[[CAST_4]]
-    # CHECK:        %[[INT_PTR_2_CAST:.+]] = arith.index_cast %[[INT_PTR_2]] : index to i32
-    # CHECK:        %[[INT_PTR_2_CAST_ADDED:.+]] = arith.addi %[[INT_PTR_2_CAST]], %{{.*}} : i32
+    # CHECK:        %[[TENSOR_DESC_0:.*]] = amdgpu.make_dma_descriptor %[[DMA_BASE0:.+]]
 
-    ### pack descriptors and invoke tensor load
-    # CHECK:        %[[D1:.*]] = vector.from_elements %{{.*}}, %[[INT_PTR_2_CAST_ADDED]], %{{.*}}, %{{.*}} : vector<4xi32>
-    # CHECK:        %[[TENSOR_DESC_1:.*]] = vector.from_elements
+    # CHECK:        %[[DMA_BASE1:.+]] = amdgpu.make_dma_base {{.*}} : {{.*}} -> !amdgpu.tdm_base<f16>
+    # CHECK:        %[[TENSOR_DESC_1:.*]] = amdgpu.make_dma_descriptor %[[DMA_BASE1:.+]]
 
     # Fused descriptors
     # CHECK:        %[[SELECTED:.*]] = arith.cmpi eq, %{{.*}}, %[[C0]] : index
-    # CHECK:        %[[D_FUSED:.*]] = arith.select %[[SELECTED]], %[[D0]], %[[D1]] : vector<4xi32>
-    # CHECK:        %[[DESC_FUSED:.*]] = arith.select %[[SELECTED]], %[[TENSOR_DESC_0]], %[[TENSOR_DESC_1]] : vector<8xi32>
+    # CHECK:        %[[DESC_FUSED:.*]] = arith.select %[[SELECTED]], %[[TENSOR_DESC_0]], %[[TENSOR_DESC_1]] : !amdgpu.tdm_descriptor
 
     ### resource provider
-    # CHECK:        amdgpu.tensor_load_to_lds  %[[D_FUSED]], %[[DESC_FUSED]], {{.*}}, {{.*}} cachepolicy {{.*}} : vector<4xi32>, vector<8xi32>
+    # CHECK:        amdgpu.tensor_load_to_lds %[[DESC_FUSED:.*]]
     # CHECK:        rocdl.s.wait.tensorcnt 0
     # CHECK:        rocdl.s.wait.dscnt 0
     # CHECK:        rocdl.s.barrier.signal id = -1