Coverage for src/flag_gems/runtime/backend/_sunrise/ops/triu.py: 0%

1import logging 

2 

3import torch 

4import triton 

5import triton.language as tl 

6 

7from flag_gems import runtime 

8from flag_gems.runtime import torch_device_fn 

9from flag_gems.utils import libentry 

10from flag_gems.utils import triton_lang_extension as ext 

11 

12logger = logging.getLogger("flag_gems").getChild(__name__.lstrip(".")) 

13 

14 

15@libentry() 

16@triton.autotune(configs=runtime.get_tuned_config("triu"), key=["M", "N"]) 

17@triton.jit(do_not_specialize=["diagonal"]) 

18def triu_kernel( 

19 X, 

20 Y, 

21 M, 

22 N, 

23 diagonal, 

24 M_BLOCK_SIZE: tl.constexpr, 

25 N_BLOCK_SIZE: tl.constexpr, 

26): 

27 pid = ext.program_id(0) 

28 row = pid * M_BLOCK_SIZE + tl.arange(0, M_BLOCK_SIZE)[:, None] 

29 m_mask = row < M 

30 X += row * N 

31 Y += row * N 

32 

33 for n_offset in range(0, N, N_BLOCK_SIZE): 

34 cols = n_offset + tl.arange(0, N_BLOCK_SIZE)[None, :] 

35 n_mask = cols < N 

36 mask = m_mask and n_mask 

37 

38 x = tl.load(X + cols, mask, other=0.0) 

39 y = tl.where(row + diagonal <= cols, x, 0.0) 

40 tl.store(Y + cols, y, mask=mask) 

41 

42 

43@libentry() 

44@triton.autotune( 

45 configs=runtime.get_tuned_config("triu_batch"), 

46 key=["batch", "MN", "N", "diagonal"], 

47) 

48@triton.jit(do_not_specialize=["diagonal"]) 

49def triu_batch_kernel( 

50 X, 

51 Y, 

52 batch, 

53 MN, 

54 N, 

55 diagonal, 

56 BATCH_BLOCK_SIZE: tl.constexpr, 

57 MN_BLOCK_SIZE: tl.constexpr, 

58): 

59 batch_id = ext.program_id(0) 

60 mn_id = ext.program_id(1) 

61 row = batch_id * BATCH_BLOCK_SIZE + tl.arange(0, BATCH_BLOCK_SIZE)[:, None] 

62 batch_mask = row < batch 

63 X += row * MN 

64 Y += row * MN 

65 

66 cols = mn_id * MN_BLOCK_SIZE + tl.arange(0, MN_BLOCK_SIZE)[None, :] 

67 mn_mask = cols < MN 

68 mask = batch_mask and mn_mask 

69 x = tl.load(X + cols, mask, other=0.0) 

70 m = cols // N 

71 n = cols % N 

72 y = tl.where(m + diagonal <= n, x, 0.0) 

73 tl.store(Y + cols, y, mask=mask) 

74 

75 

76INT32_MAX = torch.iinfo(torch.int32).max 

77 

78 

79def triu(A, diagonal=0): 

80 logger.debug("GEMS TRIU") 

81 A = A.contiguous() 

82 ori_type = A.dtype 

83 out = torch.empty(A.shape, device="ptpu").as_strided(A.shape, A.stride()) 

84 assert len(A.shape) > 1, "Input tensor must have at least 2 dimensions" 

85 M, N = A.shape[-2:] 

86 with torch_device_fn.device(A.device): 

87 if len(A.shape) == 2: 

88 grid = lambda meta: (triton.cdiv(M, meta["M_BLOCK_SIZE"]),) 

89 triu_kernel[grid](A, out, M, N, diagonal) 

90 else: 

91 batch = int(torch.numel(A) / M / N) 

92 B = A.view(batch, -1) 

93 grid = lambda meta: ( 

94 triton.cdiv(batch, meta["BATCH_BLOCK_SIZE"]), 

95 triton.cdiv(M * N, meta["MN_BLOCK_SIZE"]), 

96 ) 

97 triu_batch_kernel[grid]( 

98 B, 

99 out, 

100 batch, 

101 M * N, 

102 N, 

103 diagonal, 

104 ) 

105 out = out.view(A.shape) 

106 return out.to(ori_type)