legalizing/converting mhlo.convolution to linalg.conv

[dinkdeep]

Is there a way to convert/legalize mhlo.convolution to linalg.convolution ?
I see /mlir-hlo/tests/hlo-legalize-to-lhlo.mlir supporting a) conversion from "mhlo.convolution" to "lmhlo.convolution"
and further "lmhlo.convolution" b) conversion to "linalg.conv" in /mlir-hlo/tests/lhlo-legalize-to-linalg.mlir.
The problem with conversion b) is that the linalg.conv is not inside the linalg.generic .

#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
module {
func @conv(%arg0: memref<3x5x5x3xf32>, %arg1: memref<2x2x3x4xf32>, %arg2: memref<3x5x5x4xf32>) {
%c0 = constant 0 : index
%0 = alloc() : memref<3x5x5x4xf32>
linalg.conv(%arg0, %arg1, %0) {dilations = [1, 2], padding = dense<[[0, 1], [0, 1]]> : tensor<2x2xi64>, strides = [2, 1]} : memref<3x5x5x3xf32>, memref<2x2x3x4xf32>, memref<3x5x5x4xf32>
linalg.conv(%arg0, %arg1, %0) {dilations = [1, 1], strides = [2, 1]} : memref<3x5x5x3xf32>, memref<2x2x3x4xf32>, memref<3x5x5x4xf32>
linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%0 : memref<3x5x5x4xf32>) outs(%arg2 : memref<3x5x5x4xf32>) {
^bb0(%arg3: f32, %arg4: f32): // no predecessors
linalg.yield %arg3 : f32
}
"lmhlo.terminator"() : () -> ()
}
}

I also see the linalg-fusion-for-tensor-ops pass which can fuse hlo pointwise operators into region something like add/mul fused as std operators inside a generic op

func @float_add(%lhs: tensor<2x2xf32>,
%rhs: tensor<2x2xf32>) -> tensor<2x2xf32> {
%0 = "mhlo.add"(%lhs, %rhs) : (tensor<2x2xf32>,
tensor<2x2xf32>) -> tensor<2x2xf32>
%1 = "mhlo.multiply"(%lhs, %0) : (tensor<2x2xf32>,
tensor<2x2xf32>) -> tensor<2x2xf32>
return %1 : tensor<2x2xf32>
}

module {
func @float_add(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {
%0 = linalg.init_tensor [2, 2] : tensor<2x2xf32>
%1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<2x2xf32>, tensor<2x2xf32>) outs(%0 : tensor<2x2xf32>) {
^bb0(%arg2: f32, %arg3: f32, %arg4: f32): // no predecessors
%4 = addf %arg2, %arg3 : f32
linalg.yield %4 : f32
} -> tensor<2x2xf32>
%2 = linalg.init_tensor [2, 2] : tensor<2x2xf32>
%3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %1 : tensor<2x2xf32>, tensor<2x2xf32>) outs(%2 : tensor<2x2xf32>) {
^bb0(%arg2: f32, %arg3: f32, %arg4: f32): // no predecessors
%4 = mulf %arg2, %arg3 : f32
linalg.yield %4 : f32
} -> tensor<2x2xf32>
return %3 : tensor<2x2xf32>
}
}

#map = affine_map<(d0, d1) -> (d0, d1)>
module {
func @float_add(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {
%0 = linalg.init_tensor [2, 2] : tensor<2x2xf32>
%1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0, %arg1 : tensor<2x2xf32>, tensor<2x2xf32>) outs(%0 : tensor<2x2xf32>) {
^bb0(%arg2: f32, %arg3: f32, %arg4: f32): // no predecessors
%2 = addf %arg2, %arg3 : f32
%3 = mulf %arg2, %2 : f32
linalg.yield %3 : f32
} -> tensor<2x2xf32>
return %1 : tensor<2x2xf32>
}
}

I would like to Fuse linalg Conv Relu inside a region using how do I get something like this where Conv gets inside the block ,
Can the methodology used by pointwise operators be extended to Conv and MatMul etc.

linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%0 : memref<3x5x5x4xf32>) outs(%arg2 : memref<3x5x5x4xf32>) {
^bb0(%arg3: f32, %arg4: f32):  // no predecessors
    linalg.conv(%arg0, %arg1, %0) {dilations = [1, 1], strides = [2, 1]} : memref<3x5x5x3xf32>, memref<2x2x3x4xf32>, memref<3x5x5x4xf32>
  linalg.yield %arg3 : f32
}