Help getting started with Triton

[JCBrouwer]

Hello, I wanted to give Triton a try, so I've set out to write a parallel version of PatchMatch.

I'm running into a couple of issues with my implementation and wonder if anyone can help point me in the right direction.

My kernel code looks like this:

@triton.jit
def propagate_kernel(A_ptr, B_ptr, kNNF_ptr, output_ptr, l: int, K: int, P: int, A_h: int, A_w: int, A_c: int, **meta):
    """
    Args:
        A_ptr : pointer to image A which has shape (A_w, A_h, A_c)
        B_ptr : pointer to image B
        kNNF_ptr : pointer to k nearest neighbor field which is (A_w , A_h, K, 3)
        output_ptr : pointer to output kNNF values for this iteration which is (A_w, A_h, K)
        l (int): looking distance (where to look for candidates)
        K (int): number of nearest neighbors (default: 5)
        P (int): patch size (default: 3)
        A_h (int): height of image A
        A_w (int): width of image A
        A_c (int): number of channels in image A
    """
    
    B = meta["BLOCK_SIZE"] # 16

    # map pid to block of kNNF that it should compute
    pid = tl.program_id(axis=0)
    num_pid_h = tl.cdiv(A_h, B)
    num_pid_w = tl.cdiv(A_w, B)
    y = pid // num_pid_h
    x = pid // num_pid_w
    ys = y * B + tl.arange(0, B)
    xs = x * B + tl.arange(0, B)
    A_grid = ys[:, None] + xs[None, :]  # TODO need to add stride in ptr sum?

    # arrays for indexing certain things
    window = tl.arange(0, 3) - P // 2  # tl.arange(-P // 2, P // 2 + 1)
    channels = tl.arange(0, 3)  # A_c 
    coord_dim = tl.arange(1, 3) 
    dist_dim = tl.zeros((1,), dtype=tl.int32)
    neighbors = tl.arange(0, 5)  # K

    # extract pixel value patches from A
    A_patch_center_idxs = A_ptr + A_grid  # B, B
    A_patch_idxs = (
        A_patch_center_idxs[   :,    :, None, None, None, None]
        + window           [None, None, None,    :, None, None]
        + window           [None, None, None, None,    :, None]
        + channels         [None, None, None, None, None,    :]
    )  #                    B     B     1     P     P     A_c
    A_patches = tl.load(A_patch_idxs)

    # look for new candidate patches l pixels left, right, up, and down (9 locations)
    candidates = tl.zeros((9, B, B, 2), dtype=tl.int32)
    candidate_distances = tl.zeros((9, B, B), dtype=tl.float32)
    i = 0
    for h in range(-l, l, l):
        for w in range(-l, l, l):

            # ensure indexes stay in bounds
            candidate_ys = tl.maximum(0, tl.minimum(A_h - 1, ys + h))
            candidate_xs = tl.maximum(0, tl.minimum(A_w - 1, xs + w))

            # load candidate patch centers from kNNF
            candidate_coords = (
                kNNF_ptr
                + candidate_ys[   :, None, None, None]
                + candidate_xs[None,    :, None, None]
                + neighbors   [None, None,    :, None]
                + coord_dim   [None, None, None,    :]
            )  #               B     B     K     2

            # candidate_coords shape here is actually (B, B, K, 3) !?
            candidate_idxs = candidate_coords[:, :, :, 0] + candidate_coords[:, :, :, 1]  # Error: cannot reshape block of different shape
            # candidate_idxs = tl.sum(candidate_coords, axis=3)  # Error: Encountered unimplemented code path in sum. This is likely a bug on our side.
            B_patch_center_idxs = tl.load(candidate_idxs)

            # load corresponding patches from image B
            B_patch_idxs = (
                B_ptr
                + B_patch_center_idxs[   :,    :,    :, None, None, None]
                + window             [None, None, None,    :, None, None]
                + window             [None, None, None, None,    :, None]
                + channels           [None, None, None, None, None,    :]
            )  #                      B     B     K     P     P     A_c
            B_patches = tl.load(B_patch_idxs)

            # find distance between image A patches and candidate patches
            distances = tl.sum((B_patches - A_patches) ** 2, axis=[3, 4, 5])  # B, B, K

            # remember best candidates
            best_candidates = triton.torch.argmin(distances, axis=2)  # B, B
            candidates[i] = B_patch_center_idxs[best_candidates]
            candidate_distances[i] = distances[best_candidates]

            i += 1

    # find overal best candidates
    idxs_new = triton.torch.argmin(candidate_distances, axis=0)  # B, B
    kNNF_coord_new = candidates[idxs_new[None, :, :, None]]  # B, B, 2
    kNNF_dist_new = candidate_distances[idxs_new[None, :, :]]  # B, B

    # store 
    tl.store(output_ptr + A_grid[:, :, None] + dist_dim[None, None, :], kNNF_dist_new)
    tl.store(output_ptr + A_grid[:, :, None] + coord_dim[None, None, :], kNNF_coord_new)

First of all, I'm running into some issues setting up some arrays that I want to use to index certain parts of my tensors:

    # arrays for indexing certain things
    window = tl.arange(-P // 2, P // 2 + 1)
                                         ^
 Error: arange(): incompatible function arguments. The following argument types are supported:
    1. (arg0: int, arg1: int, arg2: triton._C.libtriton.triton.ir.builder) -> triton._C.libtriton.triton.ir.value

Invoked with: <triton._C.libtriton.triton.ir.value object at 0x7fbf23315070>, <triton._C.libtriton.triton.ir.value object at 0x7fbf23315730>, <triton._C.libtriton.triton.ir.builder object at 0x7fbdf7066db0>

Similarly:

    channels = tl.arange(0, A_c) 
                            ^
 Error: arange(): incompatible function arguments. The following argument types are supported:
    1. (arg0: int, arg1: int, arg2: triton._C.libtriton.triton.ir.builder) -> triton._C.libtriton.triton.ir.value

Invoked with: <triton._C.libtriton.triton.ir.constant_int object at 0x7f68706812b0>, <triton._C.libtriton.triton.ir.argument object at 0x7f68706773b0>, <triton._C.libtriton.triton.ir.builder object at 0x7f67443d0e50>

For the time being I've just hardcoded the literal values of the defaults in those spots. How can I create ranges based on kernel arguments?

Later on I run into some more problems:

    # look for new candidate patches l pixels left, right, up, and down (9 locations)
    candidates = tl.zeros((9, B, B, 2), dtype=tl.int32)
    candidate_distances = tl.zeros((9, B, B), dtype=tl.float32)
    i = 0
    for h in range(-l, l, l):
        for w in range(-l, l, l):

            # ensure indexes stay in bounds
            candidate_ys = tl.maximum(0, tl.minimum(A_h - 1, ys + h))
            candidate_xs = tl.maximum(0, tl.minimum(A_w - 1, xs + w))

            # load candidate patch centers from kNNF
            candidate_coords = (
                kNNF_ptr
                + candidate_ys[   :, None, None, None]
                + candidate_xs[None,    :, None, None]
                + neighbors   [None, None,    :, None]
                + coord_dim   [None, None, None,    :]
            )  #               B     B     K     2

            # candidate_coords shape here is actually (B, B, K, 3) !?

            # candidate_idxs = tl.sum(candidate_coords, axis=3)  # Error: Encountered unimplemented code path in sum. This is likely a bug on our side.
            candidate_idxs = candidate_coords[:, :, :, 0] + candidate_coords[:, :, :, 1]  # Error: cannot reshape block of different shape  
            # candidate_idxs = candidate_coords[:, :, :, [0]] + candidate_coords[:, :, :, [1]]  # Error: cannot reshape block of different shape

First of all, I noticed that the candidate_coords tensor doesn't have the shape I expect. Then, when trying to convert coordinate values from the kNNF to pointer offsets, I run into some shape problem that I'm not sure how to debug. What's going on here?

Also I have the feeling that I might not be calculating the pointer offsets correctly anyway. At the start I'm getting the offsets like this:

    # map pid to block of kNNF that it should compute
    pid = tl.program_id(axis=0)
    num_pid_h = tl.cdiv(A_h, B)
    num_pid_w = tl.cdiv(A_w, B)
    y = pid // num_pid_h
    x = pid // num_pid_w
    ys = y * B + tl.arange(0, B)
    xs = x * B + tl.arange(0, B)
    A_grid = ys[:, None] + xs[None, :]  # TODO need to add stride in ptr sum?

However, in the matmul tutorial, the stride is used to calculate the a_ptrs (what I'm calling A_grid).

    offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

I don't quite follow the grouping in the matmul tutorial, though, so I'm having trouble understanding how to translate it to my case which I'm doing without grouping for now.

To be honest, I don't know if my approach with crazy high-dimensional tensors makes sense at all for Triton (although something very similar works fine to vectorize the operations in my pytorch implementation). So I'd love to know if there's a better way of going about this.

I would appreciate any help or ideas!

The full code can be found here: https://github.com/JCBrouwer/patchmatch-triton

[ptillet]

Hey!

Thanks for your interest :) Unfortunately, Triton is under pretty active development so functional parity with pytorch is still pretty far. However, here are a few things:

• triton.arange only works on constants. Right now, this means that the arguments have to be either literal or passed as metaparameters. Pretty soon, this will mean that it will work on anything declared as constexpr, where the constexpr keyword is analogous to that found in C++
• Tensors of dimensions 3 and above are not yet supported in Triton, though this should come by version 2.0.

[JCBrouwer]

Ahh ok then, thanks for the quick response. I assume it's not easy to pass my own metaparameters, is it? I think I can probably live with the patch size and number of neighbors being fixed eventually. I guess I'll take a look at flattening a bunch of dimensions then, hopefully that'll make my pointer arithmetic a bit easier as well!