Fix use-after-free in rock-linalg-align to avoid app crash when enable debug logging

[kazhang2]

Motivation

The LinalgAlignRewriter::drainWorklist logging callbacks print the matched operation's name via op->getName(). The MemcpyRewritePattern::matchAndRewrite erase their root operation before returning success, so the success callback dereferences freed memory and crashes with an access violation (0xC0000005) inside OperationName::Impl::getName when the pass runs with -debug-only=rock-linalg-align.

Technical Details

The three logging callbacks passed to PatternApplicator::matchAndRewrite (canApply, onFailure, onSuccess) print the matched operation via op->getName(). The memcpy and reduction rewrite patterns erase their root operation (eraseOp) in MemcpyRewritePattern::matchAndRewrite

rocMLIR/mlir/lib/Dialect/Rock/Transforms/AlignTiling.cpp

Lines 1114 to 1116 in 0ceaf1c

	b.eraseOp(copy);
	return success();
	}

before returning success. When onSuccess then runs and calls op->getName(), it reads freed memory. This is only reachable in #ifndef NDEBUG builds with the debug flag set; without the flag the buggy logging never executes and the pass works correctly, which is why it went unnoticed.

Observed in practice when MIGraphX (migraphx_gpu.dll) compiles an ONNX model with ROCMLIR_DEBUG_FLAGS=--debug-only=rock-linalg-align: the process dies mid-log-line at Matched ... MemcpyRewritePattern pattern on ..., right after
** Erase : "memref.copy".

Test Plan

reproduce step
run in powershell
@'
module {
func.func @repro(%output: memref<1xf32>) attributes {rock.kernel} {
%zero = arith.constant 0.0 : f32
%regs = rock.alloc() : memref<1xf32, #gpu.address_space>
rock.fill(%regs, %zero) : memref<1xf32, #gpu.address_space>, f32
%tmp = memref.alloc() : memref<1xf32>
rock.threadwise_write_all {forceUnroll, useIndexDiffs}
%regs -> by set
: memref<1xf32, #gpu.address_space> -> memref<1xf32>
memref.copy %tmp, %output : memref<1xf32> to memref<1xf32>
return
}
}
'@ | & .\build\bin\rocmlir-opt.exe --rock-linalg-align --debug-only=rock-linalg-align

Test Result

rocmlir-opt.exe --rock-linalg-align --debug-only=rock-linalg-align
pass

Submission Checklist

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