llvm / llvm/torch-mlir

torch.aten.batch_norm is relying on incorrect annotations

Open
#663 1 comment 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
C++
Stars
1.9k
Forks
736
Avg merge
5d 22h
Merged PRs (30d)
15

Description

As discussed in https://github.com/pytorch/pytorch/issues/73050#issuecomment-1051382044, there are a few ops that don't correctly annotate that they mutate their operands. It seems like those are `aten::batch_norm` and `aten::layer_norm`.

When I revamped our ODS generator code, I tried correcting those exceptions w.r.t. the HasValueSemantics and ReadOnly traits, but it seems we were relying on the old, incorrect annotation (which I think was okay, since it only matters in the training case, which we haven't implemented yet)
https://github.com/llvm/torch-mlir/blob/a5fe0cf06308af3a372e40ae927995f7920fb55d/python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/registry.py#L258

To work on this, you just have to uncomment the code linked above and regenerate the ODS for `torch.aten.batch_norm` and see what breaks in the tests. I dug into it a little bit, and it seems like we will need some special handling in ReduceOpVariants to convert torch.aten.batch_norm to value semantics when `training == false`.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the linked registry.py code and uncomment the annotation changes for torch.aten.batch_norm and torch.aten.layer_norm. Regenerate the ODS and run the affected tests, then inspect ReduceOpVariants for the required special handling when training == false. Done means the regenerated annotations and tests work without breaking existing behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
30/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.