Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[MLIR][Arith] Add denormal attribute to binary/unary operations #112700

Merged
merged 7 commits into from
Nov 26, 2024

Conversation

chelini
Copy link
Contributor

@chelini chelini commented Oct 17, 2024

Add support for denormal in the Arith dialect (binary and unary operations).
Denormal are attached to every operation, and they can be of three different
kinds:

  1. ieee, denormal are preserved and processed as defined by IEEE 754 rules.

  2. preserve sign, a mode where denormal numbers are flushed to zero, but the
    sign of the zero (+0 or -0) is preserved.

  3. positive zero, a mode where all denormal numbers are flushed to positive zero
    (+0), ignoring the sign of the original number.

Denormal refers to both the operands and the result. Currently only lowering for
ieee is supported.

@llvmbot
Copy link
Member

llvmbot commented Oct 17, 2024

@llvm/pr-subscribers-mlir-linalg
@llvm/pr-subscribers-mlir-core
@llvm/pr-subscribers-mlir-arith
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-math

Author: lorenzo chelini (chelini)

Changes

The Flush to Zero (FTZ) modifier is used in floating-point arithmetic to set  very small numbers, known as denormal or subnormal numbers, to zero. FTZ is done to improve performance, as handling these small numbers can slow down computations. Note that this attribute does not specify if the rounding happens toward positive or negative zero since it is architecture (or vendor)-dependent.


Full diff: https://github.com/llvm/llvm-project/pull/112700.diff

3 Files Affected:

  • (modified) mlir/include/mlir/Dialect/Arith/IR/ArithBase.td (+4-2)
  • (modified) mlir/test/Dialect/Arith/ops.mlir (+15-1)
  • (modified) mlir/test/Dialect/Math/ops.mlir (+1-1)
diff --git a/mlir/include/mlir/Dialect/Arith/IR/ArithBase.td b/mlir/include/mlir/Dialect/Arith/IR/ArithBase.td
index 19a2ade2e95a0e..b7dc984817a333 100644
--- a/mlir/include/mlir/Dialect/Arith/IR/ArithBase.td
+++ b/mlir/include/mlir/Dialect/Arith/IR/ArithBase.td
@@ -108,12 +108,13 @@ def FASTMATH_NO_SIGNED_ZEROS : I32BitEnumAttrCaseBit<"nsz",      3>;
 def FASTMATH_ALLOW_RECIP     : I32BitEnumAttrCaseBit<"arcp",     4>;
 def FASTMATH_ALLOW_CONTRACT  : I32BitEnumAttrCaseBit<"contract", 5>;
 def FASTMATH_APPROX_FUNC     : I32BitEnumAttrCaseBit<"afn",      6>;
+def FASTMATH_FTZ             : I32BitEnumAttrCaseBit<"ftz",      7>;
 def FASTMATH_FAST            : I32BitEnumAttrCaseGroup<
     "fast",
     [
       FASTMATH_REASSOC,         FASTMATH_NO_NANS,     FASTMATH_NO_INFS,
       FASTMATH_NO_SIGNED_ZEROS, FASTMATH_ALLOW_RECIP, FASTMATH_ALLOW_CONTRACT,
-      FASTMATH_APPROX_FUNC]>;
+      FASTMATH_APPROX_FUNC, FASTMATH_FTZ]>;
 
 def FastMathFlags : I32BitEnumAttr<
     "FastMathFlags",
@@ -121,7 +122,8 @@ def FastMathFlags : I32BitEnumAttr<
     [
       FASTMATH_NONE,           FASTMATH_REASSOC,         FASTMATH_NO_NANS,
       FASTMATH_NO_INFS,        FASTMATH_NO_SIGNED_ZEROS, FASTMATH_ALLOW_RECIP,
-      FASTMATH_ALLOW_CONTRACT, FASTMATH_APPROX_FUNC,     FASTMATH_FAST]> {
+      FASTMATH_ALLOW_CONTRACT, FASTMATH_APPROX_FUNC,     FASTMATH_FTZ,
+      FASTMATH_FAST]> {
   let separator = ",";
   let cppNamespace = "::mlir::arith";
   let genSpecializedAttr = 0;
diff --git a/mlir/test/Dialect/Arith/ops.mlir b/mlir/test/Dialect/Arith/ops.mlir
index f684e02344a517..0209d612273e6d 100644
--- a/mlir/test/Dialect/Arith/ops.mlir
+++ b/mlir/test/Dialect/Arith/ops.mlir
@@ -1127,7 +1127,7 @@ func.func @fastmath(%arg0: f32, %arg1: f32, %arg2: i32) {
 // CHECK: {{.*}} = arith.addf %arg0, %arg1 fastmath<nnan,ninf> : f32
   %7 = arith.addf %arg0, %arg1 fastmath<nnan,ninf> : f32
 // CHECK: {{.*}} = arith.mulf %arg0, %arg1 fastmath<fast> : f32
-  %8 = arith.mulf %arg0, %arg1 fastmath<reassoc,nnan,ninf,nsz,arcp,contract,afn> : f32
+  %8 = arith.mulf %arg0, %arg1 fastmath<reassoc,nnan,ninf,nsz,arcp,contract,afn,ftz> : f32
 // CHECK: {{.*}} = arith.cmpf oeq, %arg0, %arg1 fastmath<fast> : f32
   %9 = arith.cmpf oeq, %arg0, %arg1 fastmath<fast> : f32
 
@@ -1161,3 +1161,17 @@ func.func @intflags_func(%arg0: i64, %arg1: i64) {
   %3 = arith.shli %arg0, %arg1 overflow<nsw, nuw> : i64
   return
 }
+
+// CHECK-LABEL: flush_to_zero
+// CHECK-SAME: %[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32
+func.func @flush_to_zero(%arg0: f32, %arg1: f32) {
+  // CHECK: %{{.+}} = arith.addf %[[ARG0]], %[[ARG1]] fastmath<ftz> : f32
+  // CHECK-NEXT: %{{.+}} = arith.subf %[[ARG0]], %[[ARG1]] fastmath<ftz> : f32
+  // CHECK-NEXT: %{{.+}} = arith.mulf %[[ARG0]], %[[ARG1]] fastmath<ftz> : f32
+  // CHECK-NEXT: %{{.+}} = arith.divf %[[ARG0]], %[[ARG1]] fastmath<ftz> : f32
+  %0 = arith.addf %arg0, %arg1 fastmath<ftz> : f32
+  %1 = arith.subf %arg0, %arg1 fastmath<ftz> : f32
+  %2 = arith.mulf %arg0, %arg1 fastmath<ftz> : f32
+  %3 = arith.divf %arg0, %arg1 fastmath<ftz> : f32
+  return
+}
diff --git a/mlir/test/Dialect/Math/ops.mlir b/mlir/test/Dialect/Math/ops.mlir
index 7e45d9bc6f74a9..6accd09647b8bb 100644
--- a/mlir/test/Dialect/Math/ops.mlir
+++ b/mlir/test/Dialect/Math/ops.mlir
@@ -289,7 +289,7 @@ func.func @fastmath(%f: f32, %i: i32, %v: vector<4xf32>, %t: tensor<4x4x?xf32>)
   // CHECK: math.trunc %[[F]] fastmath<fast> : f32
   %0 = math.trunc %f fastmath<fast> : f32
   // CHECK: math.powf %[[V]], %[[V]] fastmath<fast> : vector<4xf32>
-  %1 = math.powf %v, %v fastmath<reassoc,nnan,ninf,nsz,arcp,contract,afn> : vector<4xf32>
+  %1 = math.powf %v, %v fastmath<reassoc,nnan,ninf,nsz,arcp,contract,afn,ftz> : vector<4xf32>
   // CHECK: math.fma %[[T]], %[[T]], %[[T]] : tensor<4x4x?xf32>
   %2 = math.fma %t, %t, %t fastmath<none> : tensor<4x4x?xf32>
   // CHECK: math.absf %[[F]] fastmath<ninf> : f32

@chelini chelini requested a review from vzakhari October 17, 2024 12:34
@chelini chelini requested a review from grypp October 17, 2024 12:34
Copy link
Member

@kuhar kuhar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if this should be something else than a fast math flag, e.g., a rounding mode or an execution mode. In SPIR-V it's the latter and you can only set it per the whole program: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_execution_mode:~:text=4460-,DenormFlushToZero,-Any%20denormalized%20value, so I don't think we will able to lower FTZ on individual arith ops when lowering to SPIR-V.

@chelini
Copy link
Contributor Author

chelini commented Oct 18, 2024

I wonder if this should be something else than a fast math flag, e.g., a rounding mode or an execution mode. In SPIR-V it's the latter and you can only set it per the whole program: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_execution_mode:~:text=4460-,DenormFlushToZero,-Any%20denormalized%20value, so I don't think we will able to lower FTZ on individual arith ops when lowering to SPIR-V.

Hi Jakub, I would like to have this attribute on a per-op basis since this is what PTX exposes. If you are around next week at the LLVM dev meeting, @matthias-springer , @joker-eph, and I will also be there, and we can have a roundtable discussion on this topic. What do you think?

@kuhar
Copy link
Member

kuhar commented Oct 18, 2024

Sure, sounds good.

@kuhar
Copy link
Member

kuhar commented Oct 21, 2024

We had a brief offline chat with @matthias-springer and @chelini about this today. Leaving a comment just so that I forgot what we talked about. This is based on my understanding, I should verify this:

  • fast math adds assumptions but does not change semantics, should be safe to drop
  • rounding modes change semantics and cannot be simply dropped

I think what we want is a rounding mode attr with well defined behavior that in the general case cannot be dropped / ignored during lowering. If SPIR-V doesn't support something, it should be clear that we are sacrificing correctness during lowering, but this shouldn't block us having these rounding modes for targets that can support it.

@joker-eph
Copy link
Collaborator

LLVM often uses target intrinsics instead for this kind of “exotic” attributes, especially when it can’t be supported consistently across target.
Another aspect is how much the attribute would affect other transformations on this op than just lowerings.

I am not sure what’s the right thing to do here, but I wanted to mention about tradeoffs.

@dcaballe
Copy link
Contributor

I talked to Jakub during the LLVM Dev and I still have reservations about this. I don't see FTZ as significantly different from other flags in FMF that lead to approximate results. If we are concerned about the specific approximation to zero that FTZ enforces, we should consider a more open-ended “approxDenormals” FMF. Honestly, I would be surprised if functions used for “approxFunc” FMF are handling denormals properly. In my opinion, “approxDenormals” should allow for the approximation of denormals, and omitting this flag should result in more accurate results without necessarily changing the semantics. I wouldn’t walk the FMF vs rounding mode path. One could argue that “allowContract” is actually a “rounding mode” as, beyond contracting operations, what this flag states is that it’s ok to not round the intermediate results of the operations being contracted.

@chelini chelini force-pushed the lchelini/mlir/flush-to-zero branch from 98b22fd to 146b2c8 Compare November 21, 2024 18:49
@llvmbot llvmbot added mlir:core MLIR Core Infrastructure mlir:linalg labels Nov 21, 2024
@chelini chelini changed the title [mlir][Arith] Add FTZ (Flush-to-Zero) fast-math flag [MLIR][Arith] Add denormal attribute to binary/unary operations Nov 21, 2024
Copy link
Member

@kuhar kuhar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The direction looks good overall

mlir/test/Dialect/Arith/invalid.mlir Outdated Show resolved Hide resolved
mlir/include/mlir/Dialect/Arith/IR/ArithBase.td Outdated Show resolved Hide resolved
Copy link
Member

@kuhar kuhar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. I haven't reviewed the logic carefully, so please get a second approval from someone who can do that.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp Outdated Show resolved Hide resolved
@rengolin
Copy link
Member

Hi Jakub, I would like to have this attribute on a per-op basis since this is what PTX exposes.

Additional data here: in Arm, VFP is IEEE compliant while NEON is not (has denormals). So per operation would be a way to distinguish. However, the front-end would have to distinguish between scalar (VFP) IEEE compliant and vector (NEON) with denormals. Not a trivial thing to get right.

@banach-space @smithp35

Add support for denormal in the Arith dialect (binary and unary operations).
Denormal are attached to every operation, and they can be of three different
kinds:

1) ieee, denormal are preserved and processed as defined by IEEE 754 rules.

2) preserve sign, a mode where denormal numbers are flushed to zero, but the
sign of the zero (+0 or -0) is preserved.

3) positive zero, a mode where all denormal numbers are flushed to positive zero
(+0), ignoring the sign of the original number.

Denormal refers to both the operands and the result.
use name to make sure we don't print ieee
Copy link
Contributor

@dcaballe dcaballe left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! The representation and initial lowering makes sense to me. We can dig a bit more into the scalar vs Neon case separately.

@chelini chelini merged commit 4a7b56e into llvm:main Nov 26, 2024
8 checks passed
@chelini
Copy link
Contributor Author

chelini commented Nov 26, 2024

Thanks! The representation and initial lowering makes sense to me. We can dig a bit more into the scalar vs Neon case separately.

Perfect, thank you all. Let' keep discussing this. As I next step I will connect e2e with the LLVM dialect.

antiagainst added a commit to triton-lang/triton that referenced this pull request Dec 4, 2024
This pulls in the AMDGPU backend support for the
gfx950 target.

We need to fix the rewrites in `Combine.td` given that
llvm/llvm-project#112700 adds
a new attribute for denorm mode for `arith.addf`.

---------

Co-authored-by: Lei Zhang <[email protected]>
jataylo pushed a commit to jataylo/triton that referenced this pull request Dec 4, 2024
This pulls in the AMDGPU backend support for the
gfx950 target.

We need to fix the rewrites in `Combine.td` given that
llvm/llvm-project#112700 adds
a new attribute for denorm mode for `arith.addf`.

---------

Co-authored-by: Lei Zhang <[email protected]>
(cherry picked from commit 1d5e9a2)
jataylo pushed a commit to jataylo/triton that referenced this pull request Dec 4, 2024
This pulls in the AMDGPU backend support for the
gfx950 target.

We need to fix the rewrites in `Combine.td` given that
llvm/llvm-project#112700 adds
a new attribute for denorm mode for `arith.addf`.

---------

Co-authored-by: Lei Zhang <[email protected]>
(cherry picked from commit 1d5e9a2)
jataylo pushed a commit to jataylo/triton that referenced this pull request Dec 5, 2024
This pulls in the AMDGPU backend support for the
gfx950 target.

We need to fix the rewrites in `Combine.td` given that
llvm/llvm-project#112700 adds
a new attribute for denorm mode for `arith.addf`.

---------

Co-authored-by: Lei Zhang <[email protected]>
(cherry picked from commit 1d5e9a2)
jataylo pushed a commit to jataylo/triton that referenced this pull request Dec 6, 2024
This pulls in the AMDGPU backend support for the
gfx950 target.

We need to fix the rewrites in `Combine.td` given that
llvm/llvm-project#112700 adds
a new attribute for denorm mode for `arith.addf`.

---------

Co-authored-by: Lei Zhang <[email protected]>
(cherry picked from commit 1d5e9a2)
Jokeren added a commit to triton-lang/triton that referenced this pull request Dec 6, 2024
Update

Update

Update

Update

Update

Use pytest' `tmp_path` in `test_irsource.py` (#5145)

Signed-off-by: Anatoly Myachev <[email protected]>

[TEST] Make mixed matmul test deterministic (#5151)

This prevents surprises when some value may go above the tolerance
threshold

Fix `gtest_discover_tests` timeout argument (#5149)

`gtest_discover_tests` runs the built unittest executable to create a
distinct CMake target for every individual unittest in each executable.
However, this was previously noted to time out on MacOS frequently
(because MacOS scans newly built executables for viruses, or
something...) but the timeout argument was incorrectly specified.

[Triton] Remove upstream bug workaround (NFC) (#5152)

Upstream handling of splatted bools in `DenseElementsAttr` was fixed, so
the workaround can be removed when lowering `arith.constant` to
TritonGPU.

Co-authored-by: peterbell10 <[email protected]>

[Triton] Generate local MLIR reproducers when possible (#5155)

By setting a reproducer path, the pass manager will dump a standard MLIR
reproducer before each pass manager invocation. This PR also enables
additional local crash reproducer generation (to the same path set
through the env var), which tries to narrow down the specific pass that
failed, if the pass pipeline fails at any point.

Revert "[AMD][Pipeliner] Improve clustering and add prefetch (#4881)" (#5157)

This reverts commit cc25374
due to perf regressions.

[IR] Add typing for tensor descriptor types (#5147)

Currently tensor descriptors are just typed as `!tt.ptr<i8>` which is
exposing the assumption it's using a TMA descriptor. This changes it to
a custom type `!tt.tensordesc<tensor<...>>` which is lowered to a
pointer type in the LLVM IR.

I also add two new IR Ops which are used to cast between pointers and
tensordesc objects.
```mlir
tt.reinterpret_tensor_descriptor %ptr : !tt.ptr<i8> to !tt.tensordesc<...>
triton_nvidia_gpu.tensor_desc_to_tma_ptr %desc : !tt.tensordesc<...> -> !tt.ptr<i8>
```

Really both of these should be nvidia-specific but the first is exposed
in the triton IR to keep support for the by-value TMA descriptor API
around while we figure out if it's possible to update to the new style.

Load backend dialects in `IRSource` to make sure `parse_mlir_module` works for third_party backends (#5146)

The changes from #4924 do not
take into account the situation when `ttgir` level contains dialects
defined in third_party plugins (at least that's my understanding).

I'd also like to point out that the second use of `parse_mlir_module`
function (via `parse` function call) happens after the dialects are
loaded for the backend as well, which is why I thought my changes make
sense.

I hope this implementation will suit Triton, or maybe one can suggest
other options.

---------

Signed-off-by: Anatoly Myachev <[email protected]>

[BACKEND][NVIDIA] Add DotOp Hoisting Pass for WGMMA and Add Lowering for SMEM-to-MMAv3 DotOp Copy (#5003)

Hopper has two kinds of WGMMAs, "SS" (both operands in shmem) and "RS"
(LHS operand A in registers).
In cases where we apply elementwise operations on A before WGMMA, Triton
previously will copy A from global memory (GMEM) into registers (RF),
perform the elementwise ops, and then copy to shared memory (SMEM) to
perform SS WGMMA.

This PR adds an optimization for the case above to use RS GEMM. This
requires the following changes:

- In TritonGPU OptimizeDotOperands pass, add optimizations to change SS
GEMM into RS GEMM.
- Add TritonGPU -> LLVM lowering for copying from SMEM to RF in MMA v3
dotOperand layout.

NOTE: This may not see perf gain, and may even see perf loss, for
certain shapes (e.g. small-K), and additional optimizations are in a
separate [PR](openxla#19) (still more
optimizations are WIP). Please advise on the merging strategy.

Restore the CentOS 7 build (#5158)

We likely need it for the PyTorch 2.6 release

[BACKEND] Add folder for `addptr(ptr, 0) -> ptr` (#5166)

I noticed this rather obvious pattern was missing. It might come up for
example if you have an expression like:
```python
ptrs = ptr + y_stride * tl.arange(0, YBLOCK)[:, None]
```
and the `YBLOCK` is set to 1 during autotuning.

[TritonGPU] Fix incorrect mask operand used in for loop pipeliner (#5161)

When the OOB values for a `tt.load` are non-zero, the for loop pipeliner
needs to generate an `arith.select` to mask the loaded values with the
default OOB value. However, if the load memory requires a layout change,
the wrong mask operand was being passed to the `arith.select`, causing a
shape mismatch. The fix is to just use the same mask operand of the
origianl `tt.load` op.

Fixes #4739

[BACKEND] Cleanup redundant broadcast combine pattern (#5167)

Summary of changes:
- Remove `broadcast(cst) -> cst` from the triton-combine pass since it's
redundant with the existing folder.
- Reorder the triton-combine pass to come after the canonicalize pass,
to simplify pattern matching
- Cleanup patterns in triton-reorder-broadcast that called
`Op::canonicalize` in favor of `Op::getCanonicalizationPatterns`.

[AMD] NFC: Drop duplicated moveUpTranspose (#5168)

It was duplicated due to resolving merge conflicts.

[Triton] Default diagnostic handler only filters for errors (#5173)

A regular SourceMgrDiagnosticHandler is causing all remarks to be
emitted even if the user doesn't ask for it!

[AMD] Refactor instruction scheduling hints (#5144)

- Renamed instruction scheduling variants
- Enabled `buffer-ops` for `local-prefetch`
- Added  documentation regarding current variants

---------

Co-authored-by: Lei Zhang <[email protected]>

[AMD] Enable mixed precision matmul test (#5177)

This commit enables mixed precision matmul test
for AMD backend. For FP8 E4M3, we test
`fp8e4m3fnuz` given that's natively supported on
MI300 series.

Update to llvm/llvm-project@bd9145c8c213 (#5180)

This pulls in llvm/llvm-project@bd9145c8c213
to enable ASan on AMD backend.

[AMD] Implement RepOrder for AMD MMA layouts (#5126)

Implement RepOrder methods for MFMA and WMMA layouts. Both layouts have
row major rep layout. Also,
isTranspose flag in MFMA layout does not affect RepOrder, meaning
RepOrder is row major in both cases.

Co-authored-by: Ognjen Plavsic <[email protected]>

[BACKEND] Fix ProgramPoint passing in AxisInfoAnalysis (#5181)

Fixes #5122.

The `ProgramPoint`
[here](https://github.com/triton-lang/triton/blob/0bd30a2f3192204c5a50d5ffde27ad8493f6c026/lib/Analysis/AxisInfo.cpp#L1087)
is created on the stack. Then its address is
[passed](https://github.com/triton-lang/triton/blob/0bd30a2f3192204c5a50d5ffde27ad8493f6c026/lib/Analysis/AxisInfo.cpp#L1088-L1089)
to the MLIR `SparseAnalysis` code, where it is [added as a
dependency](https://github.com/llvm/llvm-project/blob/33ff9e43b4c5bdc3da31c6b11ad51d35a69bec5f/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp#L311)
and later
[dereferenced](https://github.com/llvm/llvm-project/blob/33ff9e43b4c5bdc3da31c6b11ad51d35a69bec5f/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp#L90).
By the time the `ProramPoint` is dereferenced in the
`AbstractSparseForwardDataFlowAnalysis::visit`, the
`AxisInfoAnalysis::visitForOpInductionVar` will have finished and the
`ProgramPoint` stack variable destroyed. This leads to a segfault (which
can be reproed on the base rev with the lit test added in this PR).

The code modified in this PR was originally added in #4927, in
conjunction with updating the `llvm-project` hash to `b5cc222d7429`.
However, as noted in llvm/llvm-project#110344
(the `llvm-project` PR that has made the refactoring prompting the
`AxisInfo.cpp` change in #4927):

> For dense forward data-flow analysis and other analysis (except dense
backward data-flow analysis), the program point corresponding to the
original operation can be obtained by `getProgramPointAfter(op)`

As the `AxisInfoAnalysis` (in Triton) inherits from
`SparseForwardDataFlowAnalysis` (in MLIR), in this PR we follow the
above which resolves the segfault issue (as the `ProgramPoint` is now
stored in the instance-level state of the pass).

P.S. The lit test added in this PR is not exactly minimal. However, I
did my best to minimize it starting from the 400-line repro TTGIR in

[INTERPRETER] Fix argument passing for internal parameters in function declarations (#5169)

[NFC] Use reference instead of copies in few places (#5118)

Apply fixes suggested by coverity static analysis.

Signed-off-by: Anatoly Myachev <[email protected]>

[BACKEND] Add missing precondition in optimize acc init (#5184)

We need scalar select to be able to do this optimization.

[BACKEND] Fix accumulator init optimization for integer matmuls (#5192)

[AMD][Pipeliner] Reland "Improve clustering and add prefetch" (#5175)

This unreverts commit 38c6284
to reland #4881
with the following fixes:

* Still keep `scheduleGlobalLoadLocalStore` as original--it turns
to be not totally ready to replace yet. Further iteration on it needed.
* Turn on `TRITON_HIP_STREAM_PREFETCH` if the instruction
  scheduling variant is `local-prefetch`, given it's needed there.

---------

Co-authored-by: Lei Zhang <[email protected]>

[AMD] Define an extract slice operation (#4804)

This commit introduces an extract_slice operation for AMD backend
to enable extracting slice of a tensor in registers without data
exchange.
It enables breaking down large tiles of tensors into smaller ones
for better instruction interleaving and scheduling.

This can be useful for hiding global memory latency when a global
load/store can be efficiently split into several loads/stores to be
overlapped with compute fo attention.

[BACKEND] Fix getElemsPerThread for mmav3 dot operand (#5189)

In mmav3 case the number of elements per threads should be independent
of the element type, we should only consider kWidth.
TODO: it should also be true for MMAv2 but the logic is a bit more
complicated.

Also enable larger block_m in mixed mode tests to exercise MMAv3 case

[INTERPRETER][NFC] Rename `tensor_shape` -> `block_shape` in interpreter (#5195)

`tensor_shape` is a confusing name and doesn't match block pointer's
semantic.
`block_shape` is much clearer.

[LAYOUTS] Implement LL conversion for DotOperand(Hopper) (#5193)

We also rewrite the way we implement DotOperand(Ampere) and mma Ampere
to promote code reusing. I also started using what I believe is a rather
compact pattern to write these things, where you first call `identiyND`
with the `repOrder`, which gives you an LL with the dims in the correct
order, and then you construct the final layout by specifying the tiles
by multiplying `identity1D` maps. Using this allowed me to heavily
simplify the handling of the `warps` of `DotOperand` which used to be a
tad messy.

Update README.md to remove triton conference (#5198)

It happened two months ago

[PROTON] Add `proton.state` utility (#5110)

`state` is different from `scope` in several ways:

1. State is not recursive; each operation can have only a single state.
Inner most state will overwrite the outer most state.
2. A states is a suffix, meaning that the original call path will append
a state above the name of each kernel.
3. State is compatible with both Python and shadow contexts.

[CI] remove unused inductor workflows (#5073)

These tests have completely offloaded torch inductor tests to Meta a few
months ago. They are currently disabled on GitHub.

Signed-off-by: Sébastien Han <[email protected]>

[INTERPRETER] Fix lower bound check for block pointers (#5201)

We forgot to check `offset >= 0` previously.

Now that it should match the semantic in the GPU backend

https://github.com/triton-lang/triton/blob/7bce3613755e26953518962d02315dfd343dc50c/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp#L136

[IR] Remove memdesc from `tt.trans` and implements `ttg.memdesc_trans` (#5194)

[LLs] [BE] Simplify identityND (#5199)

The auxiliary function `identityND` used to take an `order` parameter,
that comes from triton, and a set of dimensions. Now, the order in
triton is defined wrt. `dim0..dim<rank-1>`, so the dimension arg was
redundant. This was quite confusing.

We see that in all the uses of `identiyND`, we would pass the canonical
dimensions, other than in one that we simply remove as it was not
necessary.

We remove the dims arg and simply return a layout with output dims
`dim0..dim<rank-1>`.

[MXFP] Fix packing for mxfp4 type (#5197)

When packing we should have element 0 in the lower bits, until this PR
it was in higher bits.

[LAYOUTS] Unify the implementation of getShapePerCTA (#5183)

We unify it and simplify its API (it was taking an unused `shape`
parameter). While doing this, we found that the previous implementation
was incorrect at least for `AMDWmmaEncodingAttr`, as this layout was
using the shape parameter.

Interestingly enough the doc in the header file for this function noted
that the function is indeed independent of the tensor shape, even though
the function does take a shape as an input!

https://github.com/triton-lang/triton/blob/0bd30a2f3192204c5a50d5ffde27ad8493f6c026/include/triton/Dialect/TritonGPU/IR/Dialect.h#L113-L114

[BACKEND] Use the LL API to replace the using of legacy layout attribute API. (#5196)

The util function `getDistributedLayoutStr` uses the `DistributedLayout`
attribute interface, which is not flexible for third-party extensions.
Use the `getInDimSize` of the `LinearLayout`, which is better since the
legacy layout has been converted to the `LinearLayout`.

There is no new test case since it is only a change in API usage.

[CI] Fix ccache cache restoration to improve build times (#5202)

This improves a warm-cache macOS build from ~25 mins to 2 mins.

[CI] Fix `du` failling if cache restore fails (#5206)

Follow up to #5202

It's currently failing with the error
```
du: /Users/runner/.triton/**: No such file or directory
Error: Process completed with exit code 1.
```
which happens because even though the `.triton` directory exists, it is
empty. This instead uses du on `.triton` with a depth of 1.

[BACKEND][LAYOUT] Use LL for AMDMfma related layout conversions (#5210)

[BUILD] Add option to limit number of parallel link jobs (#5212)

[CI] Fix cache not saving (#5213)

1. [CI] Fix cache not saving

    Re-using the output of the cache restore step was recommended by the
`actons/cache` docs, but it doesn't work here because we actually start
from a clean cache when we run save so there is no output available to
    read.

    The annoyances of testing in the PR but main being a different
    environment.
2. Bump macOS timeout

[LAYOUTS] Implement IR support for LinearLayouts (#5170)

We also exercise this in scale_dot, where we enable support for warps of
arbitrary shape (before we just allowed `[num_warps, 1]`).

With this infra in place, it should be rather easy to move from the
legacy layouts to using LLs to represent all of our layouts.

Something I'm concerned about is the amount of recomputation that
happens when calling methods like `getSizePerThread` and the like, where
we keep recomputing the result. There might be an optimisation
opportunity here where we cache the result of all these functions.

We choose the IR representation of an LL via its canonical form + a
`repOrder` for several reasons:
- It's generally more compact
- It's easier to CSE, so it's easier to see when two layouts are in fact
  the same.
- A technical reason: the `toLinearLayout` function returns a tensor
  with dimensions `dim0, ..., dim<rank-1>`, in other words, it "forgets"
  the repetition order. Without the repetition order, we cannot recover
  the tile size of the argument. In particular, we cannot recover
  `getSizePerThread`. There is an argument to be made about whether
  `getSizePerThread` is useful on its own, or whether it is
  `getElemsPerThread` the real useful abstraction here, but for now, we
  keep both for BC.

[CI] Run tests when CI is manually triggered (#5216)

Currently you can manually call a workflow dispatch, but it won't
actually run the tests because the variable enable_integration isn't
set.

[PROTON] Introduce the Proton dialect as a third-party plugin for intra-kernel perf tooling (#5119)

This PR introduces the `Proton Dialect` to enable intra kernel profiling
and tooling for Triton. As a third-party dialect, it serves as the
building blocks to create 3rd-party perf tools (e.g., profilers,
analysis, modeling) for Triton compiler developers in a compiler-centric
way, such as an intra-kernel latency profiler to understand software
pipelining, warp specialization, and CTA fine-grained orchestration
(e.g., cuda core, tensor core, TMA). Future developments would integrate
this dialect with the existing Proton backend profiling infrastructure
to make it a powerful and general perf tool utility. As a first step,
this PR adds some basic boilerplate code and mechanics, and the
`proton.record` op for the `Proton Dialect`.

---------

Co-authored-by: Yuanwei Fang <[email protected]>
Co-authored-by: Keren Zhou <[email protected]>

[DRAFT] Completely remove `MemDesc` from the Triton dialect (#5208)

After this PR, `MemDesc` will be a type only in the TritonGPU dialect,
as will the `TensorOrMemDesc` interface.

[AMD] Prevent wrong reordering of scf operations (#5203)

The pass was reordering scf.if operations without checking the extra
dependencies coming from the region.
For now just prevent this case although this part of the code might
still be fragile.

[AMD] Cover default case in MfmaGroup (#5218)

If you build using the `CMakeLists.txt` and not `setup.py` and you build
in `Release` then you get

```
/__w/triton/triton/third_party/amd/lib/TritonAMDGPUTransforms/MfmaGroup.cpp: In function ‘std::pair<mlir::Type, mlir::Type> mlir::TypesFromMfmaId(MLIRContext*, MfmaTypeId)’:
Warning: /__w/triton/triton/third_party/amd/lib/TritonAMDGPUTransforms/MfmaGroup.cpp:240:1: warning: control reaches end of non-void function [-Wreturn-type]
```

Allow Layouts to propogate to local_load (#5219)

While working on some higher dimension tensor kernels, I noticed poor
performance due to the fact that layouts wouldn't propagate to local
loads. Since we do allow layout folding with local store and local
alloc, this seems like a bit of an oversight.

The change gives a 40% speed improvement on certain kernels for NVidia
GPUs.

This also removes asserts in lowering for higher dimensional kernels. As
far as I can tell, those restrictions aren't required in practice.

- [x] I am not making a trivial change, such as fixing a typo in a
comment.
- [x] I have written a PR description following these
[rules](https://cbea.ms/git-commit/#why-not-how).
- [x] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`.
- [x] I have added tests.
- [x] The `lit` tests I have added follow these [best
practices](https://mlir.llvm.org/getting_started/TestingGuide/#filecheck-best-practices)

[BACKEND] Fix transpose optimization missed during refactor (#5226)

[AMD] Use warp shuffle for fp8 MFMA to dot operand layout conversion (#5139)

Adding a shortcut case for fp8 MFMA to dot operand layout conversion
that avoids using shared memory, to speed up FP8 attention kernels.

[LAYOUTS] [BE] Simplify Ampere/Hopper paths introduced in #5189 (#5200)

We simplify the implementation of `getElemsPerThread` and strengthen the
preconditions of `getRepForOperand`.

More generally, we should try to minimise the calls to `isAmpere` and
`isHopper` throughout the codebase. I'll do a pass fixing many of these
once we land LLs for `ldmatrix` and Hopper.

[BACKEND] Use LL to simplify redundant elements check and fix related issues (#5225)

Make TMA tests compatible with older CUDA toolchains (#5221)

TMA fences require CUDA toolchain 12.3 or greater, but current gating
does not check the CUDA toolchain version. This causes
`test_experimental_tma.py` to fail when run with older CUDA toolchains.

With cuda-12.0:
```
55 failed, 9 passed in 18.11s
```

With cuda-12.4:
```
64 passed in 11.99s
```

With cuda-12.0:
```
9 passed, 55 skipped in 4.26s
```

With cuda-12.4:
```
64 passed in 11.96s
```

[CMake] Add C as project language (#5217)

If you build with `-DTRITON_BUILD_UT=OFF` on Mac you will get something
like

```
-- Looking for histedit.h
CMake Error at /opt/homebrew/Cellar/cmake/3.30.5/share/cmake/Modules/CheckIncludeFile.cmake:90 (try_compile):
  Unknown extension ".c" for file
-- Looking for histedit.h - not found

    /Users/runner/work/triton/triton/triton-build/CMakeFiles/CMakeScratch/TryCompile-QA06d6/CheckIncludeFile.c

  try_compile() works only for enabled languages.  Currently these are:

    CXX

  See project() command to enable other languages.
Call Stack (most recent call first):
  llvm-bd9145c8-macos-arm64/lib/cmake/llvm/FindLibEdit.cmake:28 (check_include_file)
  llvm-bd9145c8-macos-arm64/lib/cmake/llvm/LLVMConfig.cmake:177 (find_package)
  llvm-bd9145c8-macos-arm64/lib/cmake/mlir/MLIRConfig.cmake:10 (find_package)
```

because `C` isn't an enabled project language.

[AMD] Fix slow compilation due to inlining print calls (#5153)

This PR disables inline of print related functions, which speeds up
compilation of test_scan_layouts dramatically.

---------

Co-authored-by: Lei Zhang <[email protected]>

[AMD] Re-enable overflow test in test_reduce_layouts (#5233)

#5153 fixed
the issue; but we missed enabling one of the disabled
case.

[BACKEND] Fix a missed transpose optimization during refactor (#5236)

Revert "Allow Layouts to propogate to local_load" (#5237)

This is causing some performance regression. I'll investigate and reland
it.
Reverts #5219

Revert "[AMD] Use warp shuffle for MFMA to Dot operand layout conversion (FP8)" (#5240)

It is causing performance regression, revert until it can be
investigated
Reverts #5139

Updated README.md to show the steps for overriding kernel's IR (#5239)

Ensure device context before launching kernel (#3731)

If a kernel is launched on a thread which has not initialized a CUDA
context (as can happen in the linked issue), it will throw an error. A
simple fix is to call `cudaFree(0)` to establish a device context.

Fixes #3729

[LLVM] Update to llvm-project@86b69c3 (#5242)

This includes llvm/llvm-project#115627

[BUILD] Add a stable symlink to llvm in the triton cache (#5234)

Currently the llvm path changes every time the pin updates which makes
it annoying to use the included tools. e.g. I use the tablegen language
server, but currently need to update my editor config every time the
llvm pin changes.

This adds a stable symlink which for me is
`~/.triton/llvm/llvm-macos-x64`. This will always point to the most
recent version of llvm used to build triton.

As a bonus this also refactors the symlink update code which was
copy-pasted a few times.

[PIPELINER] tweak pipeline heuristic (#5247)

Don't pipeline the dot accumulator in the default heuristic.
In the finer grain control will allow user to decide.

Allow Layouts to propogate to local_load (#5219) (#5249)

recommit of #5219

While working on some higher dimension tensor kernels, I noticed poor
performance due to the fact that layouts wouldn't propagate to local
loads. Since we do allow layout folding with local store and local
alloc, this seems like a bit of an oversight.

The change gives a 40% speed improvement on certain kernels for NVidia
GPUs.

This also removes asserts in lowering for higher dimensional kernels. As
far as I can tell, those restrictions aren't required in practice.

- [x] I am not making a trivial change, such as fixing a typo in a
comment.
- [x] I have written a PR description following these
[rules](https://cbea.ms/git-commit/#why-not-how).
- [x] I have run `pre-commit run --from-ref origin/main --to-ref HEAD`.
- [x] I have added tests.
- [x] The `lit` tests I have added follow these [best
practices](https://mlir.llvm.org/getting_started/TestingGuide/#filecheck-best-practices)

<!---
The core Triton is a small number of people, and we receive many PRs
(thank
you!).  To help us review your code more quickly, **if you are a new
contributor (less than 3 PRs merged) we ask that you complete the
following
tasks and include the filled-out checklist in your PR description.**

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->

- [ ] I am not making a trivial change, such as fixing a typo in a
comment.

- [ ] I have written a PR description following these
  [rules](https://cbea.ms/git-commit/#why-not-how).

- [ ] 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](https://mlir.llvm.org/getting_started/TestingGuide/#filecheck-best-practices),
including the "tests should be minimal" section. (Usually running Python
code
    and using the instructions it generates is not minimal.)

Co-authored-by: Matthew Brookhart <[email protected]>

Windows related changes in `CMakeLists.txt` (#5186)

Upstreaming some of our Windows related changes assuming that there is
interest in this
#5094 (comment)
and hoping that it will not make it much more difficult to support this
CMake file.

---------

Signed-off-by: Anatoly Myachev <[email protected]>

[AMD] NFC: Unified header guard in third_party/amd (#5244)

This commit unified the names of header guards in third_party/amd.

[AMD] NFC: Drop v2 Suffix from Stream Pipeline (#5251)

Since StreamPipelineV2 has been the default for a while, this
commit promoted StreamPipelineV2 to the general
StreamPipeline by removing 'v2' suffix.

[NFC] Cleanup references to unused index dialect (#5257)

Also cleans up some includes clang thinks are unused.

[BUILD] Ensure parent directory exists before creating symlinks (#5258)

Fixes #5256

Tmp

[BACKEND] Fold transpose(splat_const) (#5259)

Add folding for a transpose of a splat constant.

---------

Co-authored-by: peterbell10 <[email protected]>

[LAYOUTS] Use LLs for Hopper whenever we wouldn't use ldmatrix (#5235)

The legacy path has some bugs for cases like `kWidth=1`. I'm starting to
port Hopper to use LLs to try to isolate them.

[AMD] NFC: Cleanup namespace hierachy (#5246)

Refactored namespace hierarchy by squeezing separate
namespace hierarchy together.

[AMD] Fix unhandled profile event in RoctracerProfiler (#5252)

Fixes proton unit tests when upgrading to ROCm 6.2 by
adding missing event handlers.

Magic number is replaced with the corresponding enum
value which was added by upgrading the HIP headers
#5077.

Fix Blocked FMA path in isLayoutOK (#5260)

Fixes
https://github.com/triton-lang/triton/pull/5235/files/de18e21ddf5bf03f17f779fef032d53ea87a53a0#r1858955613

[Tutorial] Remove incorrect caching from softmax tutorial (#5162)

The fused softmax implementation in the tutorial precompiles the kernel
to query the register usage of the kernel, based on the parameters used
to specialize the kernel. On top of this, it implements a simple caching
system for this step based on just the block size.

As noted in #4739, this
caching is incorrect, because it's also not keyed on the `num_stages`
constexpr argument or the shapes of the tensors. Since triton already
has its own JIT compilation cache, and this caching bit is not really
relevant to the tutorial, just remove it to get rid of the footgun.

[INSTRUMENTATION] Generalize code in `test_gpuhello.py` (#5263)

Signed-off-by: Anatoly Myachev <[email protected]>

Create an aggregate `check-triton-unit` target (#5150)

This adds a CMake target `check-triton-unit` that builds an runs all
Triton unittests written in gtest. This makes it more conveninent to
rebuild and run all unittests at once with finer granularity (instead of
`ninja; ctest`).

[NFC] Add `test_bessel` into `test_libdevice.py` (#5261)

Just a port of one of our tests. I didn't find any similar ones in
Triton itself, this should increase the test coverage.

Signed-off-by: Anatoly Myachev <[email protected]>

[NFC] Add functional regression test for cummax with bool type (#5264)

This kernel was obtained using PyTorch inductor some time ago.

Signed-off-by: Anatoly Myachev <[email protected]>

[AMD] NFC: Unified comment style (#5248)

Script:

egrep -nrI --exclude-dir "backend" "^\s*/\*+" third_party/amd

[AMD] Upgrade AMD CI docker image (#5230)

This commits updates the CI to use a new docker image that
contains ROCm 6.2.2 with ASan support and PyTorch 2.5.1.

This also switches to ubuntu's default clang toolchain instead
of using the one which comes with ROCm.

Implement `dot_scaled(mmav3)` (#5269)

As per title

[BUILD] Some CMake cleanup/modernisation (#5271)

- Prefer `find_package` over ad-hoc variable passing
- Prefer `target_` api vs global `_directories` apis
- Use `target_link_options` to specify link options instead of
`target_link_libraries`

Closes #5270

[DIALECT] Rename `triton_gpu` to `ttg` and `triton_nvidia_gpu` to `ttng` (#5266)

It may cause changes for downstream tasks but we think it's beneficial
to shorten dialect name and make them consistent. That is, we are using
`tt` to represent the `triton` dialect.

[BACKEND] Fix inline asm bug for multiple packed <32bit output (#5273)

Resolves #5272

- Fixes logic for walking result struct from LLVM InlineAsm in case of
multiple sub-32bit results
- Adds lit test

[NVIDIA][Backend] Add CoalesceAsyncCopy Pass for in-DotOpEnc Upcasting (#5222)

This is a follow-up to the dotOp hoisting optimization for WGMMA
(MMAv3). See
#5003 (comment)

In short, when upcasting operand A in registers prior to WGMMA and when
pipelining is enabled, `AsyncCopyGLobalToLocal`'s src gmem blocked
encoding will have `sizePerThread` > smem view's `vec` (along the
contiguous dimension). This will resulting in multiple `cp.async`
instructions being generated for a contiguous global data segment,
resulting in uncoalesced loads. This was previously confirmed in ncu.
See above comment for an example.

I've added a generalized fix in a new pass after the pipeliner. I've
reused the logic in the LLVM lowering for `AsyncCopyGlobalToLocal` to
calculate the max contiguous copy size. I compare that to the blockEnc's
`sizePerThread` along the inner (contiguous) dimension. If the former is
less than latter, I set the latter to former.

When A is k-major, can verify a small perf improvement and that ncu no
longer reports uncoalesced loads.
When A is m-major, this pass is a no-op because `copy size ==
sizePerThread == 16`

ptal, thanks @ThomasRaoux

[Triton] Add `tl.gather` with a naive codegen implementation (#5262)

This PR adds a `tl.gather` builtin that implements a local gather along
a single axis, with semantics matching `torch.gather`. `tl.gather`
generates a `tt.gather` op, which is piped through the compiler mostly
untouched at the moment, since the codegen is very naive.

The `tt.gather` is implemented by writing the source tensor into shared
memory and then performing a gather out of shared memory, thus it
requires scratch space to be allocated. In a follow-up, I will implement
an optimized layout rule for the op that ensures the gather axis fits
into a single warp, allowing the gather to be implemented using warp
shuffles.

There are other avenues for optimization as well: `tt.gather(tt.load)`
where the load only has one use can be lowered into a DMA from global
memory to shared, and then gather directly from shared.

[NVIDIA][Launcher] Ensure device context is valid before calling getPointer (#5276)

[CMAKE] Add `triton-tensor-layout` dep to lit tests (#5275)

Noticed this when `triton_gpu` was renamed to `ttg`.

[BACKEND] Fix and document logic for creating warp shapes in MMAv3 (#5277)

[NFC] Remove dead code for python<3.8 (#5280)

Signed-off-by: Anatoly Myachev <[email protected]>

[NFC] Remove `CMAKE_VERBOSE_MAKEFILE` var (#5282)

Warning:
```bash
  CMake Warning:
    Manually-specified variables were not used by the project:

      CMAKE_VERBOSE_MAKEFILE
```

Signed-off-by: Anatoly Myachev <[email protected]>

[AMD] Use Linear Layout convertions for AMDWmma (#5255)

Enable LL conwertions for WMMA as well as for MFMA layouts.

See also: #5210

Signed-off-by: Ilya Veselov <[email protected]>

Add tests for 3D local_load local_alloc and relax asserts (#5285)

Also switch 3D dot_operand cases to use linear layout path, This may be
suboptimal in some cases but that solves the functionality problems
which is more important. There is ongoing work from Mario that should
get the code quality to be good again soon.

[Build] Don't require Development.Embed python component (#5287)

This component is missing from the wheel building image, so we need to
make the requirement more specific.

https://github.com/triton-lang/triton/actions/runs/12081047335/job/33689420657#step:6:332

[NFC] Remove unused forOp argument from `setStageCluster` (#5288)

<git-pr-chain>

[NFC] Remove unused forOp argument from `setStageCluster`

1. 👉 #5288 👈 **YOU ARE HERE**
1. #5289
1. #5290

</git-pr-chain>

[PROTON] Don't use designated initializers in `CuptiPCSampling.cpp` as it relates to c++20 (#5291)

Hi @Jokeren,

these changes relates to your PR:
#4674, so I would like to ask
if this was done on purpose? (considering that the project declares
support for the c++17 standard).

I discovered this while trying to compile proton using MSVC. It looked
like this:
`\CuptiPCSampling.cpp(18): error C7555: use of designated initializers
requires at least '/std:c++20'`.

This might also be a good opportunity to ask you about your plans to
transition Triton to `с++20`.

---------

Signed-off-by: Anatoly Myachev <[email protected]>

Add back missing check

Replace triton_gpu with ttg

Update

Update

Update

Define `pytest-forked` and `pytest-xdist` as `tests` target deps (#5292)

This way, the dependencies needed for testing are localized in one place
- `setup.py` (instead of several), which makes maintenance easier.

Signed-off-by: Anatoly Myachev <[email protected]>

[BUILD] Skip installing test related python packages (#5294)

#5292 failed because of macOS
build. Since we don’t run any tests on macOS anyway, it’s fine to simply
skip them.

Update

Update

[TESTING] Add golden sample test for pipelining matmul with descriptors (#5289)

<git-pr-chain>

[TESTING] Add golden sample test for pipelining matmul with descriptors

1. #5288
1. 👉 #5289 👈 **YOU ARE HERE**
1. #5290

⚠️⚠️ Please **do not click the green "merge" button** unless you know
what
you're doing.  This PR is part of a chain of PRs, and clicking the merge
button will not merge it into master. ⚠️⚠️
</git-pr-chain>

Specify in `setup.py` that `setuptools>=40.8.0` is a required dependency (#5293)

Closes #5090

vancoykendall is right that the dependency is used not only during build.
However, for now I added it to `setup.py`, since the migration of
dependencies to `pyproject.toml` has not yet occurred.

Signed-off-by: Anatoly Myachev <[email protected]>

[TOOLS] Improve `generate-test-checks.py` (#5300)

- Format the doc string using the `reStructuredText` format.
- Lift the example instructions from the `.mlir.in` file to the
docstring. Previously we matched the `module` keyword twice and
encountered errors such as `assert len(output_segments) ==
len(source_segments),`. It's also fine to update the regex to something
like `\bmodule` to solve the problem, but I think lifting it from the
input file is just simpler.

[NFC][DIALECT] Remove dependency on `mlir::tensor::TensorDialect` (#5303)

[IR] Improve `ttg.memdesc` (#5296)

- Add an `allocShape` field to denote the shape a memory descriptor when
it's allocated. The value will be propagated to all its descendants
created through `subview` ops.
- Make `encoding` and `memorySpace` fields required instead of optional.
- Implement the `getAlias` function for `#ttg.shared_memory` to shorten
its length in `.mlir` files

Update

Update

[Pipeliner] Handle masking for atomic_rmw (#5231)

This commit is to support atomic_rmw in the function
predicateOp to mask operations during scheduling.

[TESTS] Forward fix for CI break (#5323)

PR #5231 was authored before the `triton_gpu` -> `ttg` rename and CI is
currently broken.

Search for `ptxas` only for cuda backend in `supports_tma` function (#5314)

For other backends, `ptxas` may not be installed.

Signed-off-by: Anatoly Myachev <[email protected]>

[LLVM] Update to llvm/llvm-project@1f20eee6dc36 (#5308)

This pulls in the AMDGPU backend support for the
gfx950 target.

We need to fix the rewrites in `Combine.td` given that
llvm/llvm-project#112700 adds
a new attribute for denorm mode for `arith.addf`.

---------

Co-authored-by: Lei Zhang <[email protected]>

[AMD][BACKEND] Add gfx950 target definitions. (#5281)

Enable new arch target since backend support has been added.

[AMD] Adjust local_store and global_load ordering (#5254)

This commit adjusts local store and global load
ordering to let local store be ahead of global
load when they are not in the same stage. It
should help GEMM kernel performance.

Re-align main and llvm-head (#5334)

We have a couple of PRs that landed in the `llvm-head` branch that are
not in `main`.

Merging those into `main` to prevent further divergence between
branches.

---------

Co-authored-by: Won-Kyu Park <[email protected]>
Co-authored-by: Lei Zhang <[email protected]>

[PIPELINER] Cleanup of LoopScheduling.cpp, introduction of AssignLatencies (#5176)

This change breaks down LoopScheduling into two sub-passes: latency
assignment and actual scheduling.
Latency assignment is a transformation that analyzes the loop and based
on the requested number of stages it assigns "latencies" to the ops that
are going to be converted to async ops by the pipeliner. Latencies are
expressed in terms of number of iterations of the loop and can be
thought as per-operation num_stages.
Scheduling transformation takes these latencies and builds a pipeliner
schedule based on it. The process of building a schedule was slightly
rewritten to simplify the code and cleanup the logic that was no longer
needed after recent refactoring.
Breaking down the schedule into latency assignment and proper scheduling
has number of purposes:
1. Code became more modular, with cleaner interfaces that helps with
maintanance
2. Both parts can be tested in separation, I have added lit tests for
both pieces. We can finally test our pipeliner infrastructure in
manageable chunks
3. It opens up opportunity to expose per-op "latencies" to the frontend,
enabling creating user-defined schedules right from the language level

Next step in the cleanup process is to clearly separate lowering and
pipelining phases.

Update

Update

Update

Update

Update

Update

Update

Update
@joker-eph
Copy link
Collaborator

I have some concerns about the tradeoffs discussed here, it does not seem to me that this got explored and spelled out entirely.
In particular, we should start from thinking about the "charter" for arith: while it does not strictly limit itself to what LLVM supports, adding things that can't be expressed in LLVM likely deserve some sort of a higher bar (it can't be a "free for all" situation where we just append every possible thing that someone could imagine).
In the things that LLVM does not support, we have grandfathered the support for tensor type (which is now adopted and heavily used by OpenAI Triton for example) and we also support some MLIR-specific thing like multi-dimensional vector. But both of these can legalize to LLVM, it's not a divergence between arith and LLVM per-se.

I tend to see the situation a bit like it is handled when adding support for this kind of things into LLVM: there is a tradeoff between the added complexity (these need to be handled by transformations, canonicalizer, etc) and the added benefits (who are the users? Is it a widespread needed feature or something exotic? What are the lowering paths? What are the alternatives? Why not using target-specific intrinsics?). It is of course a case-by-case evaluation!

Discourse is likely a better place to discuss this, and this change likely deserved and RFC in the first place before landing (sorry I should have called this clearly initially, but I thought this would be discussed much more in practice).
Happy to book an Open Meeting slot to discuss this in a high-bandwidth setting if needed as well.

It may be more prudent to revert this in the meantime.

@rengolin
Copy link
Member

rengolin commented Dec 9, 2024

It seems I have missed the discussion here. Some points...

Agreed, per operation makes more sense in the context of NEON. Though, admittedly, the only official doc that I could find for this covers ArmV7 and we are in ArmV9 era

Indeed. I mention this because, while LLVM has a definite set of targets, MLIR is more broad and widely adopted by a lot of accelerator start-ups, which usually have more than one compute unit (SIMD, systolic array, analog, etc) with different floating point behaviours.

That's interesting. Does this mean that if the scalar code is IEEE compliant and we are under a precise fp mode, the code should not be vectorized?

@dcaballe, that's exactly what the patch @banach-space mention does.

Which flags would allow changing an ieee denormal mode to ftz, for example? I think it's important that we document that.

This is again per target / extension. You may have conversion instructions, you may have soft-float helpers, and you may not. My original point here is this:

Denormal support is not a property of the operation, but the hardware micro-architecture it targets, including its runtime support.

While I said this may be thought of as a instruction-level property, it doesn't have to be. In reality, the compilers that try to compile individual instruction in separate units (think mixed AI accelerators) do so in a precise way and thus control the distribution. In the end, those instructions need to be fetched by different units and thus be at the very least in separate regions in the object file, but most commonly in separate libraries altogether.

How does this affect the implementation? I don't know. I have not spend time, nor have a use case in hand, to have a good answer.

@joker-eph, while I agree with your general point and the need for a charter, your comment is not specific enough to warrant a reversal. Perhaps you could be a bit more specific in a forum thread and we continue from there? Also, @chelini, if you could expose the usage pattern you're looking for, that'd give more power to your argument.

Generally, changes to core dialects like arith should really go through an RFC on the forum before being a PR. Usually, the simple cases are quick enough (this one took only a couple of days). This is yet another thing we need to get right in the charter for the whole project, but first we need to define what a "core dialect" means and which are they.

chsigg added a commit to triton-lang/triton that referenced this pull request Dec 9, 2024
@banach-space
Copy link
Contributor

Updated reference:

The only official doc that I could find for this covers ArmV7, and we are in the ArmV9 era.

Here’s the latest spec covering ArmV9: link.

Unfortunately, I haven’t had as much time to analyze it as I’d like. I’m hoping to focus more on this area next quarter, though I can’t promise anything definitive.

Regarding the concerns raised by @joker-eph
I suggest the following:

  • Add a note in the code marking this as "experimental" to highlight ongoing discussions.
  • Post an RFC on Discourse to gauge broader feedback on the approach taken here.

While LLVM has a liberal revert policy, these steps should help us make a more informed decision. Based on the RFC feedback, we can decide whether to revert, re-implement, or redesign.

How does this affect the implementation? I don't know. I have not spent time, nor do I have a use case in hand, to have a good answer.

I feel the same - it’s a bit of a chicken-and-egg problem at this point.

Also, @chelini, if you could expose the usage pattern you're looking for, that'd give more power to your argument.

+1 to this.

@joker-eph
Copy link
Collaborator

joker-eph commented Dec 9, 2024

@joker-eph, while I agree with your general point and the need for a charter, your comment is not specific enough to warrant a reversal. Perhaps you could be a bit more specific in a forum thread and we continue from there?

My comment is that without being able to clearly map this change to the "charter" for arith, we can't just land this as-is without proper RFC/discussion. We're back to:

Generally, changes to core dialects like arith should really go through an RFC on the forum before being a PR.

The fact that you missed this discussion is a good indication that this didn't get enough visibility.

In this case my concerns about this change and the associated tradeoffs to consider, warrants a revert IMO.
This is in line with LLVM developer policy and past practices:

  • Remember, it is normal and healthy to have patches reverted. Having a patch reverted does not necessarily mean you did anything wrong.
  • Any time you learn of a serious problem with a change, you should revert it.
  • If you are asked to revert by another contributor, please revert and discuss the merits of the request offline (unless doing so would further destabilize tip of tree).

So again: please revert, and discuss first in a RFC.

@kuhar
Copy link
Member

kuhar commented Dec 9, 2024

My position is that this change seems consistent with the existing philosophy of the arith dialect. Like I summarized on Discourse, arith is slightly higher level than both LLVM and SPIR-V and is not just a mirror of either of the two. We already have some things that are not trivially mappable to spirv (e.g., anything related to non-standard integer types, freeze semantics (adjacent to arith)). I think we generally have to be careful with these if they have a clear lowering path to llvm and spirv, but both the examples I just listed and denormal flags are fine because they will be decided by some 'frontend' instead of folds or default rewrite patterns. For this reason, I'm fine with this PR and I don't consider it going against any written documentation for arith or be otherwise precedent-setting.

@joker-eph
Copy link
Collaborator

joker-eph commented Dec 9, 2024

My position is that this change seems consistent with the existing philosophy of the arith dialect.

This is a reasonable thing we can debate in the RFC.

I'm personally not convinced: because LLVM IR and actual LLVM targets other than PTX don't seem support this right now, and the tradeoffs/alternatives with how we would manage this have to be considered, it didn't happen here so far: an RFC is the way to go to ensure this.

If we had done it, we'd be in position right now to properly document this feature on https://mlir.llvm.org/docs/ ; this PR didn't document nor discuss clearly the expectations and the plans for the support for this feature or anything, I couldn't update https://mlir.llvm.org/docs/Dialects/ArithOps/ even if I wanted to right now!
Also what's the plan with the canonicalizer/folding and all the code handling arith? Is there an actual plan to get this updated and honor this new mode of arithmetic?

@stellaraccident
Copy link
Contributor

stellaraccident commented Dec 9, 2024

+1 on continuing on discourse. -1 on a pre-emptory revert. (I very much appreciate the effort and thought leadership work that went into this and I think it is easy to look in retrospect and see that a multi-month conversation grew to something that should have been a broader discussion. But the group of stakeholders here is small enough that I think that the path forwards is a low stake discussion on Discourse about how to evolve this piece, and that it benefits nobody to lead with a revert)

@dcaballe
Copy link
Contributor

dcaballe commented Dec 9, 2024

Thanks everyone for the feedback. What started as a simple extra FMF flag evolved into a more comprehensive denormal model, which indeed would have required an RFC. Apologies for not signaling this earlier. However, the PR has been open for nearly two months and discussed multiple times offline with Jakub and others, including F2F at LLVM Dev. The proposal incorporated the feedback provided so we felt confident enough to proceed.

A few points I wanted to clarify, some already made by others:

"Free for all": I don't think describing this proposal as "free for all" is accurate. We are not proposing an Nvidia-specific modifier but a generic way to model various denormal behaviors. This is applicable to any hardware and even includes denormal modes which are not supported by our hardware.

"Exotic": I don’t think "Exotic" fits here either. Denormals (unfortunately!) are a fundamental part of the FP model, and the lack of representation is a gap in our ecosystem which is impacting everyone. Besides the emerging AI architectures that Renato mentioned, a quick search led to:

Denormals in LLVM: LLVM's denormal model is best-effort (see “denormal-fp-math” in LLVM Language Reference Manual — LLVM 20.0.0git documentation) and carries legacy baggage so I don’t think LLVM should be used as a strict role model for this particular case. It seems more reasonable to me to learn from LLVM and create something that better aligns with today's hardware requirements, which is what we were aiming for.

Per-instruction Denormal Mode: Besides the per-function mode, LLVM seems to allow finer control of the FP mode with llvm.set.fpmode and llvm.get.fpmode intrinsics (see https://llvm.org/docs/LangRef.html#llvm-set-fpmode-intrinsic). This should allow us to lower Arith to a generic per-instruction denormal model in LLVM.

Denormals and Compiler Infrastructure: As Jakub mentioned, the compiler's role with denormals is to propagate and preserve the behavior set by the front-end/framework. The compiler shouldn't be randomly changing the denormal mode. Consequently, it should be low risk to classify the denormal mode as something experimental while investigating/fixing any potential issues.

We'll follow up with an RFC on Discourse. Thanks again for your input.

@rengolin
Copy link
Member

rengolin commented Dec 9, 2024

My comment is that without being able to clearly map this change to the "charter" for arith, we can't just land this as-is without proper RFC/discussion

Unfortunately, there is no charter for arith. Once we get the governance ready (I started writing the document, will share with you once it reached the draft state), we will get to it. Right now, I don't think anyone can claim they know what the charter is.

The fact that you missed this discussion is a good indication that this didn't get enough visibility.

That's a stretch. Especially because...

the PR has been open for nearly two months and discussed multiple times offline with Jakub and others, including F2F at LLVM Dev. The proposal incorporated the feedback provided so we felt confident enough to proceed.

This is a perfectly valid course for LLVM, a bit on the cautious side. I cannot claim "because I didn't see, it wasn't properly discussed" with this history. Reverting now would warrant a stronger reason than "I didn't see" or "I wasn't convinced".

That's why we so desperately need a charter!

This is applicable to any hardware and even includes denormal modes which are not supported by our hardware.

Agree, and we gave a few examples before it was approved and merged.

Denormals (unfortunately!) are a fundamental part of the FP model, and the lack of representation is a gap in our ecosystem which is impacting everyone.

Yup.

It seems more reasonable to me to learn from LLVM and create something that better aligns with today's hardware requirements, which is what we were aiming for.

+1 to this point. I'd hate to repeat the same mistakes. My patch referenced by Andrzej was a blunt instrument to something that should have been better planned from the start, but we didn't know better. Now we do.

Consequently, it should be low risk to classify the denormal mode as something experimental while investigating/fixing any potential issues.

Exactly. No one that doesn't know what they're doing would ever pick those attributes, and the compiler should never add them without really knowing what it's doing.

I also don't see the current implementation as problematic. Maybe naive, but not wrong or in need of revert.

@joker-eph
Copy link
Collaborator

joker-eph commented Dec 9, 2024

Reverting now would warrant a stronger reason than "I didn't see" or "I wasn't convinced".

I mentioned concrete concerns, please don't reduce my design concern like this.
I believe they stand, and we have been able to process that way in LLVM and MLIR for years without the need to change the governance or write any more charters.

Such change should go through an RFC, it didn't, @dcaballe and @chelini will process with a discussion in a RFC, but that's not a reason to not revert. My concerns are warranting a revert according to every precedent in the project AFAICT.

The proposal incorporated the feedback provided so we felt confident enough to proceed.

@dcaballe, you didn't do anything wrong, quoting back the developer policy:

Remember, it is normal and healthy to have patches reverted. Having a patch reverted does not necessarily mean you did anything wrong.

@stellaraccident
Copy link
Contributor

-1 on revert

@rengolin
Copy link
Member

Such change should go through an RFC, it didn't, @dcaballe and @chelini will process with a discussion in a RFC, but that's not a reason to not revert. My concerns are warranting a revert according to every precedent in the project AFAICT.

You're mixing the arguments and picking and choosing precedent. Nothing in LLVM is this clear cut and I would not like to see it begin now. For the past few years I have supported your cause to make MLIR more mobile and have often argued in your favour in multiple forum threads. This is definitely not personal, or reductionist.

But "should" is not "must". If patches were reverted every time someone didn't do what someone else thought they should have done we would not have the project we have today. This is not a reasonable request.

@dcaballe, you didn't do anything wrong, quoting back the developer policy:

Remember, it is normal and healthy to have patches reverted. Having a patch reverted does not necessarily mean you did anything wrong.

Quoting this over and over again will not make it more valid to this particular thread. No one else in this discussion agrees with you that it applies and the patch should be reverted, and we all know by heart the LLVM policy. It is healthy to have your patches reverted, but I'm not going to start reverting patches simply because I don't agree with. It is also healthy to hash out design upstream and this is what they're doing. There's also enough precedent for that in LLVM. We'll discuss this in an RFC and if we find consensus that this is wrong, we change course. If not, we continue.

To be clear, the charter would not have stopped this PR from being merged, or forced people to submit an RFC beforehand. It would have made it easier to argument pro and against technical points in the charter, instead of what we (individually) think the other should have done.

For example, if we had a clear technical charter for arith/math, in line with llvm etc., Lorenzo et al would have started this PR towards that charter instead, and none of this discussion would have occurred, because it would have been "right" from the start. But they acted on the charter than they believed was right, others agreed, and here we are. 100% pure-bred LLVM style.

Perfect? No. But eventually consistent.

chelini added a commit that referenced this pull request Dec 10, 2024
@chelini
Copy link
Contributor Author

chelini commented Dec 10, 2024

Thank you for your feedback and reviews on this PR. I appreciate your time and effort in providing constructive comments. Reflecting on the discussions and concerns, I underestimated this change's complexity and implications. While I had discussions with Jakub and Diego and felt confident in pushing this forward, it is clear that the decision was premature. For this, I apologize; an RFC would have been a better first step.

I have reviewed the comments on this PR multiple times and have decided to revert these changes. Moving forward, we should formalize the discussion around implementing denormal handling through an RFC since canonicalization and folding concerns are grounded. I also want to avoid a precedent where a member's feedback is not considered or ignored. We don't need a consensus to revert; we need a consensus not to. The MLIR community is already fragmented, and I don't want to make this situation worse than it already is.

@chelini
Copy link
Contributor Author

chelini commented Dec 10, 2024

@chsigg I reverted this PR; let me know if you need help undoing the changes you made in Triton.

broxigarchen pushed a commit to broxigarchen/llvm-project that referenced this pull request Dec 10, 2024
vwbaker pushed a commit to openxla/triton that referenced this pull request Jan 2, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

9 participants