// -----// IR Dump Before ConvertToNVVMPass: iree-convert-to-nvvm //----- // module { func.func @main_graph$async_dispatch_0_elementwise_broadcast_Dx3584_i64xf16() attributes {gpu.known_block_size = array} { %c3584 = arith.constant 3584 : index %c256_i32 = arith.constant 256 : i32 %c8_i32 = arith.constant 8 : i32 %c14 = arith.constant 14 : index %cst = arith.constant dense<0.000000e+00> : vector<8xf16> %cst_0 = arith.constant dense : vector<8xi1> %c0 = arith.constant 0 : index %c32_i64 = arith.constant 32 : i64 %cst_1 = arith.constant dense<152064> : vector<1xi64> %cst_2 = arith.constant dense<0> : vector<1xi64> %thread_id_x = gpu.thread_id x upper_bound 32 %0 = hal.interface.constant.load layout(, #hal.pipeline.binding, #hal.pipeline.binding], flags = Indirect>) ordinal(0) : i32 %1 = hal.interface.constant.load layout(, #hal.pipeline.binding, #hal.pipeline.binding], flags = Indirect>) ordinal(1) : i32 %2 = arith.extui %0 : i32 to i64 %3 = arith.extui %1 : i32 to i64 %4 = arith.shli %3, %c32_i64 : i64 %5 = arith.ori %2, %4 : i64 %6 = arith.index_castui %5 : i64 to index %7 = util.assume.int %6 : index %8 = hal.interface.binding.subspan layout(, #hal.pipeline.binding, #hal.pipeline.binding], flags = Indirect>) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : memref<152064x3584xf16, #gpu.address_space> %assume_align = memref.assume_alignment %8, 64 : memref<152064x3584xf16, #gpu.address_space> %9 = hal.interface.binding.subspan layout(, #hal.pipeline.binding, #hal.pipeline.binding], flags = Indirect>) binding(1) alignment(64) offset(%c0) flags("ReadOnly|Indirect") : memref>{%7} %assume_align_3 = memref.assume_alignment %9, 64 : memref> %10 = hal.interface.binding.subspan layout(, #hal.pipeline.binding, #hal.pipeline.binding], flags = Indirect>) binding(2) alignment(64) offset(%c0) flags(Indirect) : memref>{%7} %assume_align_4 = memref.assume_alignment %10, 64 : memref> %11 = vector.step : vector<8xindex> %12 = arith.index_castui %thread_id_x : index to i32 %13 = arith.muli %12, %c8_i32 overflow : i32 %14 = arith.muli %7, %c14 overflow : index %workgroup_id_x = hal.interface.workgroup.id[0] upper_bound 2147483647 : index %workgroup_count_x = hal.interface.workgroup.count[0] upper_bound 2147483647 : index cf.br ^bb1(%workgroup_id_x : index) ^bb1(%15: index): // 2 preds: ^bb0, ^bb2 %16 = arith.cmpi slt, %15, %14 : index cf.cond_br %16, ^bb2, ^bb3 ^bb2: // pred: ^bb1 %17 = arith.divui %15, %c14 : index %18 = arith.remui %15, %c14 : index %19 = vector.load %assume_align_3[%17] : memref>, vector<1xi64> %20 = arith.index_castui %18 : index to i32 %21 = arith.muli %20, %c256_i32 overflow : i32 %22 = arith.addi %21, %13 : i32 %23 = vector.broadcast %22 : i32 to vector<8xi32> %24 = arith.index_castui %11 : vector<8xindex> to vector<8xi32> %25 = arith.addi %23, %24 : vector<8xi32> %26 = arith.index_castui %25 : vector<8xi32> to vector<8xindex> %27 = arith.addi %19, %cst_1 : vector<1xi64> %28 = arith.cmpi slt, %19, %cst_2 : vector<1xi64> %29 = arith.select %28, %27, %19 : vector<1xi1>, vector<1xi64> %30 = vector.extract %29[0] : i64 from vector<1xi64> %31 = arith.index_cast %30 : i64 to index %32 = arith.muli %31, %c3584 : index %33 = vector.broadcast %32 : index to vector<8xindex> %34 = arith.addi %26, %33 : vector<8xindex> %35 = vector.gather %assume_align[%c0, %c0] [%34], %cst_0, %cst : memref<152064x3584xf16, #gpu.address_space>, vector<8xindex>, vector<8xi1>, vector<8xf16> into vector<8xf16> %36 = arith.index_castui %22 : i32 to index vector.store %35, %assume_align_4[%17, %36] : memref>, vector<8xf16> %37 = arith.addi %15, %workgroup_count_x : index cf.br ^bb1(%37 : index) ^bb3: // pred: ^bb1 return } iree_codegen.dispatch_config @main_graph$async_dispatch_0_elementwise_broadcast_Dx3584_i64xf16 workgroup_size = [32, 1, 1] subgroup_size = 32 { ^bb0(%arg0: !hal.device, %arg1: index): %c1 = arith.constant 1 : index %0 = affine.min affine_map<()[s0] -> (2147483647, s0 * 14)>()[%arg1] iree_codegen.yield %0, %c1, %c1 : index, index, index } }