[AIROCMLIR-446] Lowering of Trivial Elementwise Operations

[Mr-Anyone]

Motivation

Added elementwise operator that has a one to one mapping. This allows us to test more GEMMs.

Technical Details

Added a TrivialElementwiseConverter class that performs one to one migraphx to linalg operations.

Example:

Starting from

%arg2 = migraphx.sub %arg0, %arg1 : <1x1xf32, 1x1>, <1x1xf32, 1x1> -> <1x1xf32, 1x1>

becomes:

%1 = linalg.sub ins(%expanded_0, %expanded : tensor<1x1xf32>, tensor<1x1xf32>) outs(%0 : tensor<1x1xf32>) -> tensor<1x1xf32>

Notice that there is a one to one mapping between linalg.sub and migraphx.sub.

Test Plan

Added a lit test for this initial conversion.

Test Result

Passed lit test

Submission Checklist

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

[pabloantoniom]

Add another test file for negative checks, i.e., migraphx elementwise ops that do not have a linalg conversion (yet)

[Mr-Anyone]

See linalg-to-rock-not-implemented.mlir.

[pabloantoniom]

Please also check that the operands are also set in the right place (i.e., %arg0 then %arg1, not the other way around). Dont hardcode %arg0 and %arg1 in the check, instead use variables. Look for other lit tests for examples on how to do this

[pabloantoniom]

Why MathDialect? We are not using it anywhere right?

[Mr-Anyone]

Take linalg.exp as an example, the pass fails if I don't add this line. This is what the log file shows:

[dialect-conversion:1]     //===-------------------------------------------===//
[dialect-conversion:1]     Legalizing operation : 'math.exp' (0x27df5ee0) {
[dialect-conversion:1]       %4 = "math.exp"(%arg2) <{fastmath = #arith.fastmath<none>}> : (f32) -> f32
[dialect-conversion:1]     } -> SUCCESS : operation marked legal by the target
[dialect-conversion:1]     //===-------------------------------------------===//

It seems that some math operations are emitted as intermediate steps.

[umangyadav]

That is interesting. Can you find out why and where it using math.exp ?

[Mr-Anyone]

I think that is because linalg.exp have a region that contains math.exp. This isn't shown on the pretty printed version of mlir assembly:

"builtin.module"() ({
  "func.func"() <{function_type = (tensor<16xf32>, tensor<16xf32>) -> tensor<16xf32>, sym_name = "func_exp"}> ({
  ^bb0(%arg0: tensor<16xf32>, %arg1: tensor<16xf32>):
    %0 = "tensor.expand_shape"(%arg0) <{reassociation = [[0]], static_output_shape = array<i64: 16>}> : (tensor<16xf32>) -> tensor<16xf32>
    %1 = "tensor.empty"() : () -> tensor<16xf32>
    %2 = "linalg.exp"(%0, %1) <{operandSegmentSizes = array<i32: 1, 1>}> ({
    ^bb0(%arg2: f32, %arg3: f32):
      %4 = "math.exp"(%arg2) <{fastmath = #arith.fastmath<none>}> : (f32) -> f32
      "linalg.yield"(%4) : (f32) -> ()
    }) : (tensor<16xf32>, tensor<16xf32>) -> tensor<16xf32>
    %3 = "tensor.collapse_shape"(%2) <{reassociation = [[0]]}> : (tensor<16xf32>) -> tensor<16xf32>
    "func.return"(%3) : (tensor<16xf32>) -> ()
  }) : () -> ()
}) : () -> ()

Here we can see that linalg.exp have a region that contains math.exp.

[pabloantoniom]

Rename this to ElementwiseConverter

[pabloantoniom]

Why do we need to convert named ops to generic op?

[Mr-Anyone]

I am trying to match the output of tosa pipeline. In theory, I can move this into mlir/lib/Dialect/Rock/Pipelines/Pipelines.cpp instead:

root@14298acabab2:~/rocMLIR/build-debug# cat main.mlir
func.func @func_exp(%arg0 : !migraphx.shaped<1x3x2xf32, 6x2x1>) -> !migraphx.shaped<1x3x2xf32, 6x2x1> attributes {kernel, arch="gfx950"}{
  %1 = migraphx.exp %arg0 : <1x3x2xf32, 6x2x1> -> <1x3x2xf32, 6x2x1>
  func.return %1 : !migraphx.shaped<1x3x2xf32, 6x2x1>
}
root@14298acabab2:~/rocMLIR/build-debug# ./bin/rocmlir-driver --kernel-pipeline=migraphx,highlevel main.mlir
#map = affine_map<(d0, d1, d2) -> (d1 * 2 + d2)>
#map1 = affine_map<(d0, d1) -> (0, d0, d1)>
#map2 = affine_map<(d0, d1) -> (d0, d1)>
#map3 = affine_map<(d0, d1, d2) -> (d1, d2)>
#map4 = affine_map<(d0) -> (0, d0 floordiv 2, d0 mod 2)>
#transform_map = #rock.transform_map<#map by [<Unmerge{3, 2} ["exp1", "exp2"] at [1, 2] -> ["dim0"] at [0]>, <AddDim{1} ["unit0"] at [0] -> [] at []>] bounds = [1, 3, 2] -> [6]>
#transform_map1 = #rock.transform_map<#map1 by [<Merge{1, 3} ["dim0"] at [0] -> ["col0", "col1"] at [0, 1]>, <PassThrough ["dim1"] at [1] -> ["dim1"] at [2]>] bounds = [3, 2] -> [1, 3, 2]>
#transform_map2 = #rock.transform_map<#map3 by [<Unmerge{3} ["exp1"] at [1] -> ["dim0"] at [0]>, <PassThrough ["dim1"] at [2] -> ["dim1"] at [1]>, <AddDim{1} ["unit0"] at [0] -> [] at []>] bounds = [1, 3, 2] -> [3, 2]>
#transform_map3 = #rock.transform_map<#map4 by [<Merge{1, 3, 2} ["dim0"] at [0] -> ["col0", "col1", "col2"] at [0, 1, 2]>] bounds = [6] -> [1, 3, 2]>
module {
  func.func @func_exp(%arg0: memref<6xf32>, %arg1: memref<6xf32>) attributes {arch = "gfx950", kernel} {
    %0 = rock.transform %arg0 by #transform_map : memref<6xf32> to memref<1x3x2xf32>
    %1 = rock.transform %0 by #transform_map1 : memref<1x3x2xf32> to memref<3x2xf32>
    %alloc = memref.alloc() {alignment = 64 : i64} : memref<3x2xf32>
    linalg.generic {indexing_maps = [#map2, #map2], iterator_types = ["parallel", "parallel"]} ins(%1 : memref<3x2xf32>) outs(%alloc : memref<3x2xf32>) {
    ^bb0(%in: f32, %out: f32):
      %4 = math.exp %in : f32
      linalg.yield %4 : f32
    }
    %2 = rock.transform %alloc by #transform_map2 : memref<3x2xf32> to memref<1x3x2xf32>
    %3 = rock.transform %2 by #transform_map3 : memref<1x3x2xf32> to memref<6xf32>
    memref.copy %3, %arg1 : memref<6xf32> to memref<6xf32>
    return
  }
}

[justinrosner]

Which linalg operations are being turned into linalg.generics here? This is in LinalgToRock, so we should only be converting to rock ops in this pass?

[Mr-Anyone]

There isn't a particularly strong reason, but I am just trying to match the output of the tosa lowering path. It seems that migraphx->tosa tosa->rock lowers a lot of elementwise operation into linalg.generic?

[umangyadav]

better to check output type is the same as input type before taking aType as the output for the linalg.

[Copilot]

Pull request overview

This pull request adds lowering support for trivial elementwise operations from MIGraphX dialect to Linalg dialect. The implementation introduces a generic ElementwiseConverter template class that performs one-to-one mappings between MIGraphX operations (like add, sub, mul, div, abs, etc.) and their corresponding Linalg operations. This enables testing more GEMM operations and prepares the codebase for future lowering work.

Changes:

• Added ElementwiseConverter template class for converting 14 elementwise operations (add, sub, mul, div, pow, abs, ceil, exp, floor, log, neg, sqrt, tanh, recip) from MIGraphX to Linalg
• Added Math dialect to legal dialects for the MIGraphXToLinalg conversion
• Updated pipeline to convert named linalg operations to generic form using LinalgMorphOpsPass
• Added comprehensive lit tests for operation conversions and not-yet-implemented operations

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
mlir/lib/Conversion/MIGraphXToLinalg/MIGraphXToLinalg.cpp	Implements ElementwiseConverter template class and registers 14 elementwise operation converters
mlir/lib/Conversion/MIGraphXToLinalg/MIGraphXToLinalgPass.cpp	Adds Math dialect to legal dialects for conversion
mlir/lib/Dialect/Rock/Pipelines/Pipelines.cpp	Adds LinalgMorphOpsPass to convert named linalg ops to generic form
mlir/lib/Conversion/LinalgToRock/LinalgToRockPass.cpp	Adds Math dialect to legal dialects for LinalgToRock conversion
mlir/test/Conversion/MIGraphXToLinalg/mixr-to-linalg-ops.mlir	Tests conversion of 14 elementwise operations from MIGraphX to Linalg
mlir/test/Conversion/MIGraphXToLinalg/migraphx-to-linalg-not-implemented.mlir	Tests that not-yet-implemented operations fail legalization with appropriate errors
mlir/test/fusion/pr-e2e/mixr-gemm/mixr-gemm-add.e2e.mlir	Adds end-to-end test using migraphx-linalg pipeline

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.