[AIROCMLIR-551] Completely support as_underlying_shape and as_logical_shape by Mr-Anyone · Pull Request #2265 · ROCm/rocMLIR

Mr-Anyone · 2026-03-02T21:15:08Z

Motivation

Being able to support transposed memory layout, long strides in both as_underlying_shape and as_logical_shape.

Furthermore, in as_underlying_shape, broadcasting is supported.

Technical Details

This one PR essentially implements these two commits: #2198 , this one.

as_logical_shape is implemented as the following:

Reshape from flat memory layout to memory layout
Check the stride permutation, and invert it to match the logical shape
If the memory layout tensor is "greater in dimension" than the logical dimension, it means that we have a long stride layout so we emit tensor.extract_from_slice to get the logical shape
Broadcast if needed

as_underlying_shape is implemented as the following:

Reshape to memory based on stride permutation
If there are long strides, we emit a tensor.insert_slice with a tensor.empty_op
If there is broadcasting, we error out.

Test Plan

Added e2e and lit test.

Test Result

Passed both e2e and lit test.

Submission Checklist

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

Mr-Anyone · 2026-03-02T21:15:42Z

Also, note that the code structure is really similar to the one found in the tosa pipeline if that makes review easier.

Mr-Anyone · 2026-03-02T21:24:48Z

mlir/test/Conversion/MIGraphXToLinalg/mixr-to-linalg-ops.mlir

-// CHECK-DAG: %[[expanded:.*]] = tensor.expand_shape %[[arg1]]
-// CHECK-DAG: %[[expanded_0:.*]] = tensor.expand_shape %[[arg0]]
-// CHECK-DAG: linalg.sub ins(%[[expanded_0]], %[[expanded]] {{.*}})
+// CHECK-DAG: linalg.sub ins(%[[arg0]], %[[arg1]] {{.*}})


This actually improves the previous IR. Notice that the input shape !migraphx.shaped<16xf32, 1> get converted into tensor<16xf32> which is the memory layout shape. In that case, we don't emit tensor.expand_shape.

Before:

func.func @func_sub(%arg0: tensor<16xf32>, %arg1: tensor<16xf32>) -> tensor<16xf32> { %expanded = tensor.expand_shape %arg1 [[0]] output_shape [16] : tensor<16xf32> into tensor<16xf32> %expanded_0 = tensor.expand_shape %arg0 [[0]] output_shape [16] : tensor<16xf32> into tensor<16xf32> %0 = tensor.empty() : tensor<16xf32> %1 = linalg.sub ins(%expanded_0, %expanded : tensor<16xf32>, tensor<16xf32>) outs(%0 : tensor<16xf32>) -> tensor<16xf32> %collapsed = tensor.collapse_shape %1 [[0]] : tensor<16xf32> into tensor<16xf32> return %collapsed : tensor<16xf32> }

After:

func.func @func_sub(%arg0: tensor<16xf32>, %arg1: tensor<16xf32>) -> tensor<16xf32> { %0 = tensor.empty() : tensor<16xf32> %1 = linalg.sub ins(%arg0, %arg1 : tensor<16xf32>, tensor<16xf32>) outs(%0 : tensor<16xf32>) -> tensor<16xf32> return %1 : tensor<16xf32> }

justinrosner · 2026-03-03T21:23:55Z

mlir/lib/Conversion/LinalgToRock/LinalgToRock.cpp

+    ConversionPatternRewriter &rewriter) const {
+  /// The linalg-to-rock passes emits the following expression
+  /// for expanding the strides. We are matching the following IR
+  /// clang-format off


Just curious, but why do we need to turn clang-format off here? It doesn't look like the comment is too long...

I think in this case it is ok. Usually, two lines of mlir gets merged into one line if you don't turn off clang-format

justinrosner · 2026-03-03T21:25:54Z

mlir/lib/Conversion/LinalgToRock/LinalgToRock.cpp

+  /// %empty = tensor.empty() : ....
+  /// %inserted_slice = tensor.insert_slice %actual_data into %empty ...


Is the expand strides case the only case that will produce an empty tensor + insert_slice? This seems quite broad, and potentially susceptible to pattern matching unwanted scenarios?

This is true indeed true, but my understanding is that empty tensor + insert_slice is the same as a rock.expand_strides. rock.expand_strides is the same as extract_slice (of the result) + memcpy which I think has the same semantics as above.

I think the we have to emit attribute during migraphx to linalg if we only want some case of the tensor.insert_slice to be lowered into rock. Semantically,

justinrosner · 2026-03-03T21:27:37Z

mlir/lib/Conversion/MIGraphXToLinalg/MIGraphXToLinalg.cpp

 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Statistic.h"


Is this header file still needed?

I just removed them.

justinrosner · 2026-03-03T21:32:09Z

mlir/test/Conversion/MIGraphXToLinalg/migraphx-to-linalg-non-contiguous-stride.mlir

@@ -0,0 +1,64 @@
+// RUN: rocmlir-opt -split-input-file --migraphx-to-linalg -verify-diagnostics %s | FileCheck %s
+
+// testcase taken from migraphx-to-tosa-non-contiguous-strides.mlir


nit: no need to mention this here as this comment will no longer make sense once the Tosa path is removed. This probably makes more sense to have somewhere in the PR description.

justinrosner · 2026-03-03T21:34:42Z

mlir/lib/Conversion/MIGraphXToLinalg/MIGraphXToLinalg.cpp

+  RankedTensorType transposedType =
+      dyn_cast<RankedTensorType>(result.getType());
+  if (transposedType.getShape() != ArrayRef(slicingShape)) {
+    SmallVector<int64_t, 4> starts(permutation.size(), 0);


Nit: Can we remove this unused variable?

weird, as there are no compiler warnings.

justinrosner · 2026-03-03T21:35:21Z

mlir/lib/Conversion/MIGraphXToLinalg/MIGraphXToLinalg.cpp

+  }
+
+  RankedTensorType transposedType =
+      dyn_cast<RankedTensorType>(result.getType());


We should have a null check after using dyn_cast in case we get a null value.

justinrosner · 2026-03-03T21:36:17Z

mlir/lib/Conversion/MIGraphXToLinalg/MIGraphXToLinalg.cpp

-  if (isPermutationStandardForm(permutation)) {
+  // reshape into memory layout type
+  RankedTensorType memoryType = inType.asMemoryLayoutTensor();
+  if (result.getType() != inType.asMemoryLayoutTensor()) {


Suggested change

if (result.getType() != inType.asMemoryLayoutTensor()) {

if (result.getType() != memoryType) {

justinrosner · 2026-03-03T21:38:07Z

mlir/lib/Conversion/MIGraphXToLinalg/MIGraphXToLinalg.cpp

-  migraphx::MIXRShapedType resultType = op.getResult().getType();
-  auto resultTensorType =
-      cast<RankedTensorType>(getTypeConverter()->convertType(resultType));
+  RankedTensorType inTensorType = dyn_cast<RankedTensorType>(in.getType());


Same here with the dyn_cast with no nullptr check

Mr-Anyone added 3 commits March 2, 2026 21:01

Completely support as_logical_shape/as_underlying_shape

080a3ee

Added testcase for migraphx-to-linalg

6b24412

Added Linalg to rock testcase

64bc72e

Mr-Anyone requested a review from causten as a code owner March 2, 2026 21:15

Mr-Anyone requested review from dhernandez0, justinrosner, pabloantoniom and umangyadav March 2, 2026 21:16

Mr-Anyone commented Mar 2, 2026

View reviewed changes

Early exit in as_underlying_shape

965c589

justinrosner reviewed Mar 3, 2026

View reviewed changes

Address comments round one

efbc7c1

Mr-Anyone requested a review from justinrosner March 4, 2026 20:39

		/// %empty = tensor.empty() : ....
		/// %inserted_slice = tensor.insert_slice %actual_data into %empty ...

		@@ -0,0 +1,64 @@
		// RUN: rocmlir-opt -split-input-file --migraphx-to-linalg -verify-diagnostics %s \| FileCheck %s

		// testcase taken from migraphx-to-tosa-non-contiguous-strides.mlir

	if (result.getType() != inType.asMemoryLayoutTensor()) {
	if (result.getType() != memoryType) {

Conversation

Mr-Anyone commented Mar 2, 2026

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

Mr-Anyone commented Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Mr-Anyone Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Mr-Anyone commented Mar 2, 2026 •

edited

Loading

Mr-Anyone Mar 2, 2026 •

edited

Loading