Skip to content

[AIROCMLIR-446] Lower linalg.matmul/batch_matmul into rock.gemm#2224

Open
Mr-Anyone wants to merge 11 commits intodevelopfrom
pr-template-migraphx-to-linalg-2
Open

[AIROCMLIR-446] Lower linalg.matmul/batch_matmul into rock.gemm#2224
Mr-Anyone wants to merge 11 commits intodevelopfrom
pr-template-migraphx-to-linalg-2

Conversation

@Mr-Anyone
Copy link
Member

@Mr-Anyone Mr-Anyone commented Jan 30, 2026

Motivation

Lower form Linalg into Rock.

Technical Details

This PR introduces a new LinalgToRock conversion pass that provides an alternative lowering path from MIGraphX through Linalg dialect to Rock.

The conversion handles linalg.matmul and linalg.batch_matmul operations, transforming them into rock.gemm operations.

Driver changes are added so that an E2E test can ran starting from migraphx -> linalg -> rock -> ... depending on the target platform.

Test Plan

Added e2e and a few lit test.

Test Result

Passed locally.

Submission Checklist

@Mr-Anyone Mr-Anyone requested a review from causten as a code owner January 30, 2026 18:59
@Mr-Anyone Mr-Anyone marked this pull request as draft January 30, 2026 19:00
@Mr-Anyone Mr-Anyone force-pushed the pr-template-migraphx-to-linalg-2 branch 4 times, most recently from 3a1b2cc to f525a63 Compare January 30, 2026 20:15
@Mr-Anyone Mr-Anyone marked this pull request as ready for review January 30, 2026 20:16
@@ -0,0 +1,39 @@
// RUN: rocmlir-opt -split-input-file --migraphx-to-linalg --linalg-to-rock -verify-diagnostics %s | FileCheck %s
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When writing conversion tests, the recommended practice is to test only one conversion in isolation. In this case, we want to test --linalg-to-rock. As it is right now, we are testing --migraphx-to-linalg as well, which we don't want to. This is what we do in tosa-to-rock.mlir for example.

So you should take the IR after --migraphx-to-linalg and use that as input for this test, thus removing the --migraphx-to-linalg execution.

@@ -0,0 +1,39 @@
// RUN: rocmlir-opt -split-input-file --migraphx-to-linalg --linalg-to-rock -verify-diagnostics %s | FileCheck %s
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Like the other PR, please format the test using the output from rocmlir-opt

@@ -0,0 +1,39 @@
// RUN: rocmlir-opt -split-input-file --migraphx-to-linalg --linalg-to-rock -verify-diagnostics %s | FileCheck %s
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For consistency with tosa-to-rock tests, please call this file "linalg-to-rock-matmul"

// CHECK-LABEL: func.func @dot_one(
// CHECK-SAME: %[[arg0:.*]]: tensor<6xf32>,
// CHECK-SAME: %[[arg1:.*]]: tensor<6xf32>)
// CHECK-NEXT: %[[expanded:.*]] = tensor.expand_shape %[[arg1]]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The output IR is not yet rock compatible because we dont have rock.transforms.

Looking at the tosa-to-rock tests, we have tosa-to-rock-matmul which output also does not contain rock.transforms. But in the case of tosa-to-rock we have a tosa-to-rock.mlir test, which is actually generating rock.transforms with the --rock-view-to-transform pass.

Can you add a separate test called linalg-to-rock.mlir following our tosa-to-rock approach?

// CHECK-NEXT: %[[collapsed:.*]] = tensor.collapse_shape %[[one]]
// CHECK-NEXT: return %[[collapsed]]
func.func @dot_one(%arg0 : !migraphx.shaped<1x3x2xf32, 6x2x1>, %arg1: !migraphx.shaped<1x2x3xf32, 6x3x1>)
-> !migraphx.shaped<1x3x3xf32, 9x3x1> attributes {kernel, arch="gfx950"}{
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't want tests to be arch-dependent. In other words, replace arch="gfx950" with arch-independent testing. Look at how we do it in https://github.com/ROCm/rocMLIR/blob/develop/mlir/test/Conversion/TosaToRock/tosa-to-rock.mlir for an example.

funcPm.addPass(migraphx::createMIGraphXTosaSimplifyPass());
}

void migraphx::addHighLevelMIGraphXToLinalg(PassManager &pm) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than having 2 separate functions (addHighLevelPipeline / addHighLevelMIGraphXToLinalg) we could have one that takes an extra argument to indicate if IR should be lowered to linalg or TOSA.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you should be able to make use of the lowerFromLinalg flag/option that you defined earlier.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you think a passing in a boolean value is fine? It seems a bit awkward to pass rock::BufferizeOptions because the disableRock member is never used in that function? Also, addHighLevelPipeline doesn't seem to be registered?

Value a = op.getOperand(0);
Value b = op.getOperand(1);
Value cOriginal = op.getOperand(2);
RankedTensorType outputType = cast<RankedTensorType>(cOriginal.getType());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would break if tensors are unranked which I think can happen? Please check it

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems that linalg.matmul doesn't work with unranked tensor?

main.mlir:2:25: error: custom op 'linalg.matmul' Cannot build binary Linalg operation: expects allComplex, allFloatingPoint, or allInteger, got 'tensor<*xf32>' and 'tensor<*xf32>'
  %arg2 = linalg.matmul ins(%arg0, %arg1: tensor<*xf32>, tensor<*xf32>)  outs(%arg2: tensor<*xf32>) -> tensor<*xf32>

I am not sure how to go about testing this.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fine. I would check here if we are getting static shapes at least. If the IR comes from migraphx-to-linalg they will be, but it's better if we make this pass as solid as possible, just in case we ingest IR from another source, which might contain RankedTensor with non-static shapes.

@justinrosner
Copy link
Contributor

Could you update the PR description so that it more accurately describes what the pass is doing?


PassOptions::Option<bool> lowerFromLinalg{
*this, "lower-from-linalg",
desc("Disable Rock dialect targeting when bufferizing"), init(false)};
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This description will need to be updated as to what this new flag is actually doing.

funcPm.addPass(migraphx::createMIGraphXTosaSimplifyPass());
}

void migraphx::addHighLevelMIGraphXToLinalg(PassManager &pm) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you should be able to make use of the lowerFromLinalg flag/option that you defined earlier.

Copy link
Contributor

@dhernandez0 dhernandez0 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as the other PR, can you change the title of the commit to highlight this implements rock.gemm only.


rock::StoreMethodAttr method =
rewriter.getAttr<rock::StoreMethodAttr>(rock::StoreMethod::Set);
rock::GemmOp result = rock::GemmOp::create(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

are the shapes of a and b the same for migraphx.dot, linalg.matmul and rock.gemm? Please check if that's the case (for example, if they all expect B to be [batch, K, N].

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for point this out, I just realized that this may not be always true. It is possible that linalg.matmul could be transposed if an indexing_map is set.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To the best of my knowledge, I think I have handled all the transpose case that comes with linalg. See matmul_transposed_* test cases.

Also, my understanding is that you cannot have transpose the c matrix:?

func.func @matmul_tranposed_B_3D(%arg0: tensor<1x3x2xf32>, %arg1: tensor<1x2x3xf32>) -> tensor<1x3x3xf32> attributes {arch = "gfx950", kernel} {
  %cst = arith.constant dense<0.000000e+00> : tensor<1x3x3xf32>
  %0 = linalg.batch_matmul
    indexing_maps = [affine_map<(batch, m, n, k) -> (batch, m, k)>,
                      affine_map<(batch, m, n, k) -> (batch, k, n)>,
// expected-error @+1 {{'linalg.batch_matmul' op Invalid output map result dimension.}}
                      affine_map<(batch, m, n, k) -> (batch, n, m)>]
    ins(%arg0, %arg1 : tensor<1x3x2xf32>, tensor<1x2x3xf32>) outs(%cst : tensor<1x3x3xf32>) -> tensor<1x3x3xf32>
  return %0 : tensor<1x3x3xf32>
}

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also, I think they are all [batch, K, N] and [batch, N, K]. rock.gemm seems to only support 3D or 2D.

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This pull request introduces a new LinalgToRock conversion pass that provides an alternative lowering path from MIGraphX through the Linalg dialect to Rock, enabling end-to-end compilation flows.

Changes:

  • Added LinalgToRock conversion pass that converts linalg.matmul and linalg.batch_matmul operations to rock.gemm operations
  • Extended rocmlir-driver with new pipeline options (migraphx-linalg and highlevel-linalg) to support the Linalg-based lowering path
  • Added comprehensive test coverage including unit tests and end-to-end tests

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 5 comments.

Show a summary per file
File Description
mlir/lib/Conversion/LinalgToRock/LinalgToRock.cpp Core conversion logic implementing matmul pattern rewriting from Linalg to Rock
mlir/lib/Conversion/LinalgToRock/LinalgToRockPass.cpp Pass implementation that applies the LinalgToRock conversion patterns
mlir/include/mlir/Conversion/LinalgToRock/LinalgToRock.h Header file declaring the conversion pass and pattern population functions
mlir/include/mlir/Conversion/RocMLIRPasses.h Added LinalgToRock header include
mlir/include/mlir/Conversion/RocMLIRPasses.td Added LinalgToRockPass definition
mlir/lib/Conversion/LinalgToRock/CMakeLists.txt Build configuration for the new conversion library
mlir/lib/Conversion/CMakeLists.txt Added LinalgToRock subdirectory
mlir/lib/Dialect/Rock/Pipelines/Pipelines.cpp Integrated LinalgToRock pass into the bufferization pipeline
mlir/lib/Dialect/Rock/Pipelines/CMakeLists.txt Added MLIRLinalgToRock dependency
mlir/include/mlir/Dialect/Rock/Pipelines/Pipelines.h Added lowerFromLinalg option to BufferizeOptions
mlir/lib/Dialect/MIGraphX/Pipeline/Pipeline.cpp Modified to conditionally use LinalgToRock instead of TosaToRock
mlir/lib/Dialect/MIGraphX/Pipeline/CMakeLists.txt Added MLIRMIGraphXToLinalg dependency
mlir/include/mlir/Dialect/MIGraphX/Pipeline/Pipeline.h Added lowerFromLinalg parameter to addHighLevelPipeline
mlir/tools/rocmlir-driver/rocmlir-driver.cpp Added support for new pipeline options (migraphx-linalg, highlevel-linalg)
mlir/test/Conversion/LinalgToRock/linalg-to-rock.mlir Test verifying conversion with rock-view-to-transform
mlir/test/Conversion/LinalgToRock/linalg-to-rock-matmul.mlir Test verifying basic LinalgToRock conversion
mlir/test/Conversion/LinalgToRock/e2e/linalg-to-rock-add-gemm.e2e.mlir End-to-end integration test

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

Copy link
Contributor

@pabloantoniom pabloantoniom left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't forget to update the PR title

Value a = op.getOperand(0);
Value b = op.getOperand(1);
Value cOriginal = op.getOperand(2);
RankedTensorType outputType = cast<RankedTensorType>(cOriginal.getType());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fine. I would check here if we are getting static shapes at least. If the IR comes from migraphx-to-linalg they will be, but it's better if we make this pass as solid as possible, just in case we ingest IR from another source, which might contain RankedTensor with non-static shapes.

LinalgMatOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const {
Location loc = op.getLoc();
Value a = op.getOperand(0);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit:

Suggested change
Value a = op.getOperand(0);
Value a = op.lhs();
Value b = op.rhs();
Value c = op.getOutputs()[0];

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Doesn't seem to compile?

/home/vhe/vscode/rocMLIR/mlir/lib/Conversion/LinalgToRock/LinalgToRock.cpp:41:16: error: no member named 'rhs' in 'mlir::linalg::MatmulOp'
   41 |   Value b = op.rhs();
      |             ~~ ^
/home/vhe/vscode/rocMLIR/mlir/lib/Conversion/LinalgToRock/LinalgToRock.cpp:40:16: error: no member named 'lhs' in 'mlir::linalg::BatchMatmulOp'
   40 |   Value a = op.lhs();
      |             ~~ ^

I couldn't find a method a member that has similar semantics.

@Mr-Anyone Mr-Anyone changed the title [AIROCMLIR-446] Linalg to Rock Conversion [AIROCMLIR-446] Lower linalg.matmul/batch_matmul into rock.gemm Feb 4, 2026
// CHECK-NEXT: %[[three:.*]] = rock.gemm %[[two]] = %[[one]] * %[[zero]] storeMethod = set
// CHECK-NEXT: %[[four:.*]] = rock.transform %[[three]]
// CHECK-NEXT: return %[[four]]
func.func @matmul_3D(%arg0: tensor<6xf32>, %arg1: tensor<6xf32>) -> tensor<9xf32> attributes {arch = "##TOKEN_ARCH##", kernel} {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One more interesting test would be check if output from both TosaToRock and LinalgToRock matches or not starting from MIGraphX.

e.g.

Input
migraphx.dot

  1. Lowering tosa -> rock
  2. Lowering using Linalg->Rock

Check if output from both 1 and 2 matches or not.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see linalg-to-rock-vs-tosa-to-rock.mlir. Also, rocmlir-driver is used instead because it seems that tosa-to-tensor pass is not exported to romclir-opt.

@Mr-Anyone Mr-Anyone force-pushed the pr-template-migraphx-to-linalg branch 6 times, most recently from 42811ec to 7f4ce15 Compare February 9, 2026 22:03
@Mr-Anyone Mr-Anyone force-pushed the pr-template-migraphx-to-linalg branch from 0850d14 to f12a85c Compare February 9, 2026 22:42
@Mr-Anyone Mr-Anyone force-pushed the pr-template-migraphx-to-linalg-2 branch from 9c0d44a to 271f8c8 Compare February 9, 2026 22:47
@Mr-Anyone Mr-Anyone force-pushed the pr-template-migraphx-to-linalg branch 2 times, most recently from ad1c291 to ef5d903 Compare February 10, 2026 16:09
@Mr-Anyone Mr-Anyone force-pushed the pr-template-migraphx-to-linalg-2 branch from 952693d to b501718 Compare February 10, 2026 20:02
@Mr-Anyone Mr-Anyone force-pushed the pr-template-migraphx-to-linalg-2 branch from 88d2296 to 80495d1 Compare February 10, 2026 20:53
@Mr-Anyone Mr-Anyone force-pushed the pr-template-migraphx-to-linalg branch from 97d65f9 to 405b3e8 Compare February 10, 2026 21:16
@Mr-Anyone Mr-Anyone force-pushed the pr-template-migraphx-to-linalg-2 branch from 80495d1 to 75bf455 Compare February 11, 2026 14:38
@@ -0,0 +1,12 @@
// RUN: rocmlir-gen -fut dot_add -arch %arch --clone-harness %s |\
// RUN: rocmlir-driver -host-pipeline=migraphx,highlevel -kernel-pipeline=migraphx-linalg,highlevel -targets %arch |\
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

not sure why you need multiple RUN lines if outputs are being piped to next command.

See other E2E tests.

// RUN: rocmlir-gen -fut dot_add --arch %arch --clone-harness %s | rocmlir-driver -kernel-pipeline=migraphx,highlevel -host-pipeline=migraphx,highlevel | rocmlir-gen -ph -print-results -rand none -fut dot_add_wrapper - | rocmlir-driver -host-pipeline mhal,runner -kernel-pipeline full -targets %arch | xmir-runner --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext,%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_float16_utils%shlibext,%linalg_test_lib_dir/libmlir_c_runner_utils%shlibext --entry-point-result=void | FileCheck %s

Comment on lines +132 to +133
if (auto attr = op->template getAttrOfType<StringAttr>("perf_config"))
result->setAttr("perf_config", attr);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add test for this

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see linalg-to-rock-matmul.mlir.

Comment on lines +110 to +112
if (failed(maybeAMatrixTransposed) || failed(maybeBMatrixTransposed)) {
return op.emitError("cannot determine if input matrix is transposed");
}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit :
Would be nice to have code coverage for these lines

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

see linalg-to-rock-matmul.mlir

Comment on lines +1 to +2
// RUN: sed s/##TOKEN_ARCH##/%arch/g %s | rocmlir-driver --kernel-pipeline=migraphx-linalg,highlevel| FileCheck %s
// RUN: sed s/##TOKEN_ARCH##/%arch/g %s | rocmlir-driver --kernel-pipeline=migraphx,highlevel | FileCheck %s
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a way to compare output of from both to see if they match exactly instead of relying on "CHECK" lines ?

If not, this is a good solution too.

Copy link
Member Author

@Mr-Anyone Mr-Anyone Feb 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can't think of a good way because the diff between the tosa lowering path and the linalg lowering path is slightly different (difference in the variable name, some time ordering).

Base automatically changed from pr-template-migraphx-to-linalg to develop February 11, 2026 22:46
@Mr-Anyone Mr-Anyone force-pushed the pr-template-migraphx-to-linalg-2 branch from 0d10133 to d0be693 Compare February 12, 2026 15:43
@Mr-Anyone Mr-Anyone requested a review from umangyadav February 12, 2026 17:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants