triton-lang/triton

compilation is too slow when using the tl.store

Open

#1,622 opened on May 5, 2023

 (3 comments) (0 reactions) (0 assignees)MLIR (3,126 forks)github user discovery
help wanted

Repository metrics

Stars
 (19,982 stars)
PR merge metrics
 (Avg merge 2d 18h) (185 merged PRs in 30d)

Description

Hi,I have a problem with compliation. I find that the compilation process is extremely time-comsuming(about half hour and even hours sometimes) when using the tl.store. Once I comment "tl.store", compilation will be finished soon. I do not understand why is this happening.

Note that, I run bellow code four times in loop, the test time is as follow(the unit is second): image When commenting "tl.store" in kernel function, the test time is shorter than above: 1683299588131

I provide my code and hardware information here.

Brief hareware Infomation: CPU x86_64, 64 cores GPU NVIDIA 3090 Memory 256GB

import torch
import numpy as np
import triton
import triton.language as tl
import triton.testing
import argparse 
from pathlib import Path

@triton.jit
def kernel(A,
           area_size,
           max_pixeldist_5m, 
           stride_A_M,
           stride_A_N,
           stride_om,
           stride_on,
           d2r,
           out, 
           RadioM,
           PixS,
           angles,
        BLOCK_SIZE:tl.constexpr):
    
    pid = tl.program_id(0)
    
    anchors_size_M = area_size * 2 + 1  #offset area
    anchors_size_N = anchors_size_M     #offset area
    # calculate initial memory address
    pid_m = pid // anchors_size_M
    pid_n = pid % anchors_size_N
    poi_x = pid_n + max_pixeldist_5m - 1
    poi_y = pid_m + max_pixeldist_5m - 1
    
    pixel_dist_5m = tl.arange(0, BLOCK_SIZE)  
    pixel_dist_5m = pixel_dist_5m + 1
  
    rad_cos = tl.cos((tl.arange(0, BLOCK_SIZE) * 0.5) * d2r) #angles, [0,720] is useful but must satisfy 1024
    rad_sin = tl.sin((tl.arange(0, BLOCK_SIZE) * 0.5) * d2r) # refer to above comment

    delta_y = rad_cos[:, None] * pixel_dist_5m[None, :] # shape is [1024,1024]
    delta_x = rad_sin[:, None] * pixel_dist_5m[None, :]
    

    poi_y_inter = poi_y - delta_y  #float coordinates, shape is [1024,1024]
    poi_x_inter = poi_x + delta_x
    poi_x_min = poi_x_inter.to(tl.int32) #convert  float to int, truncate the fraction
    poi_y_min = poi_y_inter.to(tl.int32)
    poi_x_max = poi_x_min + 1
    poi_y_max = poi_y_min + 1
   
    mask_col = pixel_dist_5m < max_pixeldist_5m
    mask_row = tl.arange(0, BLOCK_SIZE) < (angles + 1)
    # when interpolating, we only preverse the top-left area with shape [721, 599]
    mask = mask_row[:, None] & mask_col[None, :]

    left_top_addr = A + poi_y_min * stride_A_M + poi_x_min * stride_A_N
    left_top_val = tl.load(left_top_addr, mask=mask, other=-float('inf'))
 
    offset_out = tl.arange(0, BLOCK_SIZE)
    output_addr = out + offset_out[:, None] * stride_om + offset_out[None, :] * stride_on
    tl.store(output_addr, left_top_val)
    


def calculate(A, area_size, max_pixeldist_5m, RadioM):
    assert A.is_contiguous(), "Matrix A must be contiguous"
    assert A.is_cuda, "Maxtrix A must be on GPU"
    stride_m = A.stride(0)
    stride_n = A.stride(1)
    n_rows = (area_size * 2 + 1) ** 2
    out = torch.empty((n_rows, 1024), device=A.device, dtype=A.dtype).contiguous()
    stride_om = out.stride(0)
    stride_on = out.stride(1)
    angles= 720
    d2r = np.pi / 180
    PixS = 10
    BLOCK_SIZE = triton.next_power_of_2(max_pixeldist_5m)
    kernel[(1,)](A,   #shape is 1239 by 1239
           area_size, max_pixeldist_5m, 
           stride_m, stride_n,
           stride_om, stride_on,
           d2r, 
           out,
           angles=angles,
           RadioM=RadioM,
           PixS=PixS,
           BLOCK_SIZE=BLOCK_SIZE)
 
    return out

def main(opt) -> None:
    #load data
    data = dict(np.load(opt.input, allow_pickle=True))
    A = data['A'].astype(np.float32)
    poi = data['poi']
    area_size = int(data['area_size'])
    max_pixeldist_5m = int(data['max_pixeldist_5m'])
    
    RadioM = 1737100.0 
    device = torch.device(f"cuda:{opt.device}")
    A = torch.from_numpy(A).to(device)
    
    import time 
    for _ in range(4):
        t1 = time.time()
        output = calculate(A, area_size, max_pixeldist_5m, RadioM)
        t2 = time.time()
        print(t2-t1)
    print(output[:, :])
    return 

Hey, Anyone knows the reason? It is really wired. If neccesary, I can provide data and complete code to help you reproduce the problem.

Contributor guide