flagos-ai / flagos-ai/FlagTree
Improve integer divisibility specialization to enable matmul `cp.async` pipelining
- 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
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