flagos-ai / flagos-ai/FlagTree

Improve integer divisibility specialization to enable matmul `cp.async` pipelining

Open
#1,182 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
350
Forks
149
Avg merge
2d 4h
Merged PRs (30d)
81

Description

## Observed behavior

For `a.shape=(8192, 8200)` and `b.shape=(8200, 8192)` (matmul dimensions `[M, N, K]=[8192, 8192, 8200]`), the original compiler pipelines B through shared-memory buffers, but loads A synchronously on every K-loop iteration. After extending integer specialization to record divisibility by 8 as well as 16, a fresh compilation pipelines **both** A and B through buffers. The Tensor pointer alignment predicate remains `data_ptr % 16 == 0`; the only pointer-branch edit is renaming its result from `D_str` to `D16_str` to match the new marker names.

| Artifact | `stride_am` / `K` argument attributes | A path | B path |
| --- | --- | --- | --- |
| Original | Neither has `tt.divisibility` | `tt.load` followed by `ttg.local_alloc` inside the loop; PTX has `ld.global.b16` and `st.shared.b16` | `4x64x128xf16` shared allocation and `ttg.async_copy_global_to_local` |
| With integer D8 marker | Both have `tt.divisibility = 8` | `4x64x64xf16` shared allocation and `ttg.async_copy_global_to_local {contiguity = 8}`; PTX has 16-byte `cp.async.cg.shared.global` | Still four buffers and async copy |

In both artifacts, `a_ptr`, `b_ptr`, and `c_ptr` already have `tt.divisibility = 16`. Base-pointer specialization therefore does not explain the change in A's pipeline.

Original compilation (A is loaded inside the loop; only B has a four-buffer allocation):

```mlir
tt.func public @matmul_kernel(
%a_ptr: !tt.ptr {tt.divisibility = 16 : i32},
%b_ptr: !tt.ptr {tt.divisibility = 16 : i32},
...,
%K: i32,
%stride_am: i32,
...) {
%b = ttg.local_alloc : () -> !ttg.memdesc<4x64x128xf16, #shared, #smem, mutable>
%b_46 = ttg.async_copy_global_to_local %b_ptrs_37, %b_43
mask %accumulator_45 other %cst_2 {contiguity = 8 : i32}
%b_47 = ttg.async_commit_group tokens %b_46

scf.for ... {
%a_104 = tt.load %a_ptrs_86, %a_103, %cst_0
%a_105 = ttg.local_alloc %a_104
%b_106 = ttg.async_wait %b_93 {num = 2 : i32}
%accumulator_108 = ttng.warp_group_dot %a_105, %b_107, %accumulator_87
...
}
}
```

Fresh compilation with D8 integer specialization (both A and B have four-buffer async-copy pipelines):

```mlir
tt.func public @matmul_kernel(
%a_ptr: !tt.ptr {tt.divisibility = 16 : i32},
%b_ptr: !tt.ptr {tt.divisibility = 16 : i32},
...,
%K: i32 {tt.divisibility = 8 : i32},
%stride_am: i32 {tt.divisibility = 8 : i32},
...) {
%a = ttg.local_alloc : () -> !ttg.memdesc<4x64x64xf16, #shared, #smem, mutable>
%b = ttg.local_alloc : () -> !ttg.memdesc<4x64x128xf16, #shared, #smem, mutable>
%a_43 = ttg.memdesc_index %a[%c0_i32]
%a_46 = ttg.async_copy_global_to_local %a_ptrs_29, %a_43
mask %accumulator_45 other %cst_1 {contiguity = 8 : i32}
%a_47 = ttg.async_commit_group tokens %a_46

scf.for ... {
%a_128 = ttg.async_wait %a_117, %b_120 {num = 4 : i32}
%a_129 = ttg.memdesc_index %a[%accumulator_127]
%accumulator_131 = ttng.warp_group_dot %a_129, %b_130, %accumulator_113
...
}
}
```

## Reproduction
In `python/tutorials/03-matrix-multiplication.py`, use the unit-test configuration and FP16 tensors with the shapes above. Compile once with the original specialization and once with the integer D8 change below; compare the TTGIR and PTX for A.

```diff
diff --git a/python/src/specialize.cc b/python/src/specialize.cc
@@ -75,7 +75,8 @@
-static PyObject *D_str = nullptr;
+static PyObject *D16_str = nullptr;
+static PyObject *D8_str = nullptr;
@@ -123,7 +124,8 @@ void init_interned_strings() {
- D_str = intern_from_string("D");
+ D16_str = intern_from_string("D16");
+ D8_str = intern_from_string("D8");
@@ -250,6 +252,18 @@ std::pair specialize_tensordesc(PyObject *arg,
+// Keep scalar divisibility in its original units. Pointer analysis combines it with
+// the pointee type later; e.g. an f16 row stride of 520 is 16-byte aligned.
+static PyObject *get_integer_alignment_key(unsigned long long value, bool align) {
+ if (!align)
+ return empty_str;
+ if ((value & 15) == 0)
+ return D16_str;
+ if ((value & 7) == 0)
+ return D8_str;
+ return empty_str;
+}
@@ -269,7 +283,8 @@ std::pair handle_long_type(PyObject *backend,
- key_obj = (align && ((val & 15) == 0)) ? D_str : empty_str;
+ key_obj =
+ get_integer_alignment_key(static_cast(val), align);
@@ -285,7 +300,7 @@ std::pair handle_long_type(PyObject *backend,
- key_obj = (align && ((val_64 & 15) == 0)) ? D_str : empty_str;
+ key_obj = get_integer_alignment_key(val_64, align);
@@ -349,7 +364,7 @@ std::pair handle_tensor(PyObject *backend,
- auto key_obj = (align && ((data_ptr & 15) == 0)) ? D_str : empty_str;
+ auto key_obj = (align && ((data_ptr & 15) == 0)) ? D16_str : empty_str;

diff --git a/python/triton/backends/compiler.py b/python/triton/backends/compiler.py
@@ -111,14 +112,18 @@ class BaseBackend(metaclass=ABCMeta):
def parse_attr(desc):
assert isinstance(desc, str)
ret = []
- if "D" in desc:
+ if "D16" in desc:
ret += [["tt.divisibility", 16]]
+ if "D8" in desc:
+ ret += [["tt.divisibility", 8]]
return ret

@staticmethod
def get_int_specialization(arg, **kwargs):
- if arg % 16 == 0 and kwargs.get("align", False):
- return "D"
+ if kwargs.get("align", False):
+ for divisor in (16, 8):
+ if arg % divisor == 0:
+ return f"D{divisor}"
return ""
```

Contributor guide

Open the contributing guide

Research direction

Start in python/src/specialize.cc with the integer and tensor specialization helpers, then read parse_attr and get_int_specialization in python/triton/backends/compiler.py. Reproduce the unit-test configuration from python/tutorials/03-matrix-multiplication.py with the stated FP16 shapes, compare TTGIR and PTX for A, and verify the expected asynchronous pipeline behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
compilers
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.