[RAISE-BP] Add support for `arith.remsi|remui` as `tt.addptr` input #1570

mfrancepillois · 2024-07-05T11:27:54Z

Add minimal support for handling arith.remsi|remui as tt.addptr input.
Improve handling of unfolded arithmetic operations when evaluating the modulo property and constant values.

Closes Issue: #1436 and #1482

- Add minimal support for handling `arith.remsi|remui` as `tt.addptr` input. - Improve handling of unfolded arithmetic operations when evaluating the modulo property and constant values. Closes Issue: #1430 and #1482 Signed-off-by: Maxime France-Pillois <[email protected]>

third_party/intel/lib/TritonRaiseBlockPointer/TritonRaiseBlockPointer.cpp

test/Triton/raise-block-pointer.mlir

victor-eds

Can we have code actually exercising every case you're implementing here?

victor-eds · 2024-07-05T12:31:08Z

third_party/intel/lib/TritonRaiseBlockPointer/TritonRaiseBlockPointer.cpp

+
+// This function folds the `op` operation and returns the constant value if it
+// has successfully folded to a constant. Otherwise, it returns `std::nullopt`.
+std::optional<int64_t> getFoldedConstantValue(Operation *op) {


Can we asssert on op having exactly 1 result?

I prefer not to add an assert on the size, because it is not a error. It is not fundamentally wrong to call this function with a op with multiple results. But we cannot fold the operation to get a constant value.

Then can we early exit in that case?

I would say that assert is not for error handling (especially since it's not included in release builds), but for encoding preconditions/assumption. So if you should only ever call this function on ops with 1 results, then putting that in an assert is a good idea.

victor-eds · 2024-07-05T12:31:22Z

third_party/intel/lib/TritonRaiseBlockPointer/TritonRaiseBlockPointer.cpp

+    }
+  }
+
+  if (results.size() != 1) {


Suggested change

if (results.size() != 1) {

if (results.empty()) {

If this is not the case (see the answer to the previous point), the two codes are not semantically equivalent.

Would make sense if we early exit

third_party/intel/lib/TritonRaiseBlockPointer/TritonRaiseBlockPointer.cpp

victor-eds · 2024-07-05T12:55:24Z

third_party/intel/lib/TritonRaiseBlockPointer/TritonRaiseBlockPointer.cpp

+    // - (tl.arange(0, end)[:, None] % mod), or
+    // - (tl.arange(0, end)[None, :] % mod)


If this can cause issues, can we maybe change the order in an earlier pass?

third_party/intel/lib/TritonRaiseBlockPointer/TritonRaiseBlockPointer.cpp

victor-eds · 2024-07-05T13:00:15Z

third_party/intel/lib/TritonRaiseBlockPointer/TritonRaiseBlockPointer.cpp

+        op->emitWarning(
+            "TritonRaiseBlockPointer: allowing adding pointer state with "
+            "modulo in dim 0 to "
+            "another pointer state with offset in dim 0.\nPlease verify the "
+            "operand that contains a scalar is meant to increment pointers in "
+            "dim1. If that is not the case it WILL LEAD TO WRONG COMPILATION "
+            "RESULTS.\n\nTo avoid this warning, use expand_dims (instead of "
+            "splat) to explicitly specify which dimension contains the "


Also, if this may lead to wrong code, I'd just fail. We can think of a way to check this, but I'd sacrifice performance for correctness here.

Agree. I have removed this case.

Finally, I've some doubts about removing this special case. Indeed, this case seems to occur quite often when dealing with generated code. It is the case for 03-matrix-multiplication.py, which results in generating this pattern. So, I'm afraid that removing this special case and triggering a failure instead will reduce significantly the interest in this pass.

As safe handling of modulo requires more care. Issue #1784 has been created.
IMO, we can then keep this unsafe path out the this PR, and design a better support within the newly created issue.

etiotto · 2024-07-12T14:28:46Z

third_party/intel/lib/TritonRaiseBlockPointer/TritonRaiseBlockPointer.cpp

+    return failure();
+  }
+
+  if (state.getRank() == 1) {


Change to switch stmt

third_party/intel/lib/TritonRaiseBlockPointer/TritonRaiseBlockPointer.cpp

mfrancepillois requested review from whitneywhtsang, etiotto and a team July 5, 2024 11:27

mfrancepillois self-assigned this Jul 5, 2024

This was linked to issues Jul 5, 2024

[RAISE-BP] Add support for arith.remsi|remui as tt.addptr input #1436

Closed

[RAISE-BP] Addressed dimensions coming from non-folded arithmetic operations in dimHasModulo #1482

Closed

format

00f3d24

FMarno reviewed Jul 5, 2024

View reviewed changes

victor-eds reviewed Jul 5, 2024

View reviewed changes

Improve coding style + remove extra case that may lead to wrong codegen.

af25530

etiotto reviewed Jul 12, 2024

View reviewed changes

mfrancepillois mentioned this pull request Jul 15, 2024

[RAISE-BP] Add support for arith.remsi|remui as tt.addptr input #1436

Closed

Merge branch 'llvm-target' into maxime/raise-add-ptr-remsi-remui-support

717207e

etiotto requested review from victor-eds and FMarno July 30, 2024 13:02

mfrancepillois added 2 commits August 5, 2024 12:50

Merge branch 'llvm-target' into maxime/raise-add-ptr-remsi-remui-support

d4ff625

Improve coding style.

07518ee

mfrancepillois mentioned this pull request Aug 5, 2024

[RAISE-BP] Safely handle masks and modulo ops #1784

Open

victor-eds approved these changes Aug 5, 2024

View reviewed changes

whitneywhtsang reviewed Aug 5, 2024

View reviewed changes

third_party/intel/lib/TritonRaiseBlockPointer/TritonRaiseBlockPointer.cpp Outdated Show resolved Hide resolved

mfrancepillois added 2 commits August 5, 2024 17:31

Remove unrelevant comment

e5e6dd5

Merge branch 'llvm-target' into maxime/raise-add-ptr-remsi-remui-support

76c6cd1

etiotto approved these changes Aug 7, 2024

View reviewed changes

whitneywhtsang merged commit 57a10af into llvm-target Aug 7, 2024
4 checks passed

whitneywhtsang deleted the maxime/raise-add-ptr-remsi-remui-support branch August 7, 2024 14:17

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[RAISE-BP] Add support for `arith.remsi|remui` as `tt.addptr` input #1570

[RAISE-BP] Add support for `arith.remsi|remui` as `tt.addptr` input #1570

mfrancepillois commented Jul 5, 2024 •

edited

Loading

victor-eds left a comment

victor-eds Jul 5, 2024

mfrancepillois Jul 5, 2024

victor-eds Jul 8, 2024

FMarno Jul 30, 2024

victor-eds Jul 5, 2024

mfrancepillois Jul 5, 2024

victor-eds Jul 8, 2024

victor-eds Jul 5, 2024

victor-eds Jul 5, 2024

victor-eds Jul 5, 2024

mfrancepillois Jul 5, 2024

mfrancepillois Jul 8, 2024

mfrancepillois Aug 5, 2024

etiotto Jul 12, 2024

		// - (tl.arange(0, end)[:, None] % mod), or
		// - (tl.arange(0, end)[None, :] % mod)

[RAISE-BP] Add support for arith.remsi|remui as tt.addptr input #1570

[RAISE-BP] Add support for arith.remsi|remui as tt.addptr input #1570

Conversation

mfrancepillois commented Jul 5, 2024 • edited Loading

victor-eds left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

[RAISE-BP] Add support for `arith.remsi|remui` as `tt.addptr` input #1570

[RAISE-BP] Add support for `arith.remsi|remui` as `tt.addptr` input #1570

mfrancepillois commented Jul 5, 2024 •

edited

Loading