Change `rock::invertTransform` to use FailureOr<T> by Mr-Anyone · Pull Request #2206 · ROCm/rocMLIR

Mr-Anyone · 2026-01-15T20:10:48Z

Motivation

Change the signature of rock::invertTransform so that the user knows it can fail.

Technical Details

Refactor and change code to use rock::invertTransform.

Resolves https://github.com/ROCm/rocMLIR-internal/issues/1999

Test Plan

Build and check CI.

Test Result

Build locally.

Submission Checklist

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

Copilot

Pull request overview

This pull request changes the signature of rock::invertTransforms to return FailureOr<ArrayAttr> instead of ArrayAttr, making it explicit that the function can fail. The change updates the function declaration, implementation, and all call sites throughout the codebase.

Changes:

Modified invertTransforms to return FailureOr<ArrayAttr> instead of ArrayAttr or nullptr
Updated failure handling to return LogicalResult::failure() instead of nullptr
Updated all call sites to check for success/failure and extract values using .value()

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 14 comments.

Show a summary per file

File	Description
transformMapUtils.h	Updated function signature to return `FailureOr<ArrayAttr>`
transformMapUtils.cpp	Changed return type and failure handling in implementation
ThreadwiseGemmLowering.cpp	Updated call site with proper failure checking
ShuffleGemmForReductions.cpp	Updated call sites with mixed error handling patterns
RemoveOutputAlloc.cpp	Updated call site with proper failure checking
GridwiseGemmToBlockwise.cpp	Updated multiple call sites, some with unchecked `.value()` calls
BlockwiseLoadTileToThreadwise.cpp	Updated call sites with unchecked `.value()` calls
BlockwiseGemmToThreadwise.cpp	Updated call site with unchecked `.value()` call

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

Copilot · 2026-01-16T13:50:21Z

    ArrayAttr inputThreadSubTile2dViewInv =
-        invertTransforms(rewriter, loc, inputThreadSubTile2dView);
+        invertTransforms(rewriter, loc, inputThreadSubTile2dView).value();


Calling .value() on the result of invertTransforms without checking if the operation succeeded. If invertTransforms fails, this will cause undefined behavior. Add a failure check before calling .value() and return failure if the inversion fails.

We want to do something similar to what you did in ThreadwiseGemmLowering.cpp by calling something like succeeded on the value from invertTransforms.

justinrosner · 2026-01-16T13:52:05Z

    TransformMapAttr invertedTrMap = invertTransformMap(b, trMap, loc);
    if (!invertedTrMap)
-      return nullptr;
+      return LogicalResult::failure();


You can just return failure() here. No need for the LogicalResult part.

justinrosner · 2026-01-16T13:53:21Z

    ArrayAttr inputThreadSubTile2dViewInv =
-        invertTransforms(rewriter, loc, inputThreadSubTile2dView);
+        invertTransforms(rewriter, loc, inputThreadSubTile2dView).value();


We want to do something similar to what you did in ThreadwiseGemmLowering.cpp by calling something like succeeded on the value from invertTransforms.

pabloantoniom · 2026-01-16T16:54:48Z

+          FailureOr<ArrayAttr> inBufferViewsTrAttr =
+              invertTransforms(b, loc, inBufferViewsTr.threadSubTile);
+          if (failed(inBufferViewsTrAttr)) {
+            return failure();


Avoid this at all costs. If something fails, we should always inform the user about what failed. return failure(); will just give the user the generic message "Lowering failed." which does not help to understand what failed.

You can use something like:

return op.emitError("invertTransforms failed");

Same goes for the rest of places in this PR where we check if invertTransforms failed.

I guess the question now is if it is ok to just return failure in storeGemmInputTile? It is in GridwiseGemmToBlockwise.cpp line 800.

pabloantoniom · 2026-01-16T16:56:10Z

+
+      FailureOr<ArrayAttr> invertedThreadSubTileViews =
          invertTransforms(rewriter, loc, gemm1OutSubTileViewsTr.threadSubTile);
+      if (succeeded(invertedThreadSubTileViews)) {


Shouldn't we fail here if invertTransforms failed?

I think in this case we can exit and emit diagnostics early.

Mr-Anyone · 2026-01-16T18:59:39Z

        invertTransforms(rewriter, rOp.getLoc(), additionalOutputViews);
-    for (Attribute trMap : invertedOutViews) {
+    if (failed(maybeInvertedOutViews)) {
+        return rOp.emitError("invertTransforms failed");


To my understanding, this is where we want to attach an error message?

Mr-Anyone · 2026-01-16T19:03:51Z

+          FailureOr<ArrayAttr> inBufferViewsTrAttr =
+              invertTransforms(b, loc, inBufferViewsTr.threadSubTile);
+          if (failed(inBufferViewsTrAttr)) {
+            return failure();


I guess the question now is if it is ok to just return failure in storeGemmInputTile? It is in GridwiseGemmToBlockwise.cpp line 800.

Mr-Anyone · 2026-01-16T19:14:52Z

+
+      FailureOr<ArrayAttr> invertedThreadSubTileViews =
          invertTransforms(rewriter, loc, gemm1OutSubTileViewsTr.threadSubTile);
+      if (succeeded(invertedThreadSubTileViews)) {


I think in this case we can exit and emit diagnostics early.

justinrosner

We also want some LIT tests for this change if possible.

justinrosner · 2026-01-16T19:39:47Z

+    FailureOr<ArrayAttr> maybeInputThreadSubTile2dViewInv =
        invertTransforms(rewriter, loc, inputThreadSubTile2dView);
+    assert(succeeded(maybeInputThreadSubTile2dViewInv) &&
+           "This must work for partial reduction");


I think a comment that includes something like transforms must be invertible would make the error message more useful to users.

justinrosner · 2026-01-16T19:43:37Z

+          FailureOr<ArrayAttr> maybeInBufferViewsTrAttr =
+              invertTransforms(b, loc, inBufferViewsTr.threadSubTile);
+          if (failed(maybeInBufferViewsTrAttr)) {
+            return op.emitError("invertTransforms failed");


It might be worth it to make the comments here, and everywhere else below, more explicit to what transforms failed to invert.

justinrosner · 2026-01-16T19:44:58Z

+    FailureOr<ArrayAttr> maybeInvertedViews =
+        invertTransforms(rewriter, rOp.getLoc(), views);
+    if (failed(maybeInvertedViews) || maybeInvertedViews.value().empty()) {
      LLVM_DEBUG(llvm::dbgs() << "gemm to reduce view inversion failed.\n");


Change this debug into a op.EmitError()?

justinrosner

Small nit, but otherwise looks good!

justinrosner · 2026-01-21T18:17:50Z

      }

-      if (!gemmOutToLinalgMaps.empty()) {
+      if (!maybeGemmOutToLinalgMaps.value().empty()) {


nit: Maybe save maybeGemmOutToLinalgMaps.value() to a new variable since it's going to be used in multiple places.

pabloantoniom

LGTM, just make sure to run clang-format

Change rock::invertTransform to use FailureOr<T>

788df21

Mr-Anyone requested a review from causten as a code owner January 15, 2026 20:10

Mr-Anyone force-pushed the pr-1999 branch from c14f8f6 to 0ed30dd Compare January 15, 2026 20:28

Fixed bugs

f0fb1e4

Mr-Anyone force-pushed the pr-1999 branch from 0ed30dd to f0fb1e4 Compare January 15, 2026 20:37

Mr-Anyone added 3 commits January 15, 2026 20:42

Initialize variable

42c5a4d

Fix build

c2e93cf

ran clang-format

a91b535

justinrosner requested review from Copilot and justinrosner January 16, 2026 13:43

Copilot started reviewing on behalf of justinrosner January 16, 2026 13:45 View session

Copilot AI reviewed Jan 16, 2026

View reviewed changes

justinrosner reviewed Jan 16, 2026

View reviewed changes

Mr-Anyone added 4 commits January 16, 2026 16:15

check error

ec0604b

Check even more error

897cd28

refactor into cleaner logic

ec8c9ce

Fix format

3c3d62d

Mr-Anyone requested review from dhernandez0 and pabloantoniom January 16, 2026 16:40

pabloantoniom reviewed Jan 16, 2026

View reviewed changes

Rename variable into maybe*

b4315b6

Mr-Anyone force-pushed the pr-1999 branch from 19bd4b1 to b4315b6 Compare January 16, 2026 16:59

Address comments

e8f8575

Mr-Anyone force-pushed the pr-1999 branch from 5cf3ccb to e8f8575 Compare January 16, 2026 18:34

Ran clang-format

d82e932

Mr-Anyone force-pushed the pr-1999 branch from d11568f to d82e932 Compare January 16, 2026 18:46

Mr-Anyone added 3 commits January 16, 2026 18:51

Made some changes

f05da5a

Added missing case

59798fe

Fixed a bug by moving debug statement one line down

ce7528e

Mr-Anyone commented Jan 16, 2026

View reviewed changes

Mr-Anyone requested review from justinrosner and pabloantoniom January 16, 2026 19:24

justinrosner reviewed Jan 16, 2026

View reviewed changes

Mr-Anyone added 4 commits January 16, 2026 20:10

Address trivial comments

c7b0359

Change error messages

0fb9e20

Fixed typo

83e6b80

Added a testcase

aa1e4c9

Mr-Anyone requested a review from justinrosner January 21, 2026 17:58

justinrosner approved these changes Jan 21, 2026

View reviewed changes

Mr-Anyone force-pushed the pr-1999 branch from dd47d43 to aa1e4c9 Compare January 21, 2026 19:46

address comments

d6eee18

pabloantoniom approved these changes Jan 22, 2026

View reviewed changes

Mr-Anyone and others added 2 commits January 22, 2026 21:36

Ran clang-format

d862093

Merge branch 'develop' into pr-1999

5287c42

Mr-Anyone merged commit 9f84329 into ROCm:develop Jan 23, 2026
7 of 14 checks passed

Conversation

Mr-Anyone commented Jan 15, 2026

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI Jan 16, 2026

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!

justinrosner left a comment

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!

justinrosner left a comment

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!

pabloantoniom left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants