justinfargnoli wrote: This change looks good to me. However, I haven't worked much with the `gpu` dialect or MLIR, so I'd prefer to let someone more familiar with the code review and approve it. https://github.com/llvm/llvm-project/pull/96783