Add scatter with reduction support

[denix56]

Why
Triton doesn’t currently support scatter, which is needed for index-based writes and parity with common tensor frameworks. This PR adds scatter support with broadcastable indices
and optional reduction semantics (via combine_fn and include_self).

What

• Add scatter support in the compiler lowering and runtime.
• Allow broadcastable index in scatter (broadcast to src shape in semantic layer).
• Add reduction support via combine_fn, with include_self controlling whether original dst values participate.
• Add unit tests.

Testing

• python -m pytest -vv python/test/unit/language/test_core.py::test_scatter
• python -m pytest -vv python/test/unit/language/test_core.py::test_scatter_broadcast_index
• python -m pytest -vv python/test/unit/language/test_core.py::test_scatter_reduce

New contributor declaration

• I am not making a trivial change, such as fixing a typo in a comment.

• I have written a PR description following these
  rules.

• I have run pre-commit run --from-ref origin/main --to-ref HEAD.

• Select one of the following.

  • I have added tests.
    • /test for lit tests
    • /unittest for C++ tests
    • /python/test for end-to-end tests
  • This PR does not need a test because FILL THIS IN.

• Select one of the following.

  • I have not added any lit tests.
  • The lit tests I have added follow these best practices,
    including the "tests should be minimal" section. (Usually running Python code
    and using the instructions it generates is not minimal.)

[denix56]

I would appreciate any suggestions, especially regarding match_any implementations as I don't have access to AMD gpu for testing.

[ThomasRaoux]

I think this is out of scope for what we want in Triton unless there are real life use cases for the scatter in ML.

[denix56]

I think this is out of scope for what we want in Triton unless there are real life use cases for the scatter in ML.

I can immediately think of GNNs, K-nearest neighbors for neural networks working on point clouds, sparse tensor ops.

[CRobeck]

@denix56 is this something you would consider implementing as an extension?
https://github.com/triton-lang/triton/tree/main/examples/plugins
https://github.com/triton-lang/triton-ext
cc: @plotfi @sjw36

[denix56]

@denix56 is this something you would consider implementing as an extension? https://github.com/triton-lang/triton/tree/main/examples/plugins https://github.com/triton-lang/triton-ext cc: @plotfi @sjw36

yes, sounds like a good idea!

[CRobeck]

@denix56 is this something you would consider implementing as an extension? https://github.com/triton-lang/triton/tree/main/examples/plugins https://github.com/triton-lang/triton-ext cc: @plotfi @sjw36

yes, sounds like a good idea!

I have a draft PR that I need to finish up for custom top level DSL ops but it can be done in stages with just adding the ops at the TTGIR level first and invoking with IR over rides. The existing infrastructure should support that today.