Coverage for src/flag_gems/runtime/backend/_kunlunxin/ops/softshrink.py: 0%

1import torch 

2import triton 

3import triton.language as tl 

4 

5from flag_gems.runtime import torch_device_fn 

6 

7 

8@triton.jit 

9def softshrink_kernel(x_ptr, out_ptr, n_elements, lambd, BLOCK_SIZE: tl.constexpr): 

10 pid = tl.program_id(axis=0) 

11 block_start = pid * BLOCK_SIZE 

12 offsets = block_start + tl.arange(0, BLOCK_SIZE) 

13 mask = offsets < n_elements 

14 

15 x = tl.load(x_ptr + offsets, mask=mask, other=0) 

16 x32 = x.to(tl.float32) 

17 

18 threshold = lambd # scalar float32 

19 

20 gt = x32 > threshold 

21 lt = x32 < -threshold 

22 res32 = tl.where(gt, x32 - threshold, tl.where(lt, x32 + threshold, 0.0)) 

23 

24 # Propagate NaN: if x is NaN, keep it 

25 # res32 = tl.where(x32 != x32, x32, res32) 

26 x_bits = x32.to(tl.int32, bitcast=True) 

27 is_nan = (x_bits & 0x7FFFFFFF) > 0x7F800000 

28 res32 = tl.where(is_nan, x32, res32) 

29 

30 res = res32.to(x.dtype) 

31 tl.store(out_ptr + offsets, res, mask=mask) 

32 

33 

34def _check_supported_dtype(t: torch.Tensor): 

35 if t.dtype not in (torch.float16, torch.bfloat16, torch.float32): 

36 raise TypeError( 

37 f"Unsupported dtype {t.dtype}. Supported dtypes are float16, bfloat16, and float32." 

38 ) 

39 

40 

41def _launch_softshrink_kernel(x: torch.Tensor, out: torch.Tensor, lambd: float): 

42 n_elements = x.numel() 

43 if n_elements == 0: 

44 return 

45 BLOCK_SIZE = 1024 

46 grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) 

47 with torch_device_fn.device(x.device): 

48 softshrink_kernel[grid]( 

49 x, 

50 out, 

51 n_elements, 

52 float(lambd), 

53 BLOCK_SIZE=BLOCK_SIZE, 

54 num_warps=4, 

55 ) 

56 

57 

58def softshrink(input: torch.Tensor, lambd: float = 0.5): 

59 _check_supported_dtype(input) 

60 x = input.contiguous() 

61 out = torch.empty_like(x) 

62 _launch_softshrink_kernel(x, out, lambd) 

63 return out.reshape_as(input) 

64 

65 

66def softshrink_out(input: torch.Tensor, lambd: float = 0.5, out: torch.Tensor = None): 

67 if out is None: 

68 raise ValueError("Argument 'out' must be provided for softshrink_out.") 

69 if input.shape != out.shape: 

70 raise ValueError( 

71 f"Shape mismatch: input.shape={input.shape}, out.shape={out.shape}" 

72 ) 

73 if input.dtype != out.dtype: 

74 raise TypeError( 

75 f"Dtype mismatch: input.dtype={input.dtype}, out.dtype={out.dtype}" 

76 ) 

77 _check_supported_dtype(input) 

78 

79 x = input.contiguous() 

80 if out.is_contiguous(): 

81 out_buf = out 

82 else: 

83 out_buf = torch.empty_like(out, memory_format=torch.contiguous_format) 

84 

85 _launch_softshrink_kernel(x, out_buf, lambd) 

86 

87 if out_buf.data_ptr() != out.data_ptr(): 

88 out.copy_(out_buf) 

89 return out