Skip to content

Add decomposition for aten.miopen_batch_norm#4482

Open
rkayaith wants to merge 1 commit intollvm:mainfrom
iree-org:miopen-batch-norm
Open

Add decomposition for aten.miopen_batch_norm#4482
rkayaith wants to merge 1 commit intollvm:mainfrom
iree-org:miopen-batch-norm

Conversation

@rkayaith
Copy link
Member

@rkayaith rkayaith commented Mar 3, 2026

On ROCm 7.1+ with MIOpen 3.5+, PyTorch decomposes BatchNorm2d with channels-last layout to aten.miopen_batch_norm instead of aten._native_batch_norm_legit_functional. torch-mlir didn't support this op, causing legalization to fail. This registers the op and just decomposes it to aten.native_batch_norm, like pytorch inductor does.

Fixes #4476

Signed-off-by: Rahul Kayaith <rkayaith@gmail.com>
@rkayaith rkayaith force-pushed the miopen-batch-norm branch from b754194 to 06e640f Compare March 3, 2026 17:31
@rkayaith rkayaith marked this pull request as ready for review March 3, 2026 18:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add support for aten.miopen_batch_norm

1 participant