SplitK configs are never returned

[pfultz2]

When tuning this with mlir, it never returns splitk configs to tune:

  func.func @mlir_convolution_broadcast_add_relu(%arg0: !migraphx.shaped<1x512x7x7xf32, 25088x49x7x1>, %arg1: !migraphx.shaped<512x512x3x3xf32, 4608x9x3x1>, %arg2: !migraphx.shaped<512xf32, 1>) -> !migraphx.shaped<1x512x7x7xf32, 25088x49x7x1> attributes {arch = "", enable_splitk_for_tuning, kernel = "mixr", num_cu = 0 : i64} {
    %0 = migraphx.convolution %arg0, %arg1 {dilation = [1, 1], group = 1 : i64, padding = [1, 1, 1, 1], padding_mode = 0 : i64, stride = [1, 1]} : <1x512x7x7xf32, 25088x49x7x1>, <512x512x3x3xf32, 4608x9x3x1> -> <1x512x7x7xf32, 25088x49x7x1>
    %1 = migraphx.broadcast %arg2 {axis = 1 : i64, out_lens = [1, 512, 7, 7]} : <512xf32, 1> -> <1x512x7x7xf32, 0x1x0x0>
    %2 = migraphx.add %0, %1 : <1x512x7x7xf32, 25088x49x7x1>, <1x512x7x7xf32, 0x1x0x0> -> <1x512x7x7xf32, 25088x49x7x1>
    %3 = migraphx.relu %2 : <1x512x7x7xf32, 25088x49x7x1> -> <1x512x7x7xf32, 25088x49x7x1>
    return %3 : !migraphx.shaped<1x512x7x7xf32, 25088x49x7x1>
  }
}

However, if I remove the relu at the end it will return splitk configs:

module {
  func.func @mlir_convolution_broadcast_add(%arg0: !migraphx.shaped<1x512x7x7xf32, 25088x49x7x1>, %arg1: !migraphx.shaped<512x512x3x3xf32, 4608x9x3x1>, %arg2: !migraphx.shaped<512xf32, 1>) -> !migraphx.shaped<1x512x7x7xf32, 25088x49x7x1> attributes {arch = "", enable_splitk_for_tuning, kernel = "mixr", num_cu = 0 : i64} {
    %0 = migraphx.convolution %arg0, %arg1 {dilation = [1, 1], group = 1 : i64, padding = [1, 1, 1, 1], padding_mode = 0 : i64, stride = [1, 1]} : <1x512x7x7xf32, 25088x49x7x1>, <512x512x3x3xf32, 4608x9x3x1> -> <1x512x7x7xf32, 25088x49x7x1>
    %1 = migraphx.broadcast %arg2 {axis = 1 : i64, out_lens = [1, 512, 7, 7]} : <512xf32, 1> -> <1x512x7x7xf32, 0x1x0x0>
    %2 = migraphx.add %0, %1 : <1x512x7x7xf32, 25088x49x7x1>, <1x512x7x7xf32, 0x1x0x0> -> <1x512x7x7xf32, 25088x49x7x1>
    return %2 : !migraphx.shaped<1x512x7x7xf32, 25088x49x7x1>
  }
}

It should return splitk even with the relu, as migraphx will check the mlirIsModuleFusible and will split it into two kernels when that returns false(which is the case for the relu).

Here is the migraphx program for the above:

p = migraphx.program()
m = p.get_main_module()
x_0 = m.add_literal(migraphx.generate_argument(migraphx.shape(type="float_type", lens=[512]), 0))
x_1 = m.add_literal(migraphx.generate_argument(migraphx.shape(type="float_type", lens=[512, 512, 3, 3]), 1))
p_x0 = m.add_parameter("x0", migraphx.shape(type="float_type", lens=[1, 512, 7, 7]))
x_3 = m.add_instruction(migraphx.op("convolution", padding=[1,1,1,1]), [p_x0, x_1]) # migraphx.shape(type="float_type", lens=[1, 512, 7, 7])
x_4 = m.add_instruction(migraphx.op("broadcast", axis=1, out_lens=[1,512,7,7]), [x_0]) # migraphx.shape(type="float_type", lens=[1, 512, 7, 7], strides=[0, 1, 0, 0])
x_5 = m.add_instruction(migraphx.op("add"), [x_3, x_4]) # migraphx.shape(type="float_type", lens=[1, 512, 7, 7])
x_6 = m.add_instruction(migraphx.op("relu"), [x_5]) # migraphx.shape(type="float_type", lens=[1, 512, 7, 7])
m.add_return([x_6])