func.func @main_dispatch_18_matmul_2304x144x24_dispatch_0_matmul_2304x144x24_dispatch_0_matmul_2304x144x24() { %c8 = arith.constant 8 : index %c32 = arith.constant 32 : index %c256 = arith.constant 256 : index %c72 = arith.constant 72 : index %c144 = arith.constant 144 : index %c2304 = arith.constant 2304 : index %c0 = arith.constant 0 : index %c12_i32 = arith.constant 12 : i32 %c-128_i32 = arith.constant -128 : i32 %c127_i32 = arith.constant 127 : i32 %c0_i32 = arith.constant 0 : i32 %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor> %1 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor> %2 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor> %3 = hal.interface.binding.subspan set(0) binding(3) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor> %4 = hal.interface.binding.subspan set(0) binding(4) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor> %5 = hal.interface.binding.subspan set(0) binding(5) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor> %6 = hal.interface.binding.subspan set(0) binding(6) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor> %workgroup_id_x = hal.interface.workgroup.id[0] : index %workgroup_count_x = hal.interface.workgroup.count[0] : index %workgroup_id_y = hal.interface.workgroup.id[1] : index %workgroup_count_y = hal.interface.workgroup.count[1] : index %7 = affine.apply affine_map<()[s0] -> (s0 * 256)>()[%workgroup_id_y] %8 = affine.apply affine_map<()[s0] -> (s0 * 256)>()[%workgroup_count_y] %9 = affine.apply affine_map<()[s0] -> (s0 * 72)>()[%workgroup_id_x] %10 = affine.apply affine_map<()[s0] -> (s0 * 72)>()[%workgroup_count_x] scf.for %arg0 = %7 to %c2304 step %8 { %11 = flow.dispatch.tensor.load %0, offsets = [%arg0, 0], sizes = [256, 24], strides = [1, 1] : !flow.dispatch.tensor> -> tensor<256x24xi8> scf.for %arg1 = %9 to %c144 step %10 { %12 = flow.dispatch.tensor.load %6, offsets = [%arg0, %arg1], sizes = [256, 72], strides = [1, 1] : !flow.dispatch.tensor> -> tensor<256x72xi8> %13 = flow.dispatch.tensor.load %2, offsets = [%arg1], sizes = [72], strides = [1] : !flow.dispatch.tensor> -> tensor<72xi32> %14 = flow.dispatch.tensor.load %1, offsets = [0, %arg1], sizes = [24, 72], strides = [1, 1] : !flow.dispatch.tensor> -> tensor<24x72xi8> %15 = flow.dispatch.tensor.load %3, offsets = [%arg1], sizes = [72], strides = [1] : !flow.dispatch.tensor> -> tensor<72xi32> %16 = flow.dispatch.tensor.load %4, offsets = [%arg1], sizes = [72], strides = [1] : !flow.dispatch.tensor> -> tensor<72xi32> %17 = flow.dispatch.tensor.load %5, offsets = [%arg1], sizes = [72], strides = [1] : !flow.dispatch.tensor> -> tensor<72xi8> %18 = scf.for %arg2 = %c0 to %c256 step %c8 iter_args(%arg3 = %12) -> (tensor<256x72xi8>) { %extracted_slice = tensor.extract_slice %11[%arg2, 0] [8, 24] [1, 1] : tensor<256x24xi8> to tensor<8x24xi8> %19 = scf.for %arg4 = %c0 to %c72 step %c32 iter_args(%arg5 = %arg3) -> (tensor<256x72xi8>) { %20 = affine.min affine_map<(d0) -> (-d0 + 72, 32)>(%arg4) %extracted_slice_0 = tensor.extract_slice %13[%arg4] [%20] [1] : tensor<72xi32> to tensor %extracted_slice_1 = tensor.extract_slice %14[0, %arg4] [24, %20] [1, 1] : tensor<24x72xi8> to tensor<24x?xi8> %21 = tensor.empty(%20) : tensor<8x?xi32> %22 = linalg.fill ins(%c0_i32 : i32) outs(%21 : tensor<8x?xi32>) -> tensor<8x?xi32> %23 = linalg.matmul {lowering_config = #iree_codegen.lowering_config} ins(%extracted_slice, %extracted_slice_1 : tensor<8x24xi8>, tensor<24x?xi8>) outs(%22 : tensor<8x?xi32>) -> tensor<8x?xi32> %extracted_slice_2 = tensor.extract_slice %15[%arg4] [%20] [1] : tensor<72xi32> to tensor %extracted_slice_3 = tensor.extract_slice %16[%arg4] [%20] [1] : tensor<72xi32> to tensor %extracted_slice_4 = tensor.extract_slice %17[%arg4] [%20] [1] : tensor<72xi8> to tensor %extracted_slice_5 = tensor.extract_slice %arg5[%arg2, %arg4] [8, %20] [1, 1] : tensor<256x72xi8> to tensor<8x?xi8> %24 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%extracted_slice_0, %23, %extracted_slice_2, %extracted_slice_3, %extracted_slice_4 : tensor, tensor<8x?xi32>, tensor, tensor, tensor) outs(%extracted_slice_5 : tensor<8x?xi8>) { ^bb0(%in: i32, %in_6: i32, %in_7: i32, %in_8: i32, %in_9: i8, %out: i8): %25 = arith.muli %in_7, %c12_i32 : i32 %26 = arith.subi %in_6, %25 : i32 %27 = arith.addi %in, %26 : i32 %28 = "tosa.apply_scale"(%27, %in_8, %in_9) {double_round = true} : (i32, i32, i8) -> i32 %29 = arith.addi %28, %c-128_i32 : i32 %30 = arith.cmpi slt, %29, %c-128_i32 : i32 %31 = arith.select %30, %c-128_i32, %29 : i32 %32 = arith.cmpi sgt, %29, %c127_i32 : i32 %33 = arith.select %32, %c127_i32, %31 : i32 %34 = arith.trunci %33 : i32 to i8 linalg.yield %34 : i8 } -> tensor<8x?xi8> %inserted_slice = tensor.insert_slice %24 into %arg5[%arg2, %arg4] [8, %20] [1, 1] : tensor<8x?xi8> into tensor<256x72xi8> scf.yield %inserted_slice : tensor<256x72xi8> } scf.yield %19 : tensor<256x72xi8> } flow.dispatch.tensor.store %18, %6, offsets = [%arg0, %arg1], sizes = [256, 72], strides = [1, 1] : tensor<256x72xi8> -> !flow.dispatch.tensor> } } return }