[pytorch upstream] softmax on BMG is slow #5133
Description
Describe the issue
We run cait_m36_384 training on BMG compare to 4080s, and find softmax is slow. i.e. 0.643ms vs 0.139ms
reproduce on xpu:
import triton
import triton.language as tl
from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
triton_helpers.set_driver_to_gpu()
from torch._dynamo.testing import rand_strided
from torch._C import _xpu_getCurrentRawStream as get_raw_stream
import torch
@triton_heuristics.reduction(
size_hints={'x': 65536, 'r0_': 1024},
reduction_hint=ReductionHint.DEFAULT,
filename=__file__,
triton_meta={'signature': {'in_ptr0': '*fp16', 'in_ptr1': '*fp16', 'out_ptr2': '*fp16', 'xnumel': 'i32', 'r0_numel': 'i32', 'XBLOCK': 'constexpr', 'R0_BLOCK': 'constexpr'}, 'device': DeviceProperties(type='xpu', index=0, multi_processor_count=20, cc={'architecture': 21479031808, 'device_id': 57867, 'driver_version': '1.6.33578+13', 'gpu_eu_count': 160, 'gpu_subslice_count': 20, 'has_atomic64': True, 'has_bfloat16_conversions': True, 'has_fp16': True, 'has_fp64': True, 'has_subgroup_2d_block_io': True, 'has_subgroup_matrix_multiply_accumulate': True, 'has_subgroup_matrix_multiply_accumulate_tensor_float32': False, 'max_compute_units': 160, 'max_num_sub_groups': 64, 'max_work_group_size': 1024, 'name': 'Intel(R) Arc(TM) B580 Graphics', 'platform_name': 'Intel(R) oneAPI Unified Runtime over Level-Zero', 'sub_group_sizes': [16, 32], 'total_memory': 12168933376, 'type': 'gpu', 'vendor': 'Intel(R) Corporation', 'version': '20.1.0'}, major=None, regs_per_multiprocessor=None, max_threads_per_multi_processor=None, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]], (4,): [['tt.divisibility', 16]]}]},
inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_red_fused__softmax__unsafe_view_add_permute_4', 'mutated_arg_names': [], 'optimize_mem': False, 'no_x_dim': False, 'num_load': 4, 'num_reduction': 2, 'backend_hash': '4F98A0CCE2C4E09763E3797B17267BD555C1ED52C2FEA9E881716C4F50E8E044', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': True, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'tiling_scores': {'x': 42467360, 'r0_': 84934656}, 'kernel_num_gb': 0.084934688, 'kernel_flop': 0}
)
@triton.jit
def triton_red_fused__softmax__unsafe_view_add_permute_4(in_ptr0, in_ptr1, out_ptr2, xnumel, r0_numel, XBLOCK : tl.constexpr, R0_BLOCK : tl.constexpr):
xnumel = 36864
r0_numel = 576
rnumel = r0_numel
RBLOCK: tl.constexpr = R0_BLOCK
xoffset = tl.program_id(0) * XBLOCK
xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
xmask = tl.full([XBLOCK, R0_BLOCK], True, tl.int1)
r0_base = tl.arange(0, R0_BLOCK)[None, :]
rbase = r0_base
x0 = (xindex % 16)
x1 = xindex // 16
tmp1 = tl.load(in_ptr1 + (x0), None, eviction_policy='evict_last').to(tl.float32)
_tmp5 = tl.full([XBLOCK, R0_BLOCK], float("-inf"), tl.float32)
x5 = xindex
for r0_offset in range(0, r0_numel, R0_BLOCK):
r0_index = r0_offset + r0_base
r0_mask = r0_index < r0_numel
roffset = r0_offset
rindex = r0_index
r0_2 = r0_index
tmp0 = tl.load(in_ptr0 + (x0 + 16*r0_2 + 9216*x1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
tmp2 = tmp0 + tmp1
tmp3 = tmp2.to(tl.float32)
tmp4 = tl.broadcast_to(tmp3, [XBLOCK, R0_BLOCK])
tmp6 = triton_helpers.maximum(_tmp5, tmp4)
_tmp5 = tl.where(r0_mask, tmp6, _tmp5)
tmp5 = triton_helpers.max2(_tmp5, 1)[:, None]
_tmp13 = tl.full([XBLOCK, R0_BLOCK], 0, tl.float32)
for r0_offset in range(0, r0_numel, R0_BLOCK):
r0_index = r0_offset + r0_base
r0_mask = r0_index < r0_numel
roffset = r0_offset
rindex = r0_index
r0_2 = r0_index
tmp7 = tl.load(in_ptr0 + (x0 + 16*r0_2 + 9216*x1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
tmp8 = tmp7 + tmp1
tmp9 = tmp8.to(tl.float32)
tmp10 = tmp9 - tmp5
tmp11 = libdevice.exp(tmp10)
tmp12 = tl.broadcast_to(tmp11, [XBLOCK, R0_BLOCK])
tmp14 = _tmp13 + tmp12
_tmp13 = tl.where(r0_mask, tmp14, _tmp13)
tmp13 = tl.sum(_tmp13, 1)[:, None]
x3 = ((xindex // 16) % 576)
x4 = xindex // 9216
for r0_offset in range(0, r0_numel, R0_BLOCK):
r0_index = r0_offset + r0_base
r0_mask = r0_index < r0_numel
roffset = r0_offset
rindex = r0_index
r0_2 = r0_index
tmp15 = tl.load(in_ptr0 + (x0 + 16*r0_2 + 9216*x1), r0_mask, eviction_policy='evict_first', other=0.0).to(tl.float32)
tmp16 = tmp15 + tmp1
tmp17 = tmp16.to(tl.float32)
tmp18 = tmp17 - tmp5
tmp19 = libdevice.exp(tmp18)
tmp20 = (tmp19 / tmp13)
tmp21 = tmp20.to(tl.float32)
tl.store(out_ptr2 + (r0_2 + 576*x3 + 331776*x0 + 5308416*x4), tmp21, r0_mask)
def get_args():
arg_0 = rand_strided((1327104, 16), (16, 1), device='xpu:0', dtype=torch.float16)
arg_1 = rand_strided((16,), (1,), device='xpu:0', dtype=torch.float16)
arg_2 = rand_strided((4, 16, 576, 576), (5308416, 331776, 576, 1), device='xpu:0', dtype=torch.float16)
return arg_0, arg_1, arg_2, 36864, 576,
def call(args):
with torch.xpu._DeviceGuard(0):
torch.xpu.set_device(0)
stream0 = get_raw_stream(0)
triton_red_fused__softmax__unsafe_view_add_permute_4.run(*args, stream=stream0)
def benchmark_all_configs(args):
with torch.xpu._DeviceGuard(0):
torch.xpu.set_device(0)
return triton_red_fused__softmax__unsafe_view_add_permute_4.benchmark_all_configs(*args)
if __name__ == '__main__':
from torch._inductor.runtime.benchmarking import benchmarker
args = get_args()
ms = benchmarker.benchmark_gpu(lambda: call(args), rep=40)
num_gb = 0.084934688
gb_per_s = num_gb / (ms / 1e3)
print(f"{ms:.3f}ms {num_gb:.3f}GB {gb_per_s:.2f}GB/s")reproduce on cuda:
import triton
import triton.language as tl
from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, DeviceProperties
triton_helpers.set_driver_to_gpu()
from torch._dynamo.testing import rand_strided
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
import torch
@triton_heuristics.reduction(
size_hints={'x': 65536, 'r0_': 1024},
reduction_hint=ReductionHint.DEFAULT,
filename=__file__,
triton_meta={'signature': {'in_ptr0': '*fp16', 'in_ptr1': '*fp16', 'out_ptr2': '*fp16', 'xnumel': 'i32', 'r0_numel': 'i32', 'XBLOCK': 'constexpr', 'R0_BLOCK': 'constexpr'}, 'device': DeviceProperties(type='cuda', index=0, multi_processor_count=80, cc=89, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=1536, warp_size=32), 'constants': {}, 'configs': [{(0,): [['tt.divisibility', 16]], (1,): [['tt.divisibility', 16]], (2,): [['tt.divisibility', 16]], (3,): [['tt.divisibility', 16]], (4,): [['tt.divisibility', 16]]}]},
inductor_meta={'grid_type': 'Grid1D', 'autotune_hints': set(), 'kernel_name': 'triton_red_fused__softmax__unsafe_view_add_permute_4', 'mutated_arg_names': [], 'optimize_mem': False, 'no_x_dim': False, 'num_load': 4, 'num_reduction': 2, 'backend_hash': '7606AC1BD735D3E3F140115999815ACFE642967D9047962703B386A670225BF4', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': True, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'tiling_scores': {'x': 42467360, 'r0_': 84934656}, 'kernel_num_gb': 0.084934688, 'kernel_flop': 0}
)
@triton.jit
def triton_red_fused__softmax__unsafe_view_add_permute_4(in_ptr0, in_ptr1, out_ptr2, xnumel, r0_numel, XBLOCK : tl.constexpr, R0_BLOCK : tl.constexpr):
xnumel = 36864
r0_numel = 576
rnumel = r0_numel
RBLOCK: tl.constexpr = R0_BLOCK
xoffset = tl.program_id(0) * XBLOCK
xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
xmask = tl.full([XBLOCK, R0_BLOCK], True, tl.int1)
r0_base = tl.arange(0, R0_BLOCK)[None, :]
rbase = r0_base
x0 = (xindex % 16)
x1 = xindex // 16
tmp1 = tl.load(in_ptr1 + (x0), None, eviction_policy='evict_last').to(tl.float32)
_tmp5 = tl.full([XBLOCK, R0_BLOCK], float("-inf"), tl.float32)
x5 = xindex
for r0_offset in range(0, r0_numel, R0_BLOCK):
r0_index = r0_offset + r0_base
r0_mask = r0_index < r0_numel
roffset = r0_offset
rindex = r0_index
r0_2 = r0_index
tmp0 = tl.load(in_ptr0 + (x0 + 16*r0_2 + 9216*x1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
tmp2 = tmp0 + tmp1
tmp3 = tmp2.to(tl.float32)
tmp4 = tl.broadcast_to(tmp3, [XBLOCK, R0_BLOCK])
tmp6 = triton_helpers.maximum(_tmp5, tmp4)
_tmp5 = tl.where(r0_mask, tmp6, _tmp5)
tmp5 = triton_helpers.max2(_tmp5, 1)[:, None]
_tmp13 = tl.full([XBLOCK, R0_BLOCK], 0, tl.float32)
for r0_offset in range(0, r0_numel, R0_BLOCK):
r0_index = r0_offset + r0_base
r0_mask = r0_index < r0_numel
roffset = r0_offset
rindex = r0_index
r0_2 = r0_index
tmp7 = tl.load(in_ptr0 + (x0 + 16*r0_2 + 9216*x1), r0_mask, eviction_policy='evict_last', other=0.0).to(tl.float32)
tmp8 = tmp7 + tmp1
tmp9 = tmp8.to(tl.float32)
tmp10 = tmp9 - tmp5
tmp11 = libdevice.exp(tmp10)
tmp12 = tl.broadcast_to(tmp11, [XBLOCK, R0_BLOCK])
tmp14 = _tmp13 + tmp12
_tmp13 = tl.where(r0_mask, tmp14, _tmp13)
tmp13 = tl.sum(_tmp13, 1)[:, None]
x3 = ((xindex // 16) % 576)
x4 = xindex // 9216
for r0_offset in range(0, r0_numel, R0_BLOCK):
r0_index = r0_offset + r0_base
r0_mask = r0_index < r0_numel
roffset = r0_offset
rindex = r0_index
r0_2 = r0_index
tmp15 = tl.load(in_ptr0 + (x0 + 16*r0_2 + 9216*x1), r0_mask, eviction_policy='evict_first', other=0.0).to(tl.float32)
tmp16 = tmp15 + tmp1
tmp17 = tmp16.to(tl.float32)
tmp18 = tmp17 - tmp5
tmp19 = libdevice.exp(tmp18)
tmp20 = (tmp19 / tmp13)
tmp21 = tmp20.to(tl.float32)
tl.store(out_ptr2 + (r0_2 + 576*x3 + 331776*x0 + 5308416*x4), tmp21, r0_mask)
def get_args():
arg_0 = rand_strided((1327104, 16), (16, 1), device='cuda:0', dtype=torch.float16)
arg_1 = rand_strided((16,), (1,), device='cuda:0', dtype=torch.float16)
arg_2 = rand_strided((4, 16, 576, 576), (5308416, 331776, 576, 1), device='cuda:0', dtype=torch.float16)
return arg_0, arg_1, arg_2, 36864, 576,
def call(args):
with torch.cuda._DeviceGuard(0):
torch.cuda.set_device(0)
stream0 = get_raw_stream(0)
triton_red_fused__softmax__unsafe_view_add_permute_4.run(*args, stream=stream0)
def benchmark_all_configs(args):
with torch.cuda._DeviceGuard(0):
torch.cuda.set_device(0)
return triton_red_fused__softmax__unsafe_view_add_permute_4.benchmark_all_configs(*args)
if __name__ == '__main__':
from torch._inductor.runtime.benchmarking import benchmarker
args = get_args()
ms = benchmarker.benchmark_gpu(lambda: call(args), rep=40)
num_gb = 0.084934688
gb_per_s = num_gb / (ms / 1e3)
print(f"{ms:.3f}ms {num_gb:.3f}GB {gb_per_s:.2f}GB/s")Environment details
pytorch-triton-xpu 3.5.0+git1b0418a9
Intel(R) Arc(TM) B580 Graphics