Confusion of arange, multiple_of, max_contiguous

[sgjeon194]

Hi there!

I am very new with triton, and have started with the triton website's tutorial. During my study there was a big confusion with some api. I already checked the official document (https://triton-lang.org/main/python-api/triton.language.html), but it wasn't enough for me.

1. arange()

import torch
import triton
import triton.language as tl

@triton.jit
def kernel(debugger_ptr):
    offset_m = tl.arange(0, 8)
    tl.device_print("offset_m ", offset_m)

debugger = torch.empty(4, 2, dtype=torch.int32, device="cuda") 
kernel[(1,)](debugger, num_warps=1)

The results are like this:

offset_m 0
offset_m 1
offset_m 2
offset_m 3
offset_m 4
offset_m 5
offset_m 6
offset_m 7
offset_m 0
offset_m 1
offset_m 2
offset_m 3
offset_m 4
offset_m 5
offset_m 6
offset_m 7
offset_m 0
offset_m 1
offset_m 2
offset_m 3
offset_m 4
offset_m 5
offset_m 6
offset_m 7
offset_m 0
offset_m 1
offset_m 2
offset_m 3
offset_m 4
offset_m 5
offset_m 6
offset_m 7

This is a very simple block making me with the confusion of arange.

Because there is only 1 block with 1 warp (32 threads), the device_print works 32 times, and each prints shows an integer number 0~7, repeating 4 times.

However, according to the doc, arange is introduced like this:
Returns contiguous values within the half-open interval [start, end)

So this makes me a question, why is the device print shows me just one integer, not the whole array 0~7? I expect the output as

offset_m 0 1 2 3 4 5 6 7
offset_m 0 1 2 3 4 5 6 7
...
(repeating 32 times)

, not just 1 integer in 1 device_print call.

Is there something I'm missing?

2. max_contiguous && multiple_of

During my code reading, I bumped into this code line:

ram = tl.max_contiguous(tl.multiple_of(offset_m % M, BLOCK_M), BLOCK_M)

This code line seems a bit popular, since I found it on the pytorch library torch.mm() api, or on some other matrix multipling api. But this is quite confusing either. If I check the doc about these two,

My question is, what is the return type of each of them?

1. For multiple_of, does "check whether input are all multiple of values" means that it returns True when all of them are multiples and False when any of them are not? In that case, max_contiguous needs a boolean type parameter input?
2. Why do we need to make the complier know the first value is contiguous?
3. Why does the "ram" variable shows like an integer when I use

tl.device_print("ram ", ram)

By the way, when I check the output of the device_print("ram", ram), it looks exactly the same with the results up at the question about arange()

Thanks for reading.

[yioneko]

For Q2, there is a detailed explanation for contiguity and divisibility in source code, which might resolve your confusion:

triton/include/triton/Analysis/AxisInfo.h

Lines 38 to 105 in 65d9862

	// contiguity[d] is the length of the shortest sequence of contiguous integers
	// along dimension d.
	//
	// If we have an array of N elements with a contiguity value C, then the array
	// can be divided into a list of N/C sequences of C contiguous elements.
	// Since we have N = 2^k, C must be a power of two.
	//
	// For example, the 2D array
	//
	// [[10, 11, 12, 13, 18, 19, 20, 21],
	// [20, 21, 22, 23, 28, 29, 30, 31]]
	//
	// has contiguity [1, 4], and
	//
	// [[12, 16, 20, 24],
	// [13, 17, 21, 25],
	// [14, 18, 22, 26],
	// [15, 19, 23, 27],
	// [18, 22, 26, 30],
	// [19, 23, 27, 31]]
	//
	// has contiguity [2, 1].
	int64_t getContiguity(size_t dim) const { return contiguity[dim]; }
	const DimVectorT &getContiguity() const { return contiguity; }
	// divisibility[d] is the largest power of two that divides the first element
	// of all groups of length contiguity[d] along dimension d.
	//
	// For example,
	//
	// [[10, 11, 12, 13, 18, 19, 20, 21],
	// [20, 21, 22, 23, 28, 29, 30, 31]]
	//
	// has divisibility [1, 2], and
	//
	// [[12, 16, 20, 24],
	// [13, 17, 21, 25],
	// [14, 18, 22, 26],
	// [15, 19, 23, 27]]
	//
	// has divisibility [4, 1].
	//
	// On the other hand,
	//
	// [0, 1, 2, 0, 4, 5, 6, 7]
	//
	// has divisibility 1 because its contiguity is 1.
	int64_t getDivisibility(size_t dim) const { return divisibility[dim]; }
	const DimVectorT &getDivisibility() const { return divisibility; }
	// constancy[d] is the length of the shortest sequence of repeating integers
	// along dimension d.
	//
	// This is particularly useful to infer the contiguity of operations (e.g.
	// add) involving a constant.
	//
	// If we have an array of N elements, with a constancy value C, then the array
	// can be divided into a list of N/C sequences of C elements with the same
	// value. Since we have N = 2^k, C must be a power of two.
	//
	// For example
	//
	// [[8, 8, 8, 8, 12, 12, 12, 12],
	// [16, 16, 16, 16, 20, 20, 20, 20]]
	//
	// has constancy [1, 4].
	int64_t getConstancy(size_t dim) const { return constancy[dim]; }
	const DimVectorT &getConstancy() const { return constancy; }

I think this could be better added to document.