kernel.ttgir

#blocked = #ttg.blocked<{sizePerThread = [1, 8], threadsPerWarp = [4, 8], warpsPerCTA = [8, 1], order = [1, 0]}>
#blocked1 = #ttg.blocked<{sizePerThread = [1, 8], threadsPerWarp = [1, 32], warpsPerCTA = [8, 1], order = [1, 0]}>
#blocked2 = #ttg.blocked<{sizePerThread = [2, 1], threadsPerWarp = [32, 1], warpsPerCTA = [2, 4], order = [0, 1]}>
#linear = #ttg.linear<{register = [[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [0, 32], [0, 64]], lane = [[1, 0], [2, 0], [4, 0], [8, 0], [16, 0]], warp = [[32, 0], [64, 0], [0, 128]], block = []}>
#loc = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":28:1)
#shared = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}>
#shared1 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0, 1]}>
#shared2 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}>
#smem = #ttg.shared_memory
#tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 256, colStride = 1>
#loc50 = loc("arg_B"(#loc))
#loc51 = loc("in_ptr1"(#loc))
#loc52 = loc("in_ptr2"(#loc))
#loc53 = loc("out_ptr0"(#loc))
module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, ttg.target = "cuda:100", "ttg.threads-per-warp" = 32 : i32} {
  tt.func public @triton_tem_fused_index_mm_0(%arg_B: !tt.ptr<f16> {tt.divisibility = 16 : i32} loc("arg_B"(#loc)), %in_ptr1: !tt.ptr<i64> {tt.divisibility = 16 : i32} loc("in_ptr1"(#loc)), %in_ptr2: !tt.ptr<f16> {tt.divisibility = 16 : i32} loc("in_ptr2"(#loc)), %out_ptr0: !tt.ptr<f16> {tt.divisibility = 16 : i32} loc("out_ptr0"(#loc))) attributes {noinline = false} {
    %cst = arith.constant dense<0> : tensor<128x64xi64, #blocked> loc(#loc1)
    %cst_0 = arith.constant dense<2048> : tensor<128x64xi64, #blocked> loc(#loc1)
    %cst_1 = arith.constant dense<64> : tensor<1x64xi32, #blocked> loc(#loc1)
    %cst_2 = arith.constant dense<64> : tensor<64x1xi32, #blocked1> loc(#loc1)
    %cst_3 = arith.constant dense<true> : tensor<64x256xi1, #blocked1> loc(#loc1)
    %cst_4 = arith.constant dense<true> : tensor<128x64xi1, #blocked> loc(#loc1)
    %cst_5 = arith.constant dense<true> : tensor<128x64xi1, #blocked2> loc(#loc1)
    %c30_i32 = arith.constant 30 : i32 loc(#loc1)
    %c29_i32 = arith.constant 29 : i32 loc(#loc1)
    %c31_i32 = arith.constant 31 : i32 loc(#loc1)
    %c2_i32 = arith.constant 2 : i32 loc(#loc1)
    %cst_6 = arith.constant dense<0.000000e+00> : tensor<128x256xf32, #linear> loc(#loc1)
    %true = arith.constant true loc(#loc1)
    %c1_i32 = arith.constant 1 : i32 loc(#loc1)
    %c16_i32 = arith.constant 16 : i32 loc(#loc1)
    %c64_i32 = arith.constant 64 : i32 loc(#loc1)
    %c0_i32 = arith.constant 0 : i32 loc(#loc1)
    %c128_i32 = arith.constant 128 : i32 loc(#loc1)
    %c256_i32 = arith.constant 256 : i32 loc(#loc1)
    %c8_i32 = arith.constant 8 : i32 loc(#loc1)
    %cst_7 = arith.constant dense<2048> : tensor<128xi32, #ttg.slice<{dim = 1, parent = #blocked2}>> loc(#loc1)
    %cst_8 = arith.constant dense<2048> : tensor<256xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> loc(#loc1)
    %cst_9 = arith.constant dense<2048> : tensor<128x1xi32, #blocked1> loc(#loc1)
    %cst_10 = arith.constant dense<2048> : tensor<64x1xi32, #blocked1> loc(#loc1)
    %cst_11 = arith.constant dense<2048> : tensor<1x256xi32, #blocked1> loc(#loc1)
    %false = arith.constant false loc(#loc1)
    %pid = tt.get_program_id x : i32 loc(#loc54)
    %group_id = arith.divsi %pid, %c64_i32 : i32 loc(#loc55)
    %group_size = arith.muli %group_id, %c8_i32 : i32 loc(#loc56)
    %group_size_12 = arith.subi %c16_i32, %group_size : i32 loc(#loc57)
    %group_size_13 = arith.minsi %group_size_12, %c8_i32 : i32 loc(#loc58)
    %pid_m = arith.remsi %pid, %group_size_13 : i32 loc(#loc59)
    %pid_m_14 = arith.addi %group_size, %pid_m : i32 loc(#loc60)
    %pid_n = arith.remsi %pid, %c64_i32 : i32 loc(#loc61)
    %pid_n_15 = arith.divsi %pid_n, %group_size_13 : i32 loc(#loc62)
    %0 = arith.cmpi sge, %pid_m_14, %c0_i32 : i32 loc(#loc11)
    llvm.intr.assume %0 : i1 loc(#loc12)
    %1 = arith.cmpi sge, %pid_n_15, %c0_i32 : i32 loc(#loc13)
    llvm.intr.assume %1 : i1 loc(#loc14)
    %rm = arith.muli %pid_m_14, %c128_i32 : i32 loc(#loc63)
    %rm_16 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #ttg.slice<{dim = 1, parent = #blocked2}>> loc(#loc64)
    %rm_17 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #ttg.slice<{dim = 1, parent = #blocked1}>> loc(#loc64)
    %rm_18 = tt.splat %rm : i32 -> tensor<128xi32, #ttg.slice<{dim = 1, parent = #blocked2}>> loc(#loc63)
    %rm_19 = tt.splat %rm : i32 -> tensor<128xi32, #ttg.slice<{dim = 1, parent = #blocked1}>> loc(#loc63)
    %rm_20 = arith.addi %rm_18, %rm_16 : tensor<128xi32, #ttg.slice<{dim = 1, parent = #blocked2}>> loc(#loc63)
    %rm_21 = arith.addi %rm_19, %rm_17 : tensor<128xi32, #ttg.slice<{dim = 1, parent = #blocked1}>> loc(#loc63)
    %rn = arith.muli %pid_n_15, %c256_i32 : i32 loc(#loc65)
    %rn_22 = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor<256xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> loc(#loc66)
    %rn_23 = tt.splat %rn : i32 -> tensor<256xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> loc(#loc65)
    %rn_24 = arith.addi %rn_23, %rn_22 : tensor<256xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> loc(#loc65)
    %offs_a_m = arith.remsi %rm_20, %cst_7 {tt.contiguity = dense<128> : tensor<1xi32>, tt.divisibility = dense<128> : tensor<1xi32>} : tensor<128xi32, #ttg.slice<{dim = 1, parent = #blocked2}>> loc(#loc96)
    %offs_b_n = arith.remsi %rn_24, %cst_8 {tt.contiguity = dense<256> : tensor<1xi32>, tt.divisibility = dense<256> : tensor<1xi32>} : tensor<256xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> loc(#loc97)
    %a_k_idx_vals = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 0, parent = #blocked}>> loc(#loc69)
    %a_k_idx_vals_25 = tt.expand_dims %a_k_idx_vals {axis = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 0, parent = #blocked}>> -> tensor<1x64xi32, #blocked> loc(#loc69)
    %b_k_idx_vals = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked1}>> loc(#loc70)
    %b_k_idx_vals_26 = tt.expand_dims %b_k_idx_vals {axis = 1 : i32} : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked1}>> -> tensor<64x1xi32, #blocked1> loc(#loc70)
    %idx_m = tt.expand_dims %offs_a_m {axis = 1 : i32} : tensor<128xi32, #ttg.slice<{dim = 1, parent = #blocked2}>> -> tensor<128x1xi32, #blocked2> loc(#loc71)
    %tmp0 = tt.splat %in_ptr1 : !tt.ptr<i64> -> tensor<128x1x!tt.ptr<i64>, #blocked2> loc(#loc72)
    %tmp0_27 = tt.addptr %tmp0, %idx_m : tensor<128x1x!tt.ptr<i64>, #blocked2>, tensor<128x1xi32, #blocked2> loc(#loc72)
    %tmp0_28 = tt.broadcast %tmp0_27 : tensor<128x1x!tt.ptr<i64>, #blocked2> -> tensor<128x64x!tt.ptr<i64>, #blocked2> loc(#loc72)
    %tmp6 = tt.splat %in_ptr2 : !tt.ptr<f16> -> tensor<128x64x!tt.ptr<f16>, #blocked> loc(#loc73)
    %idx_n = tt.expand_dims %offs_b_n {axis = 0 : i32} : tensor<256xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> -> tensor<1x256xi32, #blocked1> loc(#loc74)
    %xindex = tt.broadcast %idx_n : tensor<1x256xi32, #blocked1> -> tensor<64x256xi32, #blocked1> loc(#loc75)
    %b = tt.splat %arg_B : !tt.ptr<f16> -> tensor<64x256x!tt.ptr<f16>, #blocked1> loc(#loc76)
    %acc = ttng.tmem_alloc : () -> !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable> loc(#loc77)
    ttng.tmem_store %cst_6, %acc, %true : tensor<128x256xf32, #linear> -> !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable> loc(#loc77)
    %acc_29 = ttg.local_alloc : () -> !ttg.memdesc<2x1xi64, #shared, #smem, mutable> loc(#loc78)
    %acc_30 = ttg.memdesc_index %acc_29[%c0_i32] : !ttg.memdesc<2x1xi64, #shared, #smem, mutable> -> !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc78)
    ttng.init_barrier %acc_30, 1 : !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc78)
    %acc_31 = ttg.memdesc_index %acc_29[%c1_i32] : !ttg.memdesc<2x1xi64, #shared, #smem, mutable> -> !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc78)
    ttng.init_barrier %acc_31, 1 : !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc78)
    %tmp0_32 = ttg.local_alloc : () -> !ttg.memdesc<1x128x64xi64, #shared1, #smem, mutable> loc(#loc79)
    %tmp6_33 = ttg.local_alloc : () -> !ttg.memdesc<2x128x64xf16, #shared2, #smem, mutable> loc(#loc80)
    %b_34 = ttg.local_alloc : () -> !ttg.memdesc<2x64x256xf16, #shared2, #smem, mutable> loc(#loc81)
    %tmp0_35 = ttg.memdesc_index %tmp0_32[%c0_i32] : !ttg.memdesc<1x128x64xi64, #shared1, #smem, mutable> -> !ttg.memdesc<128x64xi64, #shared1, #smem, mutable> loc(#loc79)
    %tmp0_36 = ttg.async_copy_global_to_local %tmp0_28, %tmp0_35 mask %cst_5 evictionPolicy = evict_last {contiguity = 2 : i32} : tensor<128x64x!tt.ptr<i64>, #blocked2> -> <128x64xi64, #shared1, #smem, mutable> loc(#loc79)
    %tmp0_37 = ttg.async_commit_group tokens %tmp0_36 loc(#loc79)
    %tmp0_38 = ttg.async_wait %tmp0_37 {num = 0 : i32} loc(#loc79)
    %tmp0_39 = ttg.local_load %tmp0_35 token %tmp0_38 : !ttg.memdesc<128x64xi64, #shared1, #smem, mutable> -> tensor<128x64xi64, #blocked> loc(#loc79)
    %tmp2 = arith.addi %tmp0_39, %cst_0 : tensor<128x64xi64, #blocked> loc(#loc82)
    %tmp3 = arith.cmpi slt, %tmp0_39, %cst : tensor<128x64xi64, #blocked> loc(#loc83)
    %tmp4 = arith.select %tmp3, %tmp2, %tmp0_39 : tensor<128x64xi1, #blocked>, tensor<128x64xi64, #blocked> loc(#loc84)
    %tmp6_40 = arith.muli %tmp4, %cst_0 : tensor<128x64xi64, #blocked> loc(#loc85)
    %tmp6_41 = arith.extsi %a_k_idx_vals_25 : tensor<1x64xi32, #blocked> to tensor<1x64xi64, #blocked> loc(#loc86)
    %tmp6_42 = tt.broadcast %tmp6_41 : tensor<1x64xi64, #blocked> -> tensor<128x64xi64, #blocked> loc(#loc86)
    %tmp6_43 = arith.addi %tmp6_42, %tmp6_40 : tensor<128x64xi64, #blocked> loc(#loc86)
    %tmp6_44 = tt.addptr %tmp6, %tmp6_43 : tensor<128x64x!tt.ptr<f16>, #blocked>, tensor<128x64xi64, #blocked> loc(#loc73)
    %tmp6_45 = ttg.memdesc_index %tmp6_33[%c0_i32] : !ttg.memdesc<2x128x64xf16, #shared2, #smem, mutable> -> !ttg.memdesc<128x64xf16, #shared2, #smem, mutable> loc(#loc80)
    %tmp6_46 = ttg.async_copy_global_to_local %tmp6_44, %tmp6_45 mask %cst_4 evictionPolicy = evict_last {contiguity = 8 : i32} : tensor<128x64x!tt.ptr<f16>, #blocked> -> <128x64xf16, #shared2, #smem, mutable> loc(#loc80)
    %tmp6_47 = ttg.async_commit_group tokens %tmp6_46 loc(#loc80)
    %xindex_48 = arith.muli %b_k_idx_vals_26, %cst_10 : tensor<64x1xi32, #blocked1> loc(#loc87)
    %xindex_49 = tt.broadcast %xindex_48 : tensor<64x1xi32, #blocked1> -> tensor<64x256xi32, #blocked1> loc(#loc75)
    %xindex_50 = arith.addi %xindex, %xindex_49 : tensor<64x256xi32, #blocked1> loc(#loc75)
    %b_51 = tt.addptr %b, %xindex_50 : tensor<64x256x!tt.ptr<f16>, #blocked1>, tensor<64x256xi32, #blocked1> loc(#loc76)
    %b_52 = ttg.memdesc_index %b_34[%c0_i32] : !ttg.memdesc<2x64x256xf16, #shared2, #smem, mutable> -> !ttg.memdesc<64x256xf16, #shared2, #smem, mutable> loc(#loc81)
    %b_53 = ttg.async_copy_global_to_local %b_51, %b_52 mask %cst_3 {contiguity = 8 : i32} : tensor<64x256x!tt.ptr<f16>, #blocked1> -> <64x256xf16, #shared2, #smem, mutable> loc(#loc81)
    %b_54 = ttg.async_commit_group tokens %b_53 loc(#loc81)
    %tmp0_55 = ttg.async_copy_global_to_local %tmp0_28, %tmp0_35 mask %cst_5 evictionPolicy = evict_last {contiguity = 2 : i32} : tensor<128x64x!tt.ptr<i64>, #blocked2> -> <128x64xi64, #shared1, #smem, mutable> loc(#loc79)
    %tmp0_56 = ttg.async_commit_group tokens %tmp0_55 loc(#loc79)
    %tmp6_57 = ttg.async_wait %tmp6_47, %b_54, %tmp0_56 {num = 0 : i32} loc(#loc80)
    ttng.tc_gen5_mma %tmp6_45, %b_52, %acc, %false, %true, %acc_30[%true] {is_async} : !ttg.memdesc<128x64xf16, #shared2, #smem, mutable>, !ttg.memdesc<64x256xf16, #shared2, #smem, mutable>, !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>, !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc77)
    %a_k_idx_vals_58 = arith.addi %a_k_idx_vals_25, %cst_1 : tensor<1x64xi32, #blocked> loc(#loc69)
    %b_k_idx_vals_59 = arith.addi %b_k_idx_vals_26, %cst_2 : tensor<64x1xi32, #blocked1> loc(#loc70)
    %tmp0_60 = ttg.local_load %tmp0_35 token %tmp6_57 : !ttg.memdesc<128x64xi64, #shared1, #smem, mutable> -> tensor<128x64xi64, #blocked> loc(#loc79)
    %tmp2_61 = arith.addi %tmp0_60, %cst_0 : tensor<128x64xi64, #blocked> loc(#loc82)
    %tmp3_62 = arith.cmpi slt, %tmp0_60, %cst : tensor<128x64xi64, #blocked> loc(#loc83)
    %tmp4_63 = arith.select %tmp3_62, %tmp2_61, %tmp0_60 : tensor<128x64xi1, #blocked>, tensor<128x64xi64, #blocked> loc(#loc84)
    %tmp6_64 = arith.muli %tmp4_63, %cst_0 : tensor<128x64xi64, #blocked> loc(#loc85)
    %tmp6_65 = arith.extsi %a_k_idx_vals_58 : tensor<1x64xi32, #blocked> to tensor<1x64xi64, #blocked> loc(#loc86)
    %tmp6_66 = tt.broadcast %tmp6_65 : tensor<1x64xi64, #blocked> -> tensor<128x64xi64, #blocked> loc(#loc86)
    %tmp6_67 = arith.addi %tmp6_66, %tmp6_64 : tensor<128x64xi64, #blocked> loc(#loc86)
    %tmp6_68 = tt.addptr %tmp6, %tmp6_67 : tensor<128x64x!tt.ptr<f16>, #blocked>, tensor<128x64xi64, #blocked> loc(#loc73)
    %tmp6_69 = ttg.memdesc_index %tmp6_33[%c1_i32] : !ttg.memdesc<2x128x64xf16, #shared2, #smem, mutable> -> !ttg.memdesc<128x64xf16, #shared2, #smem, mutable> loc(#loc80)
    %tmp6_70 = ttg.async_copy_global_to_local %tmp6_68, %tmp6_69 mask %cst_4 evictionPolicy = evict_last {contiguity = 8 : i32} : tensor<128x64x!tt.ptr<f16>, #blocked> -> <128x64xf16, #shared2, #smem, mutable> loc(#loc80)
    %tmp6_71 = ttg.async_commit_group tokens %tmp6_70 loc(#loc80)
    %xindex_72 = arith.muli %b_k_idx_vals_59, %cst_10 : tensor<64x1xi32, #blocked1> loc(#loc87)
    %xindex_73 = tt.broadcast %xindex_72 : tensor<64x1xi32, #blocked1> -> tensor<64x256xi32, #blocked1> loc(#loc75)
    %xindex_74 = arith.addi %xindex, %xindex_73 : tensor<64x256xi32, #blocked1> loc(#loc75)
    %b_75 = tt.addptr %b, %xindex_74 : tensor<64x256x!tt.ptr<f16>, #blocked1>, tensor<64x256xi32, #blocked1> loc(#loc76)
    %b_76 = ttg.memdesc_index %b_34[%c1_i32] : !ttg.memdesc<2x64x256xf16, #shared2, #smem, mutable> -> !ttg.memdesc<64x256xf16, #shared2, #smem, mutable> loc(#loc81)
    %b_77 = ttg.async_copy_global_to_local %b_75, %b_76 mask %cst_3 {contiguity = 8 : i32} : tensor<64x256x!tt.ptr<f16>, #blocked1> -> <64x256xf16, #shared2, #smem, mutable> loc(#loc81)
    %b_78 = ttg.async_commit_group tokens %b_77 loc(#loc81)
    %tmp0_79 = ttg.async_copy_global_to_local %tmp0_28, %tmp0_35 mask %cst_5 evictionPolicy = evict_last {contiguity = 2 : i32} : tensor<128x64x!tt.ptr<i64>, #blocked2> -> <128x64xi64, #shared1, #smem, mutable> loc(#loc79)
    %tmp0_80 = ttg.async_commit_group tokens %tmp0_79 loc(#loc79)
    %acc_81:10 = scf.for %acc_94 = %c0_i32 to %c31_i32 step %c1_i32 iter_args(%acc_95 = %c0_i32, %acc_96 = %c1_i32, %acc_97 = %c0_i32, %tmp6_98 = %tmp6_71, %b_99 = %b_78, %acc_100 = %acc_30, %acc_101 = %c0_i32, %tmp6_102 = %tmp6_45, %b_103 = %b_52, %tmp0_104 = %tmp0_80) -> (i32, i32, i32, !ttg.async.token, !ttg.async.token, !ttg.memdesc<1xi64, #shared, #smem, mutable>, i32, !ttg.memdesc<128x64xf16, #shared2, #smem, mutable>, !ttg.memdesc<64x256xf16, #shared2, #smem, mutable>, !ttg.async.token)  : i32 {
      %acc_105 = arith.cmpi slt, %acc_94, %c29_i32 : i32 loc(#loc78)
      %acc_106 = arith.cmpi slt, %acc_94, %c30_i32 : i32 loc(#loc78)
      %acc_107 = arith.addi %acc_97, %c1_i32 : i32 loc(#loc78)
      %acc_108 = arith.cmpi sge, %acc_107, %c2_i32 : i32 loc(#loc78)
      %acc_109 = arith.select %acc_108, %c0_i32, %acc_107 : i32 loc(#loc78)
      %tmp6_110 = ttg.async_wait %tmp6_98, %b_99, %tmp0_104 {num = 0 : i32} loc(#loc80)
      %tmp6_111 = ttg.memdesc_index %tmp6_33[%acc_109] : !ttg.memdesc<2x128x64xf16, #shared2, #smem, mutable> -> !ttg.memdesc<128x64xf16, #shared2, #smem, mutable> loc(#loc80)
      %b_112 = ttg.memdesc_index %b_34[%acc_109] : !ttg.memdesc<2x64x256xf16, #shared2, #smem, mutable> -> !ttg.memdesc<64x256xf16, #shared2, #smem, mutable> loc(#loc81)
      %acc_113 = ttg.memdesc_index %acc_29[%acc_96] : !ttg.memdesc<2x1xi64, #shared, #smem, mutable> -> !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc78)
      ttng.tc_gen5_mma %tmp6_111, %b_112, %acc, %true, %true, %acc_113[%true] {is_async} : !ttg.memdesc<128x64xf16, #shared2, #smem, mutable>, !ttg.memdesc<64x256xf16, #shared2, #smem, mutable>, !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>, !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc77)
      ttng.wait_barrier %acc_100, %acc_101 deps %tmp6_102, %b_103 : !ttg.memdesc<1xi64, #shared, #smem, mutable>, !ttg.memdesc<128x64xf16, #shared2, #smem, mutable>, !ttg.memdesc<64x256xf16, #shared2, #smem, mutable> loc(#loc77)
      %acc_114 = arith.xori %acc_95, %c1_i32 : i32 loc(#loc77)
      %acc_115 = arith.addi %acc_96, %c1_i32 : i32 loc(#loc77)
      %acc_116 = arith.cmpi sge, %acc_115, %c2_i32 : i32 loc(#loc77)
      %acc_117 = arith.select %acc_116, %c0_i32, %acc_115 : i32 loc(#loc77)
      %acc_118 = arith.select %acc_116, %acc_114, %acc_95 : i32 loc(#loc77)
      %acc_119 = arith.addi %acc_94, %c2_i32 : i32 loc(#loc78)
      %a_k_idx_vals_120 = arith.muli %acc_119, %c64_i32 : i32 loc(#loc88)
      %a_k_idx_vals_121 = tt.splat %a_k_idx_vals_120 : i32 -> tensor<1x64xi32, #blocked> loc(#loc69)
      %a_k_idx_vals_122 = arith.addi %a_k_idx_vals_25, %a_k_idx_vals_121 : tensor<1x64xi32, #blocked> loc(#loc69)
      %b_k_idx_vals_123 = tt.splat %a_k_idx_vals_120 : i32 -> tensor<64x1xi32, #blocked1> loc(#loc70)
      %b_k_idx_vals_124 = arith.addi %b_k_idx_vals_26, %b_k_idx_vals_123 : tensor<64x1xi32, #blocked1> loc(#loc70)
      %tmp0_125 = ttg.local_load %tmp0_35 token %tmp6_110 : !ttg.memdesc<128x64xi64, #shared1, #smem, mutable> -> tensor<128x64xi64, #blocked> loc(#loc79)
      %tmp2_126 = arith.addi %tmp0_125, %cst_0 : tensor<128x64xi64, #blocked> loc(#loc82)
      %tmp3_127 = arith.cmpi slt, %tmp0_125, %cst : tensor<128x64xi64, #blocked> loc(#loc83)
      %tmp4_128 = arith.select %tmp3_127, %tmp2_126, %tmp0_125 : tensor<128x64xi1, #blocked>, tensor<128x64xi64, #blocked> loc(#loc84)
      %tmp6_129 = arith.muli %tmp4_128, %cst_0 : tensor<128x64xi64, #blocked> loc(#loc85)
      %tmp6_130 = arith.extsi %a_k_idx_vals_122 : tensor<1x64xi32, #blocked> to tensor<1x64xi64, #blocked> loc(#loc86)
      %tmp6_131 = tt.broadcast %tmp6_130 : tensor<1x64xi64, #blocked> -> tensor<128x64xi64, #blocked> loc(#loc86)
      %tmp6_132 = arith.addi %tmp6_131, %tmp6_129 : tensor<128x64xi64, #blocked> loc(#loc86)
      %tmp6_133 = tt.addptr %tmp6, %tmp6_132 : tensor<128x64x!tt.ptr<f16>, #blocked>, tensor<128x64xi64, #blocked> loc(#loc73)
      %tmp6_134 = ttg.memdesc_index %tmp6_33[%acc_117] : !ttg.memdesc<2x128x64xf16, #shared2, #smem, mutable> -> !ttg.memdesc<128x64xf16, #shared2, #smem, mutable> loc(#loc80)
      %acc_135 = tt.splat %acc_106 : i1 -> tensor<128x64xi1, #blocked> loc(#loc78)
      %tmp6_136 = ttg.async_copy_global_to_local %tmp6_133, %tmp6_134 mask %acc_135 evictionPolicy = evict_last {contiguity = 8 : i32} : tensor<128x64x!tt.ptr<f16>, #blocked> -> <128x64xf16, #shared2, #smem, mutable> loc(#loc80)
      %tmp6_137 = ttg.async_commit_group tokens %tmp6_136 loc(#loc80)
      %xindex_138 = arith.muli %b_k_idx_vals_124, %cst_10 : tensor<64x1xi32, #blocked1> loc(#loc87)
      %xindex_139 = tt.broadcast %xindex_138 : tensor<64x1xi32, #blocked1> -> tensor<64x256xi32, #blocked1> loc(#loc75)
      %xindex_140 = arith.addi %xindex, %xindex_139 : tensor<64x256xi32, #blocked1> loc(#loc75)
      %b_141 = tt.addptr %b, %xindex_140 : tensor<64x256x!tt.ptr<f16>, #blocked1>, tensor<64x256xi32, #blocked1> loc(#loc76)
      %b_142 = ttg.memdesc_index %b_34[%acc_117] : !ttg.memdesc<2x64x256xf16, #shared2, #smem, mutable> -> !ttg.memdesc<64x256xf16, #shared2, #smem, mutable> loc(#loc81)
      %acc_143 = tt.splat %acc_106 : i1 -> tensor<64x256xi1, #blocked1> loc(#loc78)
      %b_144 = ttg.async_copy_global_to_local %b_141, %b_142 mask %acc_143 {contiguity = 8 : i32} : tensor<64x256x!tt.ptr<f16>, #blocked1> -> <64x256xf16, #shared2, #smem, mutable> loc(#loc81)
      %b_145 = ttg.async_commit_group tokens %b_144 loc(#loc81)
      %acc_146 = tt.splat %acc_105 : i1 -> tensor<128x64xi1, #blocked2> loc(#loc78)
      %tmp0_147 = ttg.async_copy_global_to_local %tmp0_28, %tmp0_35 mask %acc_146 evictionPolicy = evict_last {contiguity = 2 : i32} : tensor<128x64x!tt.ptr<i64>, #blocked2> -> <128x64xi64, #shared1, #smem, mutable> loc(#loc79)
      %tmp0_148 = ttg.async_commit_group tokens %tmp0_147 loc(#loc79)
      scf.yield %acc_118, %acc_117, %acc_109, %tmp6_137, %b_145, %acc_113, %acc_95, %tmp6_111, %b_112, %tmp0_148 : i32, i32, i32, !ttg.async.token, !ttg.async.token, !ttg.memdesc<1xi64, #shared, #smem, mutable>, i32, !ttg.memdesc<128x64xf16, #shared2, #smem, mutable>, !ttg.memdesc<64x256xf16, #shared2, #smem, mutable>, !ttg.async.token loc(#loc78)
    } loc(#loc78)
    ttng.wait_barrier %acc_81#5, %acc_81#6 deps %acc_81#7, %acc_81#8 : !ttg.memdesc<1xi64, #shared, #smem, mutable>, !ttg.memdesc<128x64xf16, #shared2, #smem, mutable>, !ttg.memdesc<64x256xf16, #shared2, #smem, mutable> loc(#loc77)
    %acc_82 = ttg.async_wait {num = 0 : i32} loc(#loc78)
    ttg.local_dealloc %b_34 : !ttg.memdesc<2x64x256xf16, #shared2, #smem, mutable> loc(#loc78)
    ttg.local_dealloc %tmp6_33 : !ttg.memdesc<2x128x64xf16, #shared2, #smem, mutable> loc(#loc78)
    ttg.local_dealloc %tmp0_32 : !ttg.memdesc<1x128x64xi64, #shared1, #smem, mutable> loc(#loc78)
    ttng.inval_barrier %acc_30 : !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc78)
    ttng.inval_barrier %acc_31 : !ttg.memdesc<1xi64, #shared, #smem, mutable> loc(#loc78)
    ttg.local_dealloc %acc_29 : !ttg.memdesc<2x1xi64, #shared, #smem, mutable> loc(#loc78)
    %idx_m_83 = tt.expand_dims %rm_21 {axis = 1 : i32} : tensor<128xi32, #ttg.slice<{dim = 1, parent = #blocked1}>> -> tensor<128x1xi32, #blocked1> loc(#loc89)
    %idx_n_84 = tt.expand_dims %rn_24 {axis = 0 : i32} : tensor<256xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> -> tensor<1x256xi32, #blocked1> loc(#loc90)
    %mask = arith.cmpi slt, %idx_m_83, %cst_9 : tensor<128x1xi32, #blocked1> loc(#loc91)
    %mask_85 = arith.cmpi slt, %idx_n_84, %cst_11 : tensor<1x256xi32, #blocked1> loc(#loc92)
    %mask_86 = tt.broadcast %mask : tensor<128x1xi1, #blocked1> -> tensor<128x256xi1, #blocked1> loc(#loc93)
    %mask_87 = tt.broadcast %mask_85 : tensor<1x256xi1, #blocked1> -> tensor<128x256xi1, #blocked1> loc(#loc93)
    %mask_88 = arith.andi %mask_86, %mask_87 : tensor<128x256xi1, #blocked1> loc(#loc93)
    %xindex_89 = arith.muli %idx_m_83, %cst_9 : tensor<128x1xi32, #blocked1> loc(#loc94)
    %xindex_90 = tt.broadcast %idx_n_84 : tensor<1x256xi32, #blocked1> -> tensor<128x256xi32, #blocked1> loc(#loc95)
    %xindex_91 = tt.broadcast %xindex_89 : tensor<128x1xi32, #blocked1> -> tensor<128x256xi32, #blocked1> loc(#loc95)
    %xindex_92 = arith.addi %xindex_90, %xindex_91 : tensor<128x256xi32, #blocked1> loc(#loc95)
    %2 = tt.splat %out_ptr0 : !tt.ptr<f16> -> tensor<128x256x!tt.ptr<f16>, #blocked1> loc(#loc48)
    %3 = tt.addptr %2, %xindex_92 : tensor<128x256x!tt.ptr<f16>, #blocked1>, tensor<128x256xi32, #blocked1> loc(#loc48)
    %acc_93 = ttng.tmem_load %acc : !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable> -> tensor<128x256xf32, #linear> loc(#loc77)
    %4 = arith.truncf %acc_93 : tensor<128x256xf32, #linear> to tensor<128x256xf16, #linear> loc(#loc49)
    %5 = ttg.convert_layout %4 : tensor<128x256xf16, #linear> -> tensor<128x256xf16, #blocked1> loc(#loc49)
    tt.store %3, %5, %mask_88 : tensor<128x256x!tt.ptr<f16>, #blocked1> loc(#loc49)
    tt.return loc(#loc)
  } loc(#loc)
} loc(#loc)
#loc1 = loc(unknown)
#loc2 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":52:11)
#loc3 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":58:16)
#loc4 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":59:31)
#loc5 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":59:22)
#loc6 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":59:18)
#loc7 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":60:35)
#loc8 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":60:13)
#loc9 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":61:14)
#loc10 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":61:13)
#loc11 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":62:15)
#loc12 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":62:5)
#loc13 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":63:15)
#loc14 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":63:5)
#loc15 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":65:10)
#loc16 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":65:28)
#loc17 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":66:10)
#loc18 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":66:28)
#loc19 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":68:53)
#loc20 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":72:53)
#loc21 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":80:24)
#loc22 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":81:24)
#loc23 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":83:17)
#loc24 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":86:24)
#loc25 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":92:24)
#loc26 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":96:17)
#loc27 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":97:18)
#loc28 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":98:21)
#loc29 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":100:16)
#loc30 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":78:5)
#loc31 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":86:16)
#loc32 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":92:16)
#loc33 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":98:13)
#loc34 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":88:16)
#loc35 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":89:16)
#loc36 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":90:16)
#loc37 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":92:59)
#loc38 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":92:51)
#loc39 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":97:26)
#loc40 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":80:43)
#loc41 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":105:13)
#loc42 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":106:13)
#loc43 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":107:13)
#loc44 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":107:27)
#loc45 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":107:12)
#loc46 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":110:22)
#loc47 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":110:14)
#loc48 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":111:14)
#loc49 = loc("/home/shunting/ws/openai_triton/licm-for-inductor/standalone_bench.py":111:5)
#loc54 = loc("pid"(#loc2))
#loc55 = loc("group_id"(#loc3))
#loc56 = loc("group_size"(#loc4))
#loc57 = loc("group_size"(#loc5))
#loc58 = loc("group_size"(#loc6))
#loc59 = loc("pid_m"(#loc7))
#loc60 = loc("pid_m"(#loc8))
#loc61 = loc("pid_n"(#loc9))
#loc62 = loc("pid_n"(#loc10))
#loc63 = loc("rm"(#loc15))
#loc64 = loc("rm"(#loc16))
#loc65 = loc("rn"(#loc17))
#loc66 = loc("rn"(#loc18))
#loc67 = loc("offs_a_m"(#loc19))
#loc68 = loc("offs_b_n"(#loc20))
#loc69 = loc("a_k_idx_vals"(#loc21))
#loc70 = loc("b_k_idx_vals"(#loc22))
#loc71 = loc("idx_m"(#loc23))
#loc72 = loc("tmp0"(#loc24))
#loc73 = loc("tmp6"(#loc25))
#loc74 = loc("idx_n"(#loc26))
#loc75 = loc("xindex"(#loc27))
#loc76 = loc("b"(#loc28))
#loc77 = loc("acc"(#loc29))
#loc78 = loc("acc"(#loc30))
#loc79 = loc("tmp0"(#loc31))
#loc80 = loc("tmp6"(#loc32))
#loc81 = loc("b"(#loc33))
#loc82 = loc("tmp2"(#loc34))
#loc83 = loc("tmp3"(#loc35))
#loc84 = loc("tmp4"(#loc36))
#loc85 = loc("tmp6"(#loc37))
#loc86 = loc("tmp6"(#loc38))
#loc87 = loc("xindex"(#loc39))
#loc88 = loc("a_k_idx_vals"(#loc40))
#loc89 = loc("idx_m"(#loc41))
#loc90 = loc("idx_n"(#loc42))
#loc91 = loc("mask"(#loc43))
#loc92 = loc("mask"(#loc44))
#loc93 = loc("mask"(#loc45))
#loc94 = loc("xindex"(#loc46))
#loc95 = loc("xindex"(#loc47))
#loc96 = loc("offs_a_m"(#loc67))
#loc97 = loc("offs_b_n"(#loc68))