diff --git a/examples/python/7.1_schedule.py b/examples/python/7.1_schedule.py index be229ee4da..4d6a7e7d0a 100755 --- a/examples/python/7.1_schedule.py +++ b/examples/python/7.1_schedule.py @@ -87,13 +87,41 @@ def _run_mxfp_gemm_preshuffle( x_scales_ps, w_scales_ps = x_scales_ps.cuda(), w_scales_ps.cuda() out = torch.zeros(x.shape[0], w_t_ps.shape[0], dtype=output_dtype).cuda() - gemm(x, x_scales_ps, w_t_ps, w_scales_ps, out) + for _ in range(100): + gemm(x, x_scales_ps, w_t_ps, w_scales_ps, out) torch.testing.assert_close( torch_out, out.cpu(), check_dtype=False, check_device=False ) +def _run_mxfp_gemm_preshuffle_transposed(gemm, shape, output_dtype=torch.float32): + """Run transposed GEMM: C^T[N,M] = B[N,K] * A[M,K]^T. + MFMA left operand ("A" role) = weight B. + MFMA right operand ("B" role, transposed by HW) = activation A. + Activation A is preshuffled (it's now in the "B" role). + The kernel stores C^T[N,M] in row-major layout. We transpose the + output before comparing with the reference C[M,N]. + """ + M_orig, N_orig, K = shape + x, w, x_scales, w_scales = generate_gemm_afp4wfp4_inputs(shape) + torch_out = torchScaledGemmMXFP4(x, w, x_scales, w_scales) + w_t = w.T.contiguous() # [N, K/2] - new "A" (left MFMA operand) + x_ps = b_preshuffle(x) # [M, K/2] - new "B" (preshuffled, right operand) + w_scales_ps = e8m0_shuffle(w_scales) # new "A" scale [N, K/32] + x_scales_ps = e8m0_shuffle(x_scales) # new "B" scale [M, K/32] + w_t, x_ps = w_t.cuda(), x_ps.cuda() + w_scales_ps, x_scales_ps = w_scales_ps.cuda(), x_scales_ps.cuda() + + out = torch.zeros(M_orig, N_orig, dtype=output_dtype).cuda() + + for _ in range(100): + gemm(w_t, w_scales_ps, x_ps, x_scales_ps, out) + torch.testing.assert_close( + torch_out, out.cpu(), check_dtype=False, check_device=False + ) + + def _get_8wave_shape_from_block(block): """Choose an 8-wave shape (4x2 or 2x4) from block M/N dims. @@ -166,7 +194,7 @@ def test_dbuf_8wave_pingpong_mxfp_gemm( def test_dbuf_8wave_pingpong_mxfp_gemm_Bshuffle( - is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256), dynamic=False + is_debug=False, shape=(6144, 8192, 8192), block=(192, 256, 256), dynamic=False ): """Double-buffered MXFP4 GEMM, 8 waves, ping-pong with stagger. A&B scales are preshuffled and read from global memory directly to VGPRs. @@ -200,44 +228,4159 @@ def test_dbuf_8wave_pingpong_mxfp_gemm_Bshuffle( def test_dbuf_8wave_pingpong_mxfp_gemm_Bshuffle_lds( - is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256), dynamic=False + is_debug=False, shape=(1792, 5376, 4096), block=(256, 192, 256), dynamic=True ): """Double-buffered MXFP4 GEMM, 8 waves, ping-pong with stagger. A&B scales are preshuffled and read from global memory directly to VGPRs. B data is preshuffled and loaded to LDS (shared memory), not directly to VGPRs. A data is read from global memory directly to LDS. """ + + mlir_256x192 = """ + #map = affine_map<()[s0] -> (s0 ceildiv 256)> + #map1 = affine_map<()[s0] -> (s0 ceildiv 192)> + #map2 = affine_map<()[s0] -> (s0 floordiv 32)> + #map3 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64) floordiv 32) * 256)> + #map4 = affine_map<()[s0] -> (s0 floordiv 2)> + #map5 = affine_map<()[s0, s1, s2] -> (s1 * 32 + s2 * 256 + s0 floordiv 8 - ((s1 * 32 + s0 floordiv 8) floordiv 256) * 256)> + #map6 = affine_map<()[s0] -> ((s0 floordiv 8) mod 8)> + #map7 = affine_map<()[s0] -> (s0 mod 8)> + #map8 = affine_map<()[s0] -> (s0 * 16)> + #map9 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64 + 8) floordiv 32) * 256 + 64)> + #map10 = affine_map<()[s0, s1, s2] -> (s1 * 32 + s2 * 256 + s0 floordiv 8 - ((s1 * 32 + s0 floordiv 8 + 64) floordiv 256) * 256 + 64)> + #map11 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64 + 16) floordiv 32) * 256 + 128)> + #map12 = affine_map<()[s0, s1, s2] -> (s1 * 32 + s2 * 256 + s0 floordiv 8 - ((s1 * 32 + s0 floordiv 8 + 128) floordiv 256) * 256 + 128)> + #map13 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64 + 24) floordiv 32) * 256 + 192)> + #map14 = affine_map<()[s0, s1, s2] -> (s1 * 32 + s2 * 256 + s0 floordiv 8 - ((s1 * 32 + s0 floordiv 8 + 192) floordiv 256) * 256 + 192)> + #map15 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64) floordiv 24) * 192)> + #map16 = affine_map<()[s0, s1] -> (s0 * 256 + (s1 floordiv 8) * 16 - ((s1 floordiv 8) floordiv 16) * 256)> + #map17 = affine_map<()[s0, s1, s2, s3] -> (s0 + s1 * 192 + (((s3 * 32 + s2 floordiv 8) mod 192) floordiv 16) * 16)> + #map18 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64 + 8) floordiv 24) * 192 + 64)> + #map19 = affine_map<()[s0, s1, s2, s3] -> (s0 + s1 * 192 + s3 * 32 + ((s2 floordiv 8) floordiv 16) * 16 - ((s3 * 32 + s2 floordiv 8 + 64) floordiv 192) * 192 + 64)> + #map20 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64 + 16) floordiv 24) * 192 + 128)> + #map21 = affine_map<()[s0, s1, s2, s3] -> (s0 + s1 * 192 + s3 * 32 + ((s2 floordiv 8) floordiv 16) * 16 - ((s3 * 32 + s2 floordiv 8 + 128) floordiv 192) * 192 + 128)> + #map22 = affine_map<()[s0] -> (s0 * 256)> + #map23 = affine_map<()[s0] -> (s0 * 128 + ((s0 mod 64) floordiv 16) * 2048 - (s0 floordiv 16) * 2048)> + #map24 = affine_map<()[s0, s1] -> (s0 + (s1 floordiv 64) * 64)> + #map25 = affine_map<()[s0] -> (s0 * 128 - (s0 floordiv 16) * 2048 + ((s0 mod 64) floordiv 16) * 2048)> + #map26 = affine_map<()[s0, s1] -> (s0 + (s1 floordiv 64) * 64 + 32)> + #map27 = affine_map<()[s0] -> (s0 * 192)> + #map28 = affine_map<()[s0, s1] -> (s0 + s1 * 96)> + #map29 = affine_map<()[s0, s1] -> (s0 + s1 * 96 + 32)> + #map30 = affine_map<()[s0, s1] -> (s0 + s1 * 96 + 64)> + #map31 = affine_map<()[s0, s1] -> (s1 * 4 + s0 floordiv 64)> + #map32 = affine_map<()[s0] -> (s0 ceildiv 256 - 1)> + #map33 = affine_map<()[s0] -> ((s0 mod 64) floordiv 16)> + #map34 = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 16 + (s0 floordiv 64) * 8192 - (s0 floordiv 16) * 2048)> + #map35 = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 16 + (s0 floordiv 64) * 8192 - (s0 floordiv 16) * 2048 + 2048)> + #map36 = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 16 + (s0 floordiv 64) * 8192 - (s0 floordiv 16) * 2048 + 4096)> + #map37 = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 16 + (s0 floordiv 64) * 8192 - (s0 floordiv 16) * 2048 + 6144)> + #map38 = affine_map<()[s0, s1, s2] -> (s0 * 12288 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048)> + #map39 = affine_map<()[s0, s1, s2] -> (s0 * 12288 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 2048)> + #map40 = affine_map<()[s0, s1, s2] -> (s0 * 12288 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 4096)> + #map41 = affine_map<()[s0, s1, s2] -> (s0 * 12288 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 6144)> + #map42 = affine_map<()[s0, s1, s2] -> (s0 * 12288 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 8192)> + #map43 = affine_map<()[s0, s1, s2] -> (s0 * 12288 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 10240)> + #map44 = affine_map<()[s0] -> ((s0 mod 64) floordiv 16 + 4)> + #map45 = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 16 + 128)> + #map46 = affine_map<()[s0, s1, s2] -> (s0 * 2048 + s1 * 256 + (s2 floordiv 8) * 16 - ((s2 floordiv 8) floordiv 16) * 256 + 2048)> + #map47 = affine_map<()[s0, s1] -> (s0 * 8192 + s1 * 128 + ((s1 mod 64) floordiv 16) * 2048 - (s1 floordiv 16) * 2048 + 8192)> + #map48 = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 8192 - (s0 floordiv 16) * 2048 + ((s0 mod 64) floordiv 16) * 2048 + 8192)> + #map49 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4)> + #map50 = affine_map<()[s0, s1] -> (s0 + s1 * 96 - (s0 floordiv 16) * 16)> + #map51 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 1)> + #map52 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 2)> + #map53 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 3)> + #map54 = affine_map<()[s0, s1] -> (s0 + s1 * 96 - (s0 floordiv 16) * 16 + 16)> + #map55 = affine_map<()[s0, s1] -> (s0 + s1 * 96 - (s0 floordiv 16) * 16 + 32)> + #map56 = affine_map<()[s0, s1] -> (s0 + s1 * 96 - (s0 floordiv 16) * 16 + 48)> + #map57 = affine_map<()[s0, s1] -> (s0 + s1 * 96 - (s0 floordiv 16) * 16 + 64)> + #map58 = affine_map<()[s0, s1] -> (s0 + s1 * 96 - (s0 floordiv 16) * 16 + 80)> + #map59 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 16)> + #map60 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 17)> + #map61 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 18)> + #map62 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 19)> + #map63 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 32)> + #map64 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 33)> + #map65 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 34)> + #map66 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 35)> + #map67 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 48)> + #map68 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 49)> + #map69 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 50)> + #map70 = affine_map<()[s0] -> ((s0 floordiv 64) * 64 + ((s0 mod 64) floordiv 16) * 4 + 51)> + #translation = #iree_codegen.translation_info + module attributes {transform.with_named_sequence} { + stream.executable private @gemm { + stream.executable.export public @gemm workgroups(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index) -> (index, index, index) { + %c1 = arith.constant 1 : index + %0 = affine.apply #map()[%arg0] + %1 = affine.apply #map1()[%arg1] + stream.return %0, %1, %c1 : index, index, index + } + builtin.module { + func.func @gemm(%arg0: !stream.binding, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: !stream.binding, %arg4: !stream.binding, %arg5: index, %arg6: index, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: index, %arg12: index) attributes {translation_info = #translation} { + %c4_i32 = arith.constant 4 : i32 + %c4294967295_i64 = arith.constant 4294967295 : i64 + %c2147483645_i64 = arith.constant 2147483645 : i64 + %c0_i32 = arith.constant 0 : i32 + %c32_i64 = arith.constant 32 : i64 + %c2147483646_i64 = arith.constant 2147483646 : i64 + %cst = arith.constant dense<0.000000e+00> : vector<4xf32> + %c2 = arith.constant 2 : index + %c1 = arith.constant 1 : index + %c0 = arith.constant 0 : index + %0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> memref + %1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> memref + %2 = stream.binding.subspan %arg2[%c0] : !stream.binding -> memref + %3 = stream.binding.subspan %arg3[%c0] : !stream.binding -> memref + %4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> memref + %block_id_x = gpu.block_id x upper_bound 2147483647 + %block_id_y = gpu.block_id y upper_bound 2147483647 + %thread_id_x = gpu.thread_id x upper_bound 256 + %thread_id_y = gpu.thread_id y upper_bound 2 + %5 = affine.apply #map2()[%arg7] + %reinterpret_cast = memref.reinterpret_cast %1 to offset: [0], sizes: [%arg5, %5], strides: [%arg9, 1] : memref to memref> + %reinterpret_cast_0 = memref.reinterpret_cast %3 to offset: [0], sizes: [%arg6, %5], strides: [%arg11, 1] : memref to memref> + %reinterpret_cast_1 = memref.reinterpret_cast %4 to offset: [0], sizes: [%arg5, %arg6], strides: [%arg12, 1] : memref to memref> + %alloc = memref.alloc() : memref<192x128xi8, #gpu.address_space> + %alloc_2 = memref.alloc() : memref<192x128xi8, #gpu.address_space> + %alloc_3 = memref.alloc() : memref<256x128xi8, #gpu.address_space> + %alloc_4 = memref.alloc() : memref<256x128xi8, #gpu.address_space> + %6 = affine.apply #map3()[%thread_id_x, %thread_id_y] + %7 = gpu.subgroup_broadcast %6, first_active_lane : index + %8 = gpu.subgroup_broadcast %c0, first_active_lane : index + %9 = affine.apply #map4()[%arg7] + %10 = affine.apply #map5()[%thread_id_x, %thread_id_y, %block_id_x] + %11 = affine.apply #map6()[%thread_id_x] + %12 = affine.apply #map7()[%thread_id_x] + %13 = arith.xori %12, %11 : index + %14 = affine.apply #map8()[%13] + %15 = arith.muli %10, %9 overflow : index + %16 = arith.addi %15, %14 overflow : index + %reinterpret_cast_5 = memref.reinterpret_cast %0 to offset: [0], sizes: [2147483646], strides: [1] : memref to memref<2147483646xi8, strided<[1]>> + %cast = memref.cast %reinterpret_cast_5 : memref<2147483646xi8, strided<[1]>> to memref> + %17 = arith.index_cast %9 : index to i14 + %18 = amdgpu.fat_raw_buffer_cast %cast validBytes(%c2147483646_i64) cacheSwizzleStride(%17) resetOffset : memref> to memref> + amdgpu.gather_to_lds %18[%16], %alloc_4[%7, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %19 = affine.apply #map9()[%thread_id_x, %thread_id_y] + %20 = gpu.subgroup_broadcast %19, first_active_lane : index + %21 = affine.apply #map10()[%thread_id_x, %thread_id_y, %block_id_x] + %22 = arith.muli %21, %9 overflow : index + %23 = arith.addi %22, %14 overflow : index + amdgpu.gather_to_lds %18[%23], %alloc_4[%20, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %24 = affine.apply #map11()[%thread_id_x, %thread_id_y] + %25 = gpu.subgroup_broadcast %24, first_active_lane : index + %26 = affine.apply #map12()[%thread_id_x, %thread_id_y, %block_id_x] + %27 = arith.muli %26, %9 overflow : index + %28 = arith.addi %27, %14 overflow : index + amdgpu.gather_to_lds %18[%28], %alloc_4[%25, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %29 = affine.apply #map13()[%thread_id_x, %thread_id_y] + %30 = gpu.subgroup_broadcast %29, first_active_lane : index + %31 = affine.apply #map14()[%thread_id_x, %thread_id_y, %block_id_x] + %32 = arith.muli %31, %9 overflow : index + %33 = arith.addi %32, %14 overflow : index + amdgpu.gather_to_lds %18[%33], %alloc_4[%30, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %34 = affine.apply #map15()[%thread_id_x, %thread_id_y] + %35 = gpu.subgroup_broadcast %34, first_active_lane : index + %36 = affine.apply #map16()[%13, %thread_id_x] + %37 = arith.index_cast %9 : index to i32 + %38 = arith.extui %37 : i32 to i64 + %39 = arith.addi %38, %c4294967295_i64 : i64 + %40 = arith.divui %39, %38 : i64 + %41 = arith.trunci %40 : i64 to i32 + %42 = arith.index_cast %36 : index to i32 + %43 = arith.extui %42 : i32 to i64 + %44 = arith.extui %41 : i32 to i64 + %45 = arith.muli %43, %44 : i64 + %46 = arith.shrui %45, %c32_i64 : i64 + %47 = arith.trunci %46 : i64 to i32 + %48 = arith.muli %47, %37 : i32 + %49 = arith.subi %42, %48 : i32 + %50 = arith.cmpi uge, %49, %37 : i32 + %51 = arith.extui %50 : i1 to i32 + %52 = arith.subi %47, %51 : i32 + %53 = arith.index_cast %52 : i32 to index + %54 = affine.apply #map17()[%53, %block_id_y, %thread_id_x, %thread_id_y] + %55 = arith.select %50, %37, %c0_i32 : i32 + %56 = arith.addi %49, %55 : i32 + %57 = arith.index_cast %56 : i32 to index + %58 = arith.muli %54, %9 overflow : index + %59 = arith.addi %58, %57 overflow : index + %reinterpret_cast_6 = memref.reinterpret_cast %2 to offset: [0], sizes: [2147483646], strides: [1] : memref to memref<2147483646xi8, strided<[1]>> + %cast_7 = memref.cast %reinterpret_cast_6 : memref<2147483646xi8, strided<[1]>> to memref> + %60 = amdgpu.fat_raw_buffer_cast %cast_7 validBytes(%c2147483646_i64) cacheSwizzleStride(%17) resetOffset : memref> to memref> + amdgpu.gather_to_lds %60[%59], %alloc_2[%35, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %61 = affine.apply #map18()[%thread_id_x, %thread_id_y] + %62 = gpu.subgroup_broadcast %61, first_active_lane : index + %63 = affine.apply #map19()[%53, %block_id_y, %thread_id_x, %thread_id_y] + %64 = arith.muli %63, %9 overflow : index + %65 = arith.addi %64, %57 overflow : index + amdgpu.gather_to_lds %60[%65], %alloc_2[%62, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %66 = affine.apply #map20()[%thread_id_x, %thread_id_y] + %67 = gpu.subgroup_broadcast %66, first_active_lane : index + %68 = affine.apply #map21()[%53, %block_id_y, %thread_id_x, %thread_id_y] + %69 = arith.muli %68, %9 overflow : index + %70 = arith.addi %69, %57 overflow : index + amdgpu.gather_to_lds %60[%70], %alloc_2[%67, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %71 = arith.index_cast %arg7 : index to i32 + %72 = arith.extui %71 : i32 to i64 + %73 = arith.addi %72, %c4294967295_i64 : i64 + %74 = arith.divui %73, %72 : i64 + %75 = arith.trunci %74 : i64 to i32 + %76 = affine.apply #map22()[%block_id_x] + %77 = affine.apply #map23()[%thread_id_x] + %78 = arith.index_cast %77 : index to i32 + %79 = arith.extui %78 : i32 to i64 + %80 = arith.extui %75 : i32 to i64 + %81 = arith.muli %79, %80 : i64 + %82 = arith.shrui %81, %c32_i64 : i64 + %83 = arith.trunci %82 : i64 to i32 + %84 = arith.muli %83, %71 : i32 + %85 = arith.subi %78, %84 : i32 + %86 = arith.cmpi uge, %85, %71 : i32 + %87 = arith.extui %86 : i1 to i32 + %88 = arith.subi %83, %87 : i32 + %89 = arith.index_cast %88 : i32 to index + %90 = affine.apply #map24()[%89, %thread_id_x] + %91 = affine.apply #map25()[%thread_id_x] + %92 = arith.index_cast %91 : index to i32 + %93 = arith.extui %92 : i32 to i64 + %94 = arith.muli %93, %80 : i64 + %95 = arith.shrui %94, %c32_i64 : i64 + %96 = arith.trunci %95 : i64 to i32 + %97 = arith.muli %96, %71 : i32 + %98 = arith.subi %92, %97 : i32 + %99 = arith.cmpi uge, %98, %71 : i32 + %100 = arith.select %99, %71, %c0_i32 : i32 + %101 = arith.addi %98, %100 : i32 + %102 = arith.index_cast %101 : i32 to index + %103 = affine.apply #map2()[%102] + %base_buffer, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %reinterpret_cast : memref> -> memref, index, index, index, index, index + %104 = arith.muli %76, %strides#0 overflow : index + %105 = arith.muli %90, %strides#0 overflow : index + %106 = arith.addi %105, %103 overflow : index + %reinterpret_cast_8 = memref.reinterpret_cast %1 to offset: [%104], sizes: [2147483646], strides: [1] : memref to memref<2147483646xi8, strided<[1], offset: ?>> + %cast_9 = memref.cast %reinterpret_cast_8 : memref<2147483646xi8, strided<[1], offset: ?>> to memref> + %107 = arith.index_cast %strides#0 : index to i14 + %108 = amdgpu.fat_raw_buffer_cast %cast_9 validBytes(%c2147483646_i64) cacheSwizzleStride(%107) resetOffset : memref> to memref> + %109 = vector.load %108[%106] : memref>, vector<4xi8> + %110 = affine.apply #map26()[%89, %thread_id_x] + %111 = arith.muli %110, %strides#0 overflow : index + %112 = arith.addi %111, %103 overflow : index + %113 = vector.load %108[%112] : memref>, vector<4xi8> + %114 = affine.apply #map27()[%block_id_y] + %115 = affine.apply #map28()[%89, %thread_id_y] + %base_buffer_10, %offset_11, %sizes_12:2, %strides_13:2 = memref.extract_strided_metadata %reinterpret_cast_0 : memref> -> memref, index, index, index, index, index + %116 = arith.muli %114, %strides_13#0 overflow : index + %117 = arith.muli %115, %strides_13#0 overflow : index + %118 = arith.addi %117, %103 overflow : index + %reinterpret_cast_14 = memref.reinterpret_cast %3 to offset: [%116], sizes: [2147483646], strides: [1] : memref to memref<2147483646xi8, strided<[1], offset: ?>> + %cast_15 = memref.cast %reinterpret_cast_14 : memref<2147483646xi8, strided<[1], offset: ?>> to memref> + %119 = arith.index_cast %strides_13#0 : index to i14 + %120 = amdgpu.fat_raw_buffer_cast %cast_15 validBytes(%c2147483646_i64) cacheSwizzleStride(%119) resetOffset : memref> to memref> + %121 = vector.load %120[%118] : memref>, vector<4xi8> + %122 = affine.apply #map29()[%89, %thread_id_y] + %123 = arith.muli %122, %strides_13#0 overflow : index + %124 = arith.addi %123, %103 overflow : index + %125 = vector.load %120[%124] : memref>, vector<4xi8> + %126 = affine.apply #map30()[%89, %thread_id_y] + %127 = arith.muli %126, %strides_13#0 overflow : index + %128 = arith.addi %127, %103 overflow : index + %129 = vector.load %120[%128] : memref>, vector<4xi8> + amdgpu.memory_counter_wait load(5) + rocdl.s.barrier + %130 = affine.apply #map31()[%thread_id_x, %thread_id_y] + %131 = arith.index_cast %130 : index to i32 + %132 = arith.cmpi sge, %131, %c4_i32 : i32 + %133 = arith.cmpi slt, %131, %c4_i32 : i32 + scf.if %132 { + rocdl.s.barrier + } + %134 = affine.apply #map32()[%arg7] + %135 = arith.remsi %134, %c2 : index + %136 = arith.subi %134, %135 : index + %reinterpret_cast_16 = memref.reinterpret_cast %alloc_4 to offset: [0], sizes: [32768], strides: [1] : memref<256x128xi8, #gpu.address_space> to memref<32768xi8, #gpu.address_space> + %137 = affine.apply #map33()[%thread_id_x] + %138 = arith.xori %137, %12 : index + %139 = affine.apply #map34()[%thread_id_x, %138] + %140 = affine.apply #map35()[%thread_id_x, %138] + %141 = affine.apply #map36()[%thread_id_x, %138] + %142 = affine.apply #map37()[%thread_id_x, %138] + %reinterpret_cast_17 = memref.reinterpret_cast %alloc_2 to offset: [0], sizes: [24576], strides: [1] : memref<192x128xi8, #gpu.address_space> to memref<24576xi8, #gpu.address_space> + %143 = affine.apply #map38()[%thread_id_y, %thread_id_x, %138] + %144 = affine.apply #map39()[%thread_id_y, %thread_id_x, %138] + %145 = affine.apply #map40()[%thread_id_y, %thread_id_x, %138] + %146 = affine.apply #map41()[%thread_id_y, %thread_id_x, %138] + %147 = affine.apply #map42()[%thread_id_y, %thread_id_x, %138] + %148 = affine.apply #map43()[%thread_id_y, %thread_id_x, %138] + %149 = affine.apply #map44()[%thread_id_x] + %150 = arith.xori %149, %12 : index + %151 = affine.apply #map34()[%thread_id_x, %150] + %152 = affine.apply #map35()[%thread_id_x, %150] + %153 = affine.apply #map36()[%thread_id_x, %150] + %154 = affine.apply #map37()[%thread_id_x, %150] + %155 = affine.apply #map38()[%thread_id_y, %thread_id_x, %150] + %156 = affine.apply #map39()[%thread_id_y, %thread_id_x, %150] + %157 = affine.apply #map40()[%thread_id_y, %thread_id_x, %150] + %158 = affine.apply #map41()[%thread_id_y, %thread_id_x, %150] + %159 = affine.apply #map42()[%thread_id_y, %thread_id_x, %150] + %160 = affine.apply #map43()[%thread_id_y, %thread_id_x, %150] + %reinterpret_cast_18 = memref.reinterpret_cast %alloc_3 to offset: [0], sizes: [32768], strides: [1] : memref<256x128xi8, #gpu.address_space> to memref<32768xi8, #gpu.address_space> + %reinterpret_cast_19 = memref.reinterpret_cast %alloc to offset: [0], sizes: [24576], strides: [1] : memref<192x128xi8, #gpu.address_space> to memref<24576xi8, #gpu.address_space> + %161:29 = scf.for %arg13 = %c0 to %136 step %c2 iter_args(%arg14 = %cst, %arg15 = %cst, %arg16 = %cst, %arg17 = %cst, %arg18 = %cst, %arg19 = %cst, %arg20 = %cst, %arg21 = %cst, %arg22 = %cst, %arg23 = %cst, %arg24 = %cst, %arg25 = %cst, %arg26 = %cst, %arg27 = %cst, %arg28 = %cst, %arg29 = %cst, %arg30 = %cst, %arg31 = %cst, %arg32 = %cst, %arg33 = %cst, %arg34 = %cst, %arg35 = %cst, %arg36 = %cst, %arg37 = %cst, %arg38 = %109, %arg39 = %113, %arg40 = %121, %arg41 = %125, %arg42 = %129) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>) { + %514 = vector.bitcast %arg39 : vector<4xi8> to vector<4xf8E8M0FNU> + %515 = vector.bitcast %arg42 : vector<4xi8> to vector<4xf8E8M0FNU> + %516 = vector.bitcast %arg41 : vector<4xi8> to vector<4xf8E8M0FNU> + %517 = vector.bitcast %arg40 : vector<4xi8> to vector<4xf8E8M0FNU> + %518 = vector.bitcast %arg38 : vector<4xi8> to vector<4xf8E8M0FNU> + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(5) + rocdl.s.barrier + %519 = affine.apply #map45()[%arg13, %13] + %520 = arith.addi %15, %519 overflow : index + amdgpu.gather_to_lds %18[%520], %alloc_3[%7, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %521 = arith.addi %22, %519 overflow : index + amdgpu.gather_to_lds %18[%521], %alloc_3[%20, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %522 = arith.addi %27, %519 overflow : index + amdgpu.gather_to_lds %18[%522], %alloc_3[%25, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %523 = arith.addi %32, %519 overflow : index + amdgpu.gather_to_lds %18[%523], %alloc_3[%30, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %524 = affine.apply #map46()[%arg13, %13, %thread_id_x] + %525 = arith.index_cast %524 : index to i32 + %526 = arith.extui %525 : i32 to i64 + %527 = arith.muli %526, %44 : i64 + %528 = arith.shrui %527, %c32_i64 : i64 + %529 = arith.trunci %528 : i64 to i32 + %530 = arith.muli %529, %37 : i32 + %531 = arith.subi %525, %530 : i32 + %532 = arith.cmpi uge, %531, %37 : i32 + %533 = arith.extui %532 : i1 to i32 + %534 = arith.subi %529, %533 : i32 + %535 = arith.index_cast %534 : i32 to index + %536 = affine.apply #map17()[%535, %block_id_y, %thread_id_x, %thread_id_y] + %537 = arith.select %532, %37, %c0_i32 : i32 + %538 = arith.addi %531, %537 : i32 + %539 = arith.index_cast %538 : i32 to index + %540 = arith.muli %536, %9 overflow : index + %541 = arith.addi %540, %539 overflow : index + amdgpu.gather_to_lds %60[%541], %alloc[%35, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %542 = affine.apply #map19()[%535, %block_id_y, %thread_id_x, %thread_id_y] + %543 = arith.muli %542, %9 overflow : index + %544 = arith.addi %543, %539 overflow : index + amdgpu.gather_to_lds %60[%544], %alloc[%62, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %545 = affine.apply #map21()[%535, %block_id_y, %thread_id_x, %thread_id_y] + %546 = arith.muli %545, %9 overflow : index + %547 = arith.addi %546, %539 overflow : index + amdgpu.gather_to_lds %60[%547], %alloc[%67, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + rocdl.sched.barrier 0 + %548 = vector.load %reinterpret_cast_16[%139] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %549 = vector.load %reinterpret_cast_16[%140] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %550 = vector.load %reinterpret_cast_16[%141] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %551 = vector.load %reinterpret_cast_16[%142] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %552 = vector.load %reinterpret_cast_17[%143] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %553 = vector.load %reinterpret_cast_17[%144] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %554 = vector.load %reinterpret_cast_17[%145] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %555 = vector.load %reinterpret_cast_17[%146] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %556 = vector.load %reinterpret_cast_17[%147] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %557 = vector.load %reinterpret_cast_17[%148] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %558 = vector.bitcast %548 : vector<16xi8> to vector<32xf4E2M1FN> + %559 = vector.bitcast %549 : vector<16xi8> to vector<32xf4E2M1FN> + %560 = vector.bitcast %550 : vector<16xi8> to vector<32xf4E2M1FN> + %561 = vector.bitcast %551 : vector<16xi8> to vector<32xf4E2M1FN> + %562 = vector.bitcast %552 : vector<16xi8> to vector<32xf4E2M1FN> + %563 = vector.bitcast %553 : vector<16xi8> to vector<32xf4E2M1FN> + %564 = vector.bitcast %554 : vector<16xi8> to vector<32xf4E2M1FN> + %565 = vector.bitcast %555 : vector<16xi8> to vector<32xf4E2M1FN> + %566 = vector.bitcast %556 : vector<16xi8> to vector<32xf4E2M1FN> + %567 = vector.bitcast %557 : vector<16xi8> to vector<32xf4E2M1FN> + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %568 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%517[0] * %562) + %arg14 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %569 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%517[1] * %563) + %arg15 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %570 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%516[0] * %564) + %arg16 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %571 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%516[1] * %565) + %arg17 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %572 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%515[0] * %566) + %arg18 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %573 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%515[1] * %567) + %arg19 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %574 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%517[0] * %562) + %arg20 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %575 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%517[1] * %563) + %arg21 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %576 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%516[0] * %564) + %arg22 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %577 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%516[1] * %565) + %arg23 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %578 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%515[0] * %566) + %arg24 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %579 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%515[1] * %567) + %arg25 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %580 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%517[0] * %562) + %arg26 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %581 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%517[1] * %563) + %arg27 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %582 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%516[0] * %564) + %arg28 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %583 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%516[1] * %565) + %arg29 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %584 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%515[0] * %566) + %arg30 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %585 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%515[1] * %567) + %arg31 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %586 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%517[0] * %562) + %arg32 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %587 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%517[1] * %563) + %arg33 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %588 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%516[0] * %564) + %arg34 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %589 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%516[1] * %565) + %arg35 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %590 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%515[0] * %566) + %arg36 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %591 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%515[1] * %567) + %arg37 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.sched.barrier 0 + %592 = affine.apply #map47()[%arg13, %thread_id_x] + %593 = arith.index_cast %592 : index to i32 + %594 = arith.extui %593 : i32 to i64 + %595 = arith.muli %594, %80 : i64 + %596 = arith.shrui %595, %c32_i64 : i64 + %597 = arith.trunci %596 : i64 to i32 + %598 = arith.muli %597, %71 : i32 + %599 = arith.subi %593, %598 : i32 + %600 = arith.cmpi uge, %599, %71 : i32 + %601 = arith.extui %600 : i1 to i32 + %602 = arith.subi %597, %601 : i32 + %603 = arith.index_cast %602 : i32 to index + %604 = affine.apply #map24()[%603, %thread_id_x] + %605 = affine.apply #map48()[%thread_id_x, %arg13] + %606 = arith.index_cast %605 : index to i32 + %607 = arith.extui %606 : i32 to i64 + %608 = arith.muli %607, %80 : i64 + %609 = arith.shrui %608, %c32_i64 : i64 + %610 = arith.trunci %609 : i64 to i32 + %611 = arith.muli %610, %71 : i32 + %612 = arith.subi %606, %611 : i32 + %613 = arith.cmpi uge, %612, %71 : i32 + %614 = arith.select %613, %71, %c0_i32 : i32 + %615 = arith.addi %612, %614 : i32 + %616 = arith.index_cast %615 : i32 to index + %617 = affine.apply #map2()[%616] + %618 = arith.muli %604, %strides#0 overflow : index + %619 = arith.addi %618, %617 overflow : index + %620 = vector.load %108[%619] : memref>, vector<4xi8> + %621 = vector.bitcast %620 : vector<4xi8> to vector<4xf8E8M0FNU> + %622 = affine.apply #map26()[%603, %thread_id_x] + %623 = arith.muli %622, %strides#0 overflow : index + %624 = arith.addi %623, %617 overflow : index + %625 = vector.load %108[%624] : memref>, vector<4xi8> + %626 = vector.bitcast %625 : vector<4xi8> to vector<4xf8E8M0FNU> + %627 = affine.apply #map28()[%603, %thread_id_y] + %628 = arith.muli %627, %strides_13#0 overflow : index + %629 = arith.addi %628, %617 overflow : index + %630 = vector.load %120[%629] : memref>, vector<4xi8> + %631 = vector.bitcast %630 : vector<4xi8> to vector<4xf8E8M0FNU> + %632 = affine.apply #map29()[%603, %thread_id_y] + %633 = arith.muli %632, %strides_13#0 overflow : index + %634 = arith.addi %633, %617 overflow : index + %635 = vector.load %120[%634] : memref>, vector<4xi8> + %636 = vector.bitcast %635 : vector<4xi8> to vector<4xf8E8M0FNU> + %637 = affine.apply #map30()[%603, %thread_id_y] + %638 = arith.muli %637, %strides_13#0 overflow : index + %639 = arith.addi %638, %617 overflow : index + %640 = vector.load %120[%639] : memref>, vector<4xi8> + %641 = vector.bitcast %640 : vector<4xi8> to vector<4xf8E8M0FNU> + %642 = vector.load %reinterpret_cast_16[%151] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %643 = vector.load %reinterpret_cast_16[%152] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %644 = vector.load %reinterpret_cast_16[%153] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %645 = vector.load %reinterpret_cast_16[%154] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %646 = vector.load %reinterpret_cast_17[%155] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %647 = vector.load %reinterpret_cast_17[%156] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %648 = vector.load %reinterpret_cast_17[%157] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %649 = vector.load %reinterpret_cast_17[%158] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %650 = vector.load %reinterpret_cast_17[%159] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %651 = vector.load %reinterpret_cast_17[%160] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %652 = vector.bitcast %642 : vector<16xi8> to vector<32xf4E2M1FN> + %653 = vector.bitcast %643 : vector<16xi8> to vector<32xf4E2M1FN> + %654 = vector.bitcast %644 : vector<16xi8> to vector<32xf4E2M1FN> + %655 = vector.bitcast %645 : vector<16xi8> to vector<32xf4E2M1FN> + %656 = vector.bitcast %646 : vector<16xi8> to vector<32xf4E2M1FN> + %657 = vector.bitcast %647 : vector<16xi8> to vector<32xf4E2M1FN> + %658 = vector.bitcast %648 : vector<16xi8> to vector<32xf4E2M1FN> + %659 = vector.bitcast %649 : vector<16xi8> to vector<32xf4E2M1FN> + %660 = vector.bitcast %650 : vector<16xi8> to vector<32xf4E2M1FN> + %661 = vector.bitcast %651 : vector<16xi8> to vector<32xf4E2M1FN> + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(7) + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %662 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %652) * (%517[2] * %656) + %568 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %663 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %652) * (%517[3] * %657) + %569 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %664 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %652) * (%516[2] * %658) + %570 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %665 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %652) * (%516[3] * %659) + %571 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %666 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %652) * (%515[2] * %660) + %572 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %667 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %652) * (%515[3] * %661) + %573 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %668 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %653) * (%517[2] * %656) + %574 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %669 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %653) * (%517[3] * %657) + %575 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %670 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %653) * (%516[2] * %658) + %576 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %671 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %653) * (%516[3] * %659) + %577 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %672 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %653) * (%515[2] * %660) + %578 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %673 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %653) * (%515[3] * %661) + %579 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %674 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %654) * (%517[2] * %656) + %580 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %675 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %654) * (%517[3] * %657) + %581 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %676 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %654) * (%516[2] * %658) + %582 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %677 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %654) * (%516[3] * %659) + %583 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %678 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %654) * (%515[2] * %660) + %584 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %679 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %654) * (%515[3] * %661) + %585 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %680 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %655) * (%517[2] * %656) + %586 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %681 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %655) * (%517[3] * %657) + %587 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %682 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %655) * (%516[2] * %658) + %588 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %683 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %655) * (%516[3] * %659) + %589 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %684 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %655) * (%515[2] * %660) + %590 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %685 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %655) * (%515[3] * %661) + %591 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + %686 = arith.addi %arg13, %c1 : index + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(5) + rocdl.s.barrier + %687 = affine.apply #map45()[%686, %13] + %688 = arith.addi %15, %687 overflow : index + amdgpu.gather_to_lds %18[%688], %alloc_4[%7, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %689 = arith.addi %22, %687 overflow : index + amdgpu.gather_to_lds %18[%689], %alloc_4[%20, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %690 = arith.addi %27, %687 overflow : index + amdgpu.gather_to_lds %18[%690], %alloc_4[%25, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %691 = arith.addi %32, %687 overflow : index + amdgpu.gather_to_lds %18[%691], %alloc_4[%30, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %692 = affine.apply #map46()[%686, %13, %thread_id_x] + %693 = arith.index_cast %692 : index to i32 + %694 = arith.extui %693 : i32 to i64 + %695 = arith.muli %694, %44 : i64 + %696 = arith.shrui %695, %c32_i64 : i64 + %697 = arith.trunci %696 : i64 to i32 + %698 = arith.muli %697, %37 : i32 + %699 = arith.subi %693, %698 : i32 + %700 = arith.cmpi uge, %699, %37 : i32 + %701 = arith.extui %700 : i1 to i32 + %702 = arith.subi %697, %701 : i32 + %703 = arith.index_cast %702 : i32 to index + %704 = affine.apply #map17()[%703, %block_id_y, %thread_id_x, %thread_id_y] + %705 = arith.select %700, %37, %c0_i32 : i32 + %706 = arith.addi %699, %705 : i32 + %707 = arith.index_cast %706 : i32 to index + %708 = arith.muli %704, %9 overflow : index + %709 = arith.addi %708, %707 overflow : index + amdgpu.gather_to_lds %60[%709], %alloc_2[%35, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %710 = affine.apply #map19()[%703, %block_id_y, %thread_id_x, %thread_id_y] + %711 = arith.muli %710, %9 overflow : index + %712 = arith.addi %711, %707 overflow : index + amdgpu.gather_to_lds %60[%712], %alloc_2[%62, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %713 = affine.apply #map21()[%703, %block_id_y, %thread_id_x, %thread_id_y] + %714 = arith.muli %713, %9 overflow : index + %715 = arith.addi %714, %707 overflow : index + amdgpu.gather_to_lds %60[%715], %alloc_2[%67, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + rocdl.sched.barrier 0 + %716 = vector.load %reinterpret_cast_18[%139] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %717 = vector.load %reinterpret_cast_18[%140] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %718 = vector.load %reinterpret_cast_18[%141] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %719 = vector.load %reinterpret_cast_18[%142] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %720 = vector.load %reinterpret_cast_19[%143] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %721 = vector.load %reinterpret_cast_19[%144] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %722 = vector.load %reinterpret_cast_19[%145] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %723 = vector.load %reinterpret_cast_19[%146] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %724 = vector.load %reinterpret_cast_19[%147] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %725 = vector.load %reinterpret_cast_19[%148] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %726 = vector.bitcast %716 : vector<16xi8> to vector<32xf4E2M1FN> + %727 = vector.bitcast %717 : vector<16xi8> to vector<32xf4E2M1FN> + %728 = vector.bitcast %718 : vector<16xi8> to vector<32xf4E2M1FN> + %729 = vector.bitcast %719 : vector<16xi8> to vector<32xf4E2M1FN> + %730 = vector.bitcast %720 : vector<16xi8> to vector<32xf4E2M1FN> + %731 = vector.bitcast %721 : vector<16xi8> to vector<32xf4E2M1FN> + %732 = vector.bitcast %722 : vector<16xi8> to vector<32xf4E2M1FN> + %733 = vector.bitcast %723 : vector<16xi8> to vector<32xf4E2M1FN> + %734 = vector.bitcast %724 : vector<16xi8> to vector<32xf4E2M1FN> + %735 = vector.bitcast %725 : vector<16xi8> to vector<32xf4E2M1FN> + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %736 = amdgpu.scaled_mfma 16x16x128 (%621[0] * %726) * (%631[0] * %730) + %662 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %737 = amdgpu.scaled_mfma 16x16x128 (%621[0] * %726) * (%631[1] * %731) + %663 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %738 = amdgpu.scaled_mfma 16x16x128 (%621[0] * %726) * (%636[0] * %732) + %664 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %739 = amdgpu.scaled_mfma 16x16x128 (%621[0] * %726) * (%636[1] * %733) + %665 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %740 = amdgpu.scaled_mfma 16x16x128 (%621[0] * %726) * (%641[0] * %734) + %666 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %741 = amdgpu.scaled_mfma 16x16x128 (%621[0] * %726) * (%641[1] * %735) + %667 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %742 = amdgpu.scaled_mfma 16x16x128 (%621[1] * %727) * (%631[0] * %730) + %668 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %743 = amdgpu.scaled_mfma 16x16x128 (%621[1] * %727) * (%631[1] * %731) + %669 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %744 = amdgpu.scaled_mfma 16x16x128 (%621[1] * %727) * (%636[0] * %732) + %670 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %745 = amdgpu.scaled_mfma 16x16x128 (%621[1] * %727) * (%636[1] * %733) + %671 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %746 = amdgpu.scaled_mfma 16x16x128 (%621[1] * %727) * (%641[0] * %734) + %672 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %747 = amdgpu.scaled_mfma 16x16x128 (%621[1] * %727) * (%641[1] * %735) + %673 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %748 = amdgpu.scaled_mfma 16x16x128 (%626[0] * %728) * (%631[0] * %730) + %674 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %749 = amdgpu.scaled_mfma 16x16x128 (%626[0] * %728) * (%631[1] * %731) + %675 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %750 = amdgpu.scaled_mfma 16x16x128 (%626[0] * %728) * (%636[0] * %732) + %676 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %751 = amdgpu.scaled_mfma 16x16x128 (%626[0] * %728) * (%636[1] * %733) + %677 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %752 = amdgpu.scaled_mfma 16x16x128 (%626[0] * %728) * (%641[0] * %734) + %678 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %753 = amdgpu.scaled_mfma 16x16x128 (%626[0] * %728) * (%641[1] * %735) + %679 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %754 = amdgpu.scaled_mfma 16x16x128 (%626[1] * %729) * (%631[0] * %730) + %680 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %755 = amdgpu.scaled_mfma 16x16x128 (%626[1] * %729) * (%631[1] * %731) + %681 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %756 = amdgpu.scaled_mfma 16x16x128 (%626[1] * %729) * (%636[0] * %732) + %682 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %757 = amdgpu.scaled_mfma 16x16x128 (%626[1] * %729) * (%636[1] * %733) + %683 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %758 = amdgpu.scaled_mfma 16x16x128 (%626[1] * %729) * (%641[0] * %734) + %684 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %759 = amdgpu.scaled_mfma 16x16x128 (%626[1] * %729) * (%641[1] * %735) + %685 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.sched.barrier 0 + %760 = affine.apply #map47()[%686, %thread_id_x] + %761 = arith.index_cast %760 : index to i32 + %762 = arith.extui %761 : i32 to i64 + %763 = arith.muli %762, %80 : i64 + %764 = arith.shrui %763, %c32_i64 : i64 + %765 = arith.trunci %764 : i64 to i32 + %766 = arith.muli %765, %71 : i32 + %767 = arith.subi %761, %766 : i32 + %768 = arith.cmpi uge, %767, %71 : i32 + %769 = arith.extui %768 : i1 to i32 + %770 = arith.subi %765, %769 : i32 + %771 = arith.index_cast %770 : i32 to index + %772 = affine.apply #map24()[%771, %thread_id_x] + %773 = affine.apply #map48()[%thread_id_x, %686] + %774 = arith.index_cast %773 : index to i32 + %775 = arith.extui %774 : i32 to i64 + %776 = arith.muli %775, %80 : i64 + %777 = arith.shrui %776, %c32_i64 : i64 + %778 = arith.trunci %777 : i64 to i32 + %779 = arith.muli %778, %71 : i32 + %780 = arith.subi %774, %779 : i32 + %781 = arith.cmpi uge, %780, %71 : i32 + %782 = arith.select %781, %71, %c0_i32 : i32 + %783 = arith.addi %780, %782 : i32 + %784 = arith.index_cast %783 : i32 to index + %785 = affine.apply #map2()[%784] + %786 = arith.muli %772, %strides#0 overflow : index + %787 = arith.addi %786, %785 overflow : index + %788 = vector.load %108[%787] : memref>, vector<4xi8> + %789 = affine.apply #map26()[%771, %thread_id_x] + %790 = arith.muli %789, %strides#0 overflow : index + %791 = arith.addi %790, %785 overflow : index + %792 = vector.load %108[%791] : memref>, vector<4xi8> + %793 = affine.apply #map28()[%771, %thread_id_y] + %794 = arith.muli %793, %strides_13#0 overflow : index + %795 = arith.addi %794, %785 overflow : index + %796 = vector.load %120[%795] : memref>, vector<4xi8> + %797 = affine.apply #map29()[%771, %thread_id_y] + %798 = arith.muli %797, %strides_13#0 overflow : index + %799 = arith.addi %798, %785 overflow : index + %800 = vector.load %120[%799] : memref>, vector<4xi8> + %801 = affine.apply #map30()[%771, %thread_id_y] + %802 = arith.muli %801, %strides_13#0 overflow : index + %803 = arith.addi %802, %785 overflow : index + %804 = vector.load %120[%803] : memref>, vector<4xi8> + %805 = vector.load %reinterpret_cast_18[%151] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %806 = vector.load %reinterpret_cast_18[%152] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %807 = vector.load %reinterpret_cast_18[%153] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %808 = vector.load %reinterpret_cast_18[%154] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %809 = vector.load %reinterpret_cast_19[%155] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %810 = vector.load %reinterpret_cast_19[%156] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %811 = vector.load %reinterpret_cast_19[%157] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %812 = vector.load %reinterpret_cast_19[%158] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %813 = vector.load %reinterpret_cast_19[%159] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %814 = vector.load %reinterpret_cast_19[%160] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %815 = vector.bitcast %805 : vector<16xi8> to vector<32xf4E2M1FN> + %816 = vector.bitcast %806 : vector<16xi8> to vector<32xf4E2M1FN> + %817 = vector.bitcast %807 : vector<16xi8> to vector<32xf4E2M1FN> + %818 = vector.bitcast %808 : vector<16xi8> to vector<32xf4E2M1FN> + %819 = vector.bitcast %809 : vector<16xi8> to vector<32xf4E2M1FN> + %820 = vector.bitcast %810 : vector<16xi8> to vector<32xf4E2M1FN> + %821 = vector.bitcast %811 : vector<16xi8> to vector<32xf4E2M1FN> + %822 = vector.bitcast %812 : vector<16xi8> to vector<32xf4E2M1FN> + %823 = vector.bitcast %813 : vector<16xi8> to vector<32xf4E2M1FN> + %824 = vector.bitcast %814 : vector<16xi8> to vector<32xf4E2M1FN> + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(7) + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %825 = amdgpu.scaled_mfma 16x16x128 (%621[2] * %815) * (%631[2] * %819) + %736 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %826 = amdgpu.scaled_mfma 16x16x128 (%621[2] * %815) * (%631[3] * %820) + %737 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %827 = amdgpu.scaled_mfma 16x16x128 (%621[2] * %815) * (%636[2] * %821) + %738 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %828 = amdgpu.scaled_mfma 16x16x128 (%621[2] * %815) * (%636[3] * %822) + %739 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %829 = amdgpu.scaled_mfma 16x16x128 (%621[2] * %815) * (%641[2] * %823) + %740 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %830 = amdgpu.scaled_mfma 16x16x128 (%621[2] * %815) * (%641[3] * %824) + %741 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %831 = amdgpu.scaled_mfma 16x16x128 (%621[3] * %816) * (%631[2] * %819) + %742 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %832 = amdgpu.scaled_mfma 16x16x128 (%621[3] * %816) * (%631[3] * %820) + %743 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %833 = amdgpu.scaled_mfma 16x16x128 (%621[3] * %816) * (%636[2] * %821) + %744 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %834 = amdgpu.scaled_mfma 16x16x128 (%621[3] * %816) * (%636[3] * %822) + %745 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %835 = amdgpu.scaled_mfma 16x16x128 (%621[3] * %816) * (%641[2] * %823) + %746 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %836 = amdgpu.scaled_mfma 16x16x128 (%621[3] * %816) * (%641[3] * %824) + %747 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %837 = amdgpu.scaled_mfma 16x16x128 (%626[2] * %817) * (%631[2] * %819) + %748 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %838 = amdgpu.scaled_mfma 16x16x128 (%626[2] * %817) * (%631[3] * %820) + %749 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %839 = amdgpu.scaled_mfma 16x16x128 (%626[2] * %817) * (%636[2] * %821) + %750 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %840 = amdgpu.scaled_mfma 16x16x128 (%626[2] * %817) * (%636[3] * %822) + %751 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %841 = amdgpu.scaled_mfma 16x16x128 (%626[2] * %817) * (%641[2] * %823) + %752 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %842 = amdgpu.scaled_mfma 16x16x128 (%626[2] * %817) * (%641[3] * %824) + %753 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %843 = amdgpu.scaled_mfma 16x16x128 (%626[3] * %818) * (%631[2] * %819) + %754 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %844 = amdgpu.scaled_mfma 16x16x128 (%626[3] * %818) * (%631[3] * %820) + %755 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %845 = amdgpu.scaled_mfma 16x16x128 (%626[3] * %818) * (%636[2] * %821) + %756 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %846 = amdgpu.scaled_mfma 16x16x128 (%626[3] * %818) * (%636[3] * %822) + %757 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %847 = amdgpu.scaled_mfma 16x16x128 (%626[3] * %818) * (%641[2] * %823) + %758 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %848 = amdgpu.scaled_mfma 16x16x128 (%626[3] * %818) * (%641[3] * %824) + %759 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + scf.yield %825, %826, %827, %828, %829, %830, %831, %832, %833, %834, %835, %836, %837, %838, %839, %840, %841, %842, %843, %844, %845, %846, %847, %848, %788, %792, %796, %800, %804 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8> + } + %162:33 = scf.for %arg13 = %136 to %134 step %c1 iter_args(%arg14 = %161#0, %arg15 = %161#1, %arg16 = %161#2, %arg17 = %161#3, %arg18 = %161#4, %arg19 = %161#5, %arg20 = %161#6, %arg21 = %161#7, %arg22 = %161#8, %arg23 = %161#9, %arg24 = %161#10, %arg25 = %161#11, %arg26 = %161#12, %arg27 = %161#13, %arg28 = %161#14, %arg29 = %161#15, %arg30 = %161#16, %arg31 = %161#17, %arg32 = %161#18, %arg33 = %161#19, %arg34 = %161#20, %arg35 = %161#21, %arg36 = %161#22, %arg37 = %161#23, %arg38 = %161#24, %arg39 = %161#25, %arg40 = %161#26, %arg41 = %161#27, %arg42 = %161#28, %arg43 = %alloc_4, %arg44 = %alloc_3, %arg45 = %alloc_2, %arg46 = %alloc) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, memref<256x128xi8, #gpu.address_space>, memref<256x128xi8, #gpu.address_space>, memref<192x128xi8, #gpu.address_space>, memref<192x128xi8, #gpu.address_space>) { + %514 = vector.bitcast %arg39 : vector<4xi8> to vector<4xf8E8M0FNU> + %515 = vector.bitcast %arg42 : vector<4xi8> to vector<4xf8E8M0FNU> + %516 = vector.bitcast %arg41 : vector<4xi8> to vector<4xf8E8M0FNU> + %517 = vector.bitcast %arg40 : vector<4xi8> to vector<4xf8E8M0FNU> + %518 = vector.bitcast %arg38 : vector<4xi8> to vector<4xf8E8M0FNU> + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(5) + rocdl.s.barrier + %519 = affine.apply #map45()[%arg13, %13] + %520 = arith.addi %15, %519 overflow : index + amdgpu.gather_to_lds %18[%520], %arg44[%7, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %521 = arith.addi %22, %519 overflow : index + amdgpu.gather_to_lds %18[%521], %arg44[%20, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %522 = arith.addi %27, %519 overflow : index + amdgpu.gather_to_lds %18[%522], %arg44[%25, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %523 = arith.addi %32, %519 overflow : index + amdgpu.gather_to_lds %18[%523], %arg44[%30, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %524 = affine.apply #map46()[%arg13, %13, %thread_id_x] + %525 = arith.index_cast %524 : index to i32 + %526 = arith.extui %525 : i32 to i64 + %527 = arith.muli %526, %44 : i64 + %528 = arith.shrui %527, %c32_i64 : i64 + %529 = arith.trunci %528 : i64 to i32 + %530 = arith.muli %529, %37 : i32 + %531 = arith.subi %525, %530 : i32 + %532 = arith.cmpi uge, %531, %37 : i32 + %533 = arith.extui %532 : i1 to i32 + %534 = arith.subi %529, %533 : i32 + %535 = arith.index_cast %534 : i32 to index + %536 = affine.apply #map17()[%535, %block_id_y, %thread_id_x, %thread_id_y] + %537 = arith.select %532, %37, %c0_i32 : i32 + %538 = arith.addi %531, %537 : i32 + %539 = arith.index_cast %538 : i32 to index + %540 = arith.muli %536, %9 overflow : index + %541 = arith.addi %540, %539 overflow : index + amdgpu.gather_to_lds %60[%541], %arg46[%35, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %542 = affine.apply #map19()[%535, %block_id_y, %thread_id_x, %thread_id_y] + %543 = arith.muli %542, %9 overflow : index + %544 = arith.addi %543, %539 overflow : index + amdgpu.gather_to_lds %60[%544], %arg46[%62, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %545 = affine.apply #map21()[%535, %block_id_y, %thread_id_x, %thread_id_y] + %546 = arith.muli %545, %9 overflow : index + %547 = arith.addi %546, %539 overflow : index + amdgpu.gather_to_lds %60[%547], %arg46[%67, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + rocdl.sched.barrier 0 + %reinterpret_cast_28 = memref.reinterpret_cast %arg43 to offset: [0], sizes: [32768], strides: [1] : memref<256x128xi8, #gpu.address_space> to memref<32768xi8, #gpu.address_space> + %548 = vector.load %reinterpret_cast_28[%139] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %549 = vector.load %reinterpret_cast_28[%140] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %550 = vector.load %reinterpret_cast_28[%141] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %551 = vector.load %reinterpret_cast_28[%142] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %reinterpret_cast_29 = memref.reinterpret_cast %arg45 to offset: [0], sizes: [24576], strides: [1] : memref<192x128xi8, #gpu.address_space> to memref<24576xi8, #gpu.address_space> + %552 = vector.load %reinterpret_cast_29[%143] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %553 = vector.load %reinterpret_cast_29[%144] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %554 = vector.load %reinterpret_cast_29[%145] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %555 = vector.load %reinterpret_cast_29[%146] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %556 = vector.load %reinterpret_cast_29[%147] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %557 = vector.load %reinterpret_cast_29[%148] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %558 = vector.bitcast %548 : vector<16xi8> to vector<32xf4E2M1FN> + %559 = vector.bitcast %549 : vector<16xi8> to vector<32xf4E2M1FN> + %560 = vector.bitcast %550 : vector<16xi8> to vector<32xf4E2M1FN> + %561 = vector.bitcast %551 : vector<16xi8> to vector<32xf4E2M1FN> + %562 = vector.bitcast %552 : vector<16xi8> to vector<32xf4E2M1FN> + %563 = vector.bitcast %553 : vector<16xi8> to vector<32xf4E2M1FN> + %564 = vector.bitcast %554 : vector<16xi8> to vector<32xf4E2M1FN> + %565 = vector.bitcast %555 : vector<16xi8> to vector<32xf4E2M1FN> + %566 = vector.bitcast %556 : vector<16xi8> to vector<32xf4E2M1FN> + %567 = vector.bitcast %557 : vector<16xi8> to vector<32xf4E2M1FN> + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %568 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%517[0] * %562) + %arg14 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %569 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%517[1] * %563) + %arg15 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %570 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%516[0] * %564) + %arg16 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %571 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%516[1] * %565) + %arg17 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %572 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%515[0] * %566) + %arg18 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %573 = amdgpu.scaled_mfma 16x16x128 (%518[0] * %558) * (%515[1] * %567) + %arg19 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %574 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%517[0] * %562) + %arg20 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %575 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%517[1] * %563) + %arg21 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %576 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%516[0] * %564) + %arg22 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %577 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%516[1] * %565) + %arg23 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %578 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%515[0] * %566) + %arg24 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %579 = amdgpu.scaled_mfma 16x16x128 (%518[1] * %559) * (%515[1] * %567) + %arg25 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %580 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%517[0] * %562) + %arg26 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %581 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%517[1] * %563) + %arg27 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %582 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%516[0] * %564) + %arg28 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %583 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%516[1] * %565) + %arg29 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %584 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%515[0] * %566) + %arg30 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %585 = amdgpu.scaled_mfma 16x16x128 (%514[0] * %560) * (%515[1] * %567) + %arg31 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %586 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%517[0] * %562) + %arg32 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %587 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%517[1] * %563) + %arg33 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %588 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%516[0] * %564) + %arg34 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %589 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%516[1] * %565) + %arg35 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %590 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%515[0] * %566) + %arg36 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %591 = amdgpu.scaled_mfma 16x16x128 (%514[1] * %561) * (%515[1] * %567) + %arg37 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.sched.barrier 0 + %592 = affine.apply #map47()[%arg13, %thread_id_x] + %593 = arith.index_cast %592 : index to i32 + %594 = arith.extui %593 : i32 to i64 + %595 = arith.muli %594, %80 : i64 + %596 = arith.shrui %595, %c32_i64 : i64 + %597 = arith.trunci %596 : i64 to i32 + %598 = arith.muli %597, %71 : i32 + %599 = arith.subi %593, %598 : i32 + %600 = arith.cmpi uge, %599, %71 : i32 + %601 = arith.extui %600 : i1 to i32 + %602 = arith.subi %597, %601 : i32 + %603 = arith.index_cast %602 : i32 to index + %604 = affine.apply #map24()[%603, %thread_id_x] + %605 = affine.apply #map48()[%thread_id_x, %arg13] + %606 = arith.index_cast %605 : index to i32 + %607 = arith.extui %606 : i32 to i64 + %608 = arith.muli %607, %80 : i64 + %609 = arith.shrui %608, %c32_i64 : i64 + %610 = arith.trunci %609 : i64 to i32 + %611 = arith.muli %610, %71 : i32 + %612 = arith.subi %606, %611 : i32 + %613 = arith.cmpi uge, %612, %71 : i32 + %614 = arith.select %613, %71, %c0_i32 : i32 + %615 = arith.addi %612, %614 : i32 + %616 = arith.index_cast %615 : i32 to index + %617 = affine.apply #map2()[%616] + %618 = arith.muli %604, %strides#0 overflow : index + %619 = arith.addi %618, %617 overflow : index + %620 = vector.load %108[%619] : memref>, vector<4xi8> + %621 = affine.apply #map26()[%603, %thread_id_x] + %622 = arith.muli %621, %strides#0 overflow : index + %623 = arith.addi %622, %617 overflow : index + %624 = vector.load %108[%623] : memref>, vector<4xi8> + %625 = affine.apply #map28()[%603, %thread_id_y] + %626 = arith.muli %625, %strides_13#0 overflow : index + %627 = arith.addi %626, %617 overflow : index + %628 = vector.load %120[%627] : memref>, vector<4xi8> + %629 = affine.apply #map29()[%603, %thread_id_y] + %630 = arith.muli %629, %strides_13#0 overflow : index + %631 = arith.addi %630, %617 overflow : index + %632 = vector.load %120[%631] : memref>, vector<4xi8> + %633 = affine.apply #map30()[%603, %thread_id_y] + %634 = arith.muli %633, %strides_13#0 overflow : index + %635 = arith.addi %634, %617 overflow : index + %636 = vector.load %120[%635] : memref>, vector<4xi8> + %637 = vector.load %reinterpret_cast_28[%151] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %638 = vector.load %reinterpret_cast_28[%152] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %639 = vector.load %reinterpret_cast_28[%153] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %640 = vector.load %reinterpret_cast_28[%154] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %641 = vector.load %reinterpret_cast_29[%155] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %642 = vector.load %reinterpret_cast_29[%156] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %643 = vector.load %reinterpret_cast_29[%157] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %644 = vector.load %reinterpret_cast_29[%158] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %645 = vector.load %reinterpret_cast_29[%159] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %646 = vector.load %reinterpret_cast_29[%160] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %647 = vector.bitcast %637 : vector<16xi8> to vector<32xf4E2M1FN> + %648 = vector.bitcast %638 : vector<16xi8> to vector<32xf4E2M1FN> + %649 = vector.bitcast %639 : vector<16xi8> to vector<32xf4E2M1FN> + %650 = vector.bitcast %640 : vector<16xi8> to vector<32xf4E2M1FN> + %651 = vector.bitcast %641 : vector<16xi8> to vector<32xf4E2M1FN> + %652 = vector.bitcast %642 : vector<16xi8> to vector<32xf4E2M1FN> + %653 = vector.bitcast %643 : vector<16xi8> to vector<32xf4E2M1FN> + %654 = vector.bitcast %644 : vector<16xi8> to vector<32xf4E2M1FN> + %655 = vector.bitcast %645 : vector<16xi8> to vector<32xf4E2M1FN> + %656 = vector.bitcast %646 : vector<16xi8> to vector<32xf4E2M1FN> + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(7) + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %657 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %647) * (%517[2] * %651) + %568 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %658 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %647) * (%517[3] * %652) + %569 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %659 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %647) * (%516[2] * %653) + %570 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %660 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %647) * (%516[3] * %654) + %571 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %661 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %647) * (%515[2] * %655) + %572 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %662 = amdgpu.scaled_mfma 16x16x128 (%518[2] * %647) * (%515[3] * %656) + %573 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %663 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %648) * (%517[2] * %651) + %574 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %664 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %648) * (%517[3] * %652) + %575 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %665 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %648) * (%516[2] * %653) + %576 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %666 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %648) * (%516[3] * %654) + %577 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %667 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %648) * (%515[2] * %655) + %578 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %668 = amdgpu.scaled_mfma 16x16x128 (%518[3] * %648) * (%515[3] * %656) + %579 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %669 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %649) * (%517[2] * %651) + %580 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %670 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %649) * (%517[3] * %652) + %581 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %671 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %649) * (%516[2] * %653) + %582 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %672 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %649) * (%516[3] * %654) + %583 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %673 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %649) * (%515[2] * %655) + %584 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %674 = amdgpu.scaled_mfma 16x16x128 (%514[2] * %649) * (%515[3] * %656) + %585 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %675 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %650) * (%517[2] * %651) + %586 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %676 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %650) * (%517[3] * %652) + %587 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %677 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %650) * (%516[2] * %653) + %588 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %678 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %650) * (%516[3] * %654) + %589 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %679 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %650) * (%515[2] * %655) + %590 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %680 = amdgpu.scaled_mfma 16x16x128 (%514[3] * %650) * (%515[3] * %656) + %591 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + scf.yield %657, %658, %659, %660, %661, %662, %663, %664, %665, %666, %667, %668, %669, %670, %671, %672, %673, %674, %675, %676, %677, %678, %679, %680, %620, %624, %628, %632, %636, %arg44, %arg43, %arg46, %arg45 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, memref<256x128xi8, #gpu.address_space>, memref<256x128xi8, #gpu.address_space>, memref<192x128xi8, #gpu.address_space>, memref<192x128xi8, #gpu.address_space> + } + %163 = vector.bitcast %162#25 : vector<4xi8> to vector<4xf8E8M0FNU> + %164 = vector.bitcast %162#28 : vector<4xi8> to vector<4xf8E8M0FNU> + %165 = vector.bitcast %162#27 : vector<4xi8> to vector<4xf8E8M0FNU> + %166 = vector.bitcast %162#26 : vector<4xi8> to vector<4xf8E8M0FNU> + %167 = vector.bitcast %162#24 : vector<4xi8> to vector<4xf8E8M0FNU> + scf.if %133 { + rocdl.s.barrier + } + amdgpu.lds_barrier + %reinterpret_cast_20 = memref.reinterpret_cast %162#31 to offset: [0], sizes: [24576], strides: [1] : memref<192x128xi8, #gpu.address_space> to memref<24576xi8, #gpu.address_space> + %168 = vector.load %reinterpret_cast_20[%143] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %169 = vector.load %reinterpret_cast_20[%155] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %170 = vector.load %reinterpret_cast_20[%144] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %171 = vector.load %reinterpret_cast_20[%156] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %172 = vector.load %reinterpret_cast_20[%145] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %173 = vector.load %reinterpret_cast_20[%157] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %174 = vector.load %reinterpret_cast_20[%146] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %175 = vector.load %reinterpret_cast_20[%158] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %176 = vector.load %reinterpret_cast_20[%147] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %177 = vector.load %reinterpret_cast_20[%159] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %178 = vector.load %reinterpret_cast_20[%148] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %179 = vector.load %reinterpret_cast_20[%160] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %reinterpret_cast_21 = memref.reinterpret_cast %162#29 to offset: [0], sizes: [32768], strides: [1] : memref<256x128xi8, #gpu.address_space> to memref<32768xi8, #gpu.address_space> + %180 = vector.load %reinterpret_cast_21[%139] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %181 = vector.load %reinterpret_cast_21[%151] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %182 = vector.load %reinterpret_cast_21[%140] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %183 = vector.load %reinterpret_cast_21[%152] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %184 = vector.load %reinterpret_cast_21[%141] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %185 = vector.load %reinterpret_cast_21[%153] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %186 = vector.load %reinterpret_cast_21[%142] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %187 = vector.load %reinterpret_cast_21[%154] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %188 = vector.bitcast %180 : vector<16xi8> to vector<32xf4E2M1FN> + %189 = vector.bitcast %181 : vector<16xi8> to vector<32xf4E2M1FN> + %190 = vector.bitcast %182 : vector<16xi8> to vector<32xf4E2M1FN> + %191 = vector.bitcast %183 : vector<16xi8> to vector<32xf4E2M1FN> + %192 = vector.bitcast %184 : vector<16xi8> to vector<32xf4E2M1FN> + %193 = vector.bitcast %185 : vector<16xi8> to vector<32xf4E2M1FN> + %194 = vector.bitcast %186 : vector<16xi8> to vector<32xf4E2M1FN> + %195 = vector.bitcast %187 : vector<16xi8> to vector<32xf4E2M1FN> + %196 = vector.bitcast %168 : vector<16xi8> to vector<32xf4E2M1FN> + %197 = vector.bitcast %169 : vector<16xi8> to vector<32xf4E2M1FN> + %198 = vector.bitcast %170 : vector<16xi8> to vector<32xf4E2M1FN> + %199 = vector.bitcast %171 : vector<16xi8> to vector<32xf4E2M1FN> + %200 = vector.bitcast %172 : vector<16xi8> to vector<32xf4E2M1FN> + %201 = vector.bitcast %173 : vector<16xi8> to vector<32xf4E2M1FN> + %202 = vector.bitcast %174 : vector<16xi8> to vector<32xf4E2M1FN> + %203 = vector.bitcast %175 : vector<16xi8> to vector<32xf4E2M1FN> + %204 = vector.bitcast %176 : vector<16xi8> to vector<32xf4E2M1FN> + %205 = vector.bitcast %177 : vector<16xi8> to vector<32xf4E2M1FN> + %206 = vector.bitcast %178 : vector<16xi8> to vector<32xf4E2M1FN> + %207 = vector.bitcast %179 : vector<16xi8> to vector<32xf4E2M1FN> + %208 = amdgpu.scaled_mfma 16x16x128 (%167[0] * %188) * (%166[0] * %196) + %162#0 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %209 = amdgpu.scaled_mfma 16x16x128 (%167[2] * %189) * (%166[2] * %197) + %208 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %210 = amdgpu.scaled_mfma 16x16x128 (%167[0] * %188) * (%166[1] * %198) + %162#1 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %211 = amdgpu.scaled_mfma 16x16x128 (%167[2] * %189) * (%166[3] * %199) + %210 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %212 = amdgpu.scaled_mfma 16x16x128 (%167[0] * %188) * (%165[0] * %200) + %162#2 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %213 = amdgpu.scaled_mfma 16x16x128 (%167[2] * %189) * (%165[2] * %201) + %212 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %214 = amdgpu.scaled_mfma 16x16x128 (%167[0] * %188) * (%165[1] * %202) + %162#3 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %215 = amdgpu.scaled_mfma 16x16x128 (%167[2] * %189) * (%165[3] * %203) + %214 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %216 = amdgpu.scaled_mfma 16x16x128 (%167[0] * %188) * (%164[0] * %204) + %162#4 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %217 = amdgpu.scaled_mfma 16x16x128 (%167[2] * %189) * (%164[2] * %205) + %216 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %218 = amdgpu.scaled_mfma 16x16x128 (%167[0] * %188) * (%164[1] * %206) + %162#5 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %219 = amdgpu.scaled_mfma 16x16x128 (%167[2] * %189) * (%164[3] * %207) + %218 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %220 = amdgpu.scaled_mfma 16x16x128 (%167[1] * %190) * (%166[0] * %196) + %162#6 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %221 = amdgpu.scaled_mfma 16x16x128 (%167[3] * %191) * (%166[2] * %197) + %220 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %222 = amdgpu.scaled_mfma 16x16x128 (%167[1] * %190) * (%166[1] * %198) + %162#7 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %223 = amdgpu.scaled_mfma 16x16x128 (%167[3] * %191) * (%166[3] * %199) + %222 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %224 = amdgpu.scaled_mfma 16x16x128 (%167[1] * %190) * (%165[0] * %200) + %162#8 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %225 = amdgpu.scaled_mfma 16x16x128 (%167[3] * %191) * (%165[2] * %201) + %224 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %226 = amdgpu.scaled_mfma 16x16x128 (%167[1] * %190) * (%165[1] * %202) + %162#9 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %227 = amdgpu.scaled_mfma 16x16x128 (%167[3] * %191) * (%165[3] * %203) + %226 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %228 = amdgpu.scaled_mfma 16x16x128 (%167[1] * %190) * (%164[0] * %204) + %162#10 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %229 = amdgpu.scaled_mfma 16x16x128 (%167[3] * %191) * (%164[2] * %205) + %228 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %230 = amdgpu.scaled_mfma 16x16x128 (%167[1] * %190) * (%164[1] * %206) + %162#11 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %231 = amdgpu.scaled_mfma 16x16x128 (%167[3] * %191) * (%164[3] * %207) + %230 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %232 = amdgpu.scaled_mfma 16x16x128 (%163[0] * %192) * (%166[0] * %196) + %162#12 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %233 = amdgpu.scaled_mfma 16x16x128 (%163[2] * %193) * (%166[2] * %197) + %232 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %234 = amdgpu.scaled_mfma 16x16x128 (%163[0] * %192) * (%166[1] * %198) + %162#13 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %235 = amdgpu.scaled_mfma 16x16x128 (%163[2] * %193) * (%166[3] * %199) + %234 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %236 = amdgpu.scaled_mfma 16x16x128 (%163[0] * %192) * (%165[0] * %200) + %162#14 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %237 = amdgpu.scaled_mfma 16x16x128 (%163[2] * %193) * (%165[2] * %201) + %236 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %238 = amdgpu.scaled_mfma 16x16x128 (%163[0] * %192) * (%165[1] * %202) + %162#15 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %239 = amdgpu.scaled_mfma 16x16x128 (%163[2] * %193) * (%165[3] * %203) + %238 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %240 = amdgpu.scaled_mfma 16x16x128 (%163[0] * %192) * (%164[0] * %204) + %162#16 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %241 = amdgpu.scaled_mfma 16x16x128 (%163[2] * %193) * (%164[2] * %205) + %240 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %242 = amdgpu.scaled_mfma 16x16x128 (%163[0] * %192) * (%164[1] * %206) + %162#17 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %243 = amdgpu.scaled_mfma 16x16x128 (%163[2] * %193) * (%164[3] * %207) + %242 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %244 = amdgpu.scaled_mfma 16x16x128 (%163[1] * %194) * (%166[0] * %196) + %162#18 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %245 = amdgpu.scaled_mfma 16x16x128 (%163[3] * %195) * (%166[2] * %197) + %244 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %246 = amdgpu.scaled_mfma 16x16x128 (%163[1] * %194) * (%166[1] * %198) + %162#19 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %247 = amdgpu.scaled_mfma 16x16x128 (%163[3] * %195) * (%166[3] * %199) + %246 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %248 = amdgpu.scaled_mfma 16x16x128 (%163[1] * %194) * (%165[0] * %200) + %162#20 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %249 = amdgpu.scaled_mfma 16x16x128 (%163[3] * %195) * (%165[2] * %201) + %248 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %250 = amdgpu.scaled_mfma 16x16x128 (%163[1] * %194) * (%165[1] * %202) + %162#21 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %251 = amdgpu.scaled_mfma 16x16x128 (%163[3] * %195) * (%165[3] * %203) + %250 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %252 = amdgpu.scaled_mfma 16x16x128 (%163[1] * %194) * (%164[0] * %204) + %162#22 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %253 = amdgpu.scaled_mfma 16x16x128 (%163[3] * %195) * (%164[2] * %205) + %252 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %254 = amdgpu.scaled_mfma 16x16x128 (%163[1] * %194) * (%164[1] * %206) + %162#23 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %255 = amdgpu.scaled_mfma 16x16x128 (%163[3] * %195) * (%164[3] * %207) + %254 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + + %ep_shuffle_offset = arith.constant 1 : i32 + %ep_shuffle_width = arith.constant 64 : i32 + %ep_one = arith.constant 1 : index + %ep_zero = arith.constant 0 : index + %ep_lane_parity = arith.andi %thread_id_x, %ep_one : index + %ep_is_even = arith.cmpi eq, %ep_lane_parity, %ep_zero : index + + %ep_base_buffer, %ep_offset, %ep_size0, %ep_size1, %ep_output_stride, %ep_col_stride = memref.extract_strided_metadata %reinterpret_cast_1 : memref> -> memref, index, index, index, index, index + %ep_wg_row_off = arith.muli %76, %ep_output_stride overflow : index + %ep_wg_offset = arith.addi %ep_wg_row_off, %114 overflow : index + %ep_reinterpret = memref.reinterpret_cast %4 to offset: [%ep_wg_offset], sizes: [1073741822], strides: [1] : memref to memref<1073741822xbf16, strided<[1], offset: ?>> + %ep_cast = memref.cast %ep_reinterpret : memref<1073741822xbf16, strided<[1], offset: ?>> to memref> + %ep_stride_i14 = arith.index_cast %ep_output_stride : index to i14 + %ep_buffer = amdgpu.fat_raw_buffer_cast %ep_cast validBytes(%c2147483645_i64) cacheSwizzleStride(%ep_stride_i14) resetOffset : memref> to memref> + + %ep_rb0_row0 = affine.apply #map49()[%thread_id_x] + %ep_rb0_row1 = affine.apply #map51()[%thread_id_x] + %ep_rb0_row2 = affine.apply #map52()[%thread_id_x] + %ep_rb0_row3 = affine.apply #map53()[%thread_id_x] + %ep_rb0_store_row_a = arith.select %ep_is_even, %ep_rb0_row0, %ep_rb0_row2 : index + %ep_rb0_store_row_b = arith.select %ep_is_even, %ep_rb0_row1, %ep_rb0_row3 : index + %ep_rb0_store_off_a = arith.muli %ep_rb0_store_row_a, %ep_output_stride overflow : index + %ep_rb0_store_off_b = arith.muli %ep_rb0_store_row_b, %ep_output_stride overflow : index + %ep_rb1_row0 = affine.apply #map59()[%thread_id_x] + %ep_rb1_row1 = affine.apply #map60()[%thread_id_x] + %ep_rb1_row2 = affine.apply #map61()[%thread_id_x] + %ep_rb1_row3 = affine.apply #map62()[%thread_id_x] + %ep_rb1_store_row_a = arith.select %ep_is_even, %ep_rb1_row0, %ep_rb1_row2 : index + %ep_rb1_store_row_b = arith.select %ep_is_even, %ep_rb1_row1, %ep_rb1_row3 : index + %ep_rb1_store_off_a = arith.muli %ep_rb1_store_row_a, %ep_output_stride overflow : index + %ep_rb1_store_off_b = arith.muli %ep_rb1_store_row_b, %ep_output_stride overflow : index + %ep_rb2_row0 = affine.apply #map63()[%thread_id_x] + %ep_rb2_row1 = affine.apply #map64()[%thread_id_x] + %ep_rb2_row2 = affine.apply #map65()[%thread_id_x] + %ep_rb2_row3 = affine.apply #map66()[%thread_id_x] + %ep_rb2_store_row_a = arith.select %ep_is_even, %ep_rb2_row0, %ep_rb2_row2 : index + %ep_rb2_store_row_b = arith.select %ep_is_even, %ep_rb2_row1, %ep_rb2_row3 : index + %ep_rb2_store_off_a = arith.muli %ep_rb2_store_row_a, %ep_output_stride overflow : index + %ep_rb2_store_off_b = arith.muli %ep_rb2_store_row_b, %ep_output_stride overflow : index + %ep_rb3_row0 = affine.apply #map67()[%thread_id_x] + %ep_rb3_row1 = affine.apply #map68()[%thread_id_x] + %ep_rb3_row2 = affine.apply #map69()[%thread_id_x] + %ep_rb3_row3 = affine.apply #map70()[%thread_id_x] + %ep_rb3_store_row_a = arith.select %ep_is_even, %ep_rb3_row0, %ep_rb3_row2 : index + %ep_rb3_store_row_b = arith.select %ep_is_even, %ep_rb3_row1, %ep_rb3_row3 : index + %ep_rb3_store_off_a = arith.muli %ep_rb3_store_row_a, %ep_output_stride overflow : index + %ep_rb3_store_off_b = arith.muli %ep_rb3_store_row_b, %ep_output_stride overflow : index + %ep_col0 = affine.apply #map50()[%thread_id_x, %thread_id_y] + %ep_col0_adj = arith.subi %ep_col0, %ep_lane_parity : index + %ep_col1 = affine.apply #map54()[%thread_id_x, %thread_id_y] + %ep_col1_adj = arith.subi %ep_col1, %ep_lane_parity : index + %ep_col2 = affine.apply #map55()[%thread_id_x, %thread_id_y] + %ep_col2_adj = arith.subi %ep_col2, %ep_lane_parity : index + %ep_col3 = affine.apply #map56()[%thread_id_x, %thread_id_y] + %ep_col3_adj = arith.subi %ep_col3, %ep_lane_parity : index + %ep_col4 = affine.apply #map57()[%thread_id_x, %thread_id_y] + %ep_col4_adj = arith.subi %ep_col4, %ep_lane_parity : index + %ep_col5 = affine.apply #map58()[%thread_id_x, %thread_id_y] + %ep_col5_adj = arith.subi %ep_col5, %ep_lane_parity : index + + %ep_v209_row0 = vector.extract %209[0] : f32 from vector<4xf32> + %ep_v209_row1 = vector.extract %209[1] : f32 from vector<4xf32> + %ep_v209_row2 = vector.extract %209[2] : f32 from vector<4xf32> + %ep_v209_row3 = vector.extract %209[3] : f32 from vector<4xf32> + %ep_v209_row0_nbr, %ep_v209_row0_valid = gpu.shuffle xor %ep_v209_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v209_row1_nbr, %ep_v209_row1_valid = gpu.shuffle xor %ep_v209_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v209_row2_nbr, %ep_v209_row2_valid = gpu.shuffle xor %ep_v209_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v209_row3_nbr, %ep_v209_row3_valid = gpu.shuffle xor %ep_v209_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v209_row0_lo = arith.select %ep_is_even, %ep_v209_row0, %ep_v209_row0_nbr : f32 + %ep_v209_row0_hi = arith.select %ep_is_even, %ep_v209_row0_nbr, %ep_v209_row0 : f32 + %ep_v209_row1_lo = arith.select %ep_is_even, %ep_v209_row1, %ep_v209_row1_nbr : f32 + %ep_v209_row1_hi = arith.select %ep_is_even, %ep_v209_row1_nbr, %ep_v209_row1 : f32 + %ep_v209_row2_lo = arith.select %ep_is_even, %ep_v209_row2, %ep_v209_row2_nbr : f32 + %ep_v209_row2_hi = arith.select %ep_is_even, %ep_v209_row2_nbr, %ep_v209_row2 : f32 + %ep_v209_row3_lo = arith.select %ep_is_even, %ep_v209_row3, %ep_v209_row3_nbr : f32 + %ep_v209_row3_hi = arith.select %ep_is_even, %ep_v209_row3_nbr, %ep_v209_row3 : f32 + %ep_v209_store_a_lo = arith.select %ep_is_even, %ep_v209_row0_lo, %ep_v209_row2_lo : f32 + %ep_v209_store_a_hi = arith.select %ep_is_even, %ep_v209_row0_hi, %ep_v209_row2_hi : f32 + %ep_v209_store_b_lo = arith.select %ep_is_even, %ep_v209_row1_lo, %ep_v209_row3_lo : f32 + %ep_v209_store_b_hi = arith.select %ep_is_even, %ep_v209_row1_hi, %ep_v209_row3_hi : f32 + %ep_v209_i32_0 = arith.constant 0 : i32 + %ep_v209_i32_1 = arith.constant 1 : i32 + %ep_v209_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v209_pair_a_0 = llvm.insertelement %ep_v209_store_a_lo, %ep_v209_undef_a[%ep_v209_i32_0 : i32] : vector<2xf32> + %ep_v209_pair_a = llvm.insertelement %ep_v209_store_a_hi, %ep_v209_pair_a_0[%ep_v209_i32_1 : i32] : vector<2xf32> + %ep_v209_store_a = llvm.fptrunc %ep_v209_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v209_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v209_pair_b_0 = llvm.insertelement %ep_v209_store_b_lo, %ep_v209_undef_b[%ep_v209_i32_0 : i32] : vector<2xf32> + %ep_v209_pair_b = llvm.insertelement %ep_v209_store_b_hi, %ep_v209_pair_b_0[%ep_v209_i32_1 : i32] : vector<2xf32> + %ep_v209_store_b = llvm.fptrunc %ep_v209_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v209_addr_a = arith.addi %ep_rb0_store_off_a, %ep_col0_adj overflow : index + %ep_v209_addr_b = arith.addi %ep_rb0_store_off_b, %ep_col0_adj overflow : index + vector.store %ep_v209_store_a, %ep_buffer[%ep_v209_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v209_store_b, %ep_buffer[%ep_v209_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v211_row0 = vector.extract %211[0] : f32 from vector<4xf32> + %ep_v211_row1 = vector.extract %211[1] : f32 from vector<4xf32> + %ep_v211_row2 = vector.extract %211[2] : f32 from vector<4xf32> + %ep_v211_row3 = vector.extract %211[3] : f32 from vector<4xf32> + %ep_v211_row0_nbr, %ep_v211_row0_valid = gpu.shuffle xor %ep_v211_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v211_row1_nbr, %ep_v211_row1_valid = gpu.shuffle xor %ep_v211_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v211_row2_nbr, %ep_v211_row2_valid = gpu.shuffle xor %ep_v211_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v211_row3_nbr, %ep_v211_row3_valid = gpu.shuffle xor %ep_v211_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v211_row0_lo = arith.select %ep_is_even, %ep_v211_row0, %ep_v211_row0_nbr : f32 + %ep_v211_row0_hi = arith.select %ep_is_even, %ep_v211_row0_nbr, %ep_v211_row0 : f32 + %ep_v211_row1_lo = arith.select %ep_is_even, %ep_v211_row1, %ep_v211_row1_nbr : f32 + %ep_v211_row1_hi = arith.select %ep_is_even, %ep_v211_row1_nbr, %ep_v211_row1 : f32 + %ep_v211_row2_lo = arith.select %ep_is_even, %ep_v211_row2, %ep_v211_row2_nbr : f32 + %ep_v211_row2_hi = arith.select %ep_is_even, %ep_v211_row2_nbr, %ep_v211_row2 : f32 + %ep_v211_row3_lo = arith.select %ep_is_even, %ep_v211_row3, %ep_v211_row3_nbr : f32 + %ep_v211_row3_hi = arith.select %ep_is_even, %ep_v211_row3_nbr, %ep_v211_row3 : f32 + %ep_v211_store_a_lo = arith.select %ep_is_even, %ep_v211_row0_lo, %ep_v211_row2_lo : f32 + %ep_v211_store_a_hi = arith.select %ep_is_even, %ep_v211_row0_hi, %ep_v211_row2_hi : f32 + %ep_v211_store_b_lo = arith.select %ep_is_even, %ep_v211_row1_lo, %ep_v211_row3_lo : f32 + %ep_v211_store_b_hi = arith.select %ep_is_even, %ep_v211_row1_hi, %ep_v211_row3_hi : f32 + %ep_v211_i32_0 = arith.constant 0 : i32 + %ep_v211_i32_1 = arith.constant 1 : i32 + %ep_v211_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v211_pair_a_0 = llvm.insertelement %ep_v211_store_a_lo, %ep_v211_undef_a[%ep_v211_i32_0 : i32] : vector<2xf32> + %ep_v211_pair_a = llvm.insertelement %ep_v211_store_a_hi, %ep_v211_pair_a_0[%ep_v211_i32_1 : i32] : vector<2xf32> + %ep_v211_store_a = llvm.fptrunc %ep_v211_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v211_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v211_pair_b_0 = llvm.insertelement %ep_v211_store_b_lo, %ep_v211_undef_b[%ep_v211_i32_0 : i32] : vector<2xf32> + %ep_v211_pair_b = llvm.insertelement %ep_v211_store_b_hi, %ep_v211_pair_b_0[%ep_v211_i32_1 : i32] : vector<2xf32> + %ep_v211_store_b = llvm.fptrunc %ep_v211_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v211_addr_a = arith.addi %ep_rb0_store_off_a, %ep_col1_adj overflow : index + %ep_v211_addr_b = arith.addi %ep_rb0_store_off_b, %ep_col1_adj overflow : index + vector.store %ep_v211_store_a, %ep_buffer[%ep_v211_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v211_store_b, %ep_buffer[%ep_v211_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v213_row0 = vector.extract %213[0] : f32 from vector<4xf32> + %ep_v213_row1 = vector.extract %213[1] : f32 from vector<4xf32> + %ep_v213_row2 = vector.extract %213[2] : f32 from vector<4xf32> + %ep_v213_row3 = vector.extract %213[3] : f32 from vector<4xf32> + %ep_v213_row0_nbr, %ep_v213_row0_valid = gpu.shuffle xor %ep_v213_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v213_row1_nbr, %ep_v213_row1_valid = gpu.shuffle xor %ep_v213_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v213_row2_nbr, %ep_v213_row2_valid = gpu.shuffle xor %ep_v213_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v213_row3_nbr, %ep_v213_row3_valid = gpu.shuffle xor %ep_v213_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v213_row0_lo = arith.select %ep_is_even, %ep_v213_row0, %ep_v213_row0_nbr : f32 + %ep_v213_row0_hi = arith.select %ep_is_even, %ep_v213_row0_nbr, %ep_v213_row0 : f32 + %ep_v213_row1_lo = arith.select %ep_is_even, %ep_v213_row1, %ep_v213_row1_nbr : f32 + %ep_v213_row1_hi = arith.select %ep_is_even, %ep_v213_row1_nbr, %ep_v213_row1 : f32 + %ep_v213_row2_lo = arith.select %ep_is_even, %ep_v213_row2, %ep_v213_row2_nbr : f32 + %ep_v213_row2_hi = arith.select %ep_is_even, %ep_v213_row2_nbr, %ep_v213_row2 : f32 + %ep_v213_row3_lo = arith.select %ep_is_even, %ep_v213_row3, %ep_v213_row3_nbr : f32 + %ep_v213_row3_hi = arith.select %ep_is_even, %ep_v213_row3_nbr, %ep_v213_row3 : f32 + %ep_v213_store_a_lo = arith.select %ep_is_even, %ep_v213_row0_lo, %ep_v213_row2_lo : f32 + %ep_v213_store_a_hi = arith.select %ep_is_even, %ep_v213_row0_hi, %ep_v213_row2_hi : f32 + %ep_v213_store_b_lo = arith.select %ep_is_even, %ep_v213_row1_lo, %ep_v213_row3_lo : f32 + %ep_v213_store_b_hi = arith.select %ep_is_even, %ep_v213_row1_hi, %ep_v213_row3_hi : f32 + %ep_v213_i32_0 = arith.constant 0 : i32 + %ep_v213_i32_1 = arith.constant 1 : i32 + %ep_v213_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v213_pair_a_0 = llvm.insertelement %ep_v213_store_a_lo, %ep_v213_undef_a[%ep_v213_i32_0 : i32] : vector<2xf32> + %ep_v213_pair_a = llvm.insertelement %ep_v213_store_a_hi, %ep_v213_pair_a_0[%ep_v213_i32_1 : i32] : vector<2xf32> + %ep_v213_store_a = llvm.fptrunc %ep_v213_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v213_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v213_pair_b_0 = llvm.insertelement %ep_v213_store_b_lo, %ep_v213_undef_b[%ep_v213_i32_0 : i32] : vector<2xf32> + %ep_v213_pair_b = llvm.insertelement %ep_v213_store_b_hi, %ep_v213_pair_b_0[%ep_v213_i32_1 : i32] : vector<2xf32> + %ep_v213_store_b = llvm.fptrunc %ep_v213_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v213_addr_a = arith.addi %ep_rb0_store_off_a, %ep_col2_adj overflow : index + %ep_v213_addr_b = arith.addi %ep_rb0_store_off_b, %ep_col2_adj overflow : index + vector.store %ep_v213_store_a, %ep_buffer[%ep_v213_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v213_store_b, %ep_buffer[%ep_v213_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v215_row0 = vector.extract %215[0] : f32 from vector<4xf32> + %ep_v215_row1 = vector.extract %215[1] : f32 from vector<4xf32> + %ep_v215_row2 = vector.extract %215[2] : f32 from vector<4xf32> + %ep_v215_row3 = vector.extract %215[3] : f32 from vector<4xf32> + %ep_v215_row0_nbr, %ep_v215_row0_valid = gpu.shuffle xor %ep_v215_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v215_row1_nbr, %ep_v215_row1_valid = gpu.shuffle xor %ep_v215_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v215_row2_nbr, %ep_v215_row2_valid = gpu.shuffle xor %ep_v215_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v215_row3_nbr, %ep_v215_row3_valid = gpu.shuffle xor %ep_v215_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v215_row0_lo = arith.select %ep_is_even, %ep_v215_row0, %ep_v215_row0_nbr : f32 + %ep_v215_row0_hi = arith.select %ep_is_even, %ep_v215_row0_nbr, %ep_v215_row0 : f32 + %ep_v215_row1_lo = arith.select %ep_is_even, %ep_v215_row1, %ep_v215_row1_nbr : f32 + %ep_v215_row1_hi = arith.select %ep_is_even, %ep_v215_row1_nbr, %ep_v215_row1 : f32 + %ep_v215_row2_lo = arith.select %ep_is_even, %ep_v215_row2, %ep_v215_row2_nbr : f32 + %ep_v215_row2_hi = arith.select %ep_is_even, %ep_v215_row2_nbr, %ep_v215_row2 : f32 + %ep_v215_row3_lo = arith.select %ep_is_even, %ep_v215_row3, %ep_v215_row3_nbr : f32 + %ep_v215_row3_hi = arith.select %ep_is_even, %ep_v215_row3_nbr, %ep_v215_row3 : f32 + %ep_v215_store_a_lo = arith.select %ep_is_even, %ep_v215_row0_lo, %ep_v215_row2_lo : f32 + %ep_v215_store_a_hi = arith.select %ep_is_even, %ep_v215_row0_hi, %ep_v215_row2_hi : f32 + %ep_v215_store_b_lo = arith.select %ep_is_even, %ep_v215_row1_lo, %ep_v215_row3_lo : f32 + %ep_v215_store_b_hi = arith.select %ep_is_even, %ep_v215_row1_hi, %ep_v215_row3_hi : f32 + %ep_v215_i32_0 = arith.constant 0 : i32 + %ep_v215_i32_1 = arith.constant 1 : i32 + %ep_v215_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v215_pair_a_0 = llvm.insertelement %ep_v215_store_a_lo, %ep_v215_undef_a[%ep_v215_i32_0 : i32] : vector<2xf32> + %ep_v215_pair_a = llvm.insertelement %ep_v215_store_a_hi, %ep_v215_pair_a_0[%ep_v215_i32_1 : i32] : vector<2xf32> + %ep_v215_store_a = llvm.fptrunc %ep_v215_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v215_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v215_pair_b_0 = llvm.insertelement %ep_v215_store_b_lo, %ep_v215_undef_b[%ep_v215_i32_0 : i32] : vector<2xf32> + %ep_v215_pair_b = llvm.insertelement %ep_v215_store_b_hi, %ep_v215_pair_b_0[%ep_v215_i32_1 : i32] : vector<2xf32> + %ep_v215_store_b = llvm.fptrunc %ep_v215_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v215_addr_a = arith.addi %ep_rb0_store_off_a, %ep_col3_adj overflow : index + %ep_v215_addr_b = arith.addi %ep_rb0_store_off_b, %ep_col3_adj overflow : index + vector.store %ep_v215_store_a, %ep_buffer[%ep_v215_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v215_store_b, %ep_buffer[%ep_v215_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v217_row0 = vector.extract %217[0] : f32 from vector<4xf32> + %ep_v217_row1 = vector.extract %217[1] : f32 from vector<4xf32> + %ep_v217_row2 = vector.extract %217[2] : f32 from vector<4xf32> + %ep_v217_row3 = vector.extract %217[3] : f32 from vector<4xf32> + %ep_v217_row0_nbr, %ep_v217_row0_valid = gpu.shuffle xor %ep_v217_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v217_row1_nbr, %ep_v217_row1_valid = gpu.shuffle xor %ep_v217_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v217_row2_nbr, %ep_v217_row2_valid = gpu.shuffle xor %ep_v217_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v217_row3_nbr, %ep_v217_row3_valid = gpu.shuffle xor %ep_v217_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v217_row0_lo = arith.select %ep_is_even, %ep_v217_row0, %ep_v217_row0_nbr : f32 + %ep_v217_row0_hi = arith.select %ep_is_even, %ep_v217_row0_nbr, %ep_v217_row0 : f32 + %ep_v217_row1_lo = arith.select %ep_is_even, %ep_v217_row1, %ep_v217_row1_nbr : f32 + %ep_v217_row1_hi = arith.select %ep_is_even, %ep_v217_row1_nbr, %ep_v217_row1 : f32 + %ep_v217_row2_lo = arith.select %ep_is_even, %ep_v217_row2, %ep_v217_row2_nbr : f32 + %ep_v217_row2_hi = arith.select %ep_is_even, %ep_v217_row2_nbr, %ep_v217_row2 : f32 + %ep_v217_row3_lo = arith.select %ep_is_even, %ep_v217_row3, %ep_v217_row3_nbr : f32 + %ep_v217_row3_hi = arith.select %ep_is_even, %ep_v217_row3_nbr, %ep_v217_row3 : f32 + %ep_v217_store_a_lo = arith.select %ep_is_even, %ep_v217_row0_lo, %ep_v217_row2_lo : f32 + %ep_v217_store_a_hi = arith.select %ep_is_even, %ep_v217_row0_hi, %ep_v217_row2_hi : f32 + %ep_v217_store_b_lo = arith.select %ep_is_even, %ep_v217_row1_lo, %ep_v217_row3_lo : f32 + %ep_v217_store_b_hi = arith.select %ep_is_even, %ep_v217_row1_hi, %ep_v217_row3_hi : f32 + %ep_v217_i32_0 = arith.constant 0 : i32 + %ep_v217_i32_1 = arith.constant 1 : i32 + %ep_v217_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v217_pair_a_0 = llvm.insertelement %ep_v217_store_a_lo, %ep_v217_undef_a[%ep_v217_i32_0 : i32] : vector<2xf32> + %ep_v217_pair_a = llvm.insertelement %ep_v217_store_a_hi, %ep_v217_pair_a_0[%ep_v217_i32_1 : i32] : vector<2xf32> + %ep_v217_store_a = llvm.fptrunc %ep_v217_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v217_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v217_pair_b_0 = llvm.insertelement %ep_v217_store_b_lo, %ep_v217_undef_b[%ep_v217_i32_0 : i32] : vector<2xf32> + %ep_v217_pair_b = llvm.insertelement %ep_v217_store_b_hi, %ep_v217_pair_b_0[%ep_v217_i32_1 : i32] : vector<2xf32> + %ep_v217_store_b = llvm.fptrunc %ep_v217_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v217_addr_a = arith.addi %ep_rb0_store_off_a, %ep_col4_adj overflow : index + %ep_v217_addr_b = arith.addi %ep_rb0_store_off_b, %ep_col4_adj overflow : index + vector.store %ep_v217_store_a, %ep_buffer[%ep_v217_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v217_store_b, %ep_buffer[%ep_v217_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v219_row0 = vector.extract %219[0] : f32 from vector<4xf32> + %ep_v219_row1 = vector.extract %219[1] : f32 from vector<4xf32> + %ep_v219_row2 = vector.extract %219[2] : f32 from vector<4xf32> + %ep_v219_row3 = vector.extract %219[3] : f32 from vector<4xf32> + %ep_v219_row0_nbr, %ep_v219_row0_valid = gpu.shuffle xor %ep_v219_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v219_row1_nbr, %ep_v219_row1_valid = gpu.shuffle xor %ep_v219_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v219_row2_nbr, %ep_v219_row2_valid = gpu.shuffle xor %ep_v219_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v219_row3_nbr, %ep_v219_row3_valid = gpu.shuffle xor %ep_v219_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v219_row0_lo = arith.select %ep_is_even, %ep_v219_row0, %ep_v219_row0_nbr : f32 + %ep_v219_row0_hi = arith.select %ep_is_even, %ep_v219_row0_nbr, %ep_v219_row0 : f32 + %ep_v219_row1_lo = arith.select %ep_is_even, %ep_v219_row1, %ep_v219_row1_nbr : f32 + %ep_v219_row1_hi = arith.select %ep_is_even, %ep_v219_row1_nbr, %ep_v219_row1 : f32 + %ep_v219_row2_lo = arith.select %ep_is_even, %ep_v219_row2, %ep_v219_row2_nbr : f32 + %ep_v219_row2_hi = arith.select %ep_is_even, %ep_v219_row2_nbr, %ep_v219_row2 : f32 + %ep_v219_row3_lo = arith.select %ep_is_even, %ep_v219_row3, %ep_v219_row3_nbr : f32 + %ep_v219_row3_hi = arith.select %ep_is_even, %ep_v219_row3_nbr, %ep_v219_row3 : f32 + %ep_v219_store_a_lo = arith.select %ep_is_even, %ep_v219_row0_lo, %ep_v219_row2_lo : f32 + %ep_v219_store_a_hi = arith.select %ep_is_even, %ep_v219_row0_hi, %ep_v219_row2_hi : f32 + %ep_v219_store_b_lo = arith.select %ep_is_even, %ep_v219_row1_lo, %ep_v219_row3_lo : f32 + %ep_v219_store_b_hi = arith.select %ep_is_even, %ep_v219_row1_hi, %ep_v219_row3_hi : f32 + %ep_v219_i32_0 = arith.constant 0 : i32 + %ep_v219_i32_1 = arith.constant 1 : i32 + %ep_v219_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v219_pair_a_0 = llvm.insertelement %ep_v219_store_a_lo, %ep_v219_undef_a[%ep_v219_i32_0 : i32] : vector<2xf32> + %ep_v219_pair_a = llvm.insertelement %ep_v219_store_a_hi, %ep_v219_pair_a_0[%ep_v219_i32_1 : i32] : vector<2xf32> + %ep_v219_store_a = llvm.fptrunc %ep_v219_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v219_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v219_pair_b_0 = llvm.insertelement %ep_v219_store_b_lo, %ep_v219_undef_b[%ep_v219_i32_0 : i32] : vector<2xf32> + %ep_v219_pair_b = llvm.insertelement %ep_v219_store_b_hi, %ep_v219_pair_b_0[%ep_v219_i32_1 : i32] : vector<2xf32> + %ep_v219_store_b = llvm.fptrunc %ep_v219_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v219_addr_a = arith.addi %ep_rb0_store_off_a, %ep_col5_adj overflow : index + %ep_v219_addr_b = arith.addi %ep_rb0_store_off_b, %ep_col5_adj overflow : index + vector.store %ep_v219_store_a, %ep_buffer[%ep_v219_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v219_store_b, %ep_buffer[%ep_v219_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v221_row0 = vector.extract %221[0] : f32 from vector<4xf32> + %ep_v221_row1 = vector.extract %221[1] : f32 from vector<4xf32> + %ep_v221_row2 = vector.extract %221[2] : f32 from vector<4xf32> + %ep_v221_row3 = vector.extract %221[3] : f32 from vector<4xf32> + %ep_v221_row0_nbr, %ep_v221_row0_valid = gpu.shuffle xor %ep_v221_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v221_row1_nbr, %ep_v221_row1_valid = gpu.shuffle xor %ep_v221_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v221_row2_nbr, %ep_v221_row2_valid = gpu.shuffle xor %ep_v221_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v221_row3_nbr, %ep_v221_row3_valid = gpu.shuffle xor %ep_v221_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v221_row0_lo = arith.select %ep_is_even, %ep_v221_row0, %ep_v221_row0_nbr : f32 + %ep_v221_row0_hi = arith.select %ep_is_even, %ep_v221_row0_nbr, %ep_v221_row0 : f32 + %ep_v221_row1_lo = arith.select %ep_is_even, %ep_v221_row1, %ep_v221_row1_nbr : f32 + %ep_v221_row1_hi = arith.select %ep_is_even, %ep_v221_row1_nbr, %ep_v221_row1 : f32 + %ep_v221_row2_lo = arith.select %ep_is_even, %ep_v221_row2, %ep_v221_row2_nbr : f32 + %ep_v221_row2_hi = arith.select %ep_is_even, %ep_v221_row2_nbr, %ep_v221_row2 : f32 + %ep_v221_row3_lo = arith.select %ep_is_even, %ep_v221_row3, %ep_v221_row3_nbr : f32 + %ep_v221_row3_hi = arith.select %ep_is_even, %ep_v221_row3_nbr, %ep_v221_row3 : f32 + %ep_v221_store_a_lo = arith.select %ep_is_even, %ep_v221_row0_lo, %ep_v221_row2_lo : f32 + %ep_v221_store_a_hi = arith.select %ep_is_even, %ep_v221_row0_hi, %ep_v221_row2_hi : f32 + %ep_v221_store_b_lo = arith.select %ep_is_even, %ep_v221_row1_lo, %ep_v221_row3_lo : f32 + %ep_v221_store_b_hi = arith.select %ep_is_even, %ep_v221_row1_hi, %ep_v221_row3_hi : f32 + %ep_v221_i32_0 = arith.constant 0 : i32 + %ep_v221_i32_1 = arith.constant 1 : i32 + %ep_v221_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v221_pair_a_0 = llvm.insertelement %ep_v221_store_a_lo, %ep_v221_undef_a[%ep_v221_i32_0 : i32] : vector<2xf32> + %ep_v221_pair_a = llvm.insertelement %ep_v221_store_a_hi, %ep_v221_pair_a_0[%ep_v221_i32_1 : i32] : vector<2xf32> + %ep_v221_store_a = llvm.fptrunc %ep_v221_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v221_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v221_pair_b_0 = llvm.insertelement %ep_v221_store_b_lo, %ep_v221_undef_b[%ep_v221_i32_0 : i32] : vector<2xf32> + %ep_v221_pair_b = llvm.insertelement %ep_v221_store_b_hi, %ep_v221_pair_b_0[%ep_v221_i32_1 : i32] : vector<2xf32> + %ep_v221_store_b = llvm.fptrunc %ep_v221_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v221_addr_a = arith.addi %ep_rb1_store_off_a, %ep_col0_adj overflow : index + %ep_v221_addr_b = arith.addi %ep_rb1_store_off_b, %ep_col0_adj overflow : index + vector.store %ep_v221_store_a, %ep_buffer[%ep_v221_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v221_store_b, %ep_buffer[%ep_v221_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v223_row0 = vector.extract %223[0] : f32 from vector<4xf32> + %ep_v223_row1 = vector.extract %223[1] : f32 from vector<4xf32> + %ep_v223_row2 = vector.extract %223[2] : f32 from vector<4xf32> + %ep_v223_row3 = vector.extract %223[3] : f32 from vector<4xf32> + %ep_v223_row0_nbr, %ep_v223_row0_valid = gpu.shuffle xor %ep_v223_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v223_row1_nbr, %ep_v223_row1_valid = gpu.shuffle xor %ep_v223_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v223_row2_nbr, %ep_v223_row2_valid = gpu.shuffle xor %ep_v223_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v223_row3_nbr, %ep_v223_row3_valid = gpu.shuffle xor %ep_v223_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v223_row0_lo = arith.select %ep_is_even, %ep_v223_row0, %ep_v223_row0_nbr : f32 + %ep_v223_row0_hi = arith.select %ep_is_even, %ep_v223_row0_nbr, %ep_v223_row0 : f32 + %ep_v223_row1_lo = arith.select %ep_is_even, %ep_v223_row1, %ep_v223_row1_nbr : f32 + %ep_v223_row1_hi = arith.select %ep_is_even, %ep_v223_row1_nbr, %ep_v223_row1 : f32 + %ep_v223_row2_lo = arith.select %ep_is_even, %ep_v223_row2, %ep_v223_row2_nbr : f32 + %ep_v223_row2_hi = arith.select %ep_is_even, %ep_v223_row2_nbr, %ep_v223_row2 : f32 + %ep_v223_row3_lo = arith.select %ep_is_even, %ep_v223_row3, %ep_v223_row3_nbr : f32 + %ep_v223_row3_hi = arith.select %ep_is_even, %ep_v223_row3_nbr, %ep_v223_row3 : f32 + %ep_v223_store_a_lo = arith.select %ep_is_even, %ep_v223_row0_lo, %ep_v223_row2_lo : f32 + %ep_v223_store_a_hi = arith.select %ep_is_even, %ep_v223_row0_hi, %ep_v223_row2_hi : f32 + %ep_v223_store_b_lo = arith.select %ep_is_even, %ep_v223_row1_lo, %ep_v223_row3_lo : f32 + %ep_v223_store_b_hi = arith.select %ep_is_even, %ep_v223_row1_hi, %ep_v223_row3_hi : f32 + %ep_v223_i32_0 = arith.constant 0 : i32 + %ep_v223_i32_1 = arith.constant 1 : i32 + %ep_v223_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v223_pair_a_0 = llvm.insertelement %ep_v223_store_a_lo, %ep_v223_undef_a[%ep_v223_i32_0 : i32] : vector<2xf32> + %ep_v223_pair_a = llvm.insertelement %ep_v223_store_a_hi, %ep_v223_pair_a_0[%ep_v223_i32_1 : i32] : vector<2xf32> + %ep_v223_store_a = llvm.fptrunc %ep_v223_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v223_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v223_pair_b_0 = llvm.insertelement %ep_v223_store_b_lo, %ep_v223_undef_b[%ep_v223_i32_0 : i32] : vector<2xf32> + %ep_v223_pair_b = llvm.insertelement %ep_v223_store_b_hi, %ep_v223_pair_b_0[%ep_v223_i32_1 : i32] : vector<2xf32> + %ep_v223_store_b = llvm.fptrunc %ep_v223_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v223_addr_a = arith.addi %ep_rb1_store_off_a, %ep_col1_adj overflow : index + %ep_v223_addr_b = arith.addi %ep_rb1_store_off_b, %ep_col1_adj overflow : index + vector.store %ep_v223_store_a, %ep_buffer[%ep_v223_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v223_store_b, %ep_buffer[%ep_v223_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v225_row0 = vector.extract %225[0] : f32 from vector<4xf32> + %ep_v225_row1 = vector.extract %225[1] : f32 from vector<4xf32> + %ep_v225_row2 = vector.extract %225[2] : f32 from vector<4xf32> + %ep_v225_row3 = vector.extract %225[3] : f32 from vector<4xf32> + %ep_v225_row0_nbr, %ep_v225_row0_valid = gpu.shuffle xor %ep_v225_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v225_row1_nbr, %ep_v225_row1_valid = gpu.shuffle xor %ep_v225_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v225_row2_nbr, %ep_v225_row2_valid = gpu.shuffle xor %ep_v225_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v225_row3_nbr, %ep_v225_row3_valid = gpu.shuffle xor %ep_v225_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v225_row0_lo = arith.select %ep_is_even, %ep_v225_row0, %ep_v225_row0_nbr : f32 + %ep_v225_row0_hi = arith.select %ep_is_even, %ep_v225_row0_nbr, %ep_v225_row0 : f32 + %ep_v225_row1_lo = arith.select %ep_is_even, %ep_v225_row1, %ep_v225_row1_nbr : f32 + %ep_v225_row1_hi = arith.select %ep_is_even, %ep_v225_row1_nbr, %ep_v225_row1 : f32 + %ep_v225_row2_lo = arith.select %ep_is_even, %ep_v225_row2, %ep_v225_row2_nbr : f32 + %ep_v225_row2_hi = arith.select %ep_is_even, %ep_v225_row2_nbr, %ep_v225_row2 : f32 + %ep_v225_row3_lo = arith.select %ep_is_even, %ep_v225_row3, %ep_v225_row3_nbr : f32 + %ep_v225_row3_hi = arith.select %ep_is_even, %ep_v225_row3_nbr, %ep_v225_row3 : f32 + %ep_v225_store_a_lo = arith.select %ep_is_even, %ep_v225_row0_lo, %ep_v225_row2_lo : f32 + %ep_v225_store_a_hi = arith.select %ep_is_even, %ep_v225_row0_hi, %ep_v225_row2_hi : f32 + %ep_v225_store_b_lo = arith.select %ep_is_even, %ep_v225_row1_lo, %ep_v225_row3_lo : f32 + %ep_v225_store_b_hi = arith.select %ep_is_even, %ep_v225_row1_hi, %ep_v225_row3_hi : f32 + %ep_v225_i32_0 = arith.constant 0 : i32 + %ep_v225_i32_1 = arith.constant 1 : i32 + %ep_v225_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v225_pair_a_0 = llvm.insertelement %ep_v225_store_a_lo, %ep_v225_undef_a[%ep_v225_i32_0 : i32] : vector<2xf32> + %ep_v225_pair_a = llvm.insertelement %ep_v225_store_a_hi, %ep_v225_pair_a_0[%ep_v225_i32_1 : i32] : vector<2xf32> + %ep_v225_store_a = llvm.fptrunc %ep_v225_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v225_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v225_pair_b_0 = llvm.insertelement %ep_v225_store_b_lo, %ep_v225_undef_b[%ep_v225_i32_0 : i32] : vector<2xf32> + %ep_v225_pair_b = llvm.insertelement %ep_v225_store_b_hi, %ep_v225_pair_b_0[%ep_v225_i32_1 : i32] : vector<2xf32> + %ep_v225_store_b = llvm.fptrunc %ep_v225_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v225_addr_a = arith.addi %ep_rb1_store_off_a, %ep_col2_adj overflow : index + %ep_v225_addr_b = arith.addi %ep_rb1_store_off_b, %ep_col2_adj overflow : index + vector.store %ep_v225_store_a, %ep_buffer[%ep_v225_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v225_store_b, %ep_buffer[%ep_v225_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v227_row0 = vector.extract %227[0] : f32 from vector<4xf32> + %ep_v227_row1 = vector.extract %227[1] : f32 from vector<4xf32> + %ep_v227_row2 = vector.extract %227[2] : f32 from vector<4xf32> + %ep_v227_row3 = vector.extract %227[3] : f32 from vector<4xf32> + %ep_v227_row0_nbr, %ep_v227_row0_valid = gpu.shuffle xor %ep_v227_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v227_row1_nbr, %ep_v227_row1_valid = gpu.shuffle xor %ep_v227_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v227_row2_nbr, %ep_v227_row2_valid = gpu.shuffle xor %ep_v227_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v227_row3_nbr, %ep_v227_row3_valid = gpu.shuffle xor %ep_v227_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v227_row0_lo = arith.select %ep_is_even, %ep_v227_row0, %ep_v227_row0_nbr : f32 + %ep_v227_row0_hi = arith.select %ep_is_even, %ep_v227_row0_nbr, %ep_v227_row0 : f32 + %ep_v227_row1_lo = arith.select %ep_is_even, %ep_v227_row1, %ep_v227_row1_nbr : f32 + %ep_v227_row1_hi = arith.select %ep_is_even, %ep_v227_row1_nbr, %ep_v227_row1 : f32 + %ep_v227_row2_lo = arith.select %ep_is_even, %ep_v227_row2, %ep_v227_row2_nbr : f32 + %ep_v227_row2_hi = arith.select %ep_is_even, %ep_v227_row2_nbr, %ep_v227_row2 : f32 + %ep_v227_row3_lo = arith.select %ep_is_even, %ep_v227_row3, %ep_v227_row3_nbr : f32 + %ep_v227_row3_hi = arith.select %ep_is_even, %ep_v227_row3_nbr, %ep_v227_row3 : f32 + %ep_v227_store_a_lo = arith.select %ep_is_even, %ep_v227_row0_lo, %ep_v227_row2_lo : f32 + %ep_v227_store_a_hi = arith.select %ep_is_even, %ep_v227_row0_hi, %ep_v227_row2_hi : f32 + %ep_v227_store_b_lo = arith.select %ep_is_even, %ep_v227_row1_lo, %ep_v227_row3_lo : f32 + %ep_v227_store_b_hi = arith.select %ep_is_even, %ep_v227_row1_hi, %ep_v227_row3_hi : f32 + %ep_v227_i32_0 = arith.constant 0 : i32 + %ep_v227_i32_1 = arith.constant 1 : i32 + %ep_v227_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v227_pair_a_0 = llvm.insertelement %ep_v227_store_a_lo, %ep_v227_undef_a[%ep_v227_i32_0 : i32] : vector<2xf32> + %ep_v227_pair_a = llvm.insertelement %ep_v227_store_a_hi, %ep_v227_pair_a_0[%ep_v227_i32_1 : i32] : vector<2xf32> + %ep_v227_store_a = llvm.fptrunc %ep_v227_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v227_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v227_pair_b_0 = llvm.insertelement %ep_v227_store_b_lo, %ep_v227_undef_b[%ep_v227_i32_0 : i32] : vector<2xf32> + %ep_v227_pair_b = llvm.insertelement %ep_v227_store_b_hi, %ep_v227_pair_b_0[%ep_v227_i32_1 : i32] : vector<2xf32> + %ep_v227_store_b = llvm.fptrunc %ep_v227_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v227_addr_a = arith.addi %ep_rb1_store_off_a, %ep_col3_adj overflow : index + %ep_v227_addr_b = arith.addi %ep_rb1_store_off_b, %ep_col3_adj overflow : index + vector.store %ep_v227_store_a, %ep_buffer[%ep_v227_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v227_store_b, %ep_buffer[%ep_v227_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v229_row0 = vector.extract %229[0] : f32 from vector<4xf32> + %ep_v229_row1 = vector.extract %229[1] : f32 from vector<4xf32> + %ep_v229_row2 = vector.extract %229[2] : f32 from vector<4xf32> + %ep_v229_row3 = vector.extract %229[3] : f32 from vector<4xf32> + %ep_v229_row0_nbr, %ep_v229_row0_valid = gpu.shuffle xor %ep_v229_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v229_row1_nbr, %ep_v229_row1_valid = gpu.shuffle xor %ep_v229_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v229_row2_nbr, %ep_v229_row2_valid = gpu.shuffle xor %ep_v229_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v229_row3_nbr, %ep_v229_row3_valid = gpu.shuffle xor %ep_v229_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v229_row0_lo = arith.select %ep_is_even, %ep_v229_row0, %ep_v229_row0_nbr : f32 + %ep_v229_row0_hi = arith.select %ep_is_even, %ep_v229_row0_nbr, %ep_v229_row0 : f32 + %ep_v229_row1_lo = arith.select %ep_is_even, %ep_v229_row1, %ep_v229_row1_nbr : f32 + %ep_v229_row1_hi = arith.select %ep_is_even, %ep_v229_row1_nbr, %ep_v229_row1 : f32 + %ep_v229_row2_lo = arith.select %ep_is_even, %ep_v229_row2, %ep_v229_row2_nbr : f32 + %ep_v229_row2_hi = arith.select %ep_is_even, %ep_v229_row2_nbr, %ep_v229_row2 : f32 + %ep_v229_row3_lo = arith.select %ep_is_even, %ep_v229_row3, %ep_v229_row3_nbr : f32 + %ep_v229_row3_hi = arith.select %ep_is_even, %ep_v229_row3_nbr, %ep_v229_row3 : f32 + %ep_v229_store_a_lo = arith.select %ep_is_even, %ep_v229_row0_lo, %ep_v229_row2_lo : f32 + %ep_v229_store_a_hi = arith.select %ep_is_even, %ep_v229_row0_hi, %ep_v229_row2_hi : f32 + %ep_v229_store_b_lo = arith.select %ep_is_even, %ep_v229_row1_lo, %ep_v229_row3_lo : f32 + %ep_v229_store_b_hi = arith.select %ep_is_even, %ep_v229_row1_hi, %ep_v229_row3_hi : f32 + %ep_v229_i32_0 = arith.constant 0 : i32 + %ep_v229_i32_1 = arith.constant 1 : i32 + %ep_v229_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v229_pair_a_0 = llvm.insertelement %ep_v229_store_a_lo, %ep_v229_undef_a[%ep_v229_i32_0 : i32] : vector<2xf32> + %ep_v229_pair_a = llvm.insertelement %ep_v229_store_a_hi, %ep_v229_pair_a_0[%ep_v229_i32_1 : i32] : vector<2xf32> + %ep_v229_store_a = llvm.fptrunc %ep_v229_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v229_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v229_pair_b_0 = llvm.insertelement %ep_v229_store_b_lo, %ep_v229_undef_b[%ep_v229_i32_0 : i32] : vector<2xf32> + %ep_v229_pair_b = llvm.insertelement %ep_v229_store_b_hi, %ep_v229_pair_b_0[%ep_v229_i32_1 : i32] : vector<2xf32> + %ep_v229_store_b = llvm.fptrunc %ep_v229_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v229_addr_a = arith.addi %ep_rb1_store_off_a, %ep_col4_adj overflow : index + %ep_v229_addr_b = arith.addi %ep_rb1_store_off_b, %ep_col4_adj overflow : index + vector.store %ep_v229_store_a, %ep_buffer[%ep_v229_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v229_store_b, %ep_buffer[%ep_v229_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v231_row0 = vector.extract %231[0] : f32 from vector<4xf32> + %ep_v231_row1 = vector.extract %231[1] : f32 from vector<4xf32> + %ep_v231_row2 = vector.extract %231[2] : f32 from vector<4xf32> + %ep_v231_row3 = vector.extract %231[3] : f32 from vector<4xf32> + %ep_v231_row0_nbr, %ep_v231_row0_valid = gpu.shuffle xor %ep_v231_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v231_row1_nbr, %ep_v231_row1_valid = gpu.shuffle xor %ep_v231_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v231_row2_nbr, %ep_v231_row2_valid = gpu.shuffle xor %ep_v231_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v231_row3_nbr, %ep_v231_row3_valid = gpu.shuffle xor %ep_v231_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v231_row0_lo = arith.select %ep_is_even, %ep_v231_row0, %ep_v231_row0_nbr : f32 + %ep_v231_row0_hi = arith.select %ep_is_even, %ep_v231_row0_nbr, %ep_v231_row0 : f32 + %ep_v231_row1_lo = arith.select %ep_is_even, %ep_v231_row1, %ep_v231_row1_nbr : f32 + %ep_v231_row1_hi = arith.select %ep_is_even, %ep_v231_row1_nbr, %ep_v231_row1 : f32 + %ep_v231_row2_lo = arith.select %ep_is_even, %ep_v231_row2, %ep_v231_row2_nbr : f32 + %ep_v231_row2_hi = arith.select %ep_is_even, %ep_v231_row2_nbr, %ep_v231_row2 : f32 + %ep_v231_row3_lo = arith.select %ep_is_even, %ep_v231_row3, %ep_v231_row3_nbr : f32 + %ep_v231_row3_hi = arith.select %ep_is_even, %ep_v231_row3_nbr, %ep_v231_row3 : f32 + %ep_v231_store_a_lo = arith.select %ep_is_even, %ep_v231_row0_lo, %ep_v231_row2_lo : f32 + %ep_v231_store_a_hi = arith.select %ep_is_even, %ep_v231_row0_hi, %ep_v231_row2_hi : f32 + %ep_v231_store_b_lo = arith.select %ep_is_even, %ep_v231_row1_lo, %ep_v231_row3_lo : f32 + %ep_v231_store_b_hi = arith.select %ep_is_even, %ep_v231_row1_hi, %ep_v231_row3_hi : f32 + %ep_v231_i32_0 = arith.constant 0 : i32 + %ep_v231_i32_1 = arith.constant 1 : i32 + %ep_v231_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v231_pair_a_0 = llvm.insertelement %ep_v231_store_a_lo, %ep_v231_undef_a[%ep_v231_i32_0 : i32] : vector<2xf32> + %ep_v231_pair_a = llvm.insertelement %ep_v231_store_a_hi, %ep_v231_pair_a_0[%ep_v231_i32_1 : i32] : vector<2xf32> + %ep_v231_store_a = llvm.fptrunc %ep_v231_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v231_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v231_pair_b_0 = llvm.insertelement %ep_v231_store_b_lo, %ep_v231_undef_b[%ep_v231_i32_0 : i32] : vector<2xf32> + %ep_v231_pair_b = llvm.insertelement %ep_v231_store_b_hi, %ep_v231_pair_b_0[%ep_v231_i32_1 : i32] : vector<2xf32> + %ep_v231_store_b = llvm.fptrunc %ep_v231_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v231_addr_a = arith.addi %ep_rb1_store_off_a, %ep_col5_adj overflow : index + %ep_v231_addr_b = arith.addi %ep_rb1_store_off_b, %ep_col5_adj overflow : index + vector.store %ep_v231_store_a, %ep_buffer[%ep_v231_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v231_store_b, %ep_buffer[%ep_v231_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v233_row0 = vector.extract %233[0] : f32 from vector<4xf32> + %ep_v233_row1 = vector.extract %233[1] : f32 from vector<4xf32> + %ep_v233_row2 = vector.extract %233[2] : f32 from vector<4xf32> + %ep_v233_row3 = vector.extract %233[3] : f32 from vector<4xf32> + %ep_v233_row0_nbr, %ep_v233_row0_valid = gpu.shuffle xor %ep_v233_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v233_row1_nbr, %ep_v233_row1_valid = gpu.shuffle xor %ep_v233_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v233_row2_nbr, %ep_v233_row2_valid = gpu.shuffle xor %ep_v233_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v233_row3_nbr, %ep_v233_row3_valid = gpu.shuffle xor %ep_v233_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v233_row0_lo = arith.select %ep_is_even, %ep_v233_row0, %ep_v233_row0_nbr : f32 + %ep_v233_row0_hi = arith.select %ep_is_even, %ep_v233_row0_nbr, %ep_v233_row0 : f32 + %ep_v233_row1_lo = arith.select %ep_is_even, %ep_v233_row1, %ep_v233_row1_nbr : f32 + %ep_v233_row1_hi = arith.select %ep_is_even, %ep_v233_row1_nbr, %ep_v233_row1 : f32 + %ep_v233_row2_lo = arith.select %ep_is_even, %ep_v233_row2, %ep_v233_row2_nbr : f32 + %ep_v233_row2_hi = arith.select %ep_is_even, %ep_v233_row2_nbr, %ep_v233_row2 : f32 + %ep_v233_row3_lo = arith.select %ep_is_even, %ep_v233_row3, %ep_v233_row3_nbr : f32 + %ep_v233_row3_hi = arith.select %ep_is_even, %ep_v233_row3_nbr, %ep_v233_row3 : f32 + %ep_v233_store_a_lo = arith.select %ep_is_even, %ep_v233_row0_lo, %ep_v233_row2_lo : f32 + %ep_v233_store_a_hi = arith.select %ep_is_even, %ep_v233_row0_hi, %ep_v233_row2_hi : f32 + %ep_v233_store_b_lo = arith.select %ep_is_even, %ep_v233_row1_lo, %ep_v233_row3_lo : f32 + %ep_v233_store_b_hi = arith.select %ep_is_even, %ep_v233_row1_hi, %ep_v233_row3_hi : f32 + %ep_v233_i32_0 = arith.constant 0 : i32 + %ep_v233_i32_1 = arith.constant 1 : i32 + %ep_v233_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v233_pair_a_0 = llvm.insertelement %ep_v233_store_a_lo, %ep_v233_undef_a[%ep_v233_i32_0 : i32] : vector<2xf32> + %ep_v233_pair_a = llvm.insertelement %ep_v233_store_a_hi, %ep_v233_pair_a_0[%ep_v233_i32_1 : i32] : vector<2xf32> + %ep_v233_store_a = llvm.fptrunc %ep_v233_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v233_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v233_pair_b_0 = llvm.insertelement %ep_v233_store_b_lo, %ep_v233_undef_b[%ep_v233_i32_0 : i32] : vector<2xf32> + %ep_v233_pair_b = llvm.insertelement %ep_v233_store_b_hi, %ep_v233_pair_b_0[%ep_v233_i32_1 : i32] : vector<2xf32> + %ep_v233_store_b = llvm.fptrunc %ep_v233_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v233_addr_a = arith.addi %ep_rb2_store_off_a, %ep_col0_adj overflow : index + %ep_v233_addr_b = arith.addi %ep_rb2_store_off_b, %ep_col0_adj overflow : index + vector.store %ep_v233_store_a, %ep_buffer[%ep_v233_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v233_store_b, %ep_buffer[%ep_v233_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v235_row0 = vector.extract %235[0] : f32 from vector<4xf32> + %ep_v235_row1 = vector.extract %235[1] : f32 from vector<4xf32> + %ep_v235_row2 = vector.extract %235[2] : f32 from vector<4xf32> + %ep_v235_row3 = vector.extract %235[3] : f32 from vector<4xf32> + %ep_v235_row0_nbr, %ep_v235_row0_valid = gpu.shuffle xor %ep_v235_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v235_row1_nbr, %ep_v235_row1_valid = gpu.shuffle xor %ep_v235_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v235_row2_nbr, %ep_v235_row2_valid = gpu.shuffle xor %ep_v235_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v235_row3_nbr, %ep_v235_row3_valid = gpu.shuffle xor %ep_v235_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v235_row0_lo = arith.select %ep_is_even, %ep_v235_row0, %ep_v235_row0_nbr : f32 + %ep_v235_row0_hi = arith.select %ep_is_even, %ep_v235_row0_nbr, %ep_v235_row0 : f32 + %ep_v235_row1_lo = arith.select %ep_is_even, %ep_v235_row1, %ep_v235_row1_nbr : f32 + %ep_v235_row1_hi = arith.select %ep_is_even, %ep_v235_row1_nbr, %ep_v235_row1 : f32 + %ep_v235_row2_lo = arith.select %ep_is_even, %ep_v235_row2, %ep_v235_row2_nbr : f32 + %ep_v235_row2_hi = arith.select %ep_is_even, %ep_v235_row2_nbr, %ep_v235_row2 : f32 + %ep_v235_row3_lo = arith.select %ep_is_even, %ep_v235_row3, %ep_v235_row3_nbr : f32 + %ep_v235_row3_hi = arith.select %ep_is_even, %ep_v235_row3_nbr, %ep_v235_row3 : f32 + %ep_v235_store_a_lo = arith.select %ep_is_even, %ep_v235_row0_lo, %ep_v235_row2_lo : f32 + %ep_v235_store_a_hi = arith.select %ep_is_even, %ep_v235_row0_hi, %ep_v235_row2_hi : f32 + %ep_v235_store_b_lo = arith.select %ep_is_even, %ep_v235_row1_lo, %ep_v235_row3_lo : f32 + %ep_v235_store_b_hi = arith.select %ep_is_even, %ep_v235_row1_hi, %ep_v235_row3_hi : f32 + %ep_v235_i32_0 = arith.constant 0 : i32 + %ep_v235_i32_1 = arith.constant 1 : i32 + %ep_v235_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v235_pair_a_0 = llvm.insertelement %ep_v235_store_a_lo, %ep_v235_undef_a[%ep_v235_i32_0 : i32] : vector<2xf32> + %ep_v235_pair_a = llvm.insertelement %ep_v235_store_a_hi, %ep_v235_pair_a_0[%ep_v235_i32_1 : i32] : vector<2xf32> + %ep_v235_store_a = llvm.fptrunc %ep_v235_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v235_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v235_pair_b_0 = llvm.insertelement %ep_v235_store_b_lo, %ep_v235_undef_b[%ep_v235_i32_0 : i32] : vector<2xf32> + %ep_v235_pair_b = llvm.insertelement %ep_v235_store_b_hi, %ep_v235_pair_b_0[%ep_v235_i32_1 : i32] : vector<2xf32> + %ep_v235_store_b = llvm.fptrunc %ep_v235_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v235_addr_a = arith.addi %ep_rb2_store_off_a, %ep_col1_adj overflow : index + %ep_v235_addr_b = arith.addi %ep_rb2_store_off_b, %ep_col1_adj overflow : index + vector.store %ep_v235_store_a, %ep_buffer[%ep_v235_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v235_store_b, %ep_buffer[%ep_v235_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v237_row0 = vector.extract %237[0] : f32 from vector<4xf32> + %ep_v237_row1 = vector.extract %237[1] : f32 from vector<4xf32> + %ep_v237_row2 = vector.extract %237[2] : f32 from vector<4xf32> + %ep_v237_row3 = vector.extract %237[3] : f32 from vector<4xf32> + %ep_v237_row0_nbr, %ep_v237_row0_valid = gpu.shuffle xor %ep_v237_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v237_row1_nbr, %ep_v237_row1_valid = gpu.shuffle xor %ep_v237_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v237_row2_nbr, %ep_v237_row2_valid = gpu.shuffle xor %ep_v237_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v237_row3_nbr, %ep_v237_row3_valid = gpu.shuffle xor %ep_v237_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v237_row0_lo = arith.select %ep_is_even, %ep_v237_row0, %ep_v237_row0_nbr : f32 + %ep_v237_row0_hi = arith.select %ep_is_even, %ep_v237_row0_nbr, %ep_v237_row0 : f32 + %ep_v237_row1_lo = arith.select %ep_is_even, %ep_v237_row1, %ep_v237_row1_nbr : f32 + %ep_v237_row1_hi = arith.select %ep_is_even, %ep_v237_row1_nbr, %ep_v237_row1 : f32 + %ep_v237_row2_lo = arith.select %ep_is_even, %ep_v237_row2, %ep_v237_row2_nbr : f32 + %ep_v237_row2_hi = arith.select %ep_is_even, %ep_v237_row2_nbr, %ep_v237_row2 : f32 + %ep_v237_row3_lo = arith.select %ep_is_even, %ep_v237_row3, %ep_v237_row3_nbr : f32 + %ep_v237_row3_hi = arith.select %ep_is_even, %ep_v237_row3_nbr, %ep_v237_row3 : f32 + %ep_v237_store_a_lo = arith.select %ep_is_even, %ep_v237_row0_lo, %ep_v237_row2_lo : f32 + %ep_v237_store_a_hi = arith.select %ep_is_even, %ep_v237_row0_hi, %ep_v237_row2_hi : f32 + %ep_v237_store_b_lo = arith.select %ep_is_even, %ep_v237_row1_lo, %ep_v237_row3_lo : f32 + %ep_v237_store_b_hi = arith.select %ep_is_even, %ep_v237_row1_hi, %ep_v237_row3_hi : f32 + %ep_v237_i32_0 = arith.constant 0 : i32 + %ep_v237_i32_1 = arith.constant 1 : i32 + %ep_v237_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v237_pair_a_0 = llvm.insertelement %ep_v237_store_a_lo, %ep_v237_undef_a[%ep_v237_i32_0 : i32] : vector<2xf32> + %ep_v237_pair_a = llvm.insertelement %ep_v237_store_a_hi, %ep_v237_pair_a_0[%ep_v237_i32_1 : i32] : vector<2xf32> + %ep_v237_store_a = llvm.fptrunc %ep_v237_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v237_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v237_pair_b_0 = llvm.insertelement %ep_v237_store_b_lo, %ep_v237_undef_b[%ep_v237_i32_0 : i32] : vector<2xf32> + %ep_v237_pair_b = llvm.insertelement %ep_v237_store_b_hi, %ep_v237_pair_b_0[%ep_v237_i32_1 : i32] : vector<2xf32> + %ep_v237_store_b = llvm.fptrunc %ep_v237_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v237_addr_a = arith.addi %ep_rb2_store_off_a, %ep_col2_adj overflow : index + %ep_v237_addr_b = arith.addi %ep_rb2_store_off_b, %ep_col2_adj overflow : index + vector.store %ep_v237_store_a, %ep_buffer[%ep_v237_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v237_store_b, %ep_buffer[%ep_v237_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v239_row0 = vector.extract %239[0] : f32 from vector<4xf32> + %ep_v239_row1 = vector.extract %239[1] : f32 from vector<4xf32> + %ep_v239_row2 = vector.extract %239[2] : f32 from vector<4xf32> + %ep_v239_row3 = vector.extract %239[3] : f32 from vector<4xf32> + %ep_v239_row0_nbr, %ep_v239_row0_valid = gpu.shuffle xor %ep_v239_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v239_row1_nbr, %ep_v239_row1_valid = gpu.shuffle xor %ep_v239_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v239_row2_nbr, %ep_v239_row2_valid = gpu.shuffle xor %ep_v239_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v239_row3_nbr, %ep_v239_row3_valid = gpu.shuffle xor %ep_v239_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v239_row0_lo = arith.select %ep_is_even, %ep_v239_row0, %ep_v239_row0_nbr : f32 + %ep_v239_row0_hi = arith.select %ep_is_even, %ep_v239_row0_nbr, %ep_v239_row0 : f32 + %ep_v239_row1_lo = arith.select %ep_is_even, %ep_v239_row1, %ep_v239_row1_nbr : f32 + %ep_v239_row1_hi = arith.select %ep_is_even, %ep_v239_row1_nbr, %ep_v239_row1 : f32 + %ep_v239_row2_lo = arith.select %ep_is_even, %ep_v239_row2, %ep_v239_row2_nbr : f32 + %ep_v239_row2_hi = arith.select %ep_is_even, %ep_v239_row2_nbr, %ep_v239_row2 : f32 + %ep_v239_row3_lo = arith.select %ep_is_even, %ep_v239_row3, %ep_v239_row3_nbr : f32 + %ep_v239_row3_hi = arith.select %ep_is_even, %ep_v239_row3_nbr, %ep_v239_row3 : f32 + %ep_v239_store_a_lo = arith.select %ep_is_even, %ep_v239_row0_lo, %ep_v239_row2_lo : f32 + %ep_v239_store_a_hi = arith.select %ep_is_even, %ep_v239_row0_hi, %ep_v239_row2_hi : f32 + %ep_v239_store_b_lo = arith.select %ep_is_even, %ep_v239_row1_lo, %ep_v239_row3_lo : f32 + %ep_v239_store_b_hi = arith.select %ep_is_even, %ep_v239_row1_hi, %ep_v239_row3_hi : f32 + %ep_v239_i32_0 = arith.constant 0 : i32 + %ep_v239_i32_1 = arith.constant 1 : i32 + %ep_v239_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v239_pair_a_0 = llvm.insertelement %ep_v239_store_a_lo, %ep_v239_undef_a[%ep_v239_i32_0 : i32] : vector<2xf32> + %ep_v239_pair_a = llvm.insertelement %ep_v239_store_a_hi, %ep_v239_pair_a_0[%ep_v239_i32_1 : i32] : vector<2xf32> + %ep_v239_store_a = llvm.fptrunc %ep_v239_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v239_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v239_pair_b_0 = llvm.insertelement %ep_v239_store_b_lo, %ep_v239_undef_b[%ep_v239_i32_0 : i32] : vector<2xf32> + %ep_v239_pair_b = llvm.insertelement %ep_v239_store_b_hi, %ep_v239_pair_b_0[%ep_v239_i32_1 : i32] : vector<2xf32> + %ep_v239_store_b = llvm.fptrunc %ep_v239_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v239_addr_a = arith.addi %ep_rb2_store_off_a, %ep_col3_adj overflow : index + %ep_v239_addr_b = arith.addi %ep_rb2_store_off_b, %ep_col3_adj overflow : index + vector.store %ep_v239_store_a, %ep_buffer[%ep_v239_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v239_store_b, %ep_buffer[%ep_v239_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v241_row0 = vector.extract %241[0] : f32 from vector<4xf32> + %ep_v241_row1 = vector.extract %241[1] : f32 from vector<4xf32> + %ep_v241_row2 = vector.extract %241[2] : f32 from vector<4xf32> + %ep_v241_row3 = vector.extract %241[3] : f32 from vector<4xf32> + %ep_v241_row0_nbr, %ep_v241_row0_valid = gpu.shuffle xor %ep_v241_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v241_row1_nbr, %ep_v241_row1_valid = gpu.shuffle xor %ep_v241_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v241_row2_nbr, %ep_v241_row2_valid = gpu.shuffle xor %ep_v241_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v241_row3_nbr, %ep_v241_row3_valid = gpu.shuffle xor %ep_v241_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v241_row0_lo = arith.select %ep_is_even, %ep_v241_row0, %ep_v241_row0_nbr : f32 + %ep_v241_row0_hi = arith.select %ep_is_even, %ep_v241_row0_nbr, %ep_v241_row0 : f32 + %ep_v241_row1_lo = arith.select %ep_is_even, %ep_v241_row1, %ep_v241_row1_nbr : f32 + %ep_v241_row1_hi = arith.select %ep_is_even, %ep_v241_row1_nbr, %ep_v241_row1 : f32 + %ep_v241_row2_lo = arith.select %ep_is_even, %ep_v241_row2, %ep_v241_row2_nbr : f32 + %ep_v241_row2_hi = arith.select %ep_is_even, %ep_v241_row2_nbr, %ep_v241_row2 : f32 + %ep_v241_row3_lo = arith.select %ep_is_even, %ep_v241_row3, %ep_v241_row3_nbr : f32 + %ep_v241_row3_hi = arith.select %ep_is_even, %ep_v241_row3_nbr, %ep_v241_row3 : f32 + %ep_v241_store_a_lo = arith.select %ep_is_even, %ep_v241_row0_lo, %ep_v241_row2_lo : f32 + %ep_v241_store_a_hi = arith.select %ep_is_even, %ep_v241_row0_hi, %ep_v241_row2_hi : f32 + %ep_v241_store_b_lo = arith.select %ep_is_even, %ep_v241_row1_lo, %ep_v241_row3_lo : f32 + %ep_v241_store_b_hi = arith.select %ep_is_even, %ep_v241_row1_hi, %ep_v241_row3_hi : f32 + %ep_v241_i32_0 = arith.constant 0 : i32 + %ep_v241_i32_1 = arith.constant 1 : i32 + %ep_v241_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v241_pair_a_0 = llvm.insertelement %ep_v241_store_a_lo, %ep_v241_undef_a[%ep_v241_i32_0 : i32] : vector<2xf32> + %ep_v241_pair_a = llvm.insertelement %ep_v241_store_a_hi, %ep_v241_pair_a_0[%ep_v241_i32_1 : i32] : vector<2xf32> + %ep_v241_store_a = llvm.fptrunc %ep_v241_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v241_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v241_pair_b_0 = llvm.insertelement %ep_v241_store_b_lo, %ep_v241_undef_b[%ep_v241_i32_0 : i32] : vector<2xf32> + %ep_v241_pair_b = llvm.insertelement %ep_v241_store_b_hi, %ep_v241_pair_b_0[%ep_v241_i32_1 : i32] : vector<2xf32> + %ep_v241_store_b = llvm.fptrunc %ep_v241_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v241_addr_a = arith.addi %ep_rb2_store_off_a, %ep_col4_adj overflow : index + %ep_v241_addr_b = arith.addi %ep_rb2_store_off_b, %ep_col4_adj overflow : index + vector.store %ep_v241_store_a, %ep_buffer[%ep_v241_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v241_store_b, %ep_buffer[%ep_v241_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v243_row0 = vector.extract %243[0] : f32 from vector<4xf32> + %ep_v243_row1 = vector.extract %243[1] : f32 from vector<4xf32> + %ep_v243_row2 = vector.extract %243[2] : f32 from vector<4xf32> + %ep_v243_row3 = vector.extract %243[3] : f32 from vector<4xf32> + %ep_v243_row0_nbr, %ep_v243_row0_valid = gpu.shuffle xor %ep_v243_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v243_row1_nbr, %ep_v243_row1_valid = gpu.shuffle xor %ep_v243_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v243_row2_nbr, %ep_v243_row2_valid = gpu.shuffle xor %ep_v243_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v243_row3_nbr, %ep_v243_row3_valid = gpu.shuffle xor %ep_v243_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v243_row0_lo = arith.select %ep_is_even, %ep_v243_row0, %ep_v243_row0_nbr : f32 + %ep_v243_row0_hi = arith.select %ep_is_even, %ep_v243_row0_nbr, %ep_v243_row0 : f32 + %ep_v243_row1_lo = arith.select %ep_is_even, %ep_v243_row1, %ep_v243_row1_nbr : f32 + %ep_v243_row1_hi = arith.select %ep_is_even, %ep_v243_row1_nbr, %ep_v243_row1 : f32 + %ep_v243_row2_lo = arith.select %ep_is_even, %ep_v243_row2, %ep_v243_row2_nbr : f32 + %ep_v243_row2_hi = arith.select %ep_is_even, %ep_v243_row2_nbr, %ep_v243_row2 : f32 + %ep_v243_row3_lo = arith.select %ep_is_even, %ep_v243_row3, %ep_v243_row3_nbr : f32 + %ep_v243_row3_hi = arith.select %ep_is_even, %ep_v243_row3_nbr, %ep_v243_row3 : f32 + %ep_v243_store_a_lo = arith.select %ep_is_even, %ep_v243_row0_lo, %ep_v243_row2_lo : f32 + %ep_v243_store_a_hi = arith.select %ep_is_even, %ep_v243_row0_hi, %ep_v243_row2_hi : f32 + %ep_v243_store_b_lo = arith.select %ep_is_even, %ep_v243_row1_lo, %ep_v243_row3_lo : f32 + %ep_v243_store_b_hi = arith.select %ep_is_even, %ep_v243_row1_hi, %ep_v243_row3_hi : f32 + %ep_v243_i32_0 = arith.constant 0 : i32 + %ep_v243_i32_1 = arith.constant 1 : i32 + %ep_v243_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v243_pair_a_0 = llvm.insertelement %ep_v243_store_a_lo, %ep_v243_undef_a[%ep_v243_i32_0 : i32] : vector<2xf32> + %ep_v243_pair_a = llvm.insertelement %ep_v243_store_a_hi, %ep_v243_pair_a_0[%ep_v243_i32_1 : i32] : vector<2xf32> + %ep_v243_store_a = llvm.fptrunc %ep_v243_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v243_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v243_pair_b_0 = llvm.insertelement %ep_v243_store_b_lo, %ep_v243_undef_b[%ep_v243_i32_0 : i32] : vector<2xf32> + %ep_v243_pair_b = llvm.insertelement %ep_v243_store_b_hi, %ep_v243_pair_b_0[%ep_v243_i32_1 : i32] : vector<2xf32> + %ep_v243_store_b = llvm.fptrunc %ep_v243_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v243_addr_a = arith.addi %ep_rb2_store_off_a, %ep_col5_adj overflow : index + %ep_v243_addr_b = arith.addi %ep_rb2_store_off_b, %ep_col5_adj overflow : index + vector.store %ep_v243_store_a, %ep_buffer[%ep_v243_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v243_store_b, %ep_buffer[%ep_v243_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v245_row0 = vector.extract %245[0] : f32 from vector<4xf32> + %ep_v245_row1 = vector.extract %245[1] : f32 from vector<4xf32> + %ep_v245_row2 = vector.extract %245[2] : f32 from vector<4xf32> + %ep_v245_row3 = vector.extract %245[3] : f32 from vector<4xf32> + %ep_v245_row0_nbr, %ep_v245_row0_valid = gpu.shuffle xor %ep_v245_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v245_row1_nbr, %ep_v245_row1_valid = gpu.shuffle xor %ep_v245_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v245_row2_nbr, %ep_v245_row2_valid = gpu.shuffle xor %ep_v245_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v245_row3_nbr, %ep_v245_row3_valid = gpu.shuffle xor %ep_v245_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v245_row0_lo = arith.select %ep_is_even, %ep_v245_row0, %ep_v245_row0_nbr : f32 + %ep_v245_row0_hi = arith.select %ep_is_even, %ep_v245_row0_nbr, %ep_v245_row0 : f32 + %ep_v245_row1_lo = arith.select %ep_is_even, %ep_v245_row1, %ep_v245_row1_nbr : f32 + %ep_v245_row1_hi = arith.select %ep_is_even, %ep_v245_row1_nbr, %ep_v245_row1 : f32 + %ep_v245_row2_lo = arith.select %ep_is_even, %ep_v245_row2, %ep_v245_row2_nbr : f32 + %ep_v245_row2_hi = arith.select %ep_is_even, %ep_v245_row2_nbr, %ep_v245_row2 : f32 + %ep_v245_row3_lo = arith.select %ep_is_even, %ep_v245_row3, %ep_v245_row3_nbr : f32 + %ep_v245_row3_hi = arith.select %ep_is_even, %ep_v245_row3_nbr, %ep_v245_row3 : f32 + %ep_v245_store_a_lo = arith.select %ep_is_even, %ep_v245_row0_lo, %ep_v245_row2_lo : f32 + %ep_v245_store_a_hi = arith.select %ep_is_even, %ep_v245_row0_hi, %ep_v245_row2_hi : f32 + %ep_v245_store_b_lo = arith.select %ep_is_even, %ep_v245_row1_lo, %ep_v245_row3_lo : f32 + %ep_v245_store_b_hi = arith.select %ep_is_even, %ep_v245_row1_hi, %ep_v245_row3_hi : f32 + %ep_v245_i32_0 = arith.constant 0 : i32 + %ep_v245_i32_1 = arith.constant 1 : i32 + %ep_v245_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v245_pair_a_0 = llvm.insertelement %ep_v245_store_a_lo, %ep_v245_undef_a[%ep_v245_i32_0 : i32] : vector<2xf32> + %ep_v245_pair_a = llvm.insertelement %ep_v245_store_a_hi, %ep_v245_pair_a_0[%ep_v245_i32_1 : i32] : vector<2xf32> + %ep_v245_store_a = llvm.fptrunc %ep_v245_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v245_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v245_pair_b_0 = llvm.insertelement %ep_v245_store_b_lo, %ep_v245_undef_b[%ep_v245_i32_0 : i32] : vector<2xf32> + %ep_v245_pair_b = llvm.insertelement %ep_v245_store_b_hi, %ep_v245_pair_b_0[%ep_v245_i32_1 : i32] : vector<2xf32> + %ep_v245_store_b = llvm.fptrunc %ep_v245_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v245_addr_a = arith.addi %ep_rb3_store_off_a, %ep_col0_adj overflow : index + %ep_v245_addr_b = arith.addi %ep_rb3_store_off_b, %ep_col0_adj overflow : index + vector.store %ep_v245_store_a, %ep_buffer[%ep_v245_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v245_store_b, %ep_buffer[%ep_v245_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v247_row0 = vector.extract %247[0] : f32 from vector<4xf32> + %ep_v247_row1 = vector.extract %247[1] : f32 from vector<4xf32> + %ep_v247_row2 = vector.extract %247[2] : f32 from vector<4xf32> + %ep_v247_row3 = vector.extract %247[3] : f32 from vector<4xf32> + %ep_v247_row0_nbr, %ep_v247_row0_valid = gpu.shuffle xor %ep_v247_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v247_row1_nbr, %ep_v247_row1_valid = gpu.shuffle xor %ep_v247_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v247_row2_nbr, %ep_v247_row2_valid = gpu.shuffle xor %ep_v247_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v247_row3_nbr, %ep_v247_row3_valid = gpu.shuffle xor %ep_v247_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v247_row0_lo = arith.select %ep_is_even, %ep_v247_row0, %ep_v247_row0_nbr : f32 + %ep_v247_row0_hi = arith.select %ep_is_even, %ep_v247_row0_nbr, %ep_v247_row0 : f32 + %ep_v247_row1_lo = arith.select %ep_is_even, %ep_v247_row1, %ep_v247_row1_nbr : f32 + %ep_v247_row1_hi = arith.select %ep_is_even, %ep_v247_row1_nbr, %ep_v247_row1 : f32 + %ep_v247_row2_lo = arith.select %ep_is_even, %ep_v247_row2, %ep_v247_row2_nbr : f32 + %ep_v247_row2_hi = arith.select %ep_is_even, %ep_v247_row2_nbr, %ep_v247_row2 : f32 + %ep_v247_row3_lo = arith.select %ep_is_even, %ep_v247_row3, %ep_v247_row3_nbr : f32 + %ep_v247_row3_hi = arith.select %ep_is_even, %ep_v247_row3_nbr, %ep_v247_row3 : f32 + %ep_v247_store_a_lo = arith.select %ep_is_even, %ep_v247_row0_lo, %ep_v247_row2_lo : f32 + %ep_v247_store_a_hi = arith.select %ep_is_even, %ep_v247_row0_hi, %ep_v247_row2_hi : f32 + %ep_v247_store_b_lo = arith.select %ep_is_even, %ep_v247_row1_lo, %ep_v247_row3_lo : f32 + %ep_v247_store_b_hi = arith.select %ep_is_even, %ep_v247_row1_hi, %ep_v247_row3_hi : f32 + %ep_v247_i32_0 = arith.constant 0 : i32 + %ep_v247_i32_1 = arith.constant 1 : i32 + %ep_v247_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v247_pair_a_0 = llvm.insertelement %ep_v247_store_a_lo, %ep_v247_undef_a[%ep_v247_i32_0 : i32] : vector<2xf32> + %ep_v247_pair_a = llvm.insertelement %ep_v247_store_a_hi, %ep_v247_pair_a_0[%ep_v247_i32_1 : i32] : vector<2xf32> + %ep_v247_store_a = llvm.fptrunc %ep_v247_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v247_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v247_pair_b_0 = llvm.insertelement %ep_v247_store_b_lo, %ep_v247_undef_b[%ep_v247_i32_0 : i32] : vector<2xf32> + %ep_v247_pair_b = llvm.insertelement %ep_v247_store_b_hi, %ep_v247_pair_b_0[%ep_v247_i32_1 : i32] : vector<2xf32> + %ep_v247_store_b = llvm.fptrunc %ep_v247_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v247_addr_a = arith.addi %ep_rb3_store_off_a, %ep_col1_adj overflow : index + %ep_v247_addr_b = arith.addi %ep_rb3_store_off_b, %ep_col1_adj overflow : index + vector.store %ep_v247_store_a, %ep_buffer[%ep_v247_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v247_store_b, %ep_buffer[%ep_v247_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v249_row0 = vector.extract %249[0] : f32 from vector<4xf32> + %ep_v249_row1 = vector.extract %249[1] : f32 from vector<4xf32> + %ep_v249_row2 = vector.extract %249[2] : f32 from vector<4xf32> + %ep_v249_row3 = vector.extract %249[3] : f32 from vector<4xf32> + %ep_v249_row0_nbr, %ep_v249_row0_valid = gpu.shuffle xor %ep_v249_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v249_row1_nbr, %ep_v249_row1_valid = gpu.shuffle xor %ep_v249_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v249_row2_nbr, %ep_v249_row2_valid = gpu.shuffle xor %ep_v249_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v249_row3_nbr, %ep_v249_row3_valid = gpu.shuffle xor %ep_v249_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v249_row0_lo = arith.select %ep_is_even, %ep_v249_row0, %ep_v249_row0_nbr : f32 + %ep_v249_row0_hi = arith.select %ep_is_even, %ep_v249_row0_nbr, %ep_v249_row0 : f32 + %ep_v249_row1_lo = arith.select %ep_is_even, %ep_v249_row1, %ep_v249_row1_nbr : f32 + %ep_v249_row1_hi = arith.select %ep_is_even, %ep_v249_row1_nbr, %ep_v249_row1 : f32 + %ep_v249_row2_lo = arith.select %ep_is_even, %ep_v249_row2, %ep_v249_row2_nbr : f32 + %ep_v249_row2_hi = arith.select %ep_is_even, %ep_v249_row2_nbr, %ep_v249_row2 : f32 + %ep_v249_row3_lo = arith.select %ep_is_even, %ep_v249_row3, %ep_v249_row3_nbr : f32 + %ep_v249_row3_hi = arith.select %ep_is_even, %ep_v249_row3_nbr, %ep_v249_row3 : f32 + %ep_v249_store_a_lo = arith.select %ep_is_even, %ep_v249_row0_lo, %ep_v249_row2_lo : f32 + %ep_v249_store_a_hi = arith.select %ep_is_even, %ep_v249_row0_hi, %ep_v249_row2_hi : f32 + %ep_v249_store_b_lo = arith.select %ep_is_even, %ep_v249_row1_lo, %ep_v249_row3_lo : f32 + %ep_v249_store_b_hi = arith.select %ep_is_even, %ep_v249_row1_hi, %ep_v249_row3_hi : f32 + %ep_v249_i32_0 = arith.constant 0 : i32 + %ep_v249_i32_1 = arith.constant 1 : i32 + %ep_v249_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v249_pair_a_0 = llvm.insertelement %ep_v249_store_a_lo, %ep_v249_undef_a[%ep_v249_i32_0 : i32] : vector<2xf32> + %ep_v249_pair_a = llvm.insertelement %ep_v249_store_a_hi, %ep_v249_pair_a_0[%ep_v249_i32_1 : i32] : vector<2xf32> + %ep_v249_store_a = llvm.fptrunc %ep_v249_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v249_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v249_pair_b_0 = llvm.insertelement %ep_v249_store_b_lo, %ep_v249_undef_b[%ep_v249_i32_0 : i32] : vector<2xf32> + %ep_v249_pair_b = llvm.insertelement %ep_v249_store_b_hi, %ep_v249_pair_b_0[%ep_v249_i32_1 : i32] : vector<2xf32> + %ep_v249_store_b = llvm.fptrunc %ep_v249_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v249_addr_a = arith.addi %ep_rb3_store_off_a, %ep_col2_adj overflow : index + %ep_v249_addr_b = arith.addi %ep_rb3_store_off_b, %ep_col2_adj overflow : index + vector.store %ep_v249_store_a, %ep_buffer[%ep_v249_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v249_store_b, %ep_buffer[%ep_v249_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v251_row0 = vector.extract %251[0] : f32 from vector<4xf32> + %ep_v251_row1 = vector.extract %251[1] : f32 from vector<4xf32> + %ep_v251_row2 = vector.extract %251[2] : f32 from vector<4xf32> + %ep_v251_row3 = vector.extract %251[3] : f32 from vector<4xf32> + %ep_v251_row0_nbr, %ep_v251_row0_valid = gpu.shuffle xor %ep_v251_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v251_row1_nbr, %ep_v251_row1_valid = gpu.shuffle xor %ep_v251_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v251_row2_nbr, %ep_v251_row2_valid = gpu.shuffle xor %ep_v251_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v251_row3_nbr, %ep_v251_row3_valid = gpu.shuffle xor %ep_v251_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v251_row0_lo = arith.select %ep_is_even, %ep_v251_row0, %ep_v251_row0_nbr : f32 + %ep_v251_row0_hi = arith.select %ep_is_even, %ep_v251_row0_nbr, %ep_v251_row0 : f32 + %ep_v251_row1_lo = arith.select %ep_is_even, %ep_v251_row1, %ep_v251_row1_nbr : f32 + %ep_v251_row1_hi = arith.select %ep_is_even, %ep_v251_row1_nbr, %ep_v251_row1 : f32 + %ep_v251_row2_lo = arith.select %ep_is_even, %ep_v251_row2, %ep_v251_row2_nbr : f32 + %ep_v251_row2_hi = arith.select %ep_is_even, %ep_v251_row2_nbr, %ep_v251_row2 : f32 + %ep_v251_row3_lo = arith.select %ep_is_even, %ep_v251_row3, %ep_v251_row3_nbr : f32 + %ep_v251_row3_hi = arith.select %ep_is_even, %ep_v251_row3_nbr, %ep_v251_row3 : f32 + %ep_v251_store_a_lo = arith.select %ep_is_even, %ep_v251_row0_lo, %ep_v251_row2_lo : f32 + %ep_v251_store_a_hi = arith.select %ep_is_even, %ep_v251_row0_hi, %ep_v251_row2_hi : f32 + %ep_v251_store_b_lo = arith.select %ep_is_even, %ep_v251_row1_lo, %ep_v251_row3_lo : f32 + %ep_v251_store_b_hi = arith.select %ep_is_even, %ep_v251_row1_hi, %ep_v251_row3_hi : f32 + %ep_v251_i32_0 = arith.constant 0 : i32 + %ep_v251_i32_1 = arith.constant 1 : i32 + %ep_v251_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v251_pair_a_0 = llvm.insertelement %ep_v251_store_a_lo, %ep_v251_undef_a[%ep_v251_i32_0 : i32] : vector<2xf32> + %ep_v251_pair_a = llvm.insertelement %ep_v251_store_a_hi, %ep_v251_pair_a_0[%ep_v251_i32_1 : i32] : vector<2xf32> + %ep_v251_store_a = llvm.fptrunc %ep_v251_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v251_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v251_pair_b_0 = llvm.insertelement %ep_v251_store_b_lo, %ep_v251_undef_b[%ep_v251_i32_0 : i32] : vector<2xf32> + %ep_v251_pair_b = llvm.insertelement %ep_v251_store_b_hi, %ep_v251_pair_b_0[%ep_v251_i32_1 : i32] : vector<2xf32> + %ep_v251_store_b = llvm.fptrunc %ep_v251_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v251_addr_a = arith.addi %ep_rb3_store_off_a, %ep_col3_adj overflow : index + %ep_v251_addr_b = arith.addi %ep_rb3_store_off_b, %ep_col3_adj overflow : index + vector.store %ep_v251_store_a, %ep_buffer[%ep_v251_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v251_store_b, %ep_buffer[%ep_v251_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v253_row0 = vector.extract %253[0] : f32 from vector<4xf32> + %ep_v253_row1 = vector.extract %253[1] : f32 from vector<4xf32> + %ep_v253_row2 = vector.extract %253[2] : f32 from vector<4xf32> + %ep_v253_row3 = vector.extract %253[3] : f32 from vector<4xf32> + %ep_v253_row0_nbr, %ep_v253_row0_valid = gpu.shuffle xor %ep_v253_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v253_row1_nbr, %ep_v253_row1_valid = gpu.shuffle xor %ep_v253_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v253_row2_nbr, %ep_v253_row2_valid = gpu.shuffle xor %ep_v253_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v253_row3_nbr, %ep_v253_row3_valid = gpu.shuffle xor %ep_v253_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v253_row0_lo = arith.select %ep_is_even, %ep_v253_row0, %ep_v253_row0_nbr : f32 + %ep_v253_row0_hi = arith.select %ep_is_even, %ep_v253_row0_nbr, %ep_v253_row0 : f32 + %ep_v253_row1_lo = arith.select %ep_is_even, %ep_v253_row1, %ep_v253_row1_nbr : f32 + %ep_v253_row1_hi = arith.select %ep_is_even, %ep_v253_row1_nbr, %ep_v253_row1 : f32 + %ep_v253_row2_lo = arith.select %ep_is_even, %ep_v253_row2, %ep_v253_row2_nbr : f32 + %ep_v253_row2_hi = arith.select %ep_is_even, %ep_v253_row2_nbr, %ep_v253_row2 : f32 + %ep_v253_row3_lo = arith.select %ep_is_even, %ep_v253_row3, %ep_v253_row3_nbr : f32 + %ep_v253_row3_hi = arith.select %ep_is_even, %ep_v253_row3_nbr, %ep_v253_row3 : f32 + %ep_v253_store_a_lo = arith.select %ep_is_even, %ep_v253_row0_lo, %ep_v253_row2_lo : f32 + %ep_v253_store_a_hi = arith.select %ep_is_even, %ep_v253_row0_hi, %ep_v253_row2_hi : f32 + %ep_v253_store_b_lo = arith.select %ep_is_even, %ep_v253_row1_lo, %ep_v253_row3_lo : f32 + %ep_v253_store_b_hi = arith.select %ep_is_even, %ep_v253_row1_hi, %ep_v253_row3_hi : f32 + %ep_v253_i32_0 = arith.constant 0 : i32 + %ep_v253_i32_1 = arith.constant 1 : i32 + %ep_v253_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v253_pair_a_0 = llvm.insertelement %ep_v253_store_a_lo, %ep_v253_undef_a[%ep_v253_i32_0 : i32] : vector<2xf32> + %ep_v253_pair_a = llvm.insertelement %ep_v253_store_a_hi, %ep_v253_pair_a_0[%ep_v253_i32_1 : i32] : vector<2xf32> + %ep_v253_store_a = llvm.fptrunc %ep_v253_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v253_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v253_pair_b_0 = llvm.insertelement %ep_v253_store_b_lo, %ep_v253_undef_b[%ep_v253_i32_0 : i32] : vector<2xf32> + %ep_v253_pair_b = llvm.insertelement %ep_v253_store_b_hi, %ep_v253_pair_b_0[%ep_v253_i32_1 : i32] : vector<2xf32> + %ep_v253_store_b = llvm.fptrunc %ep_v253_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v253_addr_a = arith.addi %ep_rb3_store_off_a, %ep_col4_adj overflow : index + %ep_v253_addr_b = arith.addi %ep_rb3_store_off_b, %ep_col4_adj overflow : index + vector.store %ep_v253_store_a, %ep_buffer[%ep_v253_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v253_store_b, %ep_buffer[%ep_v253_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + %ep_v255_row0 = vector.extract %255[0] : f32 from vector<4xf32> + %ep_v255_row1 = vector.extract %255[1] : f32 from vector<4xf32> + %ep_v255_row2 = vector.extract %255[2] : f32 from vector<4xf32> + %ep_v255_row3 = vector.extract %255[3] : f32 from vector<4xf32> + %ep_v255_row0_nbr, %ep_v255_row0_valid = gpu.shuffle xor %ep_v255_row0, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v255_row1_nbr, %ep_v255_row1_valid = gpu.shuffle xor %ep_v255_row1, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v255_row2_nbr, %ep_v255_row2_valid = gpu.shuffle xor %ep_v255_row2, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v255_row3_nbr, %ep_v255_row3_valid = gpu.shuffle xor %ep_v255_row3, %ep_shuffle_offset, %ep_shuffle_width : f32 + %ep_v255_row0_lo = arith.select %ep_is_even, %ep_v255_row0, %ep_v255_row0_nbr : f32 + %ep_v255_row0_hi = arith.select %ep_is_even, %ep_v255_row0_nbr, %ep_v255_row0 : f32 + %ep_v255_row1_lo = arith.select %ep_is_even, %ep_v255_row1, %ep_v255_row1_nbr : f32 + %ep_v255_row1_hi = arith.select %ep_is_even, %ep_v255_row1_nbr, %ep_v255_row1 : f32 + %ep_v255_row2_lo = arith.select %ep_is_even, %ep_v255_row2, %ep_v255_row2_nbr : f32 + %ep_v255_row2_hi = arith.select %ep_is_even, %ep_v255_row2_nbr, %ep_v255_row2 : f32 + %ep_v255_row3_lo = arith.select %ep_is_even, %ep_v255_row3, %ep_v255_row3_nbr : f32 + %ep_v255_row3_hi = arith.select %ep_is_even, %ep_v255_row3_nbr, %ep_v255_row3 : f32 + %ep_v255_store_a_lo = arith.select %ep_is_even, %ep_v255_row0_lo, %ep_v255_row2_lo : f32 + %ep_v255_store_a_hi = arith.select %ep_is_even, %ep_v255_row0_hi, %ep_v255_row2_hi : f32 + %ep_v255_store_b_lo = arith.select %ep_is_even, %ep_v255_row1_lo, %ep_v255_row3_lo : f32 + %ep_v255_store_b_hi = arith.select %ep_is_even, %ep_v255_row1_hi, %ep_v255_row3_hi : f32 + %ep_v255_i32_0 = arith.constant 0 : i32 + %ep_v255_i32_1 = arith.constant 1 : i32 + %ep_v255_undef_a = llvm.mlir.undef : vector<2xf32> + %ep_v255_pair_a_0 = llvm.insertelement %ep_v255_store_a_lo, %ep_v255_undef_a[%ep_v255_i32_0 : i32] : vector<2xf32> + %ep_v255_pair_a = llvm.insertelement %ep_v255_store_a_hi, %ep_v255_pair_a_0[%ep_v255_i32_1 : i32] : vector<2xf32> + %ep_v255_store_a = llvm.fptrunc %ep_v255_pair_a : vector<2xf32> to vector<2xbf16> + %ep_v255_undef_b = llvm.mlir.undef : vector<2xf32> + %ep_v255_pair_b_0 = llvm.insertelement %ep_v255_store_b_lo, %ep_v255_undef_b[%ep_v255_i32_0 : i32] : vector<2xf32> + %ep_v255_pair_b = llvm.insertelement %ep_v255_store_b_hi, %ep_v255_pair_b_0[%ep_v255_i32_1 : i32] : vector<2xf32> + %ep_v255_store_b = llvm.fptrunc %ep_v255_pair_b : vector<2xf32> to vector<2xbf16> + %ep_v255_addr_a = arith.addi %ep_rb3_store_off_a, %ep_col5_adj overflow : index + %ep_v255_addr_b = arith.addi %ep_rb3_store_off_b, %ep_col5_adj overflow : index + vector.store %ep_v255_store_a, %ep_buffer[%ep_v255_addr_a] {alignment = 4 : i64} : memref>, vector<2xbf16> + vector.store %ep_v255_store_b, %ep_buffer[%ep_v255_addr_b] {alignment = 4 : i64} : memref>, vector<2xbf16> + + return + } + } + } + func.func @isolated_benchmark$async(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view, %arg2: !hal.buffer_view, %arg3: !hal.buffer_view, %arg4: !hal.buffer_view, %arg5: index, %arg6: index, %arg7: index, %arg8: index, %arg9: index, %arg10: !hal.fence, %arg11: !hal.fence) -> !hal.buffer_view { + %0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index + %1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index + %2 = hal.tensor.import wait(%arg10) => %arg0 : !hal.buffer_view -> tensor{%0, %1} + %3 = hal.buffer_view.dim<%arg1 : !hal.buffer_view>[0] : index + %4 = hal.buffer_view.dim<%arg1 : !hal.buffer_view>[1] : index + %5 = hal.tensor.import wait(%arg10) => %arg1 : !hal.buffer_view -> tensor{%3, %4} + %6 = hal.buffer_view.dim<%arg2 : !hal.buffer_view>[0] : index + %7 = hal.buffer_view.dim<%arg2 : !hal.buffer_view>[1] : index + %8 = hal.tensor.import wait(%arg10) => %arg2 : !hal.buffer_view -> tensor{%6, %7} + %9 = hal.buffer_view.dim<%arg3 : !hal.buffer_view>[0] : index + %10 = hal.buffer_view.dim<%arg3 : !hal.buffer_view>[1] : index + %11 = hal.tensor.import wait(%arg10) => %arg3 : !hal.buffer_view -> tensor{%9, %10} + %12 = hal.buffer_view.dim<%arg4 : !hal.buffer_view>[0] : index + %13 = hal.buffer_view.dim<%arg4 : !hal.buffer_view>[1] : index + %14 = hal.tensor.import wait(%arg10) => %arg4 : !hal.buffer_view -> tensor{%12, %13} + %c0 = arith.constant 0 : index + %dim = tensor.dim %2, %c0 : tensor + %c0_0 = arith.constant 0 : index + %dim_1 = tensor.dim %8, %c0_0 : tensor + %c1 = arith.constant 1 : index + %dim_2 = tensor.dim %2, %c1 : tensor + %c2 = arith.constant 2 : index + %15 = arith.muli %dim_2, %c2 : index + %c32 = arith.constant 32 : index + %16 = arith.divsi %15, %c32 : index + %c2_3 = arith.constant 2 : index + %17 = arith.divsi %15, %c2_3 : index + %18 = flow.dispatch @gemm::@gemm[%dim, %dim_1, %15, %arg5, %arg6, %arg7, %arg8, %arg9](%2, %5, %8, %11, %14, %arg5, %arg6, %arg7, %arg8, %arg9, %dim, %dim_1, %15) : (tensor{%dim, %17}, tensor{%dim, %16}, tensor{%dim_1, %17}, tensor{%dim_1, %16}, tensor{%dim, %dim_1}, index, index, index, index, index, index, index, index) -> %14{%dim, %dim_1} + %19 = hal.tensor.barrier join(%18 : tensor) => %arg11 : !hal.fence + %c0_4 = arith.constant 0 : index + %dim_5 = tensor.dim %18, %c0_4 : tensor + %c1_6 = arith.constant 1 : index + %dim_7 = tensor.dim %18, %c1_6 : tensor + %20 = hal.tensor.export %19 : tensor{%dim_5, %dim_7} -> !hal.buffer_view + return %20 : !hal.buffer_view + } + } + """ wave_shape = _get_8wave_shape_from_block(block) gemm, options = get_tagged_mxfp4_gemm_preshuffle_scales_and_B( shape, block, wave_shape=wave_shape, b_address_space=SHARED_ADDRESS_SPACE, + output_dtype=tkl.bf16, ) options.specialize = True options.use_buffer_ops = True options.minimize_shared_allocs = False options.linearize_shared_access = True - + options.wave_runtime = True + options.dump_intermediates = "intermediates_256x192_shuffle2" + # options.print_ir_after = "all" + options.override_mlir = mlir_256x192 if dynamic: options.dynamic_symbols = [tkl.sym.M, tkl.sym.N, tkl.sym.K] for sym in options.dynamic_symbols: del options.subs[sym] schedule = get_mxfp4_dbuf_pingpong_schedule_Bshuffled_lds( - use_stagger=True, shape=shape + use_stagger=True, shape=shape, block=block ) + UNROLL_FACTOR = tkl.sym.UNROLL_FACTOR + options.subs[UNROLL_FACTOR] = 2 + options.postprocess = """ + module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { + %0 = transform.structured.match ops{["scf.for"]} in %arg0 : (!transform.any_op) -> !transform.any_op + transform.loop.unroll %0 { factor = %%UNROLL_FACTOR%% } : !transform.any_op + transform.yield + } + } + """ - options.print_ir_after = "all" if is_debug else [] options = set_default_run_config(options) gemm = wave_compile(options, gemm, schedule) - _run_mxfp_gemm_preshuffle(gemm, shape, all=True) + _run_mxfp_gemm_preshuffle(gemm, shape, all=True, output_dtype=torch.bfloat16) mode = "dynamic" if dynamic else "static" print( f"MXFP GEMM double-buffer 8-wave ping pong with scales and B shuffling and B->LDS ({mode}) test passed!" ) +def test_dbuf_8wave_pingpong_mxfp_gemm_Bshuffle_lds_transposed( + is_debug=False, shape=(1792, 5376, 4096), block=(256, 192, 256), dynamic=True +): + """Transposed MXFP4 GEMM: computes C^T = B * A^T instead of C = A * B^T. + + By transposing the computation, MFMA output elements become contiguous + along the row-major axis, enabling dwordx4 (128-bit) stores to global + memory instead of dword (32-bit) stores. + + The MFMA left operand ("A" role) is the weight matrix B. + The MFMA right operand ("B" role) is the activation matrix A. + Activation A is preshuffled (since it is now in the "B" role). + A&B scales are preshuffled and read from global memory directly to VGPRs. + """ + + mlir = """ + + #map = affine_map<()[s0] -> (s0 ceildiv 192)> + #map1 = affine_map<()[s0] -> (s0 ceildiv 256)> + #map2 = affine_map<()[s0] -> (s0 floordiv 32)> + #map3 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64) floordiv 24) * 192)> + #map4 = affine_map<()[s0] -> (s0 floordiv 2)> + #map5 = affine_map<()[s0, s1, s2] -> (s1 * 32 + s2 * 192 + s0 floordiv 8 - ((s1 * 32 + s0 floordiv 8) floordiv 192) * 192)> + #map6 = affine_map<()[s0] -> ((s0 floordiv 8) mod 8)> + #map7 = affine_map<()[s0] -> (s0 mod 8)> + #map8 = affine_map<()[s0] -> (s0 * 16)> + #map9 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64 + 8) floordiv 24) * 192 + 64)> + #map10 = affine_map<()[s0, s1, s2] -> (s1 * 32 + s2 * 192 + s0 floordiv 8 - ((s1 * 32 + s0 floordiv 8 + 64) floordiv 192) * 192 + 64)> + #map11 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64 + 16) floordiv 24) * 192 + 128)> + #map12 = affine_map<()[s0, s1, s2] -> (s1 * 32 + s2 * 192 + s0 floordiv 8 - ((s1 * 32 + s0 floordiv 8 + 128) floordiv 192) * 192 + 128)> + #map13 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64) floordiv 32) * 256)> + #map14 = affine_map<()[s0, s1] -> (s0 * 256 + (s1 floordiv 8) * 16 - ((s1 floordiv 8) floordiv 16) * 256)> + #map15 = affine_map<()[s0, s1, s2, s3] -> (s0 + s1 * 256 + (((s3 * 32 + s2 floordiv 8) mod 256) floordiv 16) * 16)> + #map16 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64 + 8) floordiv 32) * 256 + 64)> + #map17 = affine_map<()[s0, s1, s2, s3] -> (s0 + s1 * 256 + s3 * 32 + ((s2 floordiv 8) floordiv 16) * 16 - ((s3 * 32 + s2 floordiv 8 + 64) floordiv 256) * 256 + 64)> + #map18 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64 + 16) floordiv 32) * 256 + 128)> + #map19 = affine_map<()[s0, s1, s2, s3] -> (s0 + s1 * 256 + s3 * 32 + ((s2 floordiv 8) floordiv 16) * 16 - ((s3 * 32 + s2 floordiv 8 + 128) floordiv 256) * 256 + 128)> + #map20 = affine_map<()[s0, s1] -> (s1 * 32 + (s0 floordiv 64) * 8 - ((s1 * 4 + s0 floordiv 64 + 24) floordiv 32) * 256 + 192)> + #map21 = affine_map<()[s0, s1, s2, s3] -> (s0 + s1 * 256 + s3 * 32 + ((s2 floordiv 8) floordiv 16) * 16 - ((s3 * 32 + s2 floordiv 8 + 192) floordiv 256) * 256 + 192)> + #map22 = affine_map<()[s0] -> (s0 * 192)> + #map23 = affine_map<()[s0, s1] -> (s0 * 4 - (s0 floordiv 16) * 64 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + (((s0 mod 16 + (s0 floordiv 64) * 48) floordiv 32) * ((s1 floordiv 32 + 7) floordiv 8)) * 256 + ((s0 mod 64) floordiv 16) * 64)> + #map24 = affine_map<()[s0] -> (((s0 floordiv 32 + 7) floordiv 8) * 8)> + #map25 = affine_map<()[s0] -> (s0 * 4 - (s0 floordiv 16) * 64 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + ((s0 mod 64) floordiv 16) * 64)> + #map26 = affine_map<()[s0, s1] -> (s0 * 4 - (s0 floordiv 16) * 64 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + (((s0 mod 16 + (s0 floordiv 64) * 48) floordiv 32) * ((s1 floordiv 32 + 7) floordiv 8)) * 256 + ((s0 mod 64) floordiv 16) * 64 + (((s0 mod 64) floordiv 16 + 4) floordiv 8) * 252 + 2)> + #map27 = affine_map<()[s0] -> (s0 * 4 - (s0 floordiv 16) * 64 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + ((s0 mod 64) floordiv 16) * 64 + (((s0 mod 64) floordiv 16 + 4) floordiv 8) * 252 + 2)> + #map28 = affine_map<()[s0, s1] -> (s0 * 4 - (s0 floordiv 16) * 64 + (s0 floordiv 64) * 3 - ((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * 2 + (((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * ((s1 floordiv 32 + 7) floordiv 8)) * 256 + ((s0 mod 64) floordiv 16) * 64 + 1)> + #map29 = affine_map<()[s0] -> (s0 * 4 - (s0 floordiv 16) * 64 + (s0 floordiv 64) * 3 - ((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * 2 + ((s0 mod 64) floordiv 16) * 64 + 1)> + #map30 = affine_map<()[s0, s1] -> (s0 * 4 - (s0 floordiv 16) * 64 + (s0 floordiv 64) * 3 - ((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * 2 + (((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * ((s1 floordiv 32 + 7) floordiv 8)) * 256 + ((s0 mod 64) floordiv 16) * 64 + (((s0 mod 64) floordiv 16 + 4) floordiv 8) * 252 + 3)> + #map31 = affine_map<()[s0] -> (s0 * 4 - (s0 floordiv 16) * 64 + (s0 floordiv 64) * 3 - ((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * 2 + ((s0 mod 64) floordiv 16) * 64 + (((s0 mod 64) floordiv 16 + 4) floordiv 8) * 252 + 3)> + #map32 = affine_map<()[s0, s1] -> (s0 * 4 - (s0 floordiv 16) * 64 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + (((s1 floordiv 32 + 7) floordiv 8) * ((s0 mod 16 + (s0 floordiv 64) * 48) floordiv 32 + 1)) * 256 + ((s0 mod 64) floordiv 16) * 64)> + #map33 = affine_map<()[s0, s1] -> (s0 * 4 - (s0 floordiv 16) * 64 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + (((s1 floordiv 32 + 7) floordiv 8) * ((s0 mod 16 + (s0 floordiv 64) * 48) floordiv 32 + 1)) * 256 + ((s0 mod 64) floordiv 16) * 64 + (((s0 mod 64) floordiv 16 + 4) floordiv 8) * 252 + 2)> + #map34 = affine_map<()[s0] -> (s0 * 256)> + #map35 = affine_map<()[s0, s1, s2] -> (s0 * 4 - (s0 floordiv 16) * 64 + (((s1 floordiv 32 + 7) floordiv 8) * (s2 * 4)) * 256 + ((s0 mod 64) floordiv 16) * 64)> + #map36 = affine_map<()[s0] -> (s0 * 4 - (s0 floordiv 16) * 64 + ((s0 mod 64) floordiv 16) * 64)> + #map37 = affine_map<()[s0, s1, s2] -> (s0 * 4 - (s0 floordiv 16) * 64 + (((s1 floordiv 32 + 7) floordiv 8) * (s2 * 4 + 1)) * 256 + ((s0 mod 64) floordiv 16) * 64)> + #map38 = affine_map<()[s0, s1, s2] -> (s0 * 4 - (s0 floordiv 16) * 64 + (((s1 floordiv 32 + 7) floordiv 8) * (s2 * 4 + 2)) * 256 + ((s0 mod 64) floordiv 16) * 64)> + #map39 = affine_map<()[s0, s1, s2] -> (s0 * 4 - (s0 floordiv 16) * 64 + (((s1 floordiv 32 + 7) floordiv 8) * (s2 * 4 + 3)) * 256 + ((s0 mod 64) floordiv 16) * 64)> + #map40 = affine_map<()[s0, s1] -> (s1 * 4 + s0 floordiv 64)> + #map41 = affine_map<()[s0] -> (s0 ceildiv 256 - 1)> + #map42 = affine_map<()[s0] -> ((s0 mod 64) floordiv 16)> + #map43 = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 16 + (s0 floordiv 64) * 6144 - (s0 floordiv 16) * 2048)> + #map44 = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 16 + (s0 floordiv 64) * 6144 - (s0 floordiv 16) * 2048 + 2048)> + #map45 = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 16 + (s0 floordiv 64) * 6144 - (s0 floordiv 16) * 2048 + 4096)> + #map46 = affine_map<()[s0, s1, s2] -> (s0 * 16384 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048)> + #map47 = affine_map<()[s0, s1, s2] -> (s0 * 16384 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 2048)> + #map48 = affine_map<()[s0, s1, s2] -> (s0 * 16384 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 4096)> + #map49 = affine_map<()[s0, s1, s2] -> (s0 * 16384 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 6144)> + #map50 = affine_map<()[s0, s1, s2] -> (s0 * 16384 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 8192)> + #map51 = affine_map<()[s0, s1, s2] -> (s0 * 16384 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 10240)> + #map52 = affine_map<()[s0, s1, s2] -> (s0 * 16384 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 12288)> + #map53 = affine_map<()[s0, s1, s2] -> (s0 * 16384 + s1 * 128 + s2 * 16 - (s1 floordiv 16) * 2048 + 14336)> + #map54 = affine_map<()[s0] -> ((s0 mod 64) floordiv 16 + 4)> + #map55 = affine_map<()[s0, s1] -> (s0 * 128 + s1 * 16 + 128)> + #map56 = affine_map<()[s0, s1, s2] -> (s0 * 2048 + s1 * 256 + (s2 floordiv 8) * 16 - ((s2 floordiv 8) floordiv 16) * 256 + 2048)> + #map57 = affine_map<()[s0, s1, s2] -> (s0 * 4 - (s0 floordiv 16) * 64 + s2 * 256 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + (((s0 mod 16 + (s0 floordiv 64) * 48) floordiv 32) * ((s1 floordiv 32 + 7) floordiv 8)) * 256 + ((s0 mod 64) floordiv 16) * 64 + 256)> + #map58 = affine_map<()[s0, s1] -> (s0 * 4 + s1 * 256 - (s0 floordiv 16) * 64 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + ((s0 mod 64) floordiv 16) * 64 + 256)> + #map59 = affine_map<()[s0, s1, s2] -> (s0 * 4 - (s0 floordiv 16) * 64 + s2 * 256 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + (((s0 mod 16 + (s0 floordiv 64) * 48) floordiv 32) * ((s1 floordiv 32 + 7) floordiv 8)) * 256 + ((s0 mod 64) floordiv 16) * 64 + (((s0 mod 64) floordiv 16 + 12) floordiv 8) * 256 - (((s0 mod 64) floordiv 16 + 4) floordiv 8) * 4 + 2)> + #map60 = affine_map<()[s0, s1] -> (s0 * 4 + s1 * 256 - (s0 floordiv 16) * 64 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + ((s0 mod 64) floordiv 16) * 64 + (((s0 mod 64) floordiv 16 + 12) floordiv 8) * 256 - (((s0 mod 64) floordiv 16 + 4) floordiv 8) * 4 + 2)> + #map61 = affine_map<()[s0, s1, s2] -> (s0 * 4 - (s0 floordiv 16) * 64 + (s0 floordiv 64) * 3 + s2 * 256 - ((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * 2 + (((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * ((s1 floordiv 32 + 7) floordiv 8)) * 256 + ((s0 mod 64) floordiv 16) * 64 + 257)> + #map62 = affine_map<()[s0, s1] -> (s0 * 4 + s1 * 256 - (s0 floordiv 16) * 64 + (s0 floordiv 64) * 3 - ((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * 2 + ((s0 mod 64) floordiv 16) * 64 + 257)> + #map63 = affine_map<()[s0, s1, s2] -> (s0 * 4 - (s0 floordiv 16) * 64 + (s0 floordiv 64) * 3 + s2 * 256 - ((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * 2 + (((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * ((s1 floordiv 32 + 7) floordiv 8)) * 256 + ((s0 mod 64) floordiv 16) * 64 + (((s0 mod 64) floordiv 16 + 12) floordiv 8) * 256 - (((s0 mod 64) floordiv 16 + 4) floordiv 8) * 4 + 3)> + #map64 = affine_map<()[s0, s1] -> (s0 * 4 + s1 * 256 - (s0 floordiv 16) * 64 + (s0 floordiv 64) * 3 - ((s0 mod 16 + (s0 floordiv 64) * 48 + 16) floordiv 32) * 2 + ((s0 mod 64) floordiv 16) * 64 + (((s0 mod 64) floordiv 16 + 12) floordiv 8) * 256 - (((s0 mod 64) floordiv 16 + 4) floordiv 8) * 4 + 3)> + #map65 = affine_map<()[s0, s1, s2] -> (s0 * 4 - (s0 floordiv 16) * 64 + s2 * 256 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + (((s1 floordiv 32 + 7) floordiv 8) * ((s0 mod 16 + (s0 floordiv 64) * 48) floordiv 32 + 1)) * 256 + ((s0 mod 64) floordiv 16) * 64 + 256)> + #map66 = affine_map<()[s0, s1, s2] -> (s0 * 4 - (s0 floordiv 16) * 64 + s2 * 256 + ((s0 mod 16 + (s0 floordiv 64) * 48) mod 32) floordiv 16 + (((s1 floordiv 32 + 7) floordiv 8) * ((s0 mod 16 + (s0 floordiv 64) * 48) floordiv 32 + 1)) * 256 + ((s0 mod 64) floordiv 16) * 64 + (((s0 mod 64) floordiv 16 + 12) floordiv 8) * 256 - (((s0 mod 64) floordiv 16 + 4) floordiv 8) * 4 + 2)> + #map67 = affine_map<()[s0, s1, s2, s3] -> (s0 * 4 - (s0 floordiv 16) * 64 + s3 * 256 + (((s1 floordiv 32 + 7) floordiv 8) * (s2 * 4)) * 256 + ((s0 mod 64) floordiv 16) * 64 + 256)> + #map68 = affine_map<()[s0, s1] -> (s0 * 4 + s1 * 256 - (s0 floordiv 16) * 64 + ((s0 mod 64) floordiv 16) * 64 + 256)> + #map69 = affine_map<()[s0, s1, s2, s3] -> (s0 * 4 - (s0 floordiv 16) * 64 + s3 * 256 + (((s1 floordiv 32 + 7) floordiv 8) * (s2 * 4 + 1)) * 256 + ((s0 mod 64) floordiv 16) * 64 + 256)> + #map70 = affine_map<()[s0, s1, s2, s3] -> (s0 * 4 - (s0 floordiv 16) * 64 + s3 * 256 + (((s1 floordiv 32 + 7) floordiv 8) * (s2 * 4 + 2)) * 256 + ((s0 mod 64) floordiv 16) * 64 + 256)> + #map71 = affine_map<()[s0, s1, s2, s3] -> (s0 * 4 - (s0 floordiv 16) * 64 + s3 * 256 + (((s1 floordiv 32 + 7) floordiv 8) * (s2 * 4 + 3)) * 256 + ((s0 mod 64) floordiv 16) * 64 + 256)> + #map72 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4)> + #map73 = affine_map<()[s0, s1] -> (s0 + s1 * 128 - (s0 floordiv 16) * 16)> + #map74 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 1)> + #map75 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 2)> + #map76 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 3)> + #map77 = affine_map<()[s0, s1] -> (s0 + s1 * 128 - (s0 floordiv 16) * 16 + 16)> + #map78 = affine_map<()[s0, s1] -> (s0 + s1 * 128 - (s0 floordiv 16) * 16 + 32)> + #map79 = affine_map<()[s0, s1] -> (s0 + s1 * 128 - (s0 floordiv 16) * 16 + 48)> + #map80 = affine_map<()[s0, s1] -> (s0 + s1 * 128 - (s0 floordiv 16) * 16 + 64)> + #map81 = affine_map<()[s0, s1] -> (s0 + s1 * 128 - (s0 floordiv 16) * 16 + 80)> + #map82 = affine_map<()[s0, s1] -> (s0 + s1 * 128 - (s0 floordiv 16) * 16 + 96)> + #map83 = affine_map<()[s0, s1] -> (s0 + s1 * 128 - (s0 floordiv 16) * 16 + 112)> + #map84 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 16)> + #map85 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 17)> + #map86 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 18)> + #map87 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 19)> + #map88 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 32)> + #map89 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 33)> + #map90 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 34)> + #map91 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 16) * 4 + 35)> + #map92 = affine_map<()[s0] -> ((s0 mod 32) floordiv 16)> + #map93 = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + s0 - (s0 floordiv 16) * 16)> + #map94 = affine_map<()[s0, s1] -> (s1 * 128 + ((s0 mod 64) floordiv 32) * 8)> + #map_Mbase = affine_map<()[s0, s1] -> (s1 * 128 + s0 - (s0 floordiv 16) * 16)> + #map_Nbase = affine_map<()[s0] -> ((s0 floordiv 64) * 48 + ((s0 mod 64) floordiv 32) * 8)> + #map95 = affine_map<()[s0] -> (s0)> + #translation = #iree_codegen.translation_info + module attributes {transform.with_named_sequence} { + stream.executable private @gemm { + stream.executable.export public @gemm workgroups(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index, %arg7: index) -> (index, index, index) { + %c1 = arith.constant 1 : index + %0 = affine.apply #map()[%arg0] + %1 = affine.apply #map1()[%arg1] + stream.return %0, %1, %c1 : index, index, index + } + builtin.module { + func.func @gemm(%arg0: !stream.binding, %arg1: !stream.binding, %arg2: !stream.binding, %arg3: !stream.binding, %arg4: !stream.binding, %arg5: index, %arg6: index, %arg7: index, %arg8: index, %arg9: index, %arg10: index, %arg11: index, %arg12: index) attributes {translation_info = #translation} { + %c4_i32 = arith.constant 4 : i32 + %c4294967295_i64 = arith.constant 4294967295 : i64 + %c2147483645_i64 = arith.constant 2147483645 : i64 + %c0_i32 = arith.constant 0 : i32 + %c32_i64 = arith.constant 32 : i64 + %c2147483646_i64 = arith.constant 2147483646 : i64 + %cst = arith.constant dense<0.000000e+00> : vector<4xf32> + %c2 = arith.constant 2 : index + %c1 = arith.constant 1 : index + %c0 = arith.constant 0 : index + %0 = stream.binding.subspan %arg0[%c0] : !stream.binding -> memref + %1 = stream.binding.subspan %arg1[%c0] : !stream.binding -> memref + %2 = stream.binding.subspan %arg2[%c0] : !stream.binding -> memref + %3 = stream.binding.subspan %arg3[%c0] : !stream.binding -> memref + %4 = stream.binding.subspan %arg4[%c0] : !stream.binding -> memref + %block_id_x = gpu.block_id x upper_bound 2147483647 + %block_id_y = gpu.block_id y upper_bound 2147483647 + %thread_id_x = gpu.thread_id x upper_bound 256 + %thread_id_y = gpu.thread_id y upper_bound 2 + %5 = affine.apply #map2()[%arg7] + %reinterpret_cast = memref.reinterpret_cast %1 to offset: [0], sizes: [%arg5, %5], strides: [%arg9, 1] : memref to memref> + %reinterpret_cast_0 = memref.reinterpret_cast %3 to offset: [0], sizes: [%arg6, %5], strides: [%arg11, 1] : memref to memref> + %reinterpret_cast_1 = memref.reinterpret_cast %4 to offset: [0], sizes: [%arg6, %arg5], strides: [%arg5, 1] : memref to memref> + %alloc = memref.alloc() : memref<256x128xi8, #gpu.address_space> + %alloc_2 = memref.alloc() : memref<256x128xi8, #gpu.address_space> + %alloc_3 = memref.alloc() : memref<192x128xi8, #gpu.address_space> + %alloc_4 = memref.alloc() : memref<192x128xi8, #gpu.address_space> + %6 = affine.apply #map3()[%thread_id_x, %thread_id_y] + %7 = gpu.subgroup_broadcast %6, first_active_lane : index + %8 = gpu.subgroup_broadcast %c0, first_active_lane : index + %9 = affine.apply #map4()[%arg7] + %10 = affine.apply #map5()[%thread_id_x, %thread_id_y, %block_id_x] + %11 = affine.apply #map6()[%thread_id_x] + %12 = affine.apply #map7()[%thread_id_x] + %13 = arith.xori %12, %11 : index + %14 = affine.apply #map8()[%13] + %15 = arith.muli %10, %9 overflow : index + %16 = arith.addi %15, %14 overflow : index + %reinterpret_cast_5 = memref.reinterpret_cast %0 to offset: [0], sizes: [2147483646], strides: [1] : memref to memref<2147483646xi8, strided<[1]>> + %cast = memref.cast %reinterpret_cast_5 : memref<2147483646xi8, strided<[1]>> to memref> + %17 = arith.index_cast %9 : index to i14 + %18 = amdgpu.fat_raw_buffer_cast %cast validBytes(%c2147483646_i64) cacheSwizzleStride(%17) resetOffset : memref> to memref> + amdgpu.gather_to_lds %18[%16], %alloc_4[%7, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %19 = affine.apply #map9()[%thread_id_x, %thread_id_y] + %20 = gpu.subgroup_broadcast %19, first_active_lane : index + %21 = affine.apply #map10()[%thread_id_x, %thread_id_y, %block_id_x] + %22 = arith.muli %21, %9 overflow : index + %23 = arith.addi %22, %14 overflow : index + amdgpu.gather_to_lds %18[%23], %alloc_4[%20, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %24 = affine.apply #map11()[%thread_id_x, %thread_id_y] + %25 = gpu.subgroup_broadcast %24, first_active_lane : index + %26 = affine.apply #map12()[%thread_id_x, %thread_id_y, %block_id_x] + %27 = arith.muli %26, %9 overflow : index + %28 = arith.addi %27, %14 overflow : index + amdgpu.gather_to_lds %18[%28], %alloc_4[%25, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %29 = affine.apply #map13()[%thread_id_x, %thread_id_y] + %30 = gpu.subgroup_broadcast %29, first_active_lane : index + %31 = affine.apply #map14()[%13, %thread_id_x] + %32 = arith.index_cast %9 : index to i32 + %33 = arith.extui %32 : i32 to i64 + %34 = arith.addi %33, %c4294967295_i64 : i64 + %35 = arith.divui %34, %33 : i64 + %36 = arith.trunci %35 : i64 to i32 + %37 = arith.index_cast %31 : index to i32 + %38 = arith.extui %37 : i32 to i64 + %39 = arith.extui %36 : i32 to i64 + %40 = arith.muli %38, %39 : i64 + %41 = arith.shrui %40, %c32_i64 : i64 + %42 = arith.trunci %41 : i64 to i32 + %43 = arith.muli %42, %32 : i32 + %44 = arith.subi %37, %43 : i32 + %45 = arith.cmpi uge, %44, %32 : i32 + %46 = arith.extui %45 : i1 to i32 + %47 = arith.subi %42, %46 : i32 + %48 = arith.index_cast %47 : i32 to index + %49 = affine.apply #map15()[%48, %block_id_y, %thread_id_x, %thread_id_y] + %50 = arith.select %45, %32, %c0_i32 : i32 + %51 = arith.addi %44, %50 : i32 + %52 = arith.index_cast %51 : i32 to index + %53 = arith.muli %49, %9 overflow : index + %54 = arith.addi %53, %52 overflow : index + %reinterpret_cast_6 = memref.reinterpret_cast %2 to offset: [0], sizes: [2147483646], strides: [1] : memref to memref<2147483646xi8, strided<[1]>> + %cast_7 = memref.cast %reinterpret_cast_6 : memref<2147483646xi8, strided<[1]>> to memref> + %55 = amdgpu.fat_raw_buffer_cast %cast_7 validBytes(%c2147483646_i64) cacheSwizzleStride(%17) resetOffset : memref> to memref> + amdgpu.gather_to_lds %55[%54], %alloc_2[%30, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %56 = affine.apply #map16()[%thread_id_x, %thread_id_y] + %57 = gpu.subgroup_broadcast %56, first_active_lane : index + %58 = affine.apply #map17()[%48, %block_id_y, %thread_id_x, %thread_id_y] + %59 = arith.muli %58, %9 overflow : index + %60 = arith.addi %59, %52 overflow : index + amdgpu.gather_to_lds %55[%60], %alloc_2[%57, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %61 = affine.apply #map18()[%thread_id_x, %thread_id_y] + %62 = gpu.subgroup_broadcast %61, first_active_lane : index + %63 = affine.apply #map19()[%48, %block_id_y, %thread_id_x, %thread_id_y] + %64 = arith.muli %63, %9 overflow : index + %65 = arith.addi %64, %52 overflow : index + amdgpu.gather_to_lds %55[%65], %alloc_2[%62, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %66 = affine.apply #map20()[%thread_id_x, %thread_id_y] + %67 = gpu.subgroup_broadcast %66, first_active_lane : index + %68 = affine.apply #map21()[%48, %block_id_y, %thread_id_x, %thread_id_y] + %69 = arith.muli %68, %9 overflow : index + %70 = arith.addi %69, %52 overflow : index + amdgpu.gather_to_lds %55[%70], %alloc_2[%67, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %71 = affine.apply #map22()[%block_id_x] + %72 = affine.apply #map23()[%thread_id_x, %arg7] + %73 = affine.apply #map24()[%arg7] + %74 = arith.index_cast %73 : index to i32 + %75 = arith.extui %74 : i32 to i64 + %76 = arith.addi %75, %c4294967295_i64 : i64 + %77 = arith.divui %76, %75 : i64 + %78 = arith.trunci %77 : i64 to i32 + %79 = arith.index_cast %72 : index to i32 + %80 = arith.extui %79 : i32 to i64 + %81 = arith.extui %78 : i32 to i64 + %82 = arith.muli %80, %81 : i64 + %83 = arith.shrui %82, %c32_i64 : i64 + %84 = arith.trunci %83 : i64 to i32 + %85 = arith.muli %84, %74 : i32 + %86 = arith.subi %79, %85 : i32 + %87 = arith.cmpi uge, %86, %74 : i32 + %88 = arith.extui %87 : i1 to i32 + %89 = arith.subi %84, %88 : i32 + %90 = arith.index_cast %89 : i32 to index + %91 = affine.apply #map25()[%thread_id_x] + %92 = arith.index_cast %91 : index to i32 + %93 = arith.extui %92 : i32 to i64 + %94 = arith.muli %93, %81 : i64 + %95 = arith.shrui %94, %c32_i64 : i64 + %96 = arith.trunci %95 : i64 to i32 + %97 = arith.muli %96, %74 : i32 + %98 = arith.subi %92, %97 : i32 + %99 = arith.cmpi uge, %98, %74 : i32 + %100 = arith.select %99, %74, %c0_i32 : i32 + %101 = arith.addi %98, %100 : i32 + %102 = arith.index_cast %101 : i32 to index + %base_buffer, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %reinterpret_cast : memref> -> memref, index, index, index, index, index + %103 = arith.muli %71, %strides#0 overflow : index + %104 = arith.muli %90, %strides#0 overflow : index + %105 = arith.addi %104, %102 overflow : index + %reinterpret_cast_8 = memref.reinterpret_cast %1 to offset: [%103], sizes: [2147483646], strides: [1] : memref to memref<2147483646xi8, strided<[1], offset: ?>> + %cast_9 = memref.cast %reinterpret_cast_8 : memref<2147483646xi8, strided<[1], offset: ?>> to memref> + %106 = arith.index_cast %strides#0 : index to i14 + %107 = amdgpu.fat_raw_buffer_cast %cast_9 validBytes(%c2147483646_i64) cacheSwizzleStride(%106) resetOffset : memref> to memref> + %108 = vector.load %107[%105] : memref>, vector<1xi8> + %109 = affine.apply #map26()[%thread_id_x, %arg7] + %110 = arith.index_cast %109 : index to i32 + %111 = arith.extui %110 : i32 to i64 + %112 = arith.muli %111, %81 : i64 + %113 = arith.shrui %112, %c32_i64 : i64 + %114 = arith.trunci %113 : i64 to i32 + %115 = arith.muli %114, %74 : i32 + %116 = arith.subi %110, %115 : i32 + %117 = arith.cmpi uge, %116, %74 : i32 + %118 = arith.extui %117 : i1 to i32 + %119 = arith.subi %114, %118 : i32 + %120 = arith.index_cast %119 : i32 to index + %121 = affine.apply #map27()[%thread_id_x] + %122 = arith.index_cast %121 : index to i32 + %123 = arith.extui %122 : i32 to i64 + %124 = arith.muli %123, %81 : i64 + %125 = arith.shrui %124, %c32_i64 : i64 + %126 = arith.trunci %125 : i64 to i32 + %127 = arith.muli %126, %74 : i32 + %128 = arith.subi %122, %127 : i32 + %129 = arith.cmpi uge, %128, %74 : i32 + %130 = arith.select %129, %74, %c0_i32 : i32 + %131 = arith.addi %128, %130 : i32 + %132 = arith.index_cast %131 : i32 to index + %133 = arith.muli %120, %strides#0 overflow : index + %134 = arith.addi %133, %132 overflow : index + %135 = vector.load %107[%134] : memref>, vector<1xi8> + %136 = affine.apply #map28()[%thread_id_x, %arg7] + %137 = arith.index_cast %136 : index to i32 + %138 = arith.extui %137 : i32 to i64 + %139 = arith.muli %138, %81 : i64 + %140 = arith.shrui %139, %c32_i64 : i64 + %141 = arith.trunci %140 : i64 to i32 + %142 = arith.muli %141, %74 : i32 + %143 = arith.subi %137, %142 : i32 + %144 = arith.cmpi uge, %143, %74 : i32 + %145 = arith.extui %144 : i1 to i32 + %146 = arith.subi %141, %145 : i32 + %147 = arith.index_cast %146 : i32 to index + %148 = affine.apply #map29()[%thread_id_x] + %149 = arith.index_cast %148 : index to i32 + %150 = arith.extui %149 : i32 to i64 + %151 = arith.muli %150, %81 : i64 + %152 = arith.shrui %151, %c32_i64 : i64 + %153 = arith.trunci %152 : i64 to i32 + %154 = arith.muli %153, %74 : i32 + %155 = arith.subi %149, %154 : i32 + %156 = arith.cmpi uge, %155, %74 : i32 + %157 = arith.select %156, %74, %c0_i32 : i32 + %158 = arith.addi %155, %157 : i32 + %159 = arith.index_cast %158 : i32 to index + %160 = arith.muli %147, %strides#0 overflow : index + %161 = arith.addi %160, %159 overflow : index + %162 = vector.load %107[%161] : memref>, vector<1xi8> + %163 = affine.apply #map30()[%thread_id_x, %arg7] + %164 = arith.index_cast %163 : index to i32 + %165 = arith.extui %164 : i32 to i64 + %166 = arith.muli %165, %81 : i64 + %167 = arith.shrui %166, %c32_i64 : i64 + %168 = arith.trunci %167 : i64 to i32 + %169 = arith.muli %168, %74 : i32 + %170 = arith.subi %164, %169 : i32 + %171 = arith.cmpi uge, %170, %74 : i32 + %172 = arith.extui %171 : i1 to i32 + %173 = arith.subi %168, %172 : i32 + %174 = arith.index_cast %173 : i32 to index + %175 = affine.apply #map31()[%thread_id_x] + %176 = arith.index_cast %175 : index to i32 + %177 = arith.extui %176 : i32 to i64 + %178 = arith.muli %177, %81 : i64 + %179 = arith.shrui %178, %c32_i64 : i64 + %180 = arith.trunci %179 : i64 to i32 + %181 = arith.muli %180, %74 : i32 + %182 = arith.subi %176, %181 : i32 + %183 = arith.cmpi uge, %182, %74 : i32 + %184 = arith.select %183, %74, %c0_i32 : i32 + %185 = arith.addi %182, %184 : i32 + %186 = arith.index_cast %185 : i32 to index + %187 = arith.muli %174, %strides#0 overflow : index + %188 = arith.addi %187, %186 overflow : index + %189 = vector.load %107[%188] : memref>, vector<1xi8> + %190 = affine.apply #map32()[%thread_id_x, %arg7] + %191 = arith.index_cast %190 : index to i32 + %192 = arith.extui %191 : i32 to i64 + %193 = arith.muli %192, %81 : i64 + %194 = arith.shrui %193, %c32_i64 : i64 + %195 = arith.trunci %194 : i64 to i32 + %196 = arith.muli %195, %74 : i32 + %197 = arith.subi %191, %196 : i32 + %198 = arith.cmpi uge, %197, %74 : i32 + %199 = arith.extui %198 : i1 to i32 + %200 = arith.subi %195, %199 : i32 + %201 = arith.index_cast %200 : i32 to index + %202 = arith.muli %201, %strides#0 overflow : index + %203 = arith.addi %202, %102 overflow : index + %204 = vector.load %107[%203] : memref>, vector<1xi8> + %205 = affine.apply #map33()[%thread_id_x, %arg7] + %206 = arith.index_cast %205 : index to i32 + %207 = arith.extui %206 : i32 to i64 + %208 = arith.muli %207, %81 : i64 + %209 = arith.shrui %208, %c32_i64 : i64 + %210 = arith.trunci %209 : i64 to i32 + %211 = arith.muli %210, %74 : i32 + %212 = arith.subi %206, %211 : i32 + %213 = arith.cmpi uge, %212, %74 : i32 + %214 = arith.extui %213 : i1 to i32 + %215 = arith.subi %210, %214 : i32 + %216 = arith.index_cast %215 : i32 to index + %217 = arith.muli %216, %strides#0 overflow : index + %218 = arith.addi %217, %132 overflow : index + %219 = vector.load %107[%218] : memref>, vector<1xi8> + %220 = affine.apply #map34()[%block_id_y] + %221 = affine.apply #map35()[%thread_id_x, %arg7, %thread_id_y] + %222 = arith.index_cast %221 : index to i32 + %223 = arith.extui %222 : i32 to i64 + %224 = arith.muli %223, %81 : i64 + %225 = arith.shrui %224, %c32_i64 : i64 + %226 = arith.trunci %225 : i64 to i32 + %227 = arith.muli %226, %74 : i32 + %228 = arith.subi %222, %227 : i32 + %229 = arith.cmpi uge, %228, %74 : i32 + %230 = arith.extui %229 : i1 to i32 + %231 = arith.subi %226, %230 : i32 + %232 = arith.index_cast %231 : i32 to index + %233 = affine.apply #map36()[%thread_id_x] + %234 = arith.index_cast %233 : index to i32 + %235 = arith.extui %234 : i32 to i64 + %236 = arith.muli %235, %81 : i64 + %237 = arith.shrui %236, %c32_i64 : i64 + %238 = arith.trunci %237 : i64 to i32 + %239 = arith.muli %238, %74 : i32 + %240 = arith.subi %234, %239 : i32 + %241 = arith.cmpi uge, %240, %74 : i32 + %242 = arith.select %241, %74, %c0_i32 : i32 + %243 = arith.addi %240, %242 : i32 + %244 = arith.index_cast %243 : i32 to index + %base_buffer_10, %offset_11, %sizes_12:2, %strides_13:2 = memref.extract_strided_metadata %reinterpret_cast_0 : memref> -> memref, index, index, index, index, index + %245 = arith.muli %220, %strides_13#0 overflow : index + %246 = arith.muli %232, %strides_13#0 overflow : index + %247 = arith.addi %246, %244 overflow : index + %reinterpret_cast_14 = memref.reinterpret_cast %3 to offset: [%245], sizes: [2147483646], strides: [1] : memref to memref<2147483646xi8, strided<[1], offset: ?>> + %cast_15 = memref.cast %reinterpret_cast_14 : memref<2147483646xi8, strided<[1], offset: ?>> to memref> + %248 = arith.index_cast %strides_13#0 : index to i14 + %249 = amdgpu.fat_raw_buffer_cast %cast_15 validBytes(%c2147483646_i64) cacheSwizzleStride(%248) resetOffset : memref> to memref> + %250 = vector.load %249[%247] : memref>, vector<4xi8> + %251 = affine.apply #map37()[%thread_id_x, %arg7, %thread_id_y] + %252 = arith.index_cast %251 : index to i32 + %253 = arith.extui %252 : i32 to i64 + %254 = arith.muli %253, %81 : i64 + %255 = arith.shrui %254, %c32_i64 : i64 + %256 = arith.trunci %255 : i64 to i32 + %257 = arith.muli %256, %74 : i32 + %258 = arith.subi %252, %257 : i32 + %259 = arith.cmpi uge, %258, %74 : i32 + %260 = arith.extui %259 : i1 to i32 + %261 = arith.subi %256, %260 : i32 + %262 = arith.index_cast %261 : i32 to index + %263 = arith.muli %262, %strides_13#0 overflow : index + %264 = arith.addi %263, %244 overflow : index + %265 = vector.load %249[%264] : memref>, vector<4xi8> + %266 = affine.apply #map38()[%thread_id_x, %arg7, %thread_id_y] + %267 = arith.index_cast %266 : index to i32 + %268 = arith.extui %267 : i32 to i64 + %269 = arith.muli %268, %81 : i64 + %270 = arith.shrui %269, %c32_i64 : i64 + %271 = arith.trunci %270 : i64 to i32 + %272 = arith.muli %271, %74 : i32 + %273 = arith.subi %267, %272 : i32 + %274 = arith.cmpi uge, %273, %74 : i32 + %275 = arith.extui %274 : i1 to i32 + %276 = arith.subi %271, %275 : i32 + %277 = arith.index_cast %276 : i32 to index + %278 = arith.muli %277, %strides_13#0 overflow : index + %279 = arith.addi %278, %244 overflow : index + %280 = vector.load %249[%279] : memref>, vector<4xi8> + %281 = affine.apply #map39()[%thread_id_x, %arg7, %thread_id_y] + %282 = arith.index_cast %281 : index to i32 + %283 = arith.extui %282 : i32 to i64 + %284 = arith.muli %283, %81 : i64 + %285 = arith.shrui %284, %c32_i64 : i64 + %286 = arith.trunci %285 : i64 to i32 + %287 = arith.muli %286, %74 : i32 + %288 = arith.subi %282, %287 : i32 + %289 = arith.cmpi uge, %288, %74 : i32 + %290 = arith.extui %289 : i1 to i32 + %291 = arith.subi %286, %290 : i32 + %292 = arith.index_cast %291 : i32 to index + %293 = arith.muli %292, %strides_13#0 overflow : index + %294 = arith.addi %293, %244 overflow : index + %295 = vector.load %249[%294] : memref>, vector<4xi8> + amdgpu.memory_counter_wait load(0) + rocdl.s.barrier + %296 = affine.apply #map40()[%thread_id_x, %thread_id_y] + %297 = arith.index_cast %296 : index to i32 + %298 = arith.cmpi sge, %297, %c4_i32 : i32 + %299 = arith.cmpi slt, %297, %c4_i32 : i32 + scf.if %298 { + rocdl.s.barrier + } + %300 = affine.apply #map41()[%arg7] + %301 = arith.remsi %300, %c2 : index + %302 = arith.subi %300, %301 : index + %reinterpret_cast_16 = memref.reinterpret_cast %alloc_4 to offset: [0], sizes: [24576], strides: [1] : memref<192x128xi8, #gpu.address_space> to memref<24576xi8, #gpu.address_space> + %303 = affine.apply #map42()[%thread_id_x] + %304 = arith.xori %303, %12 : index + %305 = affine.apply #map43()[%thread_id_x, %304] + %306 = affine.apply #map44()[%thread_id_x, %304] + %307 = affine.apply #map45()[%thread_id_x, %304] + %reinterpret_cast_17 = memref.reinterpret_cast %alloc_2 to offset: [0], sizes: [32768], strides: [1] : memref<256x128xi8, #gpu.address_space> to memref<32768xi8, #gpu.address_space> + %308 = affine.apply #map46()[%thread_id_y, %thread_id_x, %304] + %309 = affine.apply #map47()[%thread_id_y, %thread_id_x, %304] + %310 = affine.apply #map48()[%thread_id_y, %thread_id_x, %304] + %311 = affine.apply #map49()[%thread_id_y, %thread_id_x, %304] + %312 = affine.apply #map50()[%thread_id_y, %thread_id_x, %304] + %313 = affine.apply #map51()[%thread_id_y, %thread_id_x, %304] + %314 = affine.apply #map52()[%thread_id_y, %thread_id_x, %304] + %315 = affine.apply #map53()[%thread_id_y, %thread_id_x, %304] + %316 = affine.apply #map54()[%thread_id_x] + %317 = arith.xori %316, %12 : index + %318 = affine.apply #map43()[%thread_id_x, %317] + %319 = affine.apply #map44()[%thread_id_x, %317] + %320 = affine.apply #map45()[%thread_id_x, %317] + %321 = affine.apply #map46()[%thread_id_y, %thread_id_x, %317] + %322 = affine.apply #map47()[%thread_id_y, %thread_id_x, %317] + %323 = affine.apply #map48()[%thread_id_y, %thread_id_x, %317] + %324 = affine.apply #map49()[%thread_id_y, %thread_id_x, %317] + %325 = affine.apply #map50()[%thread_id_y, %thread_id_x, %317] + %326 = affine.apply #map51()[%thread_id_y, %thread_id_x, %317] + %327 = affine.apply #map52()[%thread_id_y, %thread_id_x, %317] + %328 = affine.apply #map53()[%thread_id_y, %thread_id_x, %317] + %reinterpret_cast_18 = memref.reinterpret_cast %alloc_3 to offset: [0], sizes: [24576], strides: [1] : memref<192x128xi8, #gpu.address_space> to memref<24576xi8, #gpu.address_space> + %reinterpret_cast_19 = memref.reinterpret_cast %alloc to offset: [0], sizes: [32768], strides: [1] : memref<256x128xi8, #gpu.address_space> to memref<32768xi8, #gpu.address_space> + %329:34 = scf.for %arg13 = %c0 to %302 step %c2 iter_args(%arg14 = %cst, %arg15 = %cst, %arg16 = %cst, %arg17 = %cst, %arg18 = %cst, %arg19 = %cst, %arg20 = %cst, %arg21 = %cst, %arg22 = %cst, %arg23 = %cst, %arg24 = %cst, %arg25 = %cst, %arg26 = %cst, %arg27 = %cst, %arg28 = %cst, %arg29 = %cst, %arg30 = %cst, %arg31 = %cst, %arg32 = %cst, %arg33 = %cst, %arg34 = %cst, %arg35 = %cst, %arg36 = %cst, %arg37 = %cst, %arg38 = %108, %arg39 = %135, %arg40 = %162, %arg41 = %189, %arg42 = %204, %arg43 = %219, %arg44 = %250, %arg45 = %265, %arg46 = %280, %arg47 = %295) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>) { + %691 = vector.bitcast %arg47 : vector<4xi8> to vector<4xf8E8M0FNU> + %692 = vector.bitcast %arg46 : vector<4xi8> to vector<4xf8E8M0FNU> + %693 = vector.bitcast %arg45 : vector<4xi8> to vector<4xf8E8M0FNU> + %694 = vector.bitcast %arg44 : vector<4xi8> to vector<4xf8E8M0FNU> + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(0) + rocdl.s.barrier + %695 = affine.apply #map55()[%arg13, %13] + %696 = arith.addi %15, %695 overflow : index + amdgpu.gather_to_lds %18[%696], %alloc_3[%7, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %697 = arith.addi %22, %695 overflow : index + amdgpu.gather_to_lds %18[%697], %alloc_3[%20, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %698 = arith.addi %27, %695 overflow : index + amdgpu.gather_to_lds %18[%698], %alloc_3[%25, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %699 = affine.apply #map56()[%arg13, %13, %thread_id_x] + %700 = arith.index_cast %699 : index to i32 + %701 = arith.extui %700 : i32 to i64 + %702 = arith.muli %701, %39 : i64 + %703 = arith.shrui %702, %c32_i64 : i64 + %704 = arith.trunci %703 : i64 to i32 + %705 = arith.muli %704, %32 : i32 + %706 = arith.subi %700, %705 : i32 + %707 = arith.cmpi uge, %706, %32 : i32 + %708 = arith.extui %707 : i1 to i32 + %709 = arith.subi %704, %708 : i32 + %710 = arith.index_cast %709 : i32 to index + %711 = affine.apply #map15()[%710, %block_id_y, %thread_id_x, %thread_id_y] + %712 = arith.select %707, %32, %c0_i32 : i32 + %713 = arith.addi %706, %712 : i32 + %714 = arith.index_cast %713 : i32 to index + %715 = arith.muli %711, %9 overflow : index + %716 = arith.addi %715, %714 overflow : index + amdgpu.gather_to_lds %55[%716], %alloc[%30, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %717 = affine.apply #map17()[%710, %block_id_y, %thread_id_x, %thread_id_y] + %718 = arith.muli %717, %9 overflow : index + %719 = arith.addi %718, %714 overflow : index + amdgpu.gather_to_lds %55[%719], %alloc[%57, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %720 = affine.apply #map19()[%710, %block_id_y, %thread_id_x, %thread_id_y] + %721 = arith.muli %720, %9 overflow : index + %722 = arith.addi %721, %714 overflow : index + amdgpu.gather_to_lds %55[%722], %alloc[%62, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %723 = affine.apply #map21()[%710, %block_id_y, %thread_id_x, %thread_id_y] + %724 = arith.muli %723, %9 overflow : index + %725 = arith.addi %724, %714 overflow : index + amdgpu.gather_to_lds %55[%725], %alloc[%67, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + rocdl.sched.barrier 0 + %726 = vector.load %reinterpret_cast_16[%305] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %727 = vector.load %reinterpret_cast_16[%306] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %728 = vector.load %reinterpret_cast_16[%307] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %729 = vector.load %reinterpret_cast_17[%308] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %730 = vector.load %reinterpret_cast_17[%309] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %731 = vector.load %reinterpret_cast_17[%310] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %732 = vector.load %reinterpret_cast_17[%311] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %733 = vector.load %reinterpret_cast_17[%312] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %734 = vector.load %reinterpret_cast_17[%313] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %735 = vector.load %reinterpret_cast_17[%314] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %736 = vector.load %reinterpret_cast_17[%315] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %737 = vector.bitcast %726 : vector<16xi8> to vector<32xf4E2M1FN> + %738 = vector.bitcast %727 : vector<16xi8> to vector<32xf4E2M1FN> + %739 = vector.bitcast %728 : vector<16xi8> to vector<32xf4E2M1FN> + %740 = vector.bitcast %arg38 : vector<1xi8> to vector<1xf8E8M0FNU> + %741 = vector.bitcast %arg39 : vector<1xi8> to vector<1xf8E8M0FNU> + %742 = vector.bitcast %arg40 : vector<1xi8> to vector<1xf8E8M0FNU> + %743 = vector.bitcast %arg41 : vector<1xi8> to vector<1xf8E8M0FNU> + %744 = vector.bitcast %arg42 : vector<1xi8> to vector<1xf8E8M0FNU> + %745 = vector.bitcast %arg43 : vector<1xi8> to vector<1xf8E8M0FNU> + %746 = vector.bitcast %729 : vector<16xi8> to vector<32xf4E2M1FN> + %747 = vector.bitcast %730 : vector<16xi8> to vector<32xf4E2M1FN> + %748 = vector.bitcast %731 : vector<16xi8> to vector<32xf4E2M1FN> + %749 = vector.bitcast %732 : vector<16xi8> to vector<32xf4E2M1FN> + %750 = vector.bitcast %733 : vector<16xi8> to vector<32xf4E2M1FN> + %751 = vector.bitcast %734 : vector<16xi8> to vector<32xf4E2M1FN> + %752 = vector.bitcast %735 : vector<16xi8> to vector<32xf4E2M1FN> + %753 = vector.bitcast %736 : vector<16xi8> to vector<32xf4E2M1FN> + %754 = affine.apply #map57()[%thread_id_x, %arg7, %arg13] + %755 = arith.index_cast %754 : index to i32 + %756 = arith.extui %755 : i32 to i64 + %757 = arith.muli %756, %81 : i64 + %758 = arith.shrui %757, %c32_i64 : i64 + %759 = arith.trunci %758 : i64 to i32 + %760 = arith.muli %759, %74 : i32 + %761 = arith.subi %755, %760 : i32 + %762 = arith.cmpi uge, %761, %74 : i32 + %763 = arith.extui %762 : i1 to i32 + %764 = arith.subi %759, %763 : i32 + %765 = arith.index_cast %764 : i32 to index + %766 = affine.apply #map58()[%thread_id_x, %arg13] + %767 = arith.index_cast %766 : index to i32 + %768 = arith.extui %767 : i32 to i64 + %769 = arith.muli %768, %81 : i64 + %770 = arith.shrui %769, %c32_i64 : i64 + %771 = arith.trunci %770 : i64 to i32 + %772 = arith.muli %771, %74 : i32 + %773 = arith.subi %767, %772 : i32 + %774 = arith.cmpi uge, %773, %74 : i32 + %775 = arith.select %774, %74, %c0_i32 : i32 + %776 = arith.addi %773, %775 : i32 + %777 = arith.index_cast %776 : i32 to index + %778 = arith.muli %765, %strides#0 overflow : index + %779 = arith.addi %778, %777 overflow : index + %780 = vector.load %107[%779] : memref>, vector<1xi8> + %781 = affine.apply #map59()[%thread_id_x, %arg7, %arg13] + %782 = arith.index_cast %781 : index to i32 + %783 = arith.extui %782 : i32 to i64 + %784 = arith.muli %783, %81 : i64 + %785 = arith.shrui %784, %c32_i64 : i64 + %786 = arith.trunci %785 : i64 to i32 + %787 = arith.muli %786, %74 : i32 + %788 = arith.subi %782, %787 : i32 + %789 = arith.cmpi uge, %788, %74 : i32 + %790 = arith.extui %789 : i1 to i32 + %791 = arith.subi %786, %790 : i32 + %792 = arith.index_cast %791 : i32 to index + %793 = affine.apply #map60()[%thread_id_x, %arg13] + %794 = arith.index_cast %793 : index to i32 + %795 = arith.extui %794 : i32 to i64 + %796 = arith.muli %795, %81 : i64 + %797 = arith.shrui %796, %c32_i64 : i64 + %798 = arith.trunci %797 : i64 to i32 + %799 = arith.muli %798, %74 : i32 + %800 = arith.subi %794, %799 : i32 + %801 = arith.cmpi uge, %800, %74 : i32 + %802 = arith.select %801, %74, %c0_i32 : i32 + %803 = arith.addi %800, %802 : i32 + %804 = arith.index_cast %803 : i32 to index + %805 = arith.muli %792, %strides#0 overflow : index + %806 = arith.addi %805, %804 overflow : index + %807 = vector.load %107[%806] : memref>, vector<1xi8> + %808 = affine.apply #map61()[%thread_id_x, %arg7, %arg13] + %809 = arith.index_cast %808 : index to i32 + %810 = arith.extui %809 : i32 to i64 + %811 = arith.muli %810, %81 : i64 + %812 = arith.shrui %811, %c32_i64 : i64 + %813 = arith.trunci %812 : i64 to i32 + %814 = arith.muli %813, %74 : i32 + %815 = arith.subi %809, %814 : i32 + %816 = arith.cmpi uge, %815, %74 : i32 + %817 = arith.extui %816 : i1 to i32 + %818 = arith.subi %813, %817 : i32 + %819 = arith.index_cast %818 : i32 to index + %820 = affine.apply #map62()[%thread_id_x, %arg13] + %821 = arith.index_cast %820 : index to i32 + %822 = arith.extui %821 : i32 to i64 + %823 = arith.muli %822, %81 : i64 + %824 = arith.shrui %823, %c32_i64 : i64 + %825 = arith.trunci %824 : i64 to i32 + %826 = arith.muli %825, %74 : i32 + %827 = arith.subi %821, %826 : i32 + %828 = arith.cmpi uge, %827, %74 : i32 + %829 = arith.select %828, %74, %c0_i32 : i32 + %830 = arith.addi %827, %829 : i32 + %831 = arith.index_cast %830 : i32 to index + %832 = arith.muli %819, %strides#0 overflow : index + %833 = arith.addi %832, %831 overflow : index + %834 = vector.load %107[%833] : memref>, vector<1xi8> + %835 = affine.apply #map63()[%thread_id_x, %arg7, %arg13] + %836 = arith.index_cast %835 : index to i32 + %837 = arith.extui %836 : i32 to i64 + %838 = arith.muli %837, %81 : i64 + %839 = arith.shrui %838, %c32_i64 : i64 + %840 = arith.trunci %839 : i64 to i32 + %841 = arith.muli %840, %74 : i32 + %842 = arith.subi %836, %841 : i32 + %843 = arith.cmpi uge, %842, %74 : i32 + %844 = arith.extui %843 : i1 to i32 + %845 = arith.subi %840, %844 : i32 + %846 = arith.index_cast %845 : i32 to index + %847 = affine.apply #map64()[%thread_id_x, %arg13] + %848 = arith.index_cast %847 : index to i32 + %849 = arith.extui %848 : i32 to i64 + %850 = arith.muli %849, %81 : i64 + %851 = arith.shrui %850, %c32_i64 : i64 + %852 = arith.trunci %851 : i64 to i32 + %853 = arith.muli %852, %74 : i32 + %854 = arith.subi %848, %853 : i32 + %855 = arith.cmpi uge, %854, %74 : i32 + %856 = arith.select %855, %74, %c0_i32 : i32 + %857 = arith.addi %854, %856 : i32 + %858 = arith.index_cast %857 : i32 to index + %859 = arith.muli %846, %strides#0 overflow : index + %860 = arith.addi %859, %858 overflow : index + %861 = vector.load %107[%860] : memref>, vector<1xi8> + %862 = affine.apply #map65()[%thread_id_x, %arg7, %arg13] + %863 = arith.index_cast %862 : index to i32 + %864 = arith.extui %863 : i32 to i64 + %865 = arith.muli %864, %81 : i64 + %866 = arith.shrui %865, %c32_i64 : i64 + %867 = arith.trunci %866 : i64 to i32 + %868 = arith.muli %867, %74 : i32 + %869 = arith.subi %863, %868 : i32 + %870 = arith.cmpi uge, %869, %74 : i32 + %871 = arith.extui %870 : i1 to i32 + %872 = arith.subi %867, %871 : i32 + %873 = arith.index_cast %872 : i32 to index + %874 = arith.muli %873, %strides#0 overflow : index + %875 = arith.addi %874, %777 overflow : index + %876 = vector.load %107[%875] : memref>, vector<1xi8> + %877 = affine.apply #map66()[%thread_id_x, %arg7, %arg13] + %878 = arith.index_cast %877 : index to i32 + %879 = arith.extui %878 : i32 to i64 + %880 = arith.muli %879, %81 : i64 + %881 = arith.shrui %880, %c32_i64 : i64 + %882 = arith.trunci %881 : i64 to i32 + %883 = arith.muli %882, %74 : i32 + %884 = arith.subi %878, %883 : i32 + %885 = arith.cmpi uge, %884, %74 : i32 + %886 = arith.extui %885 : i1 to i32 + %887 = arith.subi %882, %886 : i32 + %888 = arith.index_cast %887 : i32 to index + %889 = arith.muli %888, %strides#0 overflow : index + %890 = arith.addi %889, %804 overflow : index + %891 = vector.load %107[%890] : memref>, vector<1xi8> + %892 = affine.apply #map67()[%thread_id_x, %arg7, %thread_id_y, %arg13] + %893 = arith.index_cast %892 : index to i32 + %894 = arith.extui %893 : i32 to i64 + %895 = arith.muli %894, %81 : i64 + %896 = arith.shrui %895, %c32_i64 : i64 + %897 = arith.trunci %896 : i64 to i32 + %898 = arith.muli %897, %74 : i32 + %899 = arith.subi %893, %898 : i32 + %900 = arith.cmpi uge, %899, %74 : i32 + %901 = arith.extui %900 : i1 to i32 + %902 = arith.subi %897, %901 : i32 + %903 = arith.index_cast %902 : i32 to index + %904 = affine.apply #map68()[%thread_id_x, %arg13] + %905 = arith.index_cast %904 : index to i32 + %906 = arith.extui %905 : i32 to i64 + %907 = arith.muli %906, %81 : i64 + %908 = arith.shrui %907, %c32_i64 : i64 + %909 = arith.trunci %908 : i64 to i32 + %910 = arith.muli %909, %74 : i32 + %911 = arith.subi %905, %910 : i32 + %912 = arith.cmpi uge, %911, %74 : i32 + %913 = arith.select %912, %74, %c0_i32 : i32 + %914 = arith.addi %911, %913 : i32 + %915 = arith.index_cast %914 : i32 to index + %916 = arith.muli %903, %strides_13#0 overflow : index + %917 = arith.addi %916, %915 overflow : index + %918 = vector.load %249[%917] : memref>, vector<4xi8> + %919 = vector.bitcast %918 : vector<4xi8> to vector<4xf8E8M0FNU> + %920 = affine.apply #map69()[%thread_id_x, %arg7, %thread_id_y, %arg13] + %921 = arith.index_cast %920 : index to i32 + %922 = arith.extui %921 : i32 to i64 + %923 = arith.muli %922, %81 : i64 + %924 = arith.shrui %923, %c32_i64 : i64 + %925 = arith.trunci %924 : i64 to i32 + %926 = arith.muli %925, %74 : i32 + %927 = arith.subi %921, %926 : i32 + %928 = arith.cmpi uge, %927, %74 : i32 + %929 = arith.extui %928 : i1 to i32 + %930 = arith.subi %925, %929 : i32 + %931 = arith.index_cast %930 : i32 to index + %932 = arith.muli %931, %strides_13#0 overflow : index + %933 = arith.addi %932, %915 overflow : index + %934 = vector.load %249[%933] : memref>, vector<4xi8> + %935 = vector.bitcast %934 : vector<4xi8> to vector<4xf8E8M0FNU> + %936 = affine.apply #map70()[%thread_id_x, %arg7, %thread_id_y, %arg13] + %937 = arith.index_cast %936 : index to i32 + %938 = arith.extui %937 : i32 to i64 + %939 = arith.muli %938, %81 : i64 + %940 = arith.shrui %939, %c32_i64 : i64 + %941 = arith.trunci %940 : i64 to i32 + %942 = arith.muli %941, %74 : i32 + %943 = arith.subi %937, %942 : i32 + %944 = arith.cmpi uge, %943, %74 : i32 + %945 = arith.extui %944 : i1 to i32 + %946 = arith.subi %941, %945 : i32 + %947 = arith.index_cast %946 : i32 to index + %948 = arith.muli %947, %strides_13#0 overflow : index + %949 = arith.addi %948, %915 overflow : index + %950 = vector.load %249[%949] : memref>, vector<4xi8> + %951 = vector.bitcast %950 : vector<4xi8> to vector<4xf8E8M0FNU> + %952 = affine.apply #map71()[%thread_id_x, %arg7, %thread_id_y, %arg13] + %953 = arith.index_cast %952 : index to i32 + %954 = arith.extui %953 : i32 to i64 + %955 = arith.muli %954, %81 : i64 + %956 = arith.shrui %955, %c32_i64 : i64 + %957 = arith.trunci %956 : i64 to i32 + %958 = arith.muli %957, %74 : i32 + %959 = arith.subi %953, %958 : i32 + %960 = arith.cmpi uge, %959, %74 : i32 + %961 = arith.extui %960 : i1 to i32 + %962 = arith.subi %957, %961 : i32 + %963 = arith.index_cast %962 : i32 to index + %964 = arith.muli %963, %strides_13#0 overflow : index + %965 = arith.addi %964, %915 overflow : index + %966 = vector.load %249[%965] : memref>, vector<4xi8> + %967 = vector.bitcast %966 : vector<4xi8> to vector<4xf8E8M0FNU> + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %968 = vector.extract %740[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %969 = amdgpu.scaled_mfma 16x16x128 (%968[0] * %737) * (%694[0] * %746) + %arg14 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %970 = amdgpu.scaled_mfma 16x16x128 (%968[0] * %737) * (%694[1] * %747) + %arg15 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %971 = amdgpu.scaled_mfma 16x16x128 (%968[0] * %737) * (%693[0] * %748) + %arg16 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %972 = amdgpu.scaled_mfma 16x16x128 (%968[0] * %737) * (%693[1] * %749) + %arg17 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %973 = amdgpu.scaled_mfma 16x16x128 (%968[0] * %737) * (%692[0] * %750) + %arg18 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %974 = amdgpu.scaled_mfma 16x16x128 (%968[0] * %737) * (%692[1] * %751) + %arg19 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %975 = amdgpu.scaled_mfma 16x16x128 (%968[0] * %737) * (%691[0] * %752) + %arg20 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %976 = amdgpu.scaled_mfma 16x16x128 (%968[0] * %737) * (%691[1] * %753) + %arg21 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %977 = vector.extract %742[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %978 = amdgpu.scaled_mfma 16x16x128 (%977[0] * %738) * (%694[0] * %746) + %arg22 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %979 = amdgpu.scaled_mfma 16x16x128 (%977[0] * %738) * (%694[1] * %747) + %arg23 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %980 = amdgpu.scaled_mfma 16x16x128 (%977[0] * %738) * (%693[0] * %748) + %arg24 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %981 = amdgpu.scaled_mfma 16x16x128 (%977[0] * %738) * (%693[1] * %749) + %arg25 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %982 = amdgpu.scaled_mfma 16x16x128 (%977[0] * %738) * (%692[0] * %750) + %arg26 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %983 = amdgpu.scaled_mfma 16x16x128 (%977[0] * %738) * (%692[1] * %751) + %arg27 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %984 = amdgpu.scaled_mfma 16x16x128 (%977[0] * %738) * (%691[0] * %752) + %arg28 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %985 = amdgpu.scaled_mfma 16x16x128 (%977[0] * %738) * (%691[1] * %753) + %arg29 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %986 = vector.extract %744[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %987 = amdgpu.scaled_mfma 16x16x128 (%986[0] * %739) * (%694[0] * %746) + %arg30 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %988 = amdgpu.scaled_mfma 16x16x128 (%986[0] * %739) * (%694[1] * %747) + %arg31 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %989 = amdgpu.scaled_mfma 16x16x128 (%986[0] * %739) * (%693[0] * %748) + %arg32 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %990 = amdgpu.scaled_mfma 16x16x128 (%986[0] * %739) * (%693[1] * %749) + %arg33 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %991 = amdgpu.scaled_mfma 16x16x128 (%986[0] * %739) * (%692[0] * %750) + %arg34 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %992 = amdgpu.scaled_mfma 16x16x128 (%986[0] * %739) * (%692[1] * %751) + %arg35 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %993 = amdgpu.scaled_mfma 16x16x128 (%986[0] * %739) * (%691[0] * %752) + %arg36 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %994 = amdgpu.scaled_mfma 16x16x128 (%986[0] * %739) * (%691[1] * %753) + %arg37 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.sched.barrier 0 + %995 = vector.load %reinterpret_cast_16[%318] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %996 = vector.load %reinterpret_cast_16[%319] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %997 = vector.load %reinterpret_cast_16[%320] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %998 = vector.load %reinterpret_cast_17[%321] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %999 = vector.load %reinterpret_cast_17[%322] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1000 = vector.load %reinterpret_cast_17[%323] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1001 = vector.load %reinterpret_cast_17[%324] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1002 = vector.load %reinterpret_cast_17[%325] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1003 = vector.load %reinterpret_cast_17[%326] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1004 = vector.load %reinterpret_cast_17[%327] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1005 = vector.load %reinterpret_cast_17[%328] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1006 = vector.bitcast %995 : vector<16xi8> to vector<32xf4E2M1FN> + %1007 = vector.bitcast %996 : vector<16xi8> to vector<32xf4E2M1FN> + %1008 = vector.bitcast %997 : vector<16xi8> to vector<32xf4E2M1FN> + %1009 = vector.bitcast %998 : vector<16xi8> to vector<32xf4E2M1FN> + %1010 = vector.bitcast %999 : vector<16xi8> to vector<32xf4E2M1FN> + %1011 = vector.bitcast %1000 : vector<16xi8> to vector<32xf4E2M1FN> + %1012 = vector.bitcast %1001 : vector<16xi8> to vector<32xf4E2M1FN> + %1013 = vector.bitcast %1002 : vector<16xi8> to vector<32xf4E2M1FN> + %1014 = vector.bitcast %1003 : vector<16xi8> to vector<32xf4E2M1FN> + %1015 = vector.bitcast %1004 : vector<16xi8> to vector<32xf4E2M1FN> + %1016 = vector.bitcast %1005 : vector<16xi8> to vector<32xf4E2M1FN> + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(0) + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %1017 = vector.extract %741[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1018 = amdgpu.scaled_mfma 16x16x128 (%1017[0] * %1006) * (%694[2] * %1009) + %969 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1019 = amdgpu.scaled_mfma 16x16x128 (%1017[0] * %1006) * (%694[3] * %1010) + %970 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1020 = amdgpu.scaled_mfma 16x16x128 (%1017[0] * %1006) * (%693[2] * %1011) + %971 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1021 = amdgpu.scaled_mfma 16x16x128 (%1017[0] * %1006) * (%693[3] * %1012) + %972 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1022 = amdgpu.scaled_mfma 16x16x128 (%1017[0] * %1006) * (%692[2] * %1013) + %973 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1023 = amdgpu.scaled_mfma 16x16x128 (%1017[0] * %1006) * (%692[3] * %1014) + %974 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1024 = amdgpu.scaled_mfma 16x16x128 (%1017[0] * %1006) * (%691[2] * %1015) + %975 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1025 = amdgpu.scaled_mfma 16x16x128 (%1017[0] * %1006) * (%691[3] * %1016) + %976 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1026 = vector.extract %743[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1027 = amdgpu.scaled_mfma 16x16x128 (%1026[0] * %1007) * (%694[2] * %1009) + %978 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1028 = amdgpu.scaled_mfma 16x16x128 (%1026[0] * %1007) * (%694[3] * %1010) + %979 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1029 = amdgpu.scaled_mfma 16x16x128 (%1026[0] * %1007) * (%693[2] * %1011) + %980 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1030 = amdgpu.scaled_mfma 16x16x128 (%1026[0] * %1007) * (%693[3] * %1012) + %981 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1031 = amdgpu.scaled_mfma 16x16x128 (%1026[0] * %1007) * (%692[2] * %1013) + %982 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1032 = amdgpu.scaled_mfma 16x16x128 (%1026[0] * %1007) * (%692[3] * %1014) + %983 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1033 = amdgpu.scaled_mfma 16x16x128 (%1026[0] * %1007) * (%691[2] * %1015) + %984 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1034 = amdgpu.scaled_mfma 16x16x128 (%1026[0] * %1007) * (%691[3] * %1016) + %985 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1035 = vector.extract %745[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1036 = amdgpu.scaled_mfma 16x16x128 (%1035[0] * %1008) * (%694[2] * %1009) + %987 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1037 = amdgpu.scaled_mfma 16x16x128 (%1035[0] * %1008) * (%694[3] * %1010) + %988 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1038 = amdgpu.scaled_mfma 16x16x128 (%1035[0] * %1008) * (%693[2] * %1011) + %989 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1039 = amdgpu.scaled_mfma 16x16x128 (%1035[0] * %1008) * (%693[3] * %1012) + %990 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1040 = amdgpu.scaled_mfma 16x16x128 (%1035[0] * %1008) * (%692[2] * %1013) + %991 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1041 = amdgpu.scaled_mfma 16x16x128 (%1035[0] * %1008) * (%692[3] * %1014) + %992 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1042 = amdgpu.scaled_mfma 16x16x128 (%1035[0] * %1008) * (%691[2] * %1015) + %993 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1043 = amdgpu.scaled_mfma 16x16x128 (%1035[0] * %1008) * (%691[3] * %1016) + %994 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + %1044 = arith.addi %arg13, %c1 : index + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(0) + rocdl.s.barrier + %1045 = affine.apply #map55()[%1044, %13] + %1046 = arith.addi %15, %1045 overflow : index + amdgpu.gather_to_lds %18[%1046], %alloc_4[%7, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %1047 = arith.addi %22, %1045 overflow : index + amdgpu.gather_to_lds %18[%1047], %alloc_4[%20, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %1048 = arith.addi %27, %1045 overflow : index + amdgpu.gather_to_lds %18[%1048], %alloc_4[%25, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %1049 = affine.apply #map56()[%1044, %13, %thread_id_x] + %1050 = arith.index_cast %1049 : index to i32 + %1051 = arith.extui %1050 : i32 to i64 + %1052 = arith.muli %1051, %39 : i64 + %1053 = arith.shrui %1052, %c32_i64 : i64 + %1054 = arith.trunci %1053 : i64 to i32 + %1055 = arith.muli %1054, %32 : i32 + %1056 = arith.subi %1050, %1055 : i32 + %1057 = arith.cmpi uge, %1056, %32 : i32 + %1058 = arith.extui %1057 : i1 to i32 + %1059 = arith.subi %1054, %1058 : i32 + %1060 = arith.index_cast %1059 : i32 to index + %1061 = affine.apply #map15()[%1060, %block_id_y, %thread_id_x, %thread_id_y] + %1062 = arith.select %1057, %32, %c0_i32 : i32 + %1063 = arith.addi %1056, %1062 : i32 + %1064 = arith.index_cast %1063 : i32 to index + %1065 = arith.muli %1061, %9 overflow : index + %1066 = arith.addi %1065, %1064 overflow : index + amdgpu.gather_to_lds %55[%1066], %alloc_2[%30, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %1067 = affine.apply #map17()[%1060, %block_id_y, %thread_id_x, %thread_id_y] + %1068 = arith.muli %1067, %9 overflow : index + %1069 = arith.addi %1068, %1064 overflow : index + amdgpu.gather_to_lds %55[%1069], %alloc_2[%57, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %1070 = affine.apply #map19()[%1060, %block_id_y, %thread_id_x, %thread_id_y] + %1071 = arith.muli %1070, %9 overflow : index + %1072 = arith.addi %1071, %1064 overflow : index + amdgpu.gather_to_lds %55[%1072], %alloc_2[%62, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %1073 = affine.apply #map21()[%1060, %block_id_y, %thread_id_x, %thread_id_y] + %1074 = arith.muli %1073, %9 overflow : index + %1075 = arith.addi %1074, %1064 overflow : index + amdgpu.gather_to_lds %55[%1075], %alloc_2[%67, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + rocdl.sched.barrier 0 + %1076 = vector.load %reinterpret_cast_18[%305] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %1077 = vector.load %reinterpret_cast_18[%306] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %1078 = vector.load %reinterpret_cast_18[%307] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %1079 = vector.load %reinterpret_cast_19[%308] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1080 = vector.load %reinterpret_cast_19[%309] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1081 = vector.load %reinterpret_cast_19[%310] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1082 = vector.load %reinterpret_cast_19[%311] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1083 = vector.load %reinterpret_cast_19[%312] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1084 = vector.load %reinterpret_cast_19[%313] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1085 = vector.load %reinterpret_cast_19[%314] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1086 = vector.load %reinterpret_cast_19[%315] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1087 = vector.bitcast %1076 : vector<16xi8> to vector<32xf4E2M1FN> + %1088 = vector.bitcast %1077 : vector<16xi8> to vector<32xf4E2M1FN> + %1089 = vector.bitcast %1078 : vector<16xi8> to vector<32xf4E2M1FN> + %1090 = vector.bitcast %780 : vector<1xi8> to vector<1xf8E8M0FNU> + %1091 = vector.bitcast %807 : vector<1xi8> to vector<1xf8E8M0FNU> + %1092 = vector.bitcast %834 : vector<1xi8> to vector<1xf8E8M0FNU> + %1093 = vector.bitcast %861 : vector<1xi8> to vector<1xf8E8M0FNU> + %1094 = vector.bitcast %876 : vector<1xi8> to vector<1xf8E8M0FNU> + %1095 = vector.bitcast %891 : vector<1xi8> to vector<1xf8E8M0FNU> + %1096 = vector.bitcast %1079 : vector<16xi8> to vector<32xf4E2M1FN> + %1097 = vector.bitcast %1080 : vector<16xi8> to vector<32xf4E2M1FN> + %1098 = vector.bitcast %1081 : vector<16xi8> to vector<32xf4E2M1FN> + %1099 = vector.bitcast %1082 : vector<16xi8> to vector<32xf4E2M1FN> + %1100 = vector.bitcast %1083 : vector<16xi8> to vector<32xf4E2M1FN> + %1101 = vector.bitcast %1084 : vector<16xi8> to vector<32xf4E2M1FN> + %1102 = vector.bitcast %1085 : vector<16xi8> to vector<32xf4E2M1FN> + %1103 = vector.bitcast %1086 : vector<16xi8> to vector<32xf4E2M1FN> + %1104 = affine.apply #map57()[%thread_id_x, %arg7, %1044] + %1105 = arith.index_cast %1104 : index to i32 + %1106 = arith.extui %1105 : i32 to i64 + %1107 = arith.muli %1106, %81 : i64 + %1108 = arith.shrui %1107, %c32_i64 : i64 + %1109 = arith.trunci %1108 : i64 to i32 + %1110 = arith.muli %1109, %74 : i32 + %1111 = arith.subi %1105, %1110 : i32 + %1112 = arith.cmpi uge, %1111, %74 : i32 + %1113 = arith.extui %1112 : i1 to i32 + %1114 = arith.subi %1109, %1113 : i32 + %1115 = arith.index_cast %1114 : i32 to index + %1116 = affine.apply #map58()[%thread_id_x, %1044] + %1117 = arith.index_cast %1116 : index to i32 + %1118 = arith.extui %1117 : i32 to i64 + %1119 = arith.muli %1118, %81 : i64 + %1120 = arith.shrui %1119, %c32_i64 : i64 + %1121 = arith.trunci %1120 : i64 to i32 + %1122 = arith.muli %1121, %74 : i32 + %1123 = arith.subi %1117, %1122 : i32 + %1124 = arith.cmpi uge, %1123, %74 : i32 + %1125 = arith.select %1124, %74, %c0_i32 : i32 + %1126 = arith.addi %1123, %1125 : i32 + %1127 = arith.index_cast %1126 : i32 to index + %1128 = arith.muli %1115, %strides#0 overflow : index + %1129 = arith.addi %1128, %1127 overflow : index + %1130 = vector.load %107[%1129] : memref>, vector<1xi8> + %1131 = affine.apply #map59()[%thread_id_x, %arg7, %1044] + %1132 = arith.index_cast %1131 : index to i32 + %1133 = arith.extui %1132 : i32 to i64 + %1134 = arith.muli %1133, %81 : i64 + %1135 = arith.shrui %1134, %c32_i64 : i64 + %1136 = arith.trunci %1135 : i64 to i32 + %1137 = arith.muli %1136, %74 : i32 + %1138 = arith.subi %1132, %1137 : i32 + %1139 = arith.cmpi uge, %1138, %74 : i32 + %1140 = arith.extui %1139 : i1 to i32 + %1141 = arith.subi %1136, %1140 : i32 + %1142 = arith.index_cast %1141 : i32 to index + %1143 = affine.apply #map60()[%thread_id_x, %1044] + %1144 = arith.index_cast %1143 : index to i32 + %1145 = arith.extui %1144 : i32 to i64 + %1146 = arith.muli %1145, %81 : i64 + %1147 = arith.shrui %1146, %c32_i64 : i64 + %1148 = arith.trunci %1147 : i64 to i32 + %1149 = arith.muli %1148, %74 : i32 + %1150 = arith.subi %1144, %1149 : i32 + %1151 = arith.cmpi uge, %1150, %74 : i32 + %1152 = arith.select %1151, %74, %c0_i32 : i32 + %1153 = arith.addi %1150, %1152 : i32 + %1154 = arith.index_cast %1153 : i32 to index + %1155 = arith.muli %1142, %strides#0 overflow : index + %1156 = arith.addi %1155, %1154 overflow : index + %1157 = vector.load %107[%1156] : memref>, vector<1xi8> + %1158 = affine.apply #map61()[%thread_id_x, %arg7, %1044] + %1159 = arith.index_cast %1158 : index to i32 + %1160 = arith.extui %1159 : i32 to i64 + %1161 = arith.muli %1160, %81 : i64 + %1162 = arith.shrui %1161, %c32_i64 : i64 + %1163 = arith.trunci %1162 : i64 to i32 + %1164 = arith.muli %1163, %74 : i32 + %1165 = arith.subi %1159, %1164 : i32 + %1166 = arith.cmpi uge, %1165, %74 : i32 + %1167 = arith.extui %1166 : i1 to i32 + %1168 = arith.subi %1163, %1167 : i32 + %1169 = arith.index_cast %1168 : i32 to index + %1170 = affine.apply #map62()[%thread_id_x, %1044] + %1171 = arith.index_cast %1170 : index to i32 + %1172 = arith.extui %1171 : i32 to i64 + %1173 = arith.muli %1172, %81 : i64 + %1174 = arith.shrui %1173, %c32_i64 : i64 + %1175 = arith.trunci %1174 : i64 to i32 + %1176 = arith.muli %1175, %74 : i32 + %1177 = arith.subi %1171, %1176 : i32 + %1178 = arith.cmpi uge, %1177, %74 : i32 + %1179 = arith.select %1178, %74, %c0_i32 : i32 + %1180 = arith.addi %1177, %1179 : i32 + %1181 = arith.index_cast %1180 : i32 to index + %1182 = arith.muli %1169, %strides#0 overflow : index + %1183 = arith.addi %1182, %1181 overflow : index + %1184 = vector.load %107[%1183] : memref>, vector<1xi8> + %1185 = affine.apply #map63()[%thread_id_x, %arg7, %1044] + %1186 = arith.index_cast %1185 : index to i32 + %1187 = arith.extui %1186 : i32 to i64 + %1188 = arith.muli %1187, %81 : i64 + %1189 = arith.shrui %1188, %c32_i64 : i64 + %1190 = arith.trunci %1189 : i64 to i32 + %1191 = arith.muli %1190, %74 : i32 + %1192 = arith.subi %1186, %1191 : i32 + %1193 = arith.cmpi uge, %1192, %74 : i32 + %1194 = arith.extui %1193 : i1 to i32 + %1195 = arith.subi %1190, %1194 : i32 + %1196 = arith.index_cast %1195 : i32 to index + %1197 = affine.apply #map64()[%thread_id_x, %1044] + %1198 = arith.index_cast %1197 : index to i32 + %1199 = arith.extui %1198 : i32 to i64 + %1200 = arith.muli %1199, %81 : i64 + %1201 = arith.shrui %1200, %c32_i64 : i64 + %1202 = arith.trunci %1201 : i64 to i32 + %1203 = arith.muli %1202, %74 : i32 + %1204 = arith.subi %1198, %1203 : i32 + %1205 = arith.cmpi uge, %1204, %74 : i32 + %1206 = arith.select %1205, %74, %c0_i32 : i32 + %1207 = arith.addi %1204, %1206 : i32 + %1208 = arith.index_cast %1207 : i32 to index + %1209 = arith.muli %1196, %strides#0 overflow : index + %1210 = arith.addi %1209, %1208 overflow : index + %1211 = vector.load %107[%1210] : memref>, vector<1xi8> + %1212 = affine.apply #map65()[%thread_id_x, %arg7, %1044] + %1213 = arith.index_cast %1212 : index to i32 + %1214 = arith.extui %1213 : i32 to i64 + %1215 = arith.muli %1214, %81 : i64 + %1216 = arith.shrui %1215, %c32_i64 : i64 + %1217 = arith.trunci %1216 : i64 to i32 + %1218 = arith.muli %1217, %74 : i32 + %1219 = arith.subi %1213, %1218 : i32 + %1220 = arith.cmpi uge, %1219, %74 : i32 + %1221 = arith.extui %1220 : i1 to i32 + %1222 = arith.subi %1217, %1221 : i32 + %1223 = arith.index_cast %1222 : i32 to index + %1224 = arith.muli %1223, %strides#0 overflow : index + %1225 = arith.addi %1224, %1127 overflow : index + %1226 = vector.load %107[%1225] : memref>, vector<1xi8> + %1227 = affine.apply #map66()[%thread_id_x, %arg7, %1044] + %1228 = arith.index_cast %1227 : index to i32 + %1229 = arith.extui %1228 : i32 to i64 + %1230 = arith.muli %1229, %81 : i64 + %1231 = arith.shrui %1230, %c32_i64 : i64 + %1232 = arith.trunci %1231 : i64 to i32 + %1233 = arith.muli %1232, %74 : i32 + %1234 = arith.subi %1228, %1233 : i32 + %1235 = arith.cmpi uge, %1234, %74 : i32 + %1236 = arith.extui %1235 : i1 to i32 + %1237 = arith.subi %1232, %1236 : i32 + %1238 = arith.index_cast %1237 : i32 to index + %1239 = arith.muli %1238, %strides#0 overflow : index + %1240 = arith.addi %1239, %1154 overflow : index + %1241 = vector.load %107[%1240] : memref>, vector<1xi8> + %1242 = affine.apply #map67()[%thread_id_x, %arg7, %thread_id_y, %1044] + %1243 = arith.index_cast %1242 : index to i32 + %1244 = arith.extui %1243 : i32 to i64 + %1245 = arith.muli %1244, %81 : i64 + %1246 = arith.shrui %1245, %c32_i64 : i64 + %1247 = arith.trunci %1246 : i64 to i32 + %1248 = arith.muli %1247, %74 : i32 + %1249 = arith.subi %1243, %1248 : i32 + %1250 = arith.cmpi uge, %1249, %74 : i32 + %1251 = arith.extui %1250 : i1 to i32 + %1252 = arith.subi %1247, %1251 : i32 + %1253 = arith.index_cast %1252 : i32 to index + %1254 = affine.apply #map68()[%thread_id_x, %1044] + %1255 = arith.index_cast %1254 : index to i32 + %1256 = arith.extui %1255 : i32 to i64 + %1257 = arith.muli %1256, %81 : i64 + %1258 = arith.shrui %1257, %c32_i64 : i64 + %1259 = arith.trunci %1258 : i64 to i32 + %1260 = arith.muli %1259, %74 : i32 + %1261 = arith.subi %1255, %1260 : i32 + %1262 = arith.cmpi uge, %1261, %74 : i32 + %1263 = arith.select %1262, %74, %c0_i32 : i32 + %1264 = arith.addi %1261, %1263 : i32 + %1265 = arith.index_cast %1264 : i32 to index + %1266 = arith.muli %1253, %strides_13#0 overflow : index + %1267 = arith.addi %1266, %1265 overflow : index + %1268 = vector.load %249[%1267] : memref>, vector<4xi8> + %1269 = affine.apply #map69()[%thread_id_x, %arg7, %thread_id_y, %1044] + %1270 = arith.index_cast %1269 : index to i32 + %1271 = arith.extui %1270 : i32 to i64 + %1272 = arith.muli %1271, %81 : i64 + %1273 = arith.shrui %1272, %c32_i64 : i64 + %1274 = arith.trunci %1273 : i64 to i32 + %1275 = arith.muli %1274, %74 : i32 + %1276 = arith.subi %1270, %1275 : i32 + %1277 = arith.cmpi uge, %1276, %74 : i32 + %1278 = arith.extui %1277 : i1 to i32 + %1279 = arith.subi %1274, %1278 : i32 + %1280 = arith.index_cast %1279 : i32 to index + %1281 = arith.muli %1280, %strides_13#0 overflow : index + %1282 = arith.addi %1281, %1265 overflow : index + %1283 = vector.load %249[%1282] : memref>, vector<4xi8> + %1284 = affine.apply #map70()[%thread_id_x, %arg7, %thread_id_y, %1044] + %1285 = arith.index_cast %1284 : index to i32 + %1286 = arith.extui %1285 : i32 to i64 + %1287 = arith.muli %1286, %81 : i64 + %1288 = arith.shrui %1287, %c32_i64 : i64 + %1289 = arith.trunci %1288 : i64 to i32 + %1290 = arith.muli %1289, %74 : i32 + %1291 = arith.subi %1285, %1290 : i32 + %1292 = arith.cmpi uge, %1291, %74 : i32 + %1293 = arith.extui %1292 : i1 to i32 + %1294 = arith.subi %1289, %1293 : i32 + %1295 = arith.index_cast %1294 : i32 to index + %1296 = arith.muli %1295, %strides_13#0 overflow : index + %1297 = arith.addi %1296, %1265 overflow : index + %1298 = vector.load %249[%1297] : memref>, vector<4xi8> + %1299 = affine.apply #map71()[%thread_id_x, %arg7, %thread_id_y, %1044] + %1300 = arith.index_cast %1299 : index to i32 + %1301 = arith.extui %1300 : i32 to i64 + %1302 = arith.muli %1301, %81 : i64 + %1303 = arith.shrui %1302, %c32_i64 : i64 + %1304 = arith.trunci %1303 : i64 to i32 + %1305 = arith.muli %1304, %74 : i32 + %1306 = arith.subi %1300, %1305 : i32 + %1307 = arith.cmpi uge, %1306, %74 : i32 + %1308 = arith.extui %1307 : i1 to i32 + %1309 = arith.subi %1304, %1308 : i32 + %1310 = arith.index_cast %1309 : i32 to index + %1311 = arith.muli %1310, %strides_13#0 overflow : index + %1312 = arith.addi %1311, %1265 overflow : index + %1313 = vector.load %249[%1312] : memref>, vector<4xi8> + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %1314 = vector.extract %1090[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1315 = amdgpu.scaled_mfma 16x16x128 (%1314[0] * %1087) * (%919[0] * %1096) + %1018 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1316 = amdgpu.scaled_mfma 16x16x128 (%1314[0] * %1087) * (%919[1] * %1097) + %1019 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1317 = amdgpu.scaled_mfma 16x16x128 (%1314[0] * %1087) * (%935[0] * %1098) + %1020 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1318 = amdgpu.scaled_mfma 16x16x128 (%1314[0] * %1087) * (%935[1] * %1099) + %1021 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1319 = amdgpu.scaled_mfma 16x16x128 (%1314[0] * %1087) * (%951[0] * %1100) + %1022 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1320 = amdgpu.scaled_mfma 16x16x128 (%1314[0] * %1087) * (%951[1] * %1101) + %1023 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1321 = amdgpu.scaled_mfma 16x16x128 (%1314[0] * %1087) * (%967[0] * %1102) + %1024 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1322 = amdgpu.scaled_mfma 16x16x128 (%1314[0] * %1087) * (%967[1] * %1103) + %1025 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1323 = vector.extract %1092[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1324 = amdgpu.scaled_mfma 16x16x128 (%1323[0] * %1088) * (%919[0] * %1096) + %1027 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1325 = amdgpu.scaled_mfma 16x16x128 (%1323[0] * %1088) * (%919[1] * %1097) + %1028 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1326 = amdgpu.scaled_mfma 16x16x128 (%1323[0] * %1088) * (%935[0] * %1098) + %1029 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1327 = amdgpu.scaled_mfma 16x16x128 (%1323[0] * %1088) * (%935[1] * %1099) + %1030 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1328 = amdgpu.scaled_mfma 16x16x128 (%1323[0] * %1088) * (%951[0] * %1100) + %1031 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1329 = amdgpu.scaled_mfma 16x16x128 (%1323[0] * %1088) * (%951[1] * %1101) + %1032 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1330 = amdgpu.scaled_mfma 16x16x128 (%1323[0] * %1088) * (%967[0] * %1102) + %1033 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1331 = amdgpu.scaled_mfma 16x16x128 (%1323[0] * %1088) * (%967[1] * %1103) + %1034 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1332 = vector.extract %1094[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1333 = amdgpu.scaled_mfma 16x16x128 (%1332[0] * %1089) * (%919[0] * %1096) + %1036 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1334 = amdgpu.scaled_mfma 16x16x128 (%1332[0] * %1089) * (%919[1] * %1097) + %1037 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1335 = amdgpu.scaled_mfma 16x16x128 (%1332[0] * %1089) * (%935[0] * %1098) + %1038 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1336 = amdgpu.scaled_mfma 16x16x128 (%1332[0] * %1089) * (%935[1] * %1099) + %1039 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1337 = amdgpu.scaled_mfma 16x16x128 (%1332[0] * %1089) * (%951[0] * %1100) + %1040 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1338 = amdgpu.scaled_mfma 16x16x128 (%1332[0] * %1089) * (%951[1] * %1101) + %1041 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1339 = amdgpu.scaled_mfma 16x16x128 (%1332[0] * %1089) * (%967[0] * %1102) + %1042 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1340 = amdgpu.scaled_mfma 16x16x128 (%1332[0] * %1089) * (%967[1] * %1103) + %1043 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.sched.barrier 0 + %1341 = vector.load %reinterpret_cast_18[%318] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %1342 = vector.load %reinterpret_cast_18[%319] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %1343 = vector.load %reinterpret_cast_18[%320] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %1344 = vector.load %reinterpret_cast_19[%321] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1345 = vector.load %reinterpret_cast_19[%322] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1346 = vector.load %reinterpret_cast_19[%323] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1347 = vector.load %reinterpret_cast_19[%324] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1348 = vector.load %reinterpret_cast_19[%325] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1349 = vector.load %reinterpret_cast_19[%326] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1350 = vector.load %reinterpret_cast_19[%327] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1351 = vector.load %reinterpret_cast_19[%328] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1352 = vector.bitcast %1341 : vector<16xi8> to vector<32xf4E2M1FN> + %1353 = vector.bitcast %1342 : vector<16xi8> to vector<32xf4E2M1FN> + %1354 = vector.bitcast %1343 : vector<16xi8> to vector<32xf4E2M1FN> + %1355 = vector.bitcast %1344 : vector<16xi8> to vector<32xf4E2M1FN> + %1356 = vector.bitcast %1345 : vector<16xi8> to vector<32xf4E2M1FN> + %1357 = vector.bitcast %1346 : vector<16xi8> to vector<32xf4E2M1FN> + %1358 = vector.bitcast %1347 : vector<16xi8> to vector<32xf4E2M1FN> + %1359 = vector.bitcast %1348 : vector<16xi8> to vector<32xf4E2M1FN> + %1360 = vector.bitcast %1349 : vector<16xi8> to vector<32xf4E2M1FN> + %1361 = vector.bitcast %1350 : vector<16xi8> to vector<32xf4E2M1FN> + %1362 = vector.bitcast %1351 : vector<16xi8> to vector<32xf4E2M1FN> + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(0) + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %1363 = vector.extract %1091[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1364 = amdgpu.scaled_mfma 16x16x128 (%1363[0] * %1352) * (%919[2] * %1355) + %1315 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1365 = amdgpu.scaled_mfma 16x16x128 (%1363[0] * %1352) * (%919[3] * %1356) + %1316 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1366 = amdgpu.scaled_mfma 16x16x128 (%1363[0] * %1352) * (%935[2] * %1357) + %1317 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1367 = amdgpu.scaled_mfma 16x16x128 (%1363[0] * %1352) * (%935[3] * %1358) + %1318 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1368 = amdgpu.scaled_mfma 16x16x128 (%1363[0] * %1352) * (%951[2] * %1359) + %1319 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1369 = amdgpu.scaled_mfma 16x16x128 (%1363[0] * %1352) * (%951[3] * %1360) + %1320 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1370 = amdgpu.scaled_mfma 16x16x128 (%1363[0] * %1352) * (%967[2] * %1361) + %1321 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1371 = amdgpu.scaled_mfma 16x16x128 (%1363[0] * %1352) * (%967[3] * %1362) + %1322 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1372 = vector.extract %1093[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1373 = amdgpu.scaled_mfma 16x16x128 (%1372[0] * %1353) * (%919[2] * %1355) + %1324 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1374 = amdgpu.scaled_mfma 16x16x128 (%1372[0] * %1353) * (%919[3] * %1356) + %1325 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1375 = amdgpu.scaled_mfma 16x16x128 (%1372[0] * %1353) * (%935[2] * %1357) + %1326 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1376 = amdgpu.scaled_mfma 16x16x128 (%1372[0] * %1353) * (%935[3] * %1358) + %1327 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1377 = amdgpu.scaled_mfma 16x16x128 (%1372[0] * %1353) * (%951[2] * %1359) + %1328 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1378 = amdgpu.scaled_mfma 16x16x128 (%1372[0] * %1353) * (%951[3] * %1360) + %1329 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1379 = amdgpu.scaled_mfma 16x16x128 (%1372[0] * %1353) * (%967[2] * %1361) + %1330 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1380 = amdgpu.scaled_mfma 16x16x128 (%1372[0] * %1353) * (%967[3] * %1362) + %1331 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1381 = vector.extract %1095[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1382 = amdgpu.scaled_mfma 16x16x128 (%1381[0] * %1354) * (%919[2] * %1355) + %1333 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1383 = amdgpu.scaled_mfma 16x16x128 (%1381[0] * %1354) * (%919[3] * %1356) + %1334 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1384 = amdgpu.scaled_mfma 16x16x128 (%1381[0] * %1354) * (%935[2] * %1357) + %1335 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1385 = amdgpu.scaled_mfma 16x16x128 (%1381[0] * %1354) * (%935[3] * %1358) + %1336 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1386 = amdgpu.scaled_mfma 16x16x128 (%1381[0] * %1354) * (%951[2] * %1359) + %1337 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1387 = amdgpu.scaled_mfma 16x16x128 (%1381[0] * %1354) * (%951[3] * %1360) + %1338 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1388 = amdgpu.scaled_mfma 16x16x128 (%1381[0] * %1354) * (%967[2] * %1361) + %1339 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1389 = amdgpu.scaled_mfma 16x16x128 (%1381[0] * %1354) * (%967[3] * %1362) + %1340 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + scf.yield %1364, %1365, %1366, %1367, %1368, %1369, %1370, %1371, %1373, %1374, %1375, %1376, %1377, %1378, %1379, %1380, %1382, %1383, %1384, %1385, %1386, %1387, %1388, %1389, %1130, %1157, %1184, %1211, %1226, %1241, %1268, %1283, %1298, %1313 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8> + } + %330:38 = scf.for %arg13 = %302 to %300 step %c1 iter_args(%arg14 = %329#0, %arg15 = %329#1, %arg16 = %329#2, %arg17 = %329#3, %arg18 = %329#4, %arg19 = %329#5, %arg20 = %329#6, %arg21 = %329#7, %arg22 = %329#8, %arg23 = %329#9, %arg24 = %329#10, %arg25 = %329#11, %arg26 = %329#12, %arg27 = %329#13, %arg28 = %329#14, %arg29 = %329#15, %arg30 = %329#16, %arg31 = %329#17, %arg32 = %329#18, %arg33 = %329#19, %arg34 = %329#20, %arg35 = %329#21, %arg36 = %329#22, %arg37 = %329#23, %arg38 = %329#24, %arg39 = %329#25, %arg40 = %329#26, %arg41 = %329#27, %arg42 = %329#28, %arg43 = %329#29, %arg44 = %329#30, %arg45 = %329#31, %arg46 = %329#32, %arg47 = %329#33, %arg48 = %alloc_4, %arg49 = %alloc_3, %arg50 = %alloc_2, %arg51 = %alloc) -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, memref<192x128xi8, #gpu.address_space>, memref<192x128xi8, #gpu.address_space>, memref<256x128xi8, #gpu.address_space>, memref<256x128xi8, #gpu.address_space>) { + %691 = vector.bitcast %arg47 : vector<4xi8> to vector<4xf8E8M0FNU> + %692 = vector.bitcast %arg46 : vector<4xi8> to vector<4xf8E8M0FNU> + %693 = vector.bitcast %arg45 : vector<4xi8> to vector<4xf8E8M0FNU> + %694 = vector.bitcast %arg44 : vector<4xi8> to vector<4xf8E8M0FNU> + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(0) + rocdl.s.barrier + %695 = affine.apply #map55()[%arg13, %13] + %696 = arith.addi %15, %695 overflow : index + amdgpu.gather_to_lds %18[%696], %arg49[%7, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %697 = arith.addi %22, %695 overflow : index + amdgpu.gather_to_lds %18[%697], %arg49[%20, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %698 = arith.addi %27, %695 overflow : index + amdgpu.gather_to_lds %18[%698], %arg49[%25, %8] : vector<16xi8>, memref>, memref<192x128xi8, #gpu.address_space> + %699 = affine.apply #map56()[%arg13, %13, %thread_id_x] + %700 = arith.index_cast %699 : index to i32 + %701 = arith.extui %700 : i32 to i64 + %702 = arith.muli %701, %39 : i64 + %703 = arith.shrui %702, %c32_i64 : i64 + %704 = arith.trunci %703 : i64 to i32 + %705 = arith.muli %704, %32 : i32 + %706 = arith.subi %700, %705 : i32 + %707 = arith.cmpi uge, %706, %32 : i32 + %708 = arith.extui %707 : i1 to i32 + %709 = arith.subi %704, %708 : i32 + %710 = arith.index_cast %709 : i32 to index + %711 = affine.apply #map15()[%710, %block_id_y, %thread_id_x, %thread_id_y] + %712 = arith.select %707, %32, %c0_i32 : i32 + %713 = arith.addi %706, %712 : i32 + %714 = arith.index_cast %713 : i32 to index + %715 = arith.muli %711, %9 overflow : index + %716 = arith.addi %715, %714 overflow : index + amdgpu.gather_to_lds %55[%716], %arg51[%30, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %717 = affine.apply #map17()[%710, %block_id_y, %thread_id_x, %thread_id_y] + %718 = arith.muli %717, %9 overflow : index + %719 = arith.addi %718, %714 overflow : index + amdgpu.gather_to_lds %55[%719], %arg51[%57, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %720 = affine.apply #map19()[%710, %block_id_y, %thread_id_x, %thread_id_y] + %721 = arith.muli %720, %9 overflow : index + %722 = arith.addi %721, %714 overflow : index + amdgpu.gather_to_lds %55[%722], %arg51[%62, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + %723 = affine.apply #map21()[%710, %block_id_y, %thread_id_x, %thread_id_y] + %724 = arith.muli %723, %9 overflow : index + %725 = arith.addi %724, %714 overflow : index + amdgpu.gather_to_lds %55[%725], %arg51[%67, %8] : vector<16xi8>, memref>, memref<256x128xi8, #gpu.address_space> + rocdl.sched.barrier 0 + %reinterpret_cast_28 = memref.reinterpret_cast %arg48 to offset: [0], sizes: [24576], strides: [1] : memref<192x128xi8, #gpu.address_space> to memref<24576xi8, #gpu.address_space> + %726 = vector.load %reinterpret_cast_28[%305] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %727 = vector.load %reinterpret_cast_28[%306] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %728 = vector.load %reinterpret_cast_28[%307] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %reinterpret_cast_29 = memref.reinterpret_cast %arg50 to offset: [0], sizes: [32768], strides: [1] : memref<256x128xi8, #gpu.address_space> to memref<32768xi8, #gpu.address_space> + %729 = vector.load %reinterpret_cast_29[%308] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %730 = vector.load %reinterpret_cast_29[%309] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %731 = vector.load %reinterpret_cast_29[%310] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %732 = vector.load %reinterpret_cast_29[%311] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %733 = vector.load %reinterpret_cast_29[%312] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %734 = vector.load %reinterpret_cast_29[%313] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %735 = vector.load %reinterpret_cast_29[%314] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %736 = vector.load %reinterpret_cast_29[%315] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %737 = vector.bitcast %726 : vector<16xi8> to vector<32xf4E2M1FN> + %738 = vector.bitcast %727 : vector<16xi8> to vector<32xf4E2M1FN> + %739 = vector.bitcast %728 : vector<16xi8> to vector<32xf4E2M1FN> + %740 = vector.bitcast %arg38 : vector<1xi8> to vector<1xf8E8M0FNU> + %741 = vector.bitcast %arg39 : vector<1xi8> to vector<1xf8E8M0FNU> + %742 = vector.bitcast %arg40 : vector<1xi8> to vector<1xf8E8M0FNU> + %743 = vector.bitcast %arg41 : vector<1xi8> to vector<1xf8E8M0FNU> + %744 = vector.bitcast %arg42 : vector<1xi8> to vector<1xf8E8M0FNU> + %745 = vector.bitcast %arg43 : vector<1xi8> to vector<1xf8E8M0FNU> + %746 = vector.bitcast %729 : vector<16xi8> to vector<32xf4E2M1FN> + %747 = vector.bitcast %730 : vector<16xi8> to vector<32xf4E2M1FN> + %748 = vector.bitcast %731 : vector<16xi8> to vector<32xf4E2M1FN> + %749 = vector.bitcast %732 : vector<16xi8> to vector<32xf4E2M1FN> + %750 = vector.bitcast %733 : vector<16xi8> to vector<32xf4E2M1FN> + %751 = vector.bitcast %734 : vector<16xi8> to vector<32xf4E2M1FN> + %752 = vector.bitcast %735 : vector<16xi8> to vector<32xf4E2M1FN> + %753 = vector.bitcast %736 : vector<16xi8> to vector<32xf4E2M1FN> + %754 = affine.apply #map57()[%thread_id_x, %arg7, %arg13] + %755 = arith.index_cast %754 : index to i32 + %756 = arith.extui %755 : i32 to i64 + %757 = arith.muli %756, %81 : i64 + %758 = arith.shrui %757, %c32_i64 : i64 + %759 = arith.trunci %758 : i64 to i32 + %760 = arith.muli %759, %74 : i32 + %761 = arith.subi %755, %760 : i32 + %762 = arith.cmpi uge, %761, %74 : i32 + %763 = arith.extui %762 : i1 to i32 + %764 = arith.subi %759, %763 : i32 + %765 = arith.index_cast %764 : i32 to index + %766 = affine.apply #map58()[%thread_id_x, %arg13] + %767 = arith.index_cast %766 : index to i32 + %768 = arith.extui %767 : i32 to i64 + %769 = arith.muli %768, %81 : i64 + %770 = arith.shrui %769, %c32_i64 : i64 + %771 = arith.trunci %770 : i64 to i32 + %772 = arith.muli %771, %74 : i32 + %773 = arith.subi %767, %772 : i32 + %774 = arith.cmpi uge, %773, %74 : i32 + %775 = arith.select %774, %74, %c0_i32 : i32 + %776 = arith.addi %773, %775 : i32 + %777 = arith.index_cast %776 : i32 to index + %778 = arith.muli %765, %strides#0 overflow : index + %779 = arith.addi %778, %777 overflow : index + %780 = vector.load %107[%779] : memref>, vector<1xi8> + %781 = affine.apply #map59()[%thread_id_x, %arg7, %arg13] + %782 = arith.index_cast %781 : index to i32 + %783 = arith.extui %782 : i32 to i64 + %784 = arith.muli %783, %81 : i64 + %785 = arith.shrui %784, %c32_i64 : i64 + %786 = arith.trunci %785 : i64 to i32 + %787 = arith.muli %786, %74 : i32 + %788 = arith.subi %782, %787 : i32 + %789 = arith.cmpi uge, %788, %74 : i32 + %790 = arith.extui %789 : i1 to i32 + %791 = arith.subi %786, %790 : i32 + %792 = arith.index_cast %791 : i32 to index + %793 = affine.apply #map60()[%thread_id_x, %arg13] + %794 = arith.index_cast %793 : index to i32 + %795 = arith.extui %794 : i32 to i64 + %796 = arith.muli %795, %81 : i64 + %797 = arith.shrui %796, %c32_i64 : i64 + %798 = arith.trunci %797 : i64 to i32 + %799 = arith.muli %798, %74 : i32 + %800 = arith.subi %794, %799 : i32 + %801 = arith.cmpi uge, %800, %74 : i32 + %802 = arith.select %801, %74, %c0_i32 : i32 + %803 = arith.addi %800, %802 : i32 + %804 = arith.index_cast %803 : i32 to index + %805 = arith.muli %792, %strides#0 overflow : index + %806 = arith.addi %805, %804 overflow : index + %807 = vector.load %107[%806] : memref>, vector<1xi8> + %808 = affine.apply #map61()[%thread_id_x, %arg7, %arg13] + %809 = arith.index_cast %808 : index to i32 + %810 = arith.extui %809 : i32 to i64 + %811 = arith.muli %810, %81 : i64 + %812 = arith.shrui %811, %c32_i64 : i64 + %813 = arith.trunci %812 : i64 to i32 + %814 = arith.muli %813, %74 : i32 + %815 = arith.subi %809, %814 : i32 + %816 = arith.cmpi uge, %815, %74 : i32 + %817 = arith.extui %816 : i1 to i32 + %818 = arith.subi %813, %817 : i32 + %819 = arith.index_cast %818 : i32 to index + %820 = affine.apply #map62()[%thread_id_x, %arg13] + %821 = arith.index_cast %820 : index to i32 + %822 = arith.extui %821 : i32 to i64 + %823 = arith.muli %822, %81 : i64 + %824 = arith.shrui %823, %c32_i64 : i64 + %825 = arith.trunci %824 : i64 to i32 + %826 = arith.muli %825, %74 : i32 + %827 = arith.subi %821, %826 : i32 + %828 = arith.cmpi uge, %827, %74 : i32 + %829 = arith.select %828, %74, %c0_i32 : i32 + %830 = arith.addi %827, %829 : i32 + %831 = arith.index_cast %830 : i32 to index + %832 = arith.muli %819, %strides#0 overflow : index + %833 = arith.addi %832, %831 overflow : index + %834 = vector.load %107[%833] : memref>, vector<1xi8> + %835 = affine.apply #map63()[%thread_id_x, %arg7, %arg13] + %836 = arith.index_cast %835 : index to i32 + %837 = arith.extui %836 : i32 to i64 + %838 = arith.muli %837, %81 : i64 + %839 = arith.shrui %838, %c32_i64 : i64 + %840 = arith.trunci %839 : i64 to i32 + %841 = arith.muli %840, %74 : i32 + %842 = arith.subi %836, %841 : i32 + %843 = arith.cmpi uge, %842, %74 : i32 + %844 = arith.extui %843 : i1 to i32 + %845 = arith.subi %840, %844 : i32 + %846 = arith.index_cast %845 : i32 to index + %847 = affine.apply #map64()[%thread_id_x, %arg13] + %848 = arith.index_cast %847 : index to i32 + %849 = arith.extui %848 : i32 to i64 + %850 = arith.muli %849, %81 : i64 + %851 = arith.shrui %850, %c32_i64 : i64 + %852 = arith.trunci %851 : i64 to i32 + %853 = arith.muli %852, %74 : i32 + %854 = arith.subi %848, %853 : i32 + %855 = arith.cmpi uge, %854, %74 : i32 + %856 = arith.select %855, %74, %c0_i32 : i32 + %857 = arith.addi %854, %856 : i32 + %858 = arith.index_cast %857 : i32 to index + %859 = arith.muli %846, %strides#0 overflow : index + %860 = arith.addi %859, %858 overflow : index + %861 = vector.load %107[%860] : memref>, vector<1xi8> + %862 = affine.apply #map65()[%thread_id_x, %arg7, %arg13] + %863 = arith.index_cast %862 : index to i32 + %864 = arith.extui %863 : i32 to i64 + %865 = arith.muli %864, %81 : i64 + %866 = arith.shrui %865, %c32_i64 : i64 + %867 = arith.trunci %866 : i64 to i32 + %868 = arith.muli %867, %74 : i32 + %869 = arith.subi %863, %868 : i32 + %870 = arith.cmpi uge, %869, %74 : i32 + %871 = arith.extui %870 : i1 to i32 + %872 = arith.subi %867, %871 : i32 + %873 = arith.index_cast %872 : i32 to index + %874 = arith.muli %873, %strides#0 overflow : index + %875 = arith.addi %874, %777 overflow : index + %876 = vector.load %107[%875] : memref>, vector<1xi8> + %877 = affine.apply #map66()[%thread_id_x, %arg7, %arg13] + %878 = arith.index_cast %877 : index to i32 + %879 = arith.extui %878 : i32 to i64 + %880 = arith.muli %879, %81 : i64 + %881 = arith.shrui %880, %c32_i64 : i64 + %882 = arith.trunci %881 : i64 to i32 + %883 = arith.muli %882, %74 : i32 + %884 = arith.subi %878, %883 : i32 + %885 = arith.cmpi uge, %884, %74 : i32 + %886 = arith.extui %885 : i1 to i32 + %887 = arith.subi %882, %886 : i32 + %888 = arith.index_cast %887 : i32 to index + %889 = arith.muli %888, %strides#0 overflow : index + %890 = arith.addi %889, %804 overflow : index + %891 = vector.load %107[%890] : memref>, vector<1xi8> + %892 = affine.apply #map67()[%thread_id_x, %arg7, %thread_id_y, %arg13] + %893 = arith.index_cast %892 : index to i32 + %894 = arith.extui %893 : i32 to i64 + %895 = arith.muli %894, %81 : i64 + %896 = arith.shrui %895, %c32_i64 : i64 + %897 = arith.trunci %896 : i64 to i32 + %898 = arith.muli %897, %74 : i32 + %899 = arith.subi %893, %898 : i32 + %900 = arith.cmpi uge, %899, %74 : i32 + %901 = arith.extui %900 : i1 to i32 + %902 = arith.subi %897, %901 : i32 + %903 = arith.index_cast %902 : i32 to index + %904 = affine.apply #map68()[%thread_id_x, %arg13] + %905 = arith.index_cast %904 : index to i32 + %906 = arith.extui %905 : i32 to i64 + %907 = arith.muli %906, %81 : i64 + %908 = arith.shrui %907, %c32_i64 : i64 + %909 = arith.trunci %908 : i64 to i32 + %910 = arith.muli %909, %74 : i32 + %911 = arith.subi %905, %910 : i32 + %912 = arith.cmpi uge, %911, %74 : i32 + %913 = arith.select %912, %74, %c0_i32 : i32 + %914 = arith.addi %911, %913 : i32 + %915 = arith.index_cast %914 : i32 to index + %916 = arith.muli %903, %strides_13#0 overflow : index + %917 = arith.addi %916, %915 overflow : index + %918 = vector.load %249[%917] : memref>, vector<4xi8> + %919 = affine.apply #map69()[%thread_id_x, %arg7, %thread_id_y, %arg13] + %920 = arith.index_cast %919 : index to i32 + %921 = arith.extui %920 : i32 to i64 + %922 = arith.muli %921, %81 : i64 + %923 = arith.shrui %922, %c32_i64 : i64 + %924 = arith.trunci %923 : i64 to i32 + %925 = arith.muli %924, %74 : i32 + %926 = arith.subi %920, %925 : i32 + %927 = arith.cmpi uge, %926, %74 : i32 + %928 = arith.extui %927 : i1 to i32 + %929 = arith.subi %924, %928 : i32 + %930 = arith.index_cast %929 : i32 to index + %931 = arith.muli %930, %strides_13#0 overflow : index + %932 = arith.addi %931, %915 overflow : index + %933 = vector.load %249[%932] : memref>, vector<4xi8> + %934 = affine.apply #map70()[%thread_id_x, %arg7, %thread_id_y, %arg13] + %935 = arith.index_cast %934 : index to i32 + %936 = arith.extui %935 : i32 to i64 + %937 = arith.muli %936, %81 : i64 + %938 = arith.shrui %937, %c32_i64 : i64 + %939 = arith.trunci %938 : i64 to i32 + %940 = arith.muli %939, %74 : i32 + %941 = arith.subi %935, %940 : i32 + %942 = arith.cmpi uge, %941, %74 : i32 + %943 = arith.extui %942 : i1 to i32 + %944 = arith.subi %939, %943 : i32 + %945 = arith.index_cast %944 : i32 to index + %946 = arith.muli %945, %strides_13#0 overflow : index + %947 = arith.addi %946, %915 overflow : index + %948 = vector.load %249[%947] : memref>, vector<4xi8> + %949 = affine.apply #map71()[%thread_id_x, %arg7, %thread_id_y, %arg13] + %950 = arith.index_cast %949 : index to i32 + %951 = arith.extui %950 : i32 to i64 + %952 = arith.muli %951, %81 : i64 + %953 = arith.shrui %952, %c32_i64 : i64 + %954 = arith.trunci %953 : i64 to i32 + %955 = arith.muli %954, %74 : i32 + %956 = arith.subi %950, %955 : i32 + %957 = arith.cmpi uge, %956, %74 : i32 + %958 = arith.extui %957 : i1 to i32 + %959 = arith.subi %954, %958 : i32 + %960 = arith.index_cast %959 : i32 to index + %961 = arith.muli %960, %strides_13#0 overflow : index + %962 = arith.addi %961, %915 overflow : index + %963 = vector.load %249[%962] : memref>, vector<4xi8> + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %964 = vector.extract %740[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %965 = amdgpu.scaled_mfma 16x16x128 (%964[0] * %737) * (%694[0] * %746) + %arg14 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %966 = amdgpu.scaled_mfma 16x16x128 (%964[0] * %737) * (%694[1] * %747) + %arg15 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %967 = amdgpu.scaled_mfma 16x16x128 (%964[0] * %737) * (%693[0] * %748) + %arg16 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %968 = amdgpu.scaled_mfma 16x16x128 (%964[0] * %737) * (%693[1] * %749) + %arg17 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %969 = amdgpu.scaled_mfma 16x16x128 (%964[0] * %737) * (%692[0] * %750) + %arg18 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %970 = amdgpu.scaled_mfma 16x16x128 (%964[0] * %737) * (%692[1] * %751) + %arg19 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %971 = amdgpu.scaled_mfma 16x16x128 (%964[0] * %737) * (%691[0] * %752) + %arg20 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %972 = amdgpu.scaled_mfma 16x16x128 (%964[0] * %737) * (%691[1] * %753) + %arg21 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %973 = vector.extract %742[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %974 = amdgpu.scaled_mfma 16x16x128 (%973[0] * %738) * (%694[0] * %746) + %arg22 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %975 = amdgpu.scaled_mfma 16x16x128 (%973[0] * %738) * (%694[1] * %747) + %arg23 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %976 = amdgpu.scaled_mfma 16x16x128 (%973[0] * %738) * (%693[0] * %748) + %arg24 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %977 = amdgpu.scaled_mfma 16x16x128 (%973[0] * %738) * (%693[1] * %749) + %arg25 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %978 = amdgpu.scaled_mfma 16x16x128 (%973[0] * %738) * (%692[0] * %750) + %arg26 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %979 = amdgpu.scaled_mfma 16x16x128 (%973[0] * %738) * (%692[1] * %751) + %arg27 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %980 = amdgpu.scaled_mfma 16x16x128 (%973[0] * %738) * (%691[0] * %752) + %arg28 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %981 = amdgpu.scaled_mfma 16x16x128 (%973[0] * %738) * (%691[1] * %753) + %arg29 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %982 = vector.extract %744[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %983 = amdgpu.scaled_mfma 16x16x128 (%982[0] * %739) * (%694[0] * %746) + %arg30 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %984 = amdgpu.scaled_mfma 16x16x128 (%982[0] * %739) * (%694[1] * %747) + %arg31 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %985 = amdgpu.scaled_mfma 16x16x128 (%982[0] * %739) * (%693[0] * %748) + %arg32 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %986 = amdgpu.scaled_mfma 16x16x128 (%982[0] * %739) * (%693[1] * %749) + %arg33 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %987 = amdgpu.scaled_mfma 16x16x128 (%982[0] * %739) * (%692[0] * %750) + %arg34 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %988 = amdgpu.scaled_mfma 16x16x128 (%982[0] * %739) * (%692[1] * %751) + %arg35 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %989 = amdgpu.scaled_mfma 16x16x128 (%982[0] * %739) * (%691[0] * %752) + %arg36 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %990 = amdgpu.scaled_mfma 16x16x128 (%982[0] * %739) * (%691[1] * %753) + %arg37 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.sched.barrier 0 + %991 = vector.load %reinterpret_cast_28[%318] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %992 = vector.load %reinterpret_cast_28[%319] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %993 = vector.load %reinterpret_cast_28[%320] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %994 = vector.load %reinterpret_cast_29[%321] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %995 = vector.load %reinterpret_cast_29[%322] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %996 = vector.load %reinterpret_cast_29[%323] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %997 = vector.load %reinterpret_cast_29[%324] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %998 = vector.load %reinterpret_cast_29[%325] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %999 = vector.load %reinterpret_cast_29[%326] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1000 = vector.load %reinterpret_cast_29[%327] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1001 = vector.load %reinterpret_cast_29[%328] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %1002 = vector.bitcast %991 : vector<16xi8> to vector<32xf4E2M1FN> + %1003 = vector.bitcast %992 : vector<16xi8> to vector<32xf4E2M1FN> + %1004 = vector.bitcast %993 : vector<16xi8> to vector<32xf4E2M1FN> + %1005 = vector.bitcast %994 : vector<16xi8> to vector<32xf4E2M1FN> + %1006 = vector.bitcast %995 : vector<16xi8> to vector<32xf4E2M1FN> + %1007 = vector.bitcast %996 : vector<16xi8> to vector<32xf4E2M1FN> + %1008 = vector.bitcast %997 : vector<16xi8> to vector<32xf4E2M1FN> + %1009 = vector.bitcast %998 : vector<16xi8> to vector<32xf4E2M1FN> + %1010 = vector.bitcast %999 : vector<16xi8> to vector<32xf4E2M1FN> + %1011 = vector.bitcast %1000 : vector<16xi8> to vector<32xf4E2M1FN> + %1012 = vector.bitcast %1001 : vector<16xi8> to vector<32xf4E2M1FN> + rocdl.sched.barrier 0 + amdgpu.memory_counter_wait load(0) + rocdl.s.barrier + rocdl.sched.barrier 0 + rocdl.s.setprio 1 + %1013 = vector.extract %741[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1014 = amdgpu.scaled_mfma 16x16x128 (%1013[0] * %1002) * (%694[2] * %1005) + %965 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1015 = amdgpu.scaled_mfma 16x16x128 (%1013[0] * %1002) * (%694[3] * %1006) + %966 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1016 = amdgpu.scaled_mfma 16x16x128 (%1013[0] * %1002) * (%693[2] * %1007) + %967 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1017 = amdgpu.scaled_mfma 16x16x128 (%1013[0] * %1002) * (%693[3] * %1008) + %968 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1018 = amdgpu.scaled_mfma 16x16x128 (%1013[0] * %1002) * (%692[2] * %1009) + %969 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1019 = amdgpu.scaled_mfma 16x16x128 (%1013[0] * %1002) * (%692[3] * %1010) + %970 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1020 = amdgpu.scaled_mfma 16x16x128 (%1013[0] * %1002) * (%691[2] * %1011) + %971 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1021 = amdgpu.scaled_mfma 16x16x128 (%1013[0] * %1002) * (%691[3] * %1012) + %972 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1022 = vector.extract %743[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1023 = amdgpu.scaled_mfma 16x16x128 (%1022[0] * %1003) * (%694[2] * %1005) + %974 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1024 = amdgpu.scaled_mfma 16x16x128 (%1022[0] * %1003) * (%694[3] * %1006) + %975 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1025 = amdgpu.scaled_mfma 16x16x128 (%1022[0] * %1003) * (%693[2] * %1007) + %976 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1026 = amdgpu.scaled_mfma 16x16x128 (%1022[0] * %1003) * (%693[3] * %1008) + %977 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1027 = amdgpu.scaled_mfma 16x16x128 (%1022[0] * %1003) * (%692[2] * %1009) + %978 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1028 = amdgpu.scaled_mfma 16x16x128 (%1022[0] * %1003) * (%692[3] * %1010) + %979 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1029 = amdgpu.scaled_mfma 16x16x128 (%1022[0] * %1003) * (%691[2] * %1011) + %980 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1030 = amdgpu.scaled_mfma 16x16x128 (%1022[0] * %1003) * (%691[3] * %1012) + %981 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1031 = vector.extract %745[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %1032 = amdgpu.scaled_mfma 16x16x128 (%1031[0] * %1004) * (%694[2] * %1005) + %983 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1033 = amdgpu.scaled_mfma 16x16x128 (%1031[0] * %1004) * (%694[3] * %1006) + %984 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1034 = amdgpu.scaled_mfma 16x16x128 (%1031[0] * %1004) * (%693[2] * %1007) + %985 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1035 = amdgpu.scaled_mfma 16x16x128 (%1031[0] * %1004) * (%693[3] * %1008) + %986 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1036 = amdgpu.scaled_mfma 16x16x128 (%1031[0] * %1004) * (%692[2] * %1009) + %987 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1037 = amdgpu.scaled_mfma 16x16x128 (%1031[0] * %1004) * (%692[3] * %1010) + %988 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1038 = amdgpu.scaled_mfma 16x16x128 (%1031[0] * %1004) * (%691[2] * %1011) + %989 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %1039 = amdgpu.scaled_mfma 16x16x128 (%1031[0] * %1004) * (%691[3] * %1012) + %990 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + rocdl.s.setprio 0 + rocdl.sched.barrier 0 + scf.yield %1014, %1015, %1016, %1017, %1018, %1019, %1020, %1021, %1023, %1024, %1025, %1026, %1027, %1028, %1029, %1030, %1032, %1033, %1034, %1035, %1036, %1037, %1038, %1039, %780, %807, %834, %861, %876, %891, %918, %933, %948, %963, %arg49, %arg48, %arg51, %arg50 : vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<1xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, vector<4xi8>, memref<192x128xi8, #gpu.address_space>, memref<192x128xi8, #gpu.address_space>, memref<256x128xi8, #gpu.address_space>, memref<256x128xi8, #gpu.address_space> + } + %331 = vector.bitcast %330#33 : vector<4xi8> to vector<4xf8E8M0FNU> + %332 = vector.bitcast %330#32 : vector<4xi8> to vector<4xf8E8M0FNU> + %333 = vector.bitcast %330#31 : vector<4xi8> to vector<4xf8E8M0FNU> + %334 = vector.bitcast %330#30 : vector<4xi8> to vector<4xf8E8M0FNU> + scf.if %299 { + rocdl.s.barrier + } + amdgpu.lds_barrier + %reinterpret_cast_20 = memref.reinterpret_cast %330#36 to offset: [0], sizes: [32768], strides: [1] : memref<256x128xi8, #gpu.address_space> to memref<32768xi8, #gpu.address_space> + %335 = vector.load %reinterpret_cast_20[%308] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %336 = vector.load %reinterpret_cast_20[%321] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %337 = vector.load %reinterpret_cast_20[%309] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %338 = vector.load %reinterpret_cast_20[%322] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %339 = vector.load %reinterpret_cast_20[%310] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %340 = vector.load %reinterpret_cast_20[%323] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %341 = vector.load %reinterpret_cast_20[%311] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %342 = vector.load %reinterpret_cast_20[%324] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %343 = vector.load %reinterpret_cast_20[%312] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %344 = vector.load %reinterpret_cast_20[%325] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %345 = vector.load %reinterpret_cast_20[%313] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %346 = vector.load %reinterpret_cast_20[%326] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %347 = vector.load %reinterpret_cast_20[%314] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %348 = vector.load %reinterpret_cast_20[%327] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %349 = vector.load %reinterpret_cast_20[%315] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %350 = vector.load %reinterpret_cast_20[%328] : memref<32768xi8, #gpu.address_space>, vector<16xi8> + %reinterpret_cast_21 = memref.reinterpret_cast %330#34 to offset: [0], sizes: [24576], strides: [1] : memref<192x128xi8, #gpu.address_space> to memref<24576xi8, #gpu.address_space> + %351 = vector.load %reinterpret_cast_21[%305] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %352 = vector.load %reinterpret_cast_21[%318] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %353 = vector.load %reinterpret_cast_21[%306] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %354 = vector.load %reinterpret_cast_21[%319] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %355 = vector.load %reinterpret_cast_21[%307] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %356 = vector.load %reinterpret_cast_21[%320] : memref<24576xi8, #gpu.address_space>, vector<16xi8> + %357 = vector.bitcast %351 : vector<16xi8> to vector<32xf4E2M1FN> + %358 = vector.bitcast %352 : vector<16xi8> to vector<32xf4E2M1FN> + %359 = vector.bitcast %353 : vector<16xi8> to vector<32xf4E2M1FN> + %360 = vector.bitcast %354 : vector<16xi8> to vector<32xf4E2M1FN> + %361 = vector.bitcast %355 : vector<16xi8> to vector<32xf4E2M1FN> + %362 = vector.bitcast %356 : vector<16xi8> to vector<32xf4E2M1FN> + %363 = vector.bitcast %330#24 : vector<1xi8> to vector<1xf8E8M0FNU> + %364 = vector.bitcast %330#25 : vector<1xi8> to vector<1xf8E8M0FNU> + %365 = vector.bitcast %330#26 : vector<1xi8> to vector<1xf8E8M0FNU> + %366 = vector.bitcast %330#27 : vector<1xi8> to vector<1xf8E8M0FNU> + %367 = vector.bitcast %330#28 : vector<1xi8> to vector<1xf8E8M0FNU> + %368 = vector.bitcast %330#29 : vector<1xi8> to vector<1xf8E8M0FNU> + %369 = vector.bitcast %335 : vector<16xi8> to vector<32xf4E2M1FN> + %370 = vector.bitcast %336 : vector<16xi8> to vector<32xf4E2M1FN> + %371 = vector.bitcast %337 : vector<16xi8> to vector<32xf4E2M1FN> + %372 = vector.bitcast %338 : vector<16xi8> to vector<32xf4E2M1FN> + %373 = vector.bitcast %339 : vector<16xi8> to vector<32xf4E2M1FN> + %374 = vector.bitcast %340 : vector<16xi8> to vector<32xf4E2M1FN> + %375 = vector.bitcast %341 : vector<16xi8> to vector<32xf4E2M1FN> + %376 = vector.bitcast %342 : vector<16xi8> to vector<32xf4E2M1FN> + %377 = vector.bitcast %343 : vector<16xi8> to vector<32xf4E2M1FN> + %378 = vector.bitcast %344 : vector<16xi8> to vector<32xf4E2M1FN> + %379 = vector.bitcast %345 : vector<16xi8> to vector<32xf4E2M1FN> + %380 = vector.bitcast %346 : vector<16xi8> to vector<32xf4E2M1FN> + %381 = vector.bitcast %347 : vector<16xi8> to vector<32xf4E2M1FN> + %382 = vector.bitcast %348 : vector<16xi8> to vector<32xf4E2M1FN> + %383 = vector.bitcast %349 : vector<16xi8> to vector<32xf4E2M1FN> + %384 = vector.bitcast %350 : vector<16xi8> to vector<32xf4E2M1FN> + %385 = vector.extract %363[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %386 = amdgpu.scaled_mfma 16x16x128 (%385[0] * %357) * (%334[0] * %369) + %330#0 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %387 = vector.extract %364[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %388 = amdgpu.scaled_mfma 16x16x128 (%387[0] * %358) * (%334[2] * %370) + %386 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %389 = amdgpu.scaled_mfma 16x16x128 (%385[0] * %357) * (%334[1] * %371) + %330#1 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %390 = amdgpu.scaled_mfma 16x16x128 (%387[0] * %358) * (%334[3] * %372) + %389 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %391 = amdgpu.scaled_mfma 16x16x128 (%385[0] * %357) * (%333[0] * %373) + %330#2 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %392 = amdgpu.scaled_mfma 16x16x128 (%387[0] * %358) * (%333[2] * %374) + %391 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %393 = amdgpu.scaled_mfma 16x16x128 (%385[0] * %357) * (%333[1] * %375) + %330#3 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %394 = amdgpu.scaled_mfma 16x16x128 (%387[0] * %358) * (%333[3] * %376) + %393 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %395 = amdgpu.scaled_mfma 16x16x128 (%385[0] * %357) * (%332[0] * %377) + %330#4 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %396 = amdgpu.scaled_mfma 16x16x128 (%387[0] * %358) * (%332[2] * %378) + %395 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %397 = amdgpu.scaled_mfma 16x16x128 (%385[0] * %357) * (%332[1] * %379) + %330#5 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %398 = amdgpu.scaled_mfma 16x16x128 (%387[0] * %358) * (%332[3] * %380) + %397 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %399 = amdgpu.scaled_mfma 16x16x128 (%385[0] * %357) * (%331[0] * %381) + %330#6 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %400 = amdgpu.scaled_mfma 16x16x128 (%387[0] * %358) * (%331[2] * %382) + %399 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %401 = amdgpu.scaled_mfma 16x16x128 (%385[0] * %357) * (%331[1] * %383) + %330#7 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %402 = amdgpu.scaled_mfma 16x16x128 (%387[0] * %358) * (%331[3] * %384) + %401 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %403 = vector.extract %365[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %404 = amdgpu.scaled_mfma 16x16x128 (%403[0] * %359) * (%334[0] * %369) + %330#8 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %405 = vector.extract %366[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %406 = amdgpu.scaled_mfma 16x16x128 (%405[0] * %360) * (%334[2] * %370) + %404 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %407 = amdgpu.scaled_mfma 16x16x128 (%403[0] * %359) * (%334[1] * %371) + %330#9 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %408 = amdgpu.scaled_mfma 16x16x128 (%405[0] * %360) * (%334[3] * %372) + %407 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %409 = amdgpu.scaled_mfma 16x16x128 (%403[0] * %359) * (%333[0] * %373) + %330#10 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %410 = amdgpu.scaled_mfma 16x16x128 (%405[0] * %360) * (%333[2] * %374) + %409 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %411 = amdgpu.scaled_mfma 16x16x128 (%403[0] * %359) * (%333[1] * %375) + %330#11 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %412 = amdgpu.scaled_mfma 16x16x128 (%405[0] * %360) * (%333[3] * %376) + %411 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %413 = amdgpu.scaled_mfma 16x16x128 (%403[0] * %359) * (%332[0] * %377) + %330#12 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %414 = amdgpu.scaled_mfma 16x16x128 (%405[0] * %360) * (%332[2] * %378) + %413 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %415 = amdgpu.scaled_mfma 16x16x128 (%403[0] * %359) * (%332[1] * %379) + %330#13 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %416 = amdgpu.scaled_mfma 16x16x128 (%405[0] * %360) * (%332[3] * %380) + %415 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %417 = amdgpu.scaled_mfma 16x16x128 (%403[0] * %359) * (%331[0] * %381) + %330#14 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %418 = amdgpu.scaled_mfma 16x16x128 (%405[0] * %360) * (%331[2] * %382) + %417 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %419 = amdgpu.scaled_mfma 16x16x128 (%403[0] * %359) * (%331[1] * %383) + %330#15 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %420 = amdgpu.scaled_mfma 16x16x128 (%405[0] * %360) * (%331[3] * %384) + %419 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %421 = vector.extract %367[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %422 = amdgpu.scaled_mfma 16x16x128 (%421[0] * %361) * (%334[0] * %369) + %330#16 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %423 = vector.extract %368[0] : f8E8M0FNU from vector<1xf8E8M0FNU> + %424 = amdgpu.scaled_mfma 16x16x128 (%423[0] * %362) * (%334[2] * %370) + %422 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %425 = amdgpu.scaled_mfma 16x16x128 (%421[0] * %361) * (%334[1] * %371) + %330#17 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %426 = amdgpu.scaled_mfma 16x16x128 (%423[0] * %362) * (%334[3] * %372) + %425 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %427 = amdgpu.scaled_mfma 16x16x128 (%421[0] * %361) * (%333[0] * %373) + %330#18 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %428 = amdgpu.scaled_mfma 16x16x128 (%423[0] * %362) * (%333[2] * %374) + %427 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %429 = amdgpu.scaled_mfma 16x16x128 (%421[0] * %361) * (%333[1] * %375) + %330#19 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %430 = amdgpu.scaled_mfma 16x16x128 (%423[0] * %362) * (%333[3] * %376) + %429 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %431 = amdgpu.scaled_mfma 16x16x128 (%421[0] * %361) * (%332[0] * %377) + %330#20 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %432 = amdgpu.scaled_mfma 16x16x128 (%423[0] * %362) * (%332[2] * %378) + %431 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %433 = amdgpu.scaled_mfma 16x16x128 (%421[0] * %361) * (%332[1] * %379) + %330#21 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %434 = amdgpu.scaled_mfma 16x16x128 (%423[0] * %362) * (%332[3] * %380) + %433 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %435 = amdgpu.scaled_mfma 16x16x128 (%421[0] * %361) * (%331[0] * %381) + %330#22 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %436 = amdgpu.scaled_mfma 16x16x128 (%423[0] * %362) * (%331[2] * %382) + %435 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %437 = amdgpu.scaled_mfma 16x16x128 (%421[0] * %361) * (%331[1] * %383) + %330#23 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %438 = amdgpu.scaled_mfma 16x16x128 (%423[0] * %362) * (%331[3] * %384) + %437 : f8E8M0FNU, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32> + %base_buffer_22, %offset_23, %sizes_24:2, %strides_25:2 = memref.extract_strided_metadata %reinterpret_cast_1 : memref> -> memref, index, index, index, index, index + %439 = arith.muli %220, %strides_25#0 overflow : index + %440 = arith.addi %439, %71 overflow : index + %reinterpret_cast_26 = memref.reinterpret_cast %4 to offset: [%440], sizes: [1073741822], strides: [1] : memref to memref<1073741822xbf16, strided<[1], offset: ?>> + %cast_27 = memref.cast %reinterpret_cast_26 : memref<1073741822xbf16, strided<[1], offset: ?>> to memref> + %441 = arith.index_cast %strides_25#0 : index to i14 + %442 = amdgpu.fat_raw_buffer_cast %cast_27 validBytes(%c2147483645_i64) cacheSwizzleStride(%441) resetOffset : memref> to memref> + %443 = affine.apply #map92()[%thread_id_x] + %444 = arith.cmpi eq, %443, %c0 : index + %mbase = affine.apply #map_Mbase()[%thread_id_x, %thread_id_y] + %nbase = affine.apply #map_Nbase()[%thread_id_x] + %mrow = arith.muli %mbase, %strides_25#0 overflow : index + %bn0 = arith.addi %mrow, %nbase overflow : index + %c16ep = arith.constant 16 : index + %c32ep = arith.constant 32 : index + %c48ep = arith.constant 48 : index + %c64ep = arith.constant 64 : index + %c80ep = arith.constant 80 : index + %c96ep = arith.constant 96 : index + %c112ep = arith.constant 112 : index + %bn1 = arith.addi %bn0, %c16ep overflow : index + %bn2 = arith.addi %bn0, %c32ep overflow : index + %mt0 = arith.constant 0 : index + %mt1 = arith.muli %c16ep, %strides_25#0 overflow : index + %mt2 = arith.muli %c32ep, %strides_25#0 overflow : index + %mt3 = arith.muli %c48ep, %strides_25#0 overflow : index + %mt4 = arith.muli %c64ep, %strides_25#0 overflow : index + %mt5 = arith.muli %c80ep, %strides_25#0 overflow : index + %mt6 = arith.muli %c96ep, %strides_25#0 overflow : index + %mt7 = arith.muli %c112ep, %strides_25#0 overflow : index + %466 = vector.extract %388[0] : f32 from vector<4xf32> + %467 = vector.extract %388[1] : f32 from vector<4xf32> + %468 = vector.extract %388[2] : f32 from vector<4xf32> + %469 = vector.extract %388[3] : f32 from vector<4xf32> + %470 = vector.extract %390[0] : f32 from vector<4xf32> + %471 = vector.extract %390[1] : f32 from vector<4xf32> + %472 = vector.extract %390[2] : f32 from vector<4xf32> + %473 = vector.extract %390[3] : f32 from vector<4xf32> + %474 = amdgpu.permlane_swap %466 16 : f32 + %475 = amdgpu.permlane_swap %467 16 : f32 + %476 = amdgpu.permlane_swap %468 16 : f32 + %477 = amdgpu.permlane_swap %469 16 : f32 + %478 = amdgpu.permlane_swap %470 16 : f32 + %479 = amdgpu.permlane_swap %471 16 : f32 + %480 = amdgpu.permlane_swap %472 16 : f32 + %481 = amdgpu.permlane_swap %473 16 : f32 + %482 = vector.from_elements %466, %467, %468, %469, %474, %475, %476, %477 : vector<8xf32> + %483 = arith.truncf %482 : vector<8xf32> to vector<8xbf16> + %484 = vector.from_elements %478, %479, %480, %481, %470, %471, %472, %473 : vector<8xf32> + %485 = arith.truncf %484 : vector<8xf32> to vector<8xbf16> + %486 = arith.select %444, %483, %485 : vector<8xbf16> + %487 = arith.select %444, %mt0, %mt1 : index + + %488 = arith.addi %bn0, %487 overflow : index + vector.store %486, %442[%488] : memref>, vector<8xbf16> + %489 = vector.extract %392[0] : f32 from vector<4xf32> + %490 = vector.extract %392[1] : f32 from vector<4xf32> + %491 = vector.extract %392[2] : f32 from vector<4xf32> + %492 = vector.extract %392[3] : f32 from vector<4xf32> + %493 = vector.extract %394[0] : f32 from vector<4xf32> + %494 = vector.extract %394[1] : f32 from vector<4xf32> + %495 = vector.extract %394[2] : f32 from vector<4xf32> + %496 = vector.extract %394[3] : f32 from vector<4xf32> + %497 = amdgpu.permlane_swap %489 16 : f32 + %498 = amdgpu.permlane_swap %490 16 : f32 + %499 = amdgpu.permlane_swap %491 16 : f32 + %500 = amdgpu.permlane_swap %492 16 : f32 + %501 = amdgpu.permlane_swap %493 16 : f32 + %502 = amdgpu.permlane_swap %494 16 : f32 + %503 = amdgpu.permlane_swap %495 16 : f32 + %504 = amdgpu.permlane_swap %496 16 : f32 + %505 = vector.from_elements %489, %490, %491, %492, %497, %498, %499, %500 : vector<8xf32> + %506 = arith.truncf %505 : vector<8xf32> to vector<8xbf16> + %507 = vector.from_elements %501, %502, %503, %504, %493, %494, %495, %496 : vector<8xf32> + %508 = arith.truncf %507 : vector<8xf32> to vector<8xbf16> + %509 = arith.select %444, %506, %508 : vector<8xbf16> + %510 = arith.select %444, %mt2, %mt3 : index + + %511 = arith.addi %bn0, %510 overflow : index + vector.store %509, %442[%511] : memref>, vector<8xbf16> + %512 = vector.extract %396[0] : f32 from vector<4xf32> + %513 = vector.extract %396[1] : f32 from vector<4xf32> + %514 = vector.extract %396[2] : f32 from vector<4xf32> + %515 = vector.extract %396[3] : f32 from vector<4xf32> + %516 = vector.extract %398[0] : f32 from vector<4xf32> + %517 = vector.extract %398[1] : f32 from vector<4xf32> + %518 = vector.extract %398[2] : f32 from vector<4xf32> + %519 = vector.extract %398[3] : f32 from vector<4xf32> + %520 = amdgpu.permlane_swap %512 16 : f32 + %521 = amdgpu.permlane_swap %513 16 : f32 + %522 = amdgpu.permlane_swap %514 16 : f32 + %523 = amdgpu.permlane_swap %515 16 : f32 + %524 = amdgpu.permlane_swap %516 16 : f32 + %525 = amdgpu.permlane_swap %517 16 : f32 + %526 = amdgpu.permlane_swap %518 16 : f32 + %527 = amdgpu.permlane_swap %519 16 : f32 + %528 = vector.from_elements %512, %513, %514, %515, %520, %521, %522, %523 : vector<8xf32> + %529 = arith.truncf %528 : vector<8xf32> to vector<8xbf16> + %530 = vector.from_elements %524, %525, %526, %527, %516, %517, %518, %519 : vector<8xf32> + %531 = arith.truncf %530 : vector<8xf32> to vector<8xbf16> + %532 = arith.select %444, %529, %531 : vector<8xbf16> + %533 = arith.select %444, %mt4, %mt5 : index + + %534 = arith.addi %bn0, %533 overflow : index + vector.store %532, %442[%534] : memref>, vector<8xbf16> + %535 = vector.extract %400[0] : f32 from vector<4xf32> + %536 = vector.extract %400[1] : f32 from vector<4xf32> + %537 = vector.extract %400[2] : f32 from vector<4xf32> + %538 = vector.extract %400[3] : f32 from vector<4xf32> + %539 = vector.extract %402[0] : f32 from vector<4xf32> + %540 = vector.extract %402[1] : f32 from vector<4xf32> + %541 = vector.extract %402[2] : f32 from vector<4xf32> + %542 = vector.extract %402[3] : f32 from vector<4xf32> + %543 = amdgpu.permlane_swap %535 16 : f32 + %544 = amdgpu.permlane_swap %536 16 : f32 + %545 = amdgpu.permlane_swap %537 16 : f32 + %546 = amdgpu.permlane_swap %538 16 : f32 + %547 = amdgpu.permlane_swap %539 16 : f32 + %548 = amdgpu.permlane_swap %540 16 : f32 + %549 = amdgpu.permlane_swap %541 16 : f32 + %550 = amdgpu.permlane_swap %542 16 : f32 + %551 = vector.from_elements %535, %536, %537, %538, %543, %544, %545, %546 : vector<8xf32> + %552 = arith.truncf %551 : vector<8xf32> to vector<8xbf16> + %553 = vector.from_elements %547, %548, %549, %550, %539, %540, %541, %542 : vector<8xf32> + %554 = arith.truncf %553 : vector<8xf32> to vector<8xbf16> + %555 = arith.select %444, %552, %554 : vector<8xbf16> + %556 = arith.select %444, %mt6, %mt7 : index + + %557 = arith.addi %bn0, %556 overflow : index + vector.store %555, %442[%557] : memref>, vector<8xbf16> + %558 = vector.extract %406[0] : f32 from vector<4xf32> + %559 = vector.extract %406[1] : f32 from vector<4xf32> + %560 = vector.extract %406[2] : f32 from vector<4xf32> + %561 = vector.extract %406[3] : f32 from vector<4xf32> + %562 = vector.extract %408[0] : f32 from vector<4xf32> + %563 = vector.extract %408[1] : f32 from vector<4xf32> + %564 = vector.extract %408[2] : f32 from vector<4xf32> + %565 = vector.extract %408[3] : f32 from vector<4xf32> + %566 = amdgpu.permlane_swap %558 16 : f32 + %567 = amdgpu.permlane_swap %559 16 : f32 + %568 = amdgpu.permlane_swap %560 16 : f32 + %569 = amdgpu.permlane_swap %561 16 : f32 + %570 = amdgpu.permlane_swap %562 16 : f32 + %571 = amdgpu.permlane_swap %563 16 : f32 + %572 = amdgpu.permlane_swap %564 16 : f32 + %573 = amdgpu.permlane_swap %565 16 : f32 + %574 = vector.from_elements %558, %559, %560, %561, %566, %567, %568, %569 : vector<8xf32> + %575 = arith.truncf %574 : vector<8xf32> to vector<8xbf16> + %576 = vector.from_elements %570, %571, %572, %573, %562, %563, %564, %565 : vector<8xf32> + %577 = arith.truncf %576 : vector<8xf32> to vector<8xbf16> + %578 = arith.select %444, %575, %577 : vector<8xbf16> + %579 = arith.select %444, %mt0, %mt1 : index + + %580 = arith.addi %bn1, %579 overflow : index + vector.store %578, %442[%580] : memref>, vector<8xbf16> + %581 = vector.extract %410[0] : f32 from vector<4xf32> + %582 = vector.extract %410[1] : f32 from vector<4xf32> + %583 = vector.extract %410[2] : f32 from vector<4xf32> + %584 = vector.extract %410[3] : f32 from vector<4xf32> + %585 = vector.extract %412[0] : f32 from vector<4xf32> + %586 = vector.extract %412[1] : f32 from vector<4xf32> + %587 = vector.extract %412[2] : f32 from vector<4xf32> + %588 = vector.extract %412[3] : f32 from vector<4xf32> + %589 = amdgpu.permlane_swap %581 16 : f32 + %590 = amdgpu.permlane_swap %582 16 : f32 + %591 = amdgpu.permlane_swap %583 16 : f32 + %592 = amdgpu.permlane_swap %584 16 : f32 + %593 = amdgpu.permlane_swap %585 16 : f32 + %594 = amdgpu.permlane_swap %586 16 : f32 + %595 = amdgpu.permlane_swap %587 16 : f32 + %596 = amdgpu.permlane_swap %588 16 : f32 + %597 = vector.from_elements %581, %582, %583, %584, %589, %590, %591, %592 : vector<8xf32> + %598 = arith.truncf %597 : vector<8xf32> to vector<8xbf16> + %599 = vector.from_elements %593, %594, %595, %596, %585, %586, %587, %588 : vector<8xf32> + %600 = arith.truncf %599 : vector<8xf32> to vector<8xbf16> + %601 = arith.select %444, %598, %600 : vector<8xbf16> + %602 = arith.select %444, %mt2, %mt3 : index + + %603 = arith.addi %bn1, %602 overflow : index + vector.store %601, %442[%603] : memref>, vector<8xbf16> + %604 = vector.extract %414[0] : f32 from vector<4xf32> + %605 = vector.extract %414[1] : f32 from vector<4xf32> + %606 = vector.extract %414[2] : f32 from vector<4xf32> + %607 = vector.extract %414[3] : f32 from vector<4xf32> + %608 = vector.extract %416[0] : f32 from vector<4xf32> + %609 = vector.extract %416[1] : f32 from vector<4xf32> + %610 = vector.extract %416[2] : f32 from vector<4xf32> + %611 = vector.extract %416[3] : f32 from vector<4xf32> + %612 = amdgpu.permlane_swap %604 16 : f32 + %613 = amdgpu.permlane_swap %605 16 : f32 + %614 = amdgpu.permlane_swap %606 16 : f32 + %615 = amdgpu.permlane_swap %607 16 : f32 + %616 = amdgpu.permlane_swap %608 16 : f32 + %617 = amdgpu.permlane_swap %609 16 : f32 + %618 = amdgpu.permlane_swap %610 16 : f32 + %619 = amdgpu.permlane_swap %611 16 : f32 + %620 = vector.from_elements %604, %605, %606, %607, %612, %613, %614, %615 : vector<8xf32> + %621 = arith.truncf %620 : vector<8xf32> to vector<8xbf16> + %622 = vector.from_elements %616, %617, %618, %619, %608, %609, %610, %611 : vector<8xf32> + %623 = arith.truncf %622 : vector<8xf32> to vector<8xbf16> + %624 = arith.select %444, %621, %623 : vector<8xbf16> + %625 = arith.select %444, %mt4, %mt5 : index + + %626 = arith.addi %bn1, %625 overflow : index + vector.store %624, %442[%626] : memref>, vector<8xbf16> + %627 = vector.extract %418[0] : f32 from vector<4xf32> + %628 = vector.extract %418[1] : f32 from vector<4xf32> + %629 = vector.extract %418[2] : f32 from vector<4xf32> + %630 = vector.extract %418[3] : f32 from vector<4xf32> + %631 = vector.extract %420[0] : f32 from vector<4xf32> + %632 = vector.extract %420[1] : f32 from vector<4xf32> + %633 = vector.extract %420[2] : f32 from vector<4xf32> + %634 = vector.extract %420[3] : f32 from vector<4xf32> + %635 = amdgpu.permlane_swap %627 16 : f32 + %636 = amdgpu.permlane_swap %628 16 : f32 + %637 = amdgpu.permlane_swap %629 16 : f32 + %638 = amdgpu.permlane_swap %630 16 : f32 + %639 = amdgpu.permlane_swap %631 16 : f32 + %640 = amdgpu.permlane_swap %632 16 : f32 + %641 = amdgpu.permlane_swap %633 16 : f32 + %642 = amdgpu.permlane_swap %634 16 : f32 + %643 = vector.from_elements %627, %628, %629, %630, %635, %636, %637, %638 : vector<8xf32> + %644 = arith.truncf %643 : vector<8xf32> to vector<8xbf16> + %645 = vector.from_elements %639, %640, %641, %642, %631, %632, %633, %634 : vector<8xf32> + %646 = arith.truncf %645 : vector<8xf32> to vector<8xbf16> + %647 = arith.select %444, %644, %646 : vector<8xbf16> + %648 = arith.select %444, %mt6, %mt7 : index + + %649 = arith.addi %bn1, %648 overflow : index + vector.store %647, %442[%649] : memref>, vector<8xbf16> + %650 = vector.extract %424[0] : f32 from vector<4xf32> + %651 = vector.extract %424[1] : f32 from vector<4xf32> + %652 = vector.extract %424[2] : f32 from vector<4xf32> + %653 = vector.extract %424[3] : f32 from vector<4xf32> + %654 = vector.extract %426[0] : f32 from vector<4xf32> + %655 = vector.extract %426[1] : f32 from vector<4xf32> + %656 = vector.extract %426[2] : f32 from vector<4xf32> + %657 = vector.extract %426[3] : f32 from vector<4xf32> + %658 = amdgpu.permlane_swap %650 16 : f32 + %659 = amdgpu.permlane_swap %651 16 : f32 + %660 = amdgpu.permlane_swap %652 16 : f32 + %661 = amdgpu.permlane_swap %653 16 : f32 + %662 = amdgpu.permlane_swap %654 16 : f32 + %663 = amdgpu.permlane_swap %655 16 : f32 + %664 = amdgpu.permlane_swap %656 16 : f32 + %665 = amdgpu.permlane_swap %657 16 : f32 + %666 = vector.from_elements %650, %651, %652, %653, %658, %659, %660, %661 : vector<8xf32> + %667 = arith.truncf %666 : vector<8xf32> to vector<8xbf16> + %668 = vector.from_elements %662, %663, %664, %665, %654, %655, %656, %657 : vector<8xf32> + %669 = arith.truncf %668 : vector<8xf32> to vector<8xbf16> + %670 = arith.select %444, %667, %669 : vector<8xbf16> + %671 = arith.select %444, %mt0, %mt1 : index + + %672 = arith.addi %bn2, %671 overflow : index + vector.store %670, %442[%672] : memref>, vector<8xbf16> + %673 = vector.extract %428[0] : f32 from vector<4xf32> + %674 = vector.extract %428[1] : f32 from vector<4xf32> + %675 = vector.extract %428[2] : f32 from vector<4xf32> + %676 = vector.extract %428[3] : f32 from vector<4xf32> + %677 = vector.extract %430[0] : f32 from vector<4xf32> + %678 = vector.extract %430[1] : f32 from vector<4xf32> + %679 = vector.extract %430[2] : f32 from vector<4xf32> + %680 = vector.extract %430[3] : f32 from vector<4xf32> + %681 = amdgpu.permlane_swap %673 16 : f32 + %682 = amdgpu.permlane_swap %674 16 : f32 + %683 = amdgpu.permlane_swap %675 16 : f32 + %684 = amdgpu.permlane_swap %676 16 : f32 + %685 = amdgpu.permlane_swap %677 16 : f32 + %686 = amdgpu.permlane_swap %678 16 : f32 + %687 = amdgpu.permlane_swap %679 16 : f32 + %688 = amdgpu.permlane_swap %680 16 : f32 + %689 = vector.from_elements %673, %674, %675, %676, %681, %682, %683, %684 : vector<8xf32> + %690 = arith.truncf %689 : vector<8xf32> to vector<8xbf16> + %691 = vector.from_elements %685, %686, %687, %688, %677, %678, %679, %680 : vector<8xf32> + %692 = arith.truncf %691 : vector<8xf32> to vector<8xbf16> + %693 = arith.select %444, %690, %692 : vector<8xbf16> + %694 = arith.select %444, %mt2, %mt3 : index + + %695 = arith.addi %bn2, %694 overflow : index + vector.store %693, %442[%695] : memref>, vector<8xbf16> + %696 = vector.extract %432[0] : f32 from vector<4xf32> + %697 = vector.extract %432[1] : f32 from vector<4xf32> + %698 = vector.extract %432[2] : f32 from vector<4xf32> + %699 = vector.extract %432[3] : f32 from vector<4xf32> + %700 = vector.extract %434[0] : f32 from vector<4xf32> + %701 = vector.extract %434[1] : f32 from vector<4xf32> + %702 = vector.extract %434[2] : f32 from vector<4xf32> + %703 = vector.extract %434[3] : f32 from vector<4xf32> + %704 = amdgpu.permlane_swap %696 16 : f32 + %705 = amdgpu.permlane_swap %697 16 : f32 + %706 = amdgpu.permlane_swap %698 16 : f32 + %707 = amdgpu.permlane_swap %699 16 : f32 + %708 = amdgpu.permlane_swap %700 16 : f32 + %709 = amdgpu.permlane_swap %701 16 : f32 + %710 = amdgpu.permlane_swap %702 16 : f32 + %711 = amdgpu.permlane_swap %703 16 : f32 + %712 = vector.from_elements %696, %697, %698, %699, %704, %705, %706, %707 : vector<8xf32> + %713 = arith.truncf %712 : vector<8xf32> to vector<8xbf16> + %714 = vector.from_elements %708, %709, %710, %711, %700, %701, %702, %703 : vector<8xf32> + %715 = arith.truncf %714 : vector<8xf32> to vector<8xbf16> + %716 = arith.select %444, %713, %715 : vector<8xbf16> + %717 = arith.select %444, %mt4, %mt5 : index + + %718 = arith.addi %bn2, %717 overflow : index + vector.store %716, %442[%718] : memref>, vector<8xbf16> + %719 = vector.extract %436[0] : f32 from vector<4xf32> + %720 = vector.extract %436[1] : f32 from vector<4xf32> + %721 = vector.extract %436[2] : f32 from vector<4xf32> + %722 = vector.extract %436[3] : f32 from vector<4xf32> + %723 = vector.extract %438[0] : f32 from vector<4xf32> + %724 = vector.extract %438[1] : f32 from vector<4xf32> + %725 = vector.extract %438[2] : f32 from vector<4xf32> + %726 = vector.extract %438[3] : f32 from vector<4xf32> + %727 = amdgpu.permlane_swap %719 16 : f32 + %728 = amdgpu.permlane_swap %720 16 : f32 + %729 = amdgpu.permlane_swap %721 16 : f32 + %730 = amdgpu.permlane_swap %722 16 : f32 + %731 = amdgpu.permlane_swap %723 16 : f32 + %732 = amdgpu.permlane_swap %724 16 : f32 + %733 = amdgpu.permlane_swap %725 16 : f32 + %734 = amdgpu.permlane_swap %726 16 : f32 + %735 = vector.from_elements %719, %720, %721, %722, %727, %728, %729, %730 : vector<8xf32> + %736 = arith.truncf %735 : vector<8xf32> to vector<8xbf16> + %737 = vector.from_elements %731, %732, %733, %734, %723, %724, %725, %726 : vector<8xf32> + %738 = arith.truncf %737 : vector<8xf32> to vector<8xbf16> + %739 = arith.select %444, %736, %738 : vector<8xbf16> + %740 = arith.select %444, %mt6, %mt7 : index + + %741 = arith.addi %bn2, %740 overflow : index + vector.store %739, %442[%741] : memref>, vector<8xbf16> + return + } + } + } + func.func @isolated_benchmark$async(%arg0: !hal.buffer_view, %arg1: !hal.buffer_view, %arg2: !hal.buffer_view, %arg3: !hal.buffer_view, %arg4: !hal.buffer_view, %arg5: index, %arg6: index, %arg7: index, %arg8: index, %arg9: index, %arg10: !hal.fence, %arg11: !hal.fence) -> !hal.buffer_view { + %0 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[0] : index + %1 = hal.buffer_view.dim<%arg0 : !hal.buffer_view>[1] : index + %2 = hal.tensor.import wait(%arg10) => %arg0 : !hal.buffer_view -> tensor{%0, %1} + %3 = hal.buffer_view.dim<%arg1 : !hal.buffer_view>[0] : index + %4 = hal.buffer_view.dim<%arg1 : !hal.buffer_view>[1] : index + %5 = hal.tensor.import wait(%arg10) => %arg1 : !hal.buffer_view -> tensor{%3, %4} + %6 = hal.buffer_view.dim<%arg2 : !hal.buffer_view>[0] : index + %7 = hal.buffer_view.dim<%arg2 : !hal.buffer_view>[1] : index + %8 = hal.tensor.import wait(%arg10) => %arg2 : !hal.buffer_view -> tensor{%6, %7} + %9 = hal.buffer_view.dim<%arg3 : !hal.buffer_view>[0] : index + %10 = hal.buffer_view.dim<%arg3 : !hal.buffer_view>[1] : index + %11 = hal.tensor.import wait(%arg10) => %arg3 : !hal.buffer_view -> tensor{%9, %10} + %12 = hal.buffer_view.dim<%arg4 : !hal.buffer_view>[0] : index + %13 = hal.buffer_view.dim<%arg4 : !hal.buffer_view>[1] : index + %14 = hal.tensor.import wait(%arg10) => %arg4 : !hal.buffer_view -> tensor{%12, %13} + %c0 = arith.constant 0 : index + %dim = tensor.dim %2, %c0 : tensor + %c0_0 = arith.constant 0 : index + %dim_1 = tensor.dim %8, %c0_0 : tensor + %c1 = arith.constant 1 : index + %dim_2 = tensor.dim %2, %c1 : tensor + %c2 = arith.constant 2 : index + %15 = arith.muli %dim_2, %c2 : index + %c32 = arith.constant 32 : index + %16 = arith.divsi %15, %c32 : index + %c2_3 = arith.constant 2 : index + %17 = arith.divsi %15, %c2_3 : index + %18 = flow.dispatch @gemm::@gemm[%dim, %dim_1, %15, %arg5, %arg6, %arg7, %arg8, %arg9](%2, %5, %8, %11, %14, %arg5, %arg6, %arg7, %arg8, %arg9, %dim, %dim_1, %15) : (tensor{%dim, %17}, tensor{%dim, %16}, tensor{%dim_1, %17}, tensor{%dim_1, %16}, tensor{%dim_1, %dim}, index, index, index, index, index, index, index, index) -> %14{%dim_1, %dim} + %19 = hal.tensor.barrier join(%18 : tensor) => %arg11 : !hal.fence + %c0_4 = arith.constant 0 : index + %dim_5 = tensor.dim %18, %c0_4 : tensor + %c1_6 = arith.constant 1 : index + %dim_7 = tensor.dim %18, %c1_6 : tensor + %20 = hal.tensor.export %19 : tensor{%dim_5, %dim_7} -> !hal.buffer_view + return %20 : !hal.buffer_view + } + } + """ + M_orig, N_orig, K = shape + shape_t = (N_orig, M_orig, K) + block_t = (block[1], block[0], block[2]) + + wave_shape = _get_8wave_shape_from_block(block_t) + gemm, options = get_tagged_mxfp4_gemm_preshuffle_scales_and_B( + shape_t, + block_t, + wave_shape=wave_shape, + b_address_space=SHARED_ADDRESS_SPACE, + output_dtype=tkl.bf16, + ) + options.specialize = True + options.use_buffer_ops = True + options.minimize_shared_allocs = False + options.linearize_shared_access = True + options.wave_runtime = True + options.dump_intermediates = "intermediates_transposed" + options.override_mlir = mlir + if dynamic: + options.dynamic_symbols = [tkl.sym.M, tkl.sym.N, tkl.sym.K] + for sym in options.dynamic_symbols: + del options.subs[sym] + schedule = get_mxfp4_dbuf_pingpong_schedule_Bshuffled_lds( + use_stagger=True, shape=shape_t, block=block_t + ) + UNROLL_FACTOR = tkl.sym.UNROLL_FACTOR + options.subs[UNROLL_FACTOR] = 2 + options.postprocess = """ + module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { + %0 = transform.structured.match ops{["scf.for"]} in %arg0 : (!transform.any_op) -> !transform.any_op + transform.loop.unroll %0 { factor = %%UNROLL_FACTOR%% } : !transform.any_op + transform.yield + } + } + """ + + options = set_default_run_config(options) + gemm = wave_compile(options, gemm, schedule) + print(gemm.asm) + + _run_mxfp_gemm_preshuffle_transposed(gemm, shape, output_dtype=torch.bfloat16) + mode = "dynamic" if dynamic else "static" + print( + f"MXFP GEMM transposed (C^T=B*A^T) 8-wave ping pong B->LDS ({mode}) test passed!" + ) + + def test_dbuf_8wave_mixed_pingpong_mxfp_gemm( is_debug=False, shape=(1024, 1024, 8192), block=(256, 256, 256) ): @@ -350,7 +4493,7 @@ def test_dbuf_4wave_mxfp_preshuffle_b_gemm( def test_dbuf_4wave_mxfp_asymmetric_gemm_cpp( - is_debug=False, shape=(1024, 1024, 8192), block=(128, 256, 256) + is_debug=False, shape=(1024, 1024, 8192), block=(256, 192, 256) ): """Asymmetric MXFP4 GEMM using C++ WaveASM backend (no preshuffle).""" gemm, options = get_tagged_mxfp4_gemm( diff --git a/wave_lang/kernel/compiler/wave_codegen/emitter.py b/wave_lang/kernel/compiler/wave_codegen/emitter.py index c97d311582..c99cc317c4 100644 --- a/wave_lang/kernel/compiler/wave_codegen/emitter.py +++ b/wave_lang/kernel/compiler/wave_codegen/emitter.py @@ -633,10 +633,13 @@ def add_emitter_subs( _emulate_ceildiv = bool(int(environ.get("WAVE_EMULATE_CEILDIV", 0))) _use_affine_expr = bool(int(environ.get("WAVE_USE_AFFINE_EXPR", 1))) +_magic_number_enabled = bool(int(environ.get("WAVE_MAGIC_NUMBER_DIV", 1))) _Rational = namedtuple("_Rational", ["numerator", "denominator"]) _ApplyExpr = namedtuple("_ApplyExpr", ["expr", "args"]) +_magic_number_cache: dict = {} + def gen_sympy_index(dynamics: dict[IndexSymbol, Value], expr: sympy.Expr) -> Value: use_affine_expr = _use_affine_expr @@ -714,6 +717,26 @@ def get_const_val(arg): return None + def _same_value(a, b) -> bool: + """Structural equality for emitter values (_ApplyExpr / OpResult). + + Detects when two independently-built IR values represent the same + computation. Used in _add to avoid denominator explosion when + adding _Rationals that share a common denominator. + """ + if isinstance(a, _ApplyExpr) and isinstance(b, _ApplyExpr): + if a.expr != b.expr or len(a.args) != len(b.args): + return False + return all(_same_value(x, y) for x, y in zip(a.args, b.args)) + if isinstance(a, (Value, OpResult)) and isinstance(b, (Value, OpResult)): + if a is b: + return True + a_val = get_const_val(a) + b_val = get_const_val(b) + if a_val is not None and b_val is not None: + return a_val == b_val + return False + overflow_flags = arith_d.IntegerOverflowFlags.nsw | arith_d.IntegerOverflowFlags.nuw def muli(lhs, rhs): @@ -778,16 +801,105 @@ def muli_expr(lhs, rhs): return op_expr(lhs, rhs, lambda a, b: a * b) + def _is_dynamic_divisor(val) -> bool: + """Check if a value is NOT a compile-time constant.""" + if isinstance(val, _ApplyExpr): + if all( + isinstance(a, OpResult) and get_const_val(a) is not None + for a in val.args + ): + return False + val = _get_ir_value(val) + if isinstance(val, OpResult): + return get_const_val(val) is None + return True + + def _mulhi_u32(n_i32, m_i32): + """Unsigned 32-bit multiply-high: (n * m) >> 32, via 64-bit multiply.""" + i64 = IntegerType.get_signless(64) + c32_i64 = arith_d.constant(i64, 32) + n_i64 = arith_d.extui(i64, n_i32) + m_i64 = arith_d.extui(i64, m_i32) + prod_i64 = arith_d.muli(n_i64, m_i64) + hi_i64 = arith_d.shrui(prod_i64, c32_i64) + i32 = IntegerType.get_signless(32) + return arith_d.trunci(i32, hi_i64) + + def _precompute_magic_number(divisor_index: Value): + """ + Compute magic = ceil(2^32 / d) from a dynamic divisor. + Returns (magic_i32, d_i32) both as i32 Values. + """ + i32 = IntegerType.get_signless(32) + i64 = IntegerType.get_signless(64) + d_i32 = arith_d.index_cast(i32, divisor_index) + d_i64 = arith_d.extui(i64, d_i32) + c1_i64 = arith_d.constant(i64, 1) + c32_i64 = arith_d.constant(i64, 32) + pow32 = arith_d.shli(c1_i64, c32_i64) + d_minus_1_i64 = arith_d.subi(d_i64, c1_i64) + numer_i64 = arith_d.addi(pow32, d_minus_1_i64) + magic_i64 = arith_d.divui(numer_i64, d_i64) + magic_i32 = arith_d.trunci(i32, magic_i64) + return magic_i32, d_i32 + + def _get_or_create_magic(divisor: Value): + """Get cached (magic_i32, d_i32) or compute and cache them.""" + key = id(divisor) + if key in _magic_number_cache: + return _magic_number_cache[key] + magic_i32, d_i32 = _precompute_magic_number(divisor) + _magic_number_cache[key] = (magic_i32, d_i32) + return magic_i32, d_i32 + + def _magic_div_and_rem(lhs_val, rhs_val): + """ + Compute (quotient, remainder) of lhs_val // rhs_val using + magic number multiplication: q = mulhi(n, magic), with a + one-step correction for exactness. + Returns (quotient_index, remainder_index). + """ + i32 = IntegerType.get_signless(32) + magic_i32, d_i32 = _get_or_create_magic(rhs_val) + n_i32 = arith_d.index_cast(i32, lhs_val) + q_i32 = _mulhi_u32(n_i32, magic_i32) + qd_i32 = arith_d.muli(q_i32, d_i32) + r_i32 = arith_d.subi(n_i32, qd_i32) + # Correction: ceil(2^32/d) can overestimate quotient by 1. + # Detect via unsigned remainder >= divisor (wraps on overestimate). + too_big = arith_d.cmpi(arith_d.CmpIPredicate.uge, r_i32, d_i32) + c1_i32 = arith_d.constant(i32, 1) + c0_i32 = arith_d.constant(i32, 0) + corr = arith_d.select(too_big, c1_i32, c0_i32) + q_final = arith_d.subi(q_i32, corr) + d_or_zero = arith_d.select(too_big, d_i32, c0_i32) + r_final = arith_d.addi(r_i32, d_or_zero) + q_index = arith_d.index_cast(IndexType.get(), q_final) + r_index = arith_d.index_cast(IndexType.get(), r_final) + return q_index, r_index + def rem_expr(lhs, rhs): if not use_affine_expr or not check_index_types(lhs, rhs): return arith_d.remsi(*_broadcast(lhs, rhs)) + if _magic_number_enabled and _is_dynamic_divisor(rhs): + lhs_val = _get_ir_value(lhs) if isinstance(lhs, _ApplyExpr) else lhs + rhs_val = _get_ir_value(rhs) if isinstance(rhs, _ApplyExpr) else rhs + _, r = _magic_div_and_rem(lhs_val, rhs_val) + return r + return op_expr(lhs, rhs, lambda a, b: a % b) def floordiv_expr(lhs, rhs): if not use_affine_expr or not check_index_types(lhs, rhs): return arith_d.divsi(*_broadcast(lhs, rhs)) + if _magic_number_enabled and _is_dynamic_divisor(rhs): + lhs_val = _get_ir_value(lhs) if isinstance(lhs, _ApplyExpr) else lhs + rhs_val = _get_ir_value(rhs) if isinstance(rhs, _ApplyExpr) else rhs + q, _ = _magic_div_and_rem(lhs_val, rhs_val) + return q + return op_expr(lhs, rhs, lambda a, b: AffineExpr.get_floor_div(a, b)) def ceildiv_expr(lhs, rhs): @@ -820,6 +932,9 @@ def _add(lhs, rhs): numerator = add_expr(numerator, rhs.numerator) return _Rational(numerator, rhs.denominator) elif is_rational_lhs and is_rational_rhs: + if _same_value(lhs.denominator, rhs.denominator): + numerator = add_expr(lhs.numerator, rhs.numerator) + return _Rational(numerator, lhs.denominator) lhs_numerator = muli_expr(lhs.numerator, rhs.denominator) rhs_numerator = muli_expr(rhs.numerator, lhs.denominator) numerator = add_expr(lhs_numerator, rhs_numerator) diff --git a/wave_lang/kernel/wave/gather_to_shared.py b/wave_lang/kernel/wave/gather_to_shared.py index 23e3345ca4..ee1c633d54 100644 --- a/wave_lang/kernel/wave/gather_to_shared.py +++ b/wave_lang/kernel/wave/gather_to_shared.py @@ -27,8 +27,10 @@ Write, get_custom, ) +from ..wave.assumptions import get_divisibility_subs from ..wave.constraints import ( Constraint, + DistributionConstraint, TilingConstraint, WaveConstraint, WorkgroupConstraint, @@ -49,6 +51,7 @@ remove_thread_indexing, remove_global_indexing, ) +from .generate_bounds_exprs import is_divisible from .utils.general_utils import is_gather from .utils.graph_utils import DCE from .utils.symbol_utils import subs_idxc @@ -542,6 +545,18 @@ def gather_to_shared( constraints, read.index, vector_shapes, symbolic_shape ) + # Remove bounds where divisibility assumptions prove full tile alignment + fwd, _ = get_divisibility_subs(constraints) + if bounds and fwd: + for c in constraints: + if ( + isinstance(c, DistributionConstraint) + and c.dim in bounds + and is_divisible(c.dim, c.tile_size, fwd) + ): + del bounds[c.dim] + bounds = bounds or None + logger.info(f"bounds={bounds}") fastest_dim_bound = bounds.get(symbolic_shape[-1], None) if bounds else None diff --git a/wave_lang/kernel/wave/schedules/gemm_mxfp4_double_buffer.py b/wave_lang/kernel/wave/schedules/gemm_mxfp4_double_buffer.py index e68a5ae823..06678f714b 100755 --- a/wave_lang/kernel/wave/schedules/gemm_mxfp4_double_buffer.py +++ b/wave_lang/kernel/wave/schedules/gemm_mxfp4_double_buffer.py @@ -44,8 +44,9 @@ def mxfp4_dbuf_schedule(): global_to_shared_a = tkw.filter_nodes(all_read_a, node_type=tkw.GatherToLDS) shared_load_a = tkw.filter_nodes(all_read_a, node_type=tkw.Read) - # Matrix A scale + # Matrix A scale (global -> VGPR reads) all_read_a_scale = tkw.get_node_by_tag("read_a_scale") + read_a_scale = tkw.filter_nodes(all_read_a_scale, node_type=tkw.Read) global_to_shared_a_scale = tkw.filter_nodes( all_read_a_scale, node_type=tkw.GatherToLDS ) @@ -56,8 +57,9 @@ def mxfp4_dbuf_schedule(): global_to_shared_b = tkw.filter_nodes(all_read_b, node_type=tkw.GatherToLDS) shared_load_b = tkw.filter_nodes(all_read_b, node_type=tkw.Read) - # Matrix B scale + # Matrix B scale (global -> VGPR reads) all_read_b_scale = tkw.get_node_by_tag("read_b_scale") + read_b_scale = tkw.filter_nodes(all_read_b_scale, node_type=tkw.Read) global_to_shared_b_scale = tkw.filter_nodes( all_read_b_scale, node_type=tkw.GatherToLDS ) @@ -276,16 +278,20 @@ def mxfp4_dbuf_schedule(): global_to_shared_a = tkw.filter_nodes(all_read_a, node_type=tkw.GatherToLDS) shared_load_a = tkw.filter_nodes(all_read_a, node_type=tkw.Read) - # Matrix A scale + # Matrix A scale (global -> VGPR reads) all_read_a_scale = tkw.get_node_by_tag("read_a_scale") + read_a_scale = tkw.filter_nodes(all_read_a_scale, node_type=tkw.Read) + read_a_scale = tkw.filter_nodes(all_read_a_scale, node_type=tkw.Read) # Matrix B data all_read_b = tkw.get_node_by_tag("read_b") global_to_shared_b = tkw.filter_nodes(all_read_b, node_type=tkw.GatherToLDS) shared_load_b = tkw.filter_nodes(all_read_b, node_type=tkw.Read) - # Matrix B scale + # Matrix B scale (global -> VGPR reads) all_read_b_scale = tkw.get_node_by_tag("read_b_scale") + read_b_scale = tkw.filter_nodes(all_read_b_scale, node_type=tkw.Read) + read_b_scale = tkw.filter_nodes(all_read_b_scale, node_type=tkw.Read) # Bitcast operations (needed alongside compute) bitcast_a = tkw.get_node_by_tag("bitcast_a") @@ -308,8 +314,8 @@ def mxfp4_dbuf_schedule(): ( global_to_shared_a, global_to_shared_b, - all_read_a_scale, - all_read_b_scale, + read_a_scale, + read_b_scale, ), (), (), @@ -342,11 +348,11 @@ def mxfp4_dbuf_schedule(): loop_shared_load_b = tkw.filter_nodes( shared_load_b, subgraph=pipeline_loop.KERNEL ) - loop_all_read_a_scale = tkw.filter_nodes( - all_read_a_scale, subgraph=pipeline_loop.KERNEL + loop_read_a_scale = tkw.filter_nodes( + read_a_scale, subgraph=pipeline_loop.KERNEL ) - loop_all_read_b_scale = tkw.filter_nodes( - all_read_b_scale, subgraph=pipeline_loop.KERNEL + loop_read_b_scale = tkw.filter_nodes( + read_b_scale, subgraph=pipeline_loop.KERNEL ) loop_bitcast_a = tkw.filter_nodes(bitcast_a, subgraph=pipeline_loop.KERNEL) loop_bitcast_a_scale = tkw.filter_nodes( @@ -367,11 +373,11 @@ def mxfp4_dbuf_schedule(): loop_shared_load_b_0, loop_shared_load_b_1 = tkw.partition_by_dim( loop_shared_load_b, dim=K, num_partitions=2 ) - loop_all_read_a_scale_0, loop_all_read_a_scale_1 = tkw.partition_by_dim( - loop_all_read_a_scale, dim=K, num_partitions=2 + loop_read_a_scale_0, loop_read_a_scale_1 = tkw.partition_by_dim( + loop_read_a_scale, dim=K, num_partitions=2 ) - loop_all_read_b_scale_0, loop_all_read_b_scale_1 = tkw.partition_by_dim( - loop_all_read_b_scale, dim=K, num_partitions=2 + loop_read_b_scale_0, loop_read_b_scale_1 = tkw.partition_by_dim( + loop_read_b_scale, dim=K, num_partitions=2 ) loop_bitcast_a_0, loop_bitcast_a_1 = tkw.partition_by_dim( @@ -409,8 +415,8 @@ def mxfp4_dbuf_schedule(): loop_bitcast_a_scale_0, loop_bitcast_b_0, loop_bitcast_b_scale_0, - loop_all_read_a_scale_0, # prefetch A & B scales for next iteration - loop_all_read_b_scale_0, + loop_read_a_scale_0, # prefetch A & B scales for next iteration + loop_read_b_scale_0, tkw.SchedulingBarrier([]), ] ) @@ -446,8 +452,8 @@ def mxfp4_dbuf_schedule(): loop_bitcast_a_scale_1, loop_bitcast_b_1, loop_bitcast_b_scale_1, - loop_all_read_a_scale_1, - loop_all_read_b_scale_1, + loop_read_a_scale_1, + loop_read_b_scale_1, tkw.SchedulingBarrier([]), tkw.WorkgroupBarrier(), tkw.SchedulingBarrier([]), @@ -507,14 +513,16 @@ def mxfp4_dbuf_schedule(): global_to_shared_a = tkw.filter_nodes(all_read_a, node_type=tkw.GatherToLDS) shared_load_a = tkw.filter_nodes(all_read_a, node_type=tkw.Read) - # Matrix A scale + # Matrix A scale (global -> VGPR reads) all_read_a_scale = tkw.get_node_by_tag("read_a_scale") + read_a_scale = tkw.filter_nodes(all_read_a_scale, node_type=tkw.Read) # Matrix B data all_read_b = tkw.get_node_by_tag("read_b") - # Matrix B scale + # Matrix B scale (global -> VGPR reads) all_read_b_scale = tkw.get_node_by_tag("read_b_scale") + read_b_scale = tkw.filter_nodes(all_read_b_scale, node_type=tkw.Read) # Bitcast operations (needed alongside compute) bitcast_a = tkw.get_node_by_tag("bitcast_a") @@ -683,7 +691,7 @@ def mxfp4_dbuf_schedule(): def get_mxfp4_dbuf_pingpong_schedule_Bshuffled_lds( - use_stagger: bool = True, shape: tuple = None + use_stagger: bool = True, shape: tuple = None, block: tuple = None ): """Return a double-buffered MXFP4 schedule for wave_compile(). Same as get_mxfp4_dbuf_pingpong_schedule_Bshuffled(), but B data is read @@ -710,7 +718,7 @@ def mxfp4_dbuf_schedule(): global_to_shared_a = tkw.filter_nodes(all_read_a, node_type=tkw.GatherToLDS) shared_load_a = tkw.filter_nodes(all_read_a, node_type=tkw.Read) - # Matrix A scale + # Matrix A scale (global -> VGPR reads) all_read_a_scale = tkw.get_node_by_tag("read_a_scale") # Matrix B data - GatherToLDS (global->shared) + Read (shared load) @@ -718,8 +726,9 @@ def mxfp4_dbuf_schedule(): global_to_shared_b = tkw.filter_nodes(all_read_b, node_type=tkw.GatherToLDS) shared_load_b = tkw.filter_nodes(all_read_b, node_type=tkw.Read) - # Matrix B scale + # Matrix B scale (global -> VGPR reads) all_read_b_scale = tkw.get_node_by_tag("read_b_scale") + read_b_scale = tkw.filter_nodes(all_read_b_scale, node_type=tkw.Read) # Bitcast operations (needed alongside compute) bitcast_a = tkw.get_node_by_tag("bitcast_a") @@ -807,79 +816,258 @@ def mxfp4_dbuf_schedule(): loop_bitcast_b, dim=K, num_partitions=2 ) + # TODO: FIX THIS, make this dynamic based on the block size + # Count only actual global->VGPR Read nodes from the scheduled scale groups. + # loop_a_scale_reads = tkw.filter_nodes(loop_all_read_a_scale, node_type=tkw.Read) + # loop_b_scale_reads = tkw.filter_nodes(loop_all_read_b_scale, node_type=tkw.Read) + # number_outstanding_loads_to_vgpr = len(loop_a_scale_reads) + len( + # loop_b_scale_reads + # ) + number_outstanding_loads_to_vgpr = 0 + safe = 0 + if block is not None and block == (256, 192, 256): + safe = 2 + number_outstanding_loads_to_vgpr = 5 ## HARDCODED FOR NOW TODO: FIX THIS + + if block is not None and block == (256, 160, 256): + safe = 5 + number_outstanding_loads_to_vgpr = 12 + + if block is not None and (block == (64, 64, 256) or block == (128, 128, 256)): + safe = 0 + number_outstanding_loads_to_vgpr = 3 + + if block is not None and block == (128, 64, 256): + safe = 0 + number_outstanding_loads_to_vgpr = 3 + + if block is not None and block == (128, 256, 256): + safe = 0 + number_outstanding_loads_to_vgpr = 5 + if block is not None and block == (256, 128, 256): + safe = 0 + number_outstanding_loads_to_vgpr = 4 + if block is not None and block == (256, 256, 256): + safe = 0 + number_outstanding_loads_to_vgpr = 6 + # If the bus gets congested and cluster memory dependency are affected, we must add a second barrier to fix the timing and prevent incorrect output results. # In case a second a second workgroup barrier is needed, another schedule is created to hide the latency of that second barrier, by scheduling safe ds_read ops before the second barrier (see get_mxfp4_dbuf_mixed_pingpong_schedule). - use_extra_barrier = True - # Build cluster 0: first K-partition loads + bitcasts + GatherToLDS - cluster_0_ops = [ - tkw.SchedulingBarrier([]), - tkw.MemoryCounterWait(load=0), - tkw.WorkgroupBarrier(), - ] - if use_extra_barrier: - cluster_0_ops.append(tkw.WorkgroupBarrier()) - cluster_0_ops.extend( - [ - loop_global_to_shared, - tkw.SchedulingBarrier([]), - loop_shared_load_a_0, - loop_shared_load_b_0, - loop_bitcast_a_0, - loop_bitcast_a_scale, - loop_bitcast_b_0, - loop_bitcast_b_scale, - loop_all_read_a_scale, # prefetch A & B scales for next iteration - loop_all_read_b_scale, + use_extra_barrier = False + + if block is not None and block == (256, 192, 256): + + cluster_0_ops = [ tkw.SchedulingBarrier([]), + tkw.MemoryCounterWait(load=number_outstanding_loads_to_vgpr), + tkw.WorkgroupBarrier(), ] - ) - if use_stagger: + if use_extra_barrier: + cluster_0_ops.append(tkw.WorkgroupBarrier()) cluster_0_ops.extend( [ - tkw.WorkgroupBarrier(), + loop_global_to_shared, + tkw.SchedulingBarrier([]), + loop_shared_load_a_0, + loop_shared_load_b_0, + loop_bitcast_a_0, + loop_bitcast_a_scale, + loop_bitcast_b_0, + loop_bitcast_b_scale, tkw.SchedulingBarrier([]), ] ) + if use_stagger: + cluster_0_ops.extend( + [ + tkw.WorkgroupBarrier(), + tkw.SchedulingBarrier([]), + ] + ) - clusters = [ - # Cluster 0: First K-partition shared loads/bitcasts + async GatherToLDS - tkw.cluster(cluster_0_ops), - # Cluster 1: First K-partition scaled_mma (high priority) - tkw.cluster( - [ - tkw.SetWavePrio(1), - loop_scaled_mma_0, - tkw.SetWavePrio(0), - tkw.SchedulingBarrier([]), - tkw.WorkgroupBarrier(), - tkw.SchedulingBarrier([]), - ], - ), - # Cluster 2: Second K-partition shared loads/bitcasts - tkw.cluster( + clusters = [ + tkw.cluster(cluster_0_ops), + tkw.cluster( + [ + tkw.SetWavePrio(1), + loop_scaled_mma_0, + tkw.SetWavePrio(0), + tkw.SchedulingBarrier([]), + tkw.WorkgroupBarrier(), + tkw.SchedulingBarrier([]), + ], + ), + tkw.cluster( + [ + tkw.SchedulingBarrier([]), + loop_all_read_a_scale, + loop_all_read_b_scale, + loop_shared_load_a_1, + loop_shared_load_b_1, + loop_bitcast_a_1, + loop_bitcast_b_1, + tkw.SchedulingBarrier([]), + tkw.MemoryCounterWait( + load=(number_outstanding_loads_to_vgpr + safe) + ), + tkw.WorkgroupBarrier(), + tkw.SchedulingBarrier([]), + ], + ), + tkw.cluster( + [ + tkw.SetWavePrio(1), + loop_scaled_mma_1, + tkw.SetWavePrio(0), + tkw.SchedulingBarrier([]), + ], + ), + ] + tkw.insert_before( + pipeline_loop.KERNEL, + tkw.MemoryCounterWait(load=number_outstanding_loads_to_vgpr), + ) + + elif block is not None and block == (256, 160, 256): + cluster_0_ops = [ + tkw.SchedulingBarrier([]), + tkw.MemoryCounterWait(load=number_outstanding_loads_to_vgpr), + tkw.WorkgroupBarrier(), + ] + if use_extra_barrier: + cluster_0_ops.append(tkw.WorkgroupBarrier()) + cluster_0_ops.extend( [ + loop_global_to_shared, tkw.SchedulingBarrier([]), + loop_shared_load_a_0, + loop_shared_load_b_0, + loop_bitcast_a_0, + loop_bitcast_a_scale, + loop_bitcast_b_0, + loop_bitcast_b_scale, loop_shared_load_a_1, loop_shared_load_b_1, - loop_bitcast_a_1, - loop_bitcast_b_1, - tkw.SchedulingBarrier([]), - tkw.WorkgroupBarrier(), tkw.SchedulingBarrier([]), - ], - ), - # Cluster 3: Second K-partition scaled_mma (high priority) - tkw.cluster( + ] + ) + if use_stagger: + cluster_0_ops.extend( + [ + tkw.WorkgroupBarrier(), + tkw.SchedulingBarrier([]), + ] + ) + + clusters = [ + tkw.cluster(cluster_0_ops), + tkw.cluster( + [ + tkw.SetWavePrio(1), + loop_scaled_mma_0, + tkw.SetWavePrio(0), + tkw.SchedulingBarrier([]), + tkw.WorkgroupBarrier(), + tkw.SchedulingBarrier([]), + ], + ), + tkw.cluster( + [ + tkw.SchedulingBarrier([]), + loop_bitcast_a_1, + loop_bitcast_b_1, + tkw.SchedulingBarrier([]), + tkw.MemoryCounterWait( + load=(number_outstanding_loads_to_vgpr + safe) + ), + tkw.WorkgroupBarrier(), + tkw.SchedulingBarrier([]), + ], + ), + tkw.cluster( + [ + tkw.SetWavePrio(1), + loop_scaled_mma_1, + tkw.SetWavePrio(0), + tkw.SchedulingBarrier([]), + ], + ), + ] + tkw.insert_before( + pipeline_loop.KERNEL, + tkw.MemoryCounterWait(load=number_outstanding_loads_to_vgpr), + ) + + else: + + cluster_0_ops = [ + tkw.SchedulingBarrier([]), + tkw.MemoryCounterWait(load=0), + tkw.WorkgroupBarrier(), + ] + if use_extra_barrier: + cluster_0_ops.append(tkw.WorkgroupBarrier()) + cluster_0_ops.extend( [ - tkw.SetWavePrio(1), - loop_scaled_mma_1, - tkw.SetWavePrio(0), + loop_global_to_shared, tkw.SchedulingBarrier([]), - ], - ), - ] + loop_shared_load_a_0, + loop_shared_load_b_0, + loop_bitcast_a_0, + loop_bitcast_a_scale, + loop_bitcast_b_0, + loop_bitcast_b_scale, + loop_all_read_a_scale, + loop_all_read_b_scale, + tkw.SchedulingBarrier([]), + ] + ) + if use_stagger: + cluster_0_ops.extend( + [ + tkw.WorkgroupBarrier(), + tkw.SchedulingBarrier([]), + ] + ) + + clusters = [ + tkw.cluster(cluster_0_ops), + tkw.cluster( + [ + tkw.SetWavePrio(1), + loop_scaled_mma_0, + tkw.SetWavePrio(0), + tkw.SchedulingBarrier([]), + tkw.WorkgroupBarrier(), + tkw.SchedulingBarrier([]), + ], + ), + tkw.cluster( + [ + tkw.SchedulingBarrier([]), + loop_shared_load_a_1, + loop_shared_load_b_1, + loop_bitcast_a_1, + loop_bitcast_b_1, + tkw.SchedulingBarrier([]), + tkw.MemoryCounterWait(load=(number_outstanding_loads_to_vgpr)), + tkw.WorkgroupBarrier(), + tkw.SchedulingBarrier([]), + ], + ), + tkw.cluster( + [ + tkw.SetWavePrio(1), + loop_scaled_mma_1, + tkw.SetWavePrio(0), + tkw.SchedulingBarrier([]), + ], + ), + ] + tkw.insert_before(pipeline_loop.KERNEL, tkw.MemoryCounterWait(load=0)) # Insert barriers at loop boundaries + tkw.insert_before(pipeline_loop.KERNEL, tkw.WorkgroupBarrier()) tkw.insert_after(pipeline_loop.KERNEL, tkw.SharedMemoryBarrier()) diff --git a/wave_lang/kernel/wave/templates/tagged_mxfp4_gemm.py b/wave_lang/kernel/wave/templates/tagged_mxfp4_gemm.py index b7b013460b..564667e263 100755 --- a/wave_lang/kernel/wave/templates/tagged_mxfp4_gemm.py +++ b/wave_lang/kernel/wave/templates/tagged_mxfp4_gemm.py @@ -149,6 +149,7 @@ def _get_tagged_mxfp4_gemm_preshuffle_scales_impl( b_preshuffled: bool = False, reorder_workgroups: bool = False, group_size_n=32, + output_dtype=tkl.f32, ): """Shared implementation: preshuffle scales only, or scales + B data. @@ -179,6 +180,19 @@ def _get_tagged_mxfp4_gemm_preshuffle_scales_impl( constraints += [tkw.WaveConstraint(N, BLOCK_N / wave_shape[1])] constraints += [tkw.HardwareConstraint(threads_per_wave=64, mma_type=mfma_variant)] + # Divisibility assumptions for M, N, K (no effect for static shapes). + constraints += [tkw.Assumption(Eq(M % 32, 0))] + constraints += [tkw.Assumption(Eq(N % 32, 0))] + constraints += [tkw.Assumption(Eq(K % 256, 0))] + + # Include assumption that K is divisible by BLOCK_K to allow gather_to_shared ops to omit masking predicates. + constraints += [tkw.Assumption(Eq(K % BLOCK_K, 0))] + constraints += [tkw.Assumption(Eq(M % BLOCK_M, 0))] + constraints += [tkw.Assumption(Eq(N % BLOCK_N, 0))] + + # K is always large enough for software pipelining. + constraints += [tkw.Assumption(K > BLOCK_K * 6)] + if reorder_workgroups: new_wg0, new_wg1 = _reorder_mxfp4_workgroups( M, N, BLOCK_M, BLOCK_N, GROUP_SIZE_N @@ -249,7 +263,7 @@ def gemm( a_scale: tkl.Memory[M, K / 32, GLOBAL_ADDRESS_SPACE, tkl.i8], b: tkl.Memory[N, K / 2, B_ADDRESS_SPACE, tkl.i8], b_scale: tkl.Memory[N, K / 32, GLOBAL_ADDRESS_SPACE, tkl.i8], - c: tkl.Memory[M, N, C_ADDRESS_SPACE, tkl.f32], + c: tkl.Memory[M, N, C_ADDRESS_SPACE, output_dtype], ): c_reg = tkl.Register[M, N, tkl.f32](0.0) @@ -273,6 +287,9 @@ def repeat( ) return acc + if output_dtype == tkl.bf16: + repeat = tkw.cast(repeat, tkl.bf16) + tkw.write(repeat, c) hyperparams = { @@ -290,7 +307,7 @@ def repeat( M: shape[0], N: shape[1], K: shape[2], - K_SCALE_SHUFFLED: (((shape[2] // 32) + 7) // 8) * 8, + K_SCALE_SHUFFLED: (((K // 32) + 7) // 8) * 8, } if b_preshuffled: hyperparams[K_PACKED] = K // 2 @@ -348,6 +365,7 @@ def get_tagged_mxfp4_gemm_preshuffle_scales_and_B( mfma_variant: ScaledMMAType = ScaledMMAType.F32_16x16x128_F8F6F4, a_address_space: tkl.AddressSpace = SHARED_ADDRESS_SPACE, b_address_space: tkl.AddressSpace | None = None, + output_dtype=tkl.f32, ): """Return a tagged MXFP4 scaled GEMM kernel with preshuffled B and B_scale. @@ -374,6 +392,7 @@ def get_tagged_mxfp4_gemm_preshuffle_scales_and_B( a_address_space, b_address_space, b_preshuffled=True, + output_dtype=output_dtype, )