[AIROCMLIR-446] Lower linalg.matmul/batch_matmul into rock.gemm#2224
[AIROCMLIR-446] Lower linalg.matmul/batch_matmul into rock.gemm#2224
linalg.matmul/batch_matmul into rock.gemm#2224Conversation
3a1b2cc to
f525a63
Compare
| @@ -0,0 +1,39 @@ | |||
| // RUN: rocmlir-opt -split-input-file --migraphx-to-linalg --linalg-to-rock -verify-diagnostics %s | FileCheck %s | |||
There was a problem hiding this comment.
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 | |||
There was a problem hiding this comment.
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 | |||
There was a problem hiding this comment.
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]] |
There was a problem hiding this comment.
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"}{ |
There was a problem hiding this comment.
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) { |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
I think you should be able to make use of the lowerFromLinalg flag/option that you defined earlier.
There was a problem hiding this comment.
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()); |
There was a problem hiding this comment.
This would break if tensors are unranked which I think can happen? Please check it
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
|
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)}; |
There was a problem hiding this comment.
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) { |
There was a problem hiding this comment.
I think you should be able to make use of the lowerFromLinalg flag/option that you defined earlier.
8759718 to
6d3702b
Compare
f525a63 to
9ae19b8
Compare
|
|
||
| rock::StoreMethodAttr method = | ||
| rewriter.getAttr<rock::StoreMethodAttr>(rock::StoreMethod::Set); | ||
| rock::GemmOp result = rock::GemmOp::create( |
There was a problem hiding this comment.
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].
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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>
}
There was a problem hiding this comment.
Also, I think they are all [batch, K, N] and [batch, N, K]. rock.gemm seems to only support 3D or 2D.
There was a problem hiding this comment.
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.matmulandlinalg.batch_matmuloperations torock.gemmoperations - Extended rocmlir-driver with new pipeline options (
migraphx-linalgandhighlevel-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.
pabloantoniom
left a comment
There was a problem hiding this comment.
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()); |
There was a problem hiding this comment.
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); |
There was a problem hiding this comment.
nit:
| Value a = op.getOperand(0); | |
| Value a = op.lhs(); | |
| Value b = op.rhs(); | |
| Value c = op.getOutputs()[0]; |
There was a problem hiding this comment.
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.
linalg.matmul/batch_matmul into rock.gemm
mlir/test/Conversion/LinalgToRock/e2e/linalg-to-rock-add-gemm.e2e.mlir
Outdated
Show resolved
Hide resolved
| // 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} { |
There was a problem hiding this comment.
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
- Lowering tosa -> rock
- Lowering using Linalg->Rock
Check if output from both 1 and 2 matches or not.
There was a problem hiding this comment.
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.
42811ec to
7f4ce15
Compare
0850d14 to
f12a85c
Compare
9c0d44a to
271f8c8
Compare
ad1c291 to
ef5d903
Compare
952693d to
b501718
Compare
88d2296 to
80495d1
Compare
97d65f9 to
405b3e8
Compare
80495d1 to
75bf455
Compare
| @@ -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 |\ | |||
There was a problem hiding this comment.
not sure why you need multiple RUN lines if outputs are being piped to next command.
See other E2E tests.
| if (auto attr = op->template getAttrOfType<StringAttr>("perf_config")) | ||
| result->setAttr("perf_config", attr); |
There was a problem hiding this comment.
see linalg-to-rock-matmul.mlir.
| if (failed(maybeAMatrixTransposed) || failed(maybeBMatrixTransposed)) { | ||
| return op.emitError("cannot determine if input matrix is transposed"); | ||
| } |
There was a problem hiding this comment.
nit :
Would be nice to have code coverage for these lines
There was a problem hiding this comment.
see linalg-to-rock-matmul.mlir
| // 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 |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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).
0d10133 to
d0be693
Compare
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