[AIROCMIR-446] Lower migraphx.greater/equal into linalg.generic

[Mr-Anyone]

Motivation

Lower migraphx.greater/equal into linalg.generic. This allow the user to use migraphx.greater and migraphx.equal operations.

Technical Details

migraphx.greater/equal doesn't map directly into linalg.elementwise. In this case, we emit a linalg.generic instead. Also, note that the lowering is exactly the same as tosa path.

Test Plan

Added a lit test.

Test Result

Passed lit test.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Mr-Anyone]

I am not sure if this is a bug in the tosa lowering pipeline, but consider the following:

**root@14298acabab2:~/rocMLIR/build-release# cat main.mlir
func.func @func_greater_signed(%arg0: !migraphx.shaped<1xsi32, 1>, %arg1: !migraphx.shaped<1xsi32, 1>) ->  !migraphx.shaped<1xsi32, 1> attributes {kernel, arch="gfx950"}{
  %result = migraphx.greater %arg0, %arg1: <1xsi32, 1>, <1xsi32, 1> -> <1xsi32, 1>
  func.return %result : !migraphx.shaped<1xsi32, 1>
}

func.func @func_greater_unsigned(%arg0: !migraphx.shaped<1xui32, 1>, %arg1: !migraphx.shaped<1xui32, 1>) ->  !migraphx.shaped<1xui32, 1> attributes {kernel, arch="gfx950"}{
  %result = migraphx.greater %arg0, %arg1: <1xui32, 1>, <1xui32, 1> -> <1xui32, 1>
  func.return %result : !migraphx.shaped<1xui32, 1>
}
root@14298acabab2:~/rocMLIR/build-release# ./bin/rocmlir-driver --kernel-pipeline=migraphx,highlevel main.mlir
#map = affine_map<() -> (0)>
#map1 = affine_map<() -> ()>
#map2 = affine_map<(d0) -> ()>
#transform_map = #rock.transform_map<#map by [<ConstDim{0, 1} [] at [] -> ["const0"] at [0]>] bounds = [] -> [1]>
#transform_map1 = #rock.transform_map<#map2 by [<AddDim{1} ["unit0"] at [0] -> [] at []>] bounds = [1] -> []>
module {
  func.func @func_greater_signed(%arg0: memref<1xi32>, %arg1: memref<1xi32>, %arg2: memref<1xi32>) attributes {arch = "gfx950", kernel} {
    %0 = rock.transform %arg0 by #transform_map : memref<1xi32> to memref<i32>
    %1 = rock.transform %arg1 by #transform_map : memref<1xi32> to memref<i32>
    %alloc = memref.alloc() {alignment = 64 : i64} : memref<i32>
    linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = []} ins(%0, %1 : memref<i32>, memref<i32>) outs(%alloc : memref<i32>) {
    ^bb0(%in: i32, %in_0: i32, %out: i32):
      %3 = arith.cmpi sgt, %in, %in_0 : i32
      %4 = arith.extui %3 : i1 to i32
      linalg.yield %4 : i32
    }
    %2 = rock.transform %alloc by #transform_map1 : memref<i32> to memref<1xi32>
    memref.copy %2, %arg2 : memref<1xi32> to memref<1xi32>
    return
  }
  func.func @func_greater_unsigned(%arg0: memref<1xi32>, %arg1: memref<1xi32>, %arg2: memref<1xi32>) attributes {arch = "gfx950", kernel} {
    %0 = rock.transform %arg0 by #transform_map : memref<1xi32> to memref<i32>
    %1 = rock.transform %arg1 by #transform_map : memref<1xi32> to memref<i32>
    %alloc = memref.alloc() {alignment = 64 : i64} : memref<i32>
    linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = []} ins(%0, %1 : memref<i32>, memref<i32>) outs(%alloc : memref<i32>) {
    ^bb0(%in: i32, %in_0: i32, %out: i32):
      %3 = arith.cmpi sgt, %in, %in_0 : i32
      %4 = arith.extui %3 : i1 to i32
      linalg.yield %4 : i32
    }
    %2 = rock.transform %alloc by #transform_map1 : memref<i32> to memref<1xi32>
    memref.copy %2, %arg2 : memref<1xi32> to memref<1xi32>
    return
  }
}

The MLIR file below defines two functions that are identical except for the signedness of their input and output types: one operates on i32 (signed), and the other on ui32 (unsigned). Notice that the tosa pipeline emits arith.cmpi sgt for both cases which assume the operand is signed.