Created
May 6, 2024 19:21
-
-
Save eellison/36f0c3e6025360315dd67932461ff11b to your computer and use it in GitHub Desktop.
repro.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# AOT ID: ['0_inference'] | |
from ctypes import c_void_p, c_long | |
import torch | |
import math | |
import random | |
import os | |
import tempfile | |
from math import inf, nan | |
from torch._inductor.hooks import run_intermediate_hooks | |
from torch._inductor.utils import maybe_profile | |
from torch._inductor.codegen.memory_planning import _align as align | |
from torch import device, empty_strided | |
from torch._inductor.codecache import AsyncCompile | |
from torch._inductor.select_algorithm import extern_kernels | |
from torch._inductor.codegen.multi_kernel import MultiKernelCall | |
aten = torch.ops.aten | |
inductor_ops = torch.ops.inductor | |
_quantized = torch.ops._quantized | |
assert_size_stride = torch._C._dynamo.guards.assert_size_stride | |
empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu | |
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda | |
alloc_from_pool = torch.ops.inductor._alloc_from_pool | |
reinterpret_tensor = torch.ops.inductor._reinterpret_tensor | |
async_compile = AsyncCompile() | |
# kernel path: /tmp/torchinductor_tarasg/rf/crfz34h2i76un6i6igtlap7xy6f3ukco3czymaj4tdhrdcoredec.py | |
# Source Nodes: [padded1], Original ATen: [test_inductor_ops.jagged_to_padded_dense] | |
# padded1 => jagged_to_padded_dense | |
triton_poi_fused_jagged_to_padded_dense_0 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.pointwise( | |
size_hints=[131072], | |
filename=__file__, | |
triton_meta={'signature': {0: '*i32', 1: '*i64', 2: '*i64', 3: 'i32', 4: 'i32', 5: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 5), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_jagged_to_padded_dense_0', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True}, | |
min_elem_per_thread=0 | |
) | |
@triton.jit | |
def triton_(in_ptr0, in_ptr1, out_ptr0, ks0, ks1, xnumel, XBLOCK : tl.constexpr): | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:] | |
xmask = xindex < xnumel | |
x1 = (xindex // ks0) | |
x0 = xindex % ks0 | |
x2 = xindex | |
tmp0 = tl.load(in_ptr0 + (x1), xmask, eviction_policy='evict_last') | |
tmp4 = tl.load(in_ptr0 + (1 + x1), xmask, eviction_policy='evict_last') | |
tmp1 = tmp0 + ks1 + 1 | |
tmp2 = tmp0 < 0 | |
tmp3 = tl.where(tmp2, tmp1, tmp0) | |
tl.device_assert(((0 <= tmp3) & (tmp3 < 1 + ks1)) | ~xmask, "index out of bounds: 0 <= tmp3 < 1 + ks1") | |
tmp5 = tmp3 + x0 | |
tmp6 = tmp5 < tmp4 | |
tmp7 = tl.load(in_ptr1 + (tmp3 + x0), tmp6 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp8 = tl.full(tmp7.shape, 4611686018427387904, tmp7.dtype) | |
tmp9 = tl.where(tmp6, tmp7, tmp8) | |
tl.store(out_ptr0 + (x2), tmp9, xmask) | |
''', device_str='cuda') | |
import triton | |
import triton.language as tl | |
from torch._inductor.runtime.triton_heuristics import grid, split_scan_grid, start_graph, end_graph | |
from torch._C import _cuda_getCurrentRawStream as get_raw_stream | |
# kernel path: /tmp/torchinductor_tarasg/xb/cxboeyvzkjclnnuaecctdhoa7tuhbvk2j2wobcybczuumws5c5vh.py | |
# Source Nodes: [unique_0hot], Original ATen: [aten.cat] | |
# unique_0hot => cat | |
triton_poi_fused_cat_1 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.pointwise( | |
size_hints=[131072], | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: 'i32', 3: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 3), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_cat_1', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True}, | |
min_elem_per_thread=0 | |
) | |
@triton.jit | |
def triton_(in_ptr0, out_ptr0, ks0, xnumel, XBLOCK : tl.constexpr): | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:] | |
xmask = xindex < xnumel | |
x0 = xindex % ks0 | |
x2 = xindex | |
tmp0 = x0 | |
tmp1 = tl.full([1], 0, tl.int64) | |
tmp2 = tmp0 >= tmp1 | |
tmp3 = tl.full([1], 1, tl.int64) | |
tmp4 = tmp0 < tmp3 | |
tmp5 = tl.full(tmp1.shape, 0.0, tmp1.dtype) | |
tmp6 = tl.where(tmp4, tmp1, tmp5) | |
tmp7 = tmp0 >= tmp3 | |
tmp8 = ks0 | |
tmp9 = tmp0 < tmp8 | |
tmp10 = tl.load(in_ptr0 + (x2), tmp7 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp11 = tl.load(in_ptr0 + ((-1) + x2), tmp7 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp12 = tmp10 - tmp11 | |
tmp13 = tmp12 > tmp1 | |
tmp14 = tmp13.to(tl.int8) | |
tmp15 = tl.full([1], 1, tl.int8) | |
tmp16 = tmp15 - tmp14 | |
tmp17 = tmp16.to(tl.int64) | |
tmp18 = tl.full(tmp17.shape, 0.0, tmp17.dtype) | |
tmp19 = tl.where(tmp7, tmp17, tmp18) | |
tmp20 = tl.where(tmp4, tmp6, tmp19) | |
tl.store(out_ptr0 + (x2), tmp20, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/kh/ckhdqw5p3h5ibw7ycw4thl5uf7m4rtyghbdkarofakutjtmfxj7v.py | |
# Source Nodes: [mask], Original ATen: [aten.zeros_like] | |
# mask => full_1 | |
triton_poi_fused_zeros_like_2 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.pointwise( | |
size_hints=[131072], | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_zeros_like_2', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True}, | |
min_elem_per_thread=0 | |
) | |
@triton.jit | |
def triton_(out_ptr0, xnumel, XBLOCK : tl.constexpr): | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:] | |
xmask = xindex < xnumel | |
x0 = xindex | |
tmp0 = tl.full([1], 0, tl.int64) | |
tl.store(out_ptr0 + (x0), tmp0, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/nl/cnlzhvarzel643wxe6ed6obbhv64tfuqtadbiq2iab6whtwb6nif.py | |
# Source Nodes: [adj_diff, diff, gt, mask, n_unique, setitem], Original ATen: [aten._to_copy, aten.gt, aten.index_put, aten.lift_fresh, aten.sub, aten.sum, aten.zeros_like] | |
# adj_diff => convert_element_type | |
# diff => sub_1 | |
# gt => gt | |
# mask => full_1 | |
# n_unique => sum_1 | |
# setitem => full_default_1, index_put | |
triton_red_fused__to_copy_gt_index_put_lift_fresh_sub_sum_zeros_like_3 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.reduction( | |
size_hints=[1024, 128], | |
reduction_hint=ReductionHint.INNER, | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: 'i32', 3: 'i32', 4: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 3), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_red_fused__to_copy_gt_index_put_lift_fresh_sub_sum_zeros_like_3', 'mutated_arg_names': ['out_ptr1'], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True} | |
) | |
@triton.jit | |
def triton_(in_ptr0, out_ptr1, ks0, xnumel, rnumel, XBLOCK : tl.constexpr, RBLOCK : tl.constexpr): | |
xnumel = 1024 | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:, None] | |
xmask = xindex < xnumel | |
rbase = tl.arange(0, RBLOCK)[None, :] | |
x0 = xindex | |
_tmp8 = tl.full([XBLOCK, RBLOCK], 0, tl.int64) | |
for roffset in range(0, rnumel, RBLOCK): | |
rindex = roffset + rbase | |
rmask = rindex < rnumel | |
r1 = rindex | |
tmp0 = tl.load(in_ptr0 + (1 + r1 + x0 + (ks0*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp1 = tl.load(in_ptr0 + (r1 + x0 + (ks0*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp2 = tmp0 - tmp1 | |
tmp3 = tl.full([1, 1], 0, tl.int64) | |
tmp4 = tmp2 > tmp3 | |
tmp5 = tmp4.to(tl.int8) | |
tmp6 = tmp5.to(tl.int64) | |
tmp7 = tl.broadcast_to(tmp6, [XBLOCK, RBLOCK]) | |
tmp9 = _tmp8 + tmp7 | |
_tmp8 = tl.where(rmask & xmask, tmp9, _tmp8) | |
tmp8 = tl.sum(_tmp8, 1)[:, None] | |
tmp10 = tmp8 + ks0 + 1 | |
tmp11 = tmp8 < 0 | |
tmp12 = tl.where(tmp11, tmp10, tmp8) | |
tl.device_assert(((0 <= tmp12) & (tmp12 < 1 + ks0)) | ~xmask, "index out of bounds: 0 <= tmp12 < 1 + ks0") | |
tmp13 = tl.full([1, 1], 1, tl.int64) | |
tl.store(out_ptr1 + (tmp12 + x0 + (ks0*x0)), tmp13, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/6m/c6m53wyb7krckmgwrwcecvcyjtuu664zyttgnscwbulfnkdmjjjb.py | |
# Source Nodes: [cumsum, gather_out, setitem_1], Original ATen: [aten.cumsum, aten.gather, aten.index_put, aten.lift_fresh] | |
# cumsum => cumsum | |
# gather_out => gather | |
# setitem_1 => full_default_2, index_put_1 | |
triton_red_fused_cumsum_gather_index_put_lift_fresh_4 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton.jit | |
def _triton_helper_fn_add0(arg0_0, arg1_0): | |
tmp0 = arg0_0 + arg1_0 | |
return tmp0 | |
@triton_heuristics.reduction( | |
size_hints=[1024, 128], | |
reduction_hint=ReductionHint.INNER, | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: '*i64', 3: '*i64', 4: '*i64', 5: 'i32', 6: 'i32', 7: 'i32', 8: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4, 7), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_red_fused_cumsum_gather_index_put_lift_fresh_4', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True} | |
) | |
@triton.jit | |
def triton_(in_ptr0, in_ptr1, in_ptr2, out_ptr0, out_ptr1, ks0, ks1, xnumel, rnumel, XBLOCK : tl.constexpr, RBLOCK : tl.constexpr): | |
xnumel = 1024 | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:, None] | |
xmask = xindex < xnumel | |
rbase = tl.arange(0, RBLOCK)[None, :] | |
x0 = xindex | |
tmp3 = tl.full([XBLOCK, 1], -1, tl.int64) | |
for roffset in range(0, rnumel, RBLOCK): | |
rindex = roffset + rbase | |
rmask = rindex < rnumel | |
r1 = rindex | |
tmp0 = tl.load(in_ptr0 + (r1 + x0 + (ks0*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp10 = tl.load(in_ptr1 + (r1 + x0 + (ks0*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp1 = tmp0.to(tl.int64) | |
tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK]) | |
tmp4, = tl.associative_scan((tmp2,), 1, _triton_helper_fn_add0) | |
tmp5, = tl.reduce((tmp2,), 1, _triton_helper_fn_add0) | |
tmp6 = tmp3 + tmp5[:, None] | |
tmp7 = tmp3 + tmp4 | |
tmp8 = tl.where(roffset > 0, tmp7, tmp4) | |
tmp3 = tl.where(roffset > 0, tmp6, tmp5[:, None]) | |
tmp9 = (tmp8 != 0) | |
tmp11 = tmp10 + ks0 + 1 | |
tmp12 = tmp10 < 0 | |
tmp13 = tl.where(tmp12, tmp11, tmp10) | |
tl.device_assert(((0 <= tmp13) & (tmp13 < 1 + ks0)) | ~(xmask & rmask), "index out of bounds: 0 <= tmp13 < 1 + ks0") | |
tmp14 = tl.load(in_ptr2 + (tmp13 + x0 + (ks0*x0)), rmask & xmask, eviction_policy='evict_last') | |
tmp15 = tl.full([1, 1], 4611686018427387904, tl.int64) | |
tmp16 = tl.where(tmp9, tmp15, tmp14) | |
tl.store(out_ptr0 + (r1 + x0 + (ks0*x0)), tmp8, rmask & xmask) | |
tl.store(out_ptr1 + (r1 + (2*x0) + (ks0*x0) + (ks1*x0)), tmp16, rmask & xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/gg/cggya3gpn3ejtsdq2ezgyaqi55k6325s5kndciveu5yjwvohcztz.py | |
# Source Nodes: [cumsum_1, gather_out_1, setitem_4], Original ATen: [aten.cumsum, aten.gather, aten.index_put, aten.lift_fresh] | |
# cumsum_1 => cumsum_1 | |
# gather_out_1 => gather_1 | |
# setitem_4 => full_default_6, index_put_4 | |
triton_red_fused_cumsum_gather_index_put_lift_fresh_5 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton.jit | |
def _triton_helper_fn_add0(arg0_0, arg1_0): | |
tmp0 = arg0_0 + arg1_0 | |
return tmp0 | |
@triton_heuristics.reduction( | |
size_hints=[1024, 128], | |
reduction_hint=ReductionHint.INNER, | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: '*i64', 3: '*i64', 4: '*i64', 5: 'i32', 6: 'i32', 7: 'i32', 8: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 7), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_red_fused_cumsum_gather_index_put_lift_fresh_5', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True} | |
) | |
@triton.jit | |
def triton_(in_ptr0, in_ptr1, in_ptr2, out_ptr0, out_ptr1, ks0, ks1, xnumel, rnumel, XBLOCK : tl.constexpr, RBLOCK : tl.constexpr): | |
xnumel = 1024 | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:, None] | |
xmask = xindex < xnumel | |
rbase = tl.arange(0, RBLOCK)[None, :] | |
x0 = xindex | |
tmp3 = tl.full([XBLOCK, 1], -1, tl.int64) | |
for roffset in range(0, rnumel, RBLOCK): | |
rindex = roffset + rbase | |
rmask = rindex < rnumel | |
r1 = rindex | |
tmp0 = tl.load(in_ptr0 + (r1 + x0 + (ks0*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp10 = tl.load(in_ptr1 + (r1 + x0 + (ks0*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp1 = tmp0.to(tl.int64) | |
tmp2 = tl.broadcast_to(tmp1, [XBLOCK, RBLOCK]) | |
tmp4, = tl.associative_scan((tmp2,), 1, _triton_helper_fn_add0) | |
tmp5, = tl.reduce((tmp2,), 1, _triton_helper_fn_add0) | |
tmp6 = tmp3 + tmp5[:, None] | |
tmp7 = tmp3 + tmp4 | |
tmp8 = tl.where(roffset > 0, tmp7, tmp4) | |
tmp3 = tl.where(roffset > 0, tmp6, tmp5[:, None]) | |
tmp9 = (tmp8 != 0) | |
tmp11 = tmp10 + ks0 + 1 | |
tmp12 = tmp10 < 0 | |
tmp13 = tl.where(tmp12, tmp11, tmp10) | |
tl.device_assert(((0 <= tmp13) & (tmp13 < 1 + ks0)) | ~(xmask & rmask), "index out of bounds: 0 <= tmp13 < 1 + ks0") | |
tmp14 = tl.load(in_ptr2 + (tmp13 + x0 + (ks0*x0)), rmask & xmask, eviction_policy='evict_last') | |
tmp15 = tl.full([1, 1], 4611686018427387904, tl.int64) | |
tmp16 = tl.where(tmp9, tmp15, tmp14) | |
tl.store(out_ptr0 + (r1 + x0 + (ks0*x0)), tmp8, rmask & xmask) | |
tl.store(out_ptr1 + (r1 + (2*x0) + (ks0*x0) + (ks1*x0)), tmp16, rmask & xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/et/cetmg7225yixhvwn5ik2slfjataizwcl2zguwayk7ysvmj6xdjdc.py | |
# Source Nodes: [setitem_2, unique_indices_1], Original ATen: [aten.index_put, aten.lift_fresh, aten.sub] | |
# setitem_2 => full_default_3, index_put_2 | |
# unique_indices_1 => sub_4 | |
triton_poi_fused_index_put_lift_fresh_sub_6 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.pointwise( | |
size_hints=[131072], | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: '*i64', 3: 'i32', 4: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 4), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_index_put_lift_fresh_sub_6', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True}, | |
min_elem_per_thread=0 | |
) | |
@triton.jit | |
def triton_(in_ptr0, in_ptr1, out_ptr0, ks0, xnumel, XBLOCK : tl.constexpr): | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:] | |
xmask = xindex < xnumel | |
x0 = xindex % ks0 | |
x1 = (xindex // ks0) | |
x2 = xindex | |
tmp0 = tl.load(in_ptr0 + (x0 + x1 + (ks0*x1)), xmask, eviction_policy='evict_last') | |
tmp2 = tl.load(in_ptr1 + (1 + x0 + x1 + (ks0*x1)), xmask, eviction_policy='evict_last') | |
tmp3 = tl.load(in_ptr1 + (x0 + x1 + (ks0*x1)), xmask, eviction_policy='evict_last') | |
tmp1 = (tmp0 != 0) | |
tmp4 = tmp2 - tmp3 | |
tmp5 = tl.full([1], 0, tl.int64) | |
tmp6 = tl.where(tmp1, tmp5, tmp4) | |
tl.store(out_ptr0 + (x2), tmp6, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/pf/cpfuizamjfs7clym7hsjrhozara6z4tylhydowigp4ttpalrmdz6.py | |
# Source Nodes: [joined_padded_counts], Original ATen: [aten.cat] | |
# joined_padded_counts => cat_5 | |
triton_poi_fused_cat_7 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.pointwise( | |
size_hints=[262144], | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: '*i64', 3: 'i32', 4: 'i32', 5: 'i32', 6: 'i32', 7: 'i32', 8: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 8), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_cat_7', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True}, | |
min_elem_per_thread=0 | |
) | |
@triton.jit | |
def triton_(in_ptr0, in_ptr1, out_ptr0, ks0, ks1, ks2, ks3, ks4, xnumel, XBLOCK : tl.constexpr): | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:] | |
xmask = xindex < xnumel | |
x0 = xindex % ks0 | |
x1 = (xindex // ks0) | |
x2 = xindex | |
tmp0 = x0 | |
tmp1 = tl.full([1], 0, tl.int64) | |
tmp2 = tmp0 >= tmp1 | |
tmp3 = ks1 | |
tmp4 = tmp0 < tmp3 | |
tmp5 = tl.broadcast_to(ks2, [XBLOCK]) | |
tmp6 = tmp0 < tmp5 | |
tmp7 = tmp6 & tmp4 | |
tmp8 = tl.load(in_ptr0 + (x0 + (ks2*x1)), tmp7 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp9 = tl.full(tmp8.shape, 0.0, tmp8.dtype) | |
tmp10 = tl.where(tmp7, tmp8, tmp9) | |
tmp11 = tmp0 >= tmp5 | |
tmp12 = tl.broadcast_to(ks1, [XBLOCK]) | |
tmp13 = tmp0 < tmp12 | |
tmp14 = tmp11 & tmp4 | |
tmp15 = tl.load(in_ptr0 + ((-1) + ks2 + (ks2*x1)), tmp14 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp16 = tl.full(tmp15.shape, 0.0, tmp15.dtype) | |
tmp17 = tl.where(tmp14, tmp15, tmp16) | |
tmp18 = tl.where(tmp6, tmp10, tmp17) | |
tmp19 = tl.full(tmp18.shape, 0.0, tmp18.dtype) | |
tmp20 = tl.where(tmp4, tmp18, tmp19) | |
tmp21 = tmp0 >= tmp3 | |
tmp22 = ks0 | |
tmp23 = tmp0 < tmp22 | |
tmp24 = (-1) + x0 + ((-1)*ks2) | |
tmp25 = tmp24 >= tmp1 | |
tmp26 = tl.broadcast_to(ks3, [XBLOCK]) | |
tmp27 = tmp24 < tmp26 | |
tmp28 = tmp27 & tmp21 | |
tmp29 = tl.load(in_ptr1 + ((-1) + x0 + ((-1)*ks2) + (ks3*x1)), tmp28 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp30 = tl.full(tmp29.shape, 0.0, tmp29.dtype) | |
tmp31 = tl.where(tmp28, tmp29, tmp30) | |
tmp32 = tmp24 >= tmp26 | |
tmp33 = tl.broadcast_to(ks4, [XBLOCK]) | |
tmp34 = tmp24 < tmp33 | |
tmp35 = tmp32 & tmp21 | |
tmp36 = tl.load(in_ptr1 + ((-1) + ks3 + (ks3*x1)), tmp35 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp37 = tl.full(tmp36.shape, 0.0, tmp36.dtype) | |
tmp38 = tl.where(tmp35, tmp36, tmp37) | |
tmp39 = tl.where(tmp27, tmp31, tmp38) | |
tmp40 = tl.full(tmp39.shape, 0.0, tmp39.dtype) | |
tmp41 = tl.where(tmp21, tmp39, tmp40) | |
tmp42 = tl.where(tmp4, tmp20, tmp41) | |
tl.store(out_ptr0 + (x2), tmp42, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/ce/ccedtkpzbs66fozel5to6qjqjdjurk5kdgejw6n63xvxefut46ie.py | |
# Source Nodes: [min_1], Original ATen: [aten.min] | |
# min_1 => min_1 | |
triton_poi_fused_min_8 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.pointwise( | |
size_hints=[262144], | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: '*i64', 3: 'i32', 4: 'i32', 5: 'i32', 6: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 6), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_min_8', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True}, | |
min_elem_per_thread=0 | |
) | |
@triton.jit | |
def triton_(in_ptr0, in_ptr1, out_ptr0, ks0, ks1, ks2, xnumel, XBLOCK : tl.constexpr): | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:] | |
xmask = xindex < xnumel | |
x1 = (xindex // ks0) | |
x2 = xindex | |
x0 = xindex % ks0 | |
tmp0 = x1 | |
tmp1 = tl.full([1], 0, tl.int64) | |
tmp2 = tmp0 >= tmp1 | |
tmp3 = tl.full([1], 1024, tl.int64) | |
tmp4 = tmp0 < tmp3 | |
tmp5 = tl.load(in_ptr0 + (x2), tmp4 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp6 = tmp5 + ks1 + ks2 + 2 | |
tmp7 = tmp5 < 0 | |
tmp8 = tl.where(tmp7, tmp6, tmp5) | |
tl.device_assert(((0 <= tmp8) & (tmp8 < 2 + ks1 + ks2)) | ~(tmp4 & xmask), "index out of bounds: 0 <= tmp8 < 2 + ks1 + ks2") | |
tmp9 = tl.load(in_ptr1 + (tmp8 + (2*x1) + (ks1*x1) + (ks2*x1)), tmp4 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype) | |
tmp11 = tl.where(tmp4, tmp9, tmp10) | |
tmp12 = tmp0 >= tmp3 | |
tmp13 = tl.full([1], 2048, tl.int64) | |
tmp14 = tmp0 < tmp13 | |
tmp15 = tl.load(in_ptr0 + ((-2048) + ((-1024)*ks1) + ((-1024)*ks2) + (2*x1) + (ks1*x1) + (ks2*x1) + ((x0 + ((3 + ks1 + ks2) % ks0)) % ks0)), tmp12 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp16 = tmp15 + ks1 + ks2 + 2 | |
tmp17 = tmp15 < 0 | |
tmp18 = tl.where(tmp17, tmp16, tmp15) | |
tl.device_assert(((0 <= tmp18) & (tmp18 < 2 + ks1 + ks2)) | ~(tmp12 & xmask), "index out of bounds: 0 <= tmp18 < 2 + ks1 + ks2") | |
tmp19 = tl.load(in_ptr1 + ((-2048) + tmp18 + ((-1024)*ks1) + ((-1024)*ks2) + (2*x1) + (ks1*x1) + (ks2*x1)), tmp12 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp20 = tl.full(tmp19.shape, 0.0, tmp19.dtype) | |
tmp21 = tl.where(tmp12, tmp19, tmp20) | |
tmp22 = tl.where(tmp4, tmp11, tmp21) | |
tmp23 = 1024 + x1 | |
tmp24 = tmp23 >= tmp1 | |
tmp25 = tmp23 < tmp3 | |
tmp26 = tl.load(in_ptr0 + (2048 + x2 + (1024*ks1) + (1024*ks2)), tmp25 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp27 = tmp26 + ks1 + ks2 + 2 | |
tmp28 = tmp26 < 0 | |
tmp29 = tl.where(tmp28, tmp27, tmp26) | |
tl.device_assert(((0 <= tmp29) & (tmp29 < 2 + ks1 + ks2)) | ~(xmask & tmp25), "index out of bounds: 0 <= tmp29 < 2 + ks1 + ks2") | |
tmp30 = tl.load(in_ptr1 + (2048 + tmp29 + (2*x1) + (1024*ks1) + (1024*ks2) + (ks1*x1) + (ks2*x1)), tmp25 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp31 = tl.full(tmp30.shape, 0.0, tmp30.dtype) | |
tmp32 = tl.where(tmp25, tmp30, tmp31) | |
tmp33 = tmp23 >= tmp3 | |
tmp34 = tmp23 < tmp13 | |
tmp35 = tl.load(in_ptr0 + ((2*x1) + (ks1*x1) + (ks2*x1) + ((x0 + ((3 + ks1 + ks2) % ks0)) % ks0)), tmp33 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp36 = tmp35 + ks1 + ks2 + 2 | |
tmp37 = tmp35 < 0 | |
tmp38 = tl.where(tmp37, tmp36, tmp35) | |
tl.device_assert(((0 <= tmp38) & (tmp38 < 2 + ks1 + ks2)) | ~(xmask & tmp33), "index out of bounds: 0 <= tmp38 < 2 + ks1 + ks2") | |
tmp39 = tl.load(in_ptr1 + (tmp38 + (2*x1) + (ks1*x1) + (ks2*x1)), tmp33 & xmask, eviction_policy='evict_last', other=0.0) | |
tmp40 = tl.full(tmp39.shape, 0.0, tmp39.dtype) | |
tmp41 = tl.where(tmp33, tmp39, tmp40) | |
tmp42 = tl.where(tmp25, tmp32, tmp41) | |
tmp43 = triton_helpers.minimum(tmp22, tmp42) | |
tl.store(out_ptr0 + (x2), tmp43, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/j5/cj5tpoxiug34serot5wjm52idc2zg37fxasdgotwwdomnlu45xch.py | |
# Source Nodes: [diff_4, eq, joined_padded_lengths, padded_diff, padding_mask, stack_5, stack_7, to_4], Original ATen: [aten._to_copy, aten.eq, aten.gt, aten.mul, aten.stack, aten.sub, aten.sum] | |
# diff_4 => sub_11 | |
# eq => eq | |
# joined_padded_lengths => sum_3 | |
# padded_diff => mul | |
# padding_mask => gt_2 | |
# stack_5 => cat_12 | |
# stack_7 => cat_9 | |
# to_4 => convert_element_type_4 | |
triton_red_fused__to_copy_eq_gt_mul_stack_sub_sum_9 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.reduction( | |
size_hints=[1024, 256], | |
reduction_hint=ReductionHint.DEFAULT, | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: '*i64', 3: '*i8', 4: '*i64', 5: '*i64', 6: '*i64', 7: '*i64', 8: 'i32', 9: 'i32', 10: 'i32', 11: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4, 6, 10), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_red_fused__to_copy_eq_gt_mul_stack_sub_sum_9', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True} | |
) | |
@triton.jit | |
def triton_(in_ptr0, in_ptr1, in_ptr2, out_ptr0, out_ptr2, out_ptr3, out_ptr4, out_ptr5, ks0, ks1, xnumel, rnumel, XBLOCK : tl.constexpr, RBLOCK : tl.constexpr): | |
xnumel = 1024 | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:, None] | |
xmask = xindex < xnumel | |
rbase = tl.arange(0, RBLOCK)[None, :] | |
x0 = xindex | |
_tmp15 = tl.full([XBLOCK, RBLOCK], 0, tl.int64) | |
for roffset in range(0, rnumel, RBLOCK): | |
rindex = roffset + rbase | |
rmask = rindex < rnumel | |
r1 = rindex | |
tmp0 = tl.load(in_ptr0 + (1 + r1 + (2*x0) + (ks0*x0) + (ks1*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp1 = tl.load(in_ptr0 + (r1 + (2*x0) + (ks0*x0) + (ks1*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp5 = tl.load(in_ptr1 + (r1 + (2*x0) + (ks0*x0) + (ks1*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp2 = tmp0 - tmp1 | |
tmp3 = tl.full([1, 1], 0, tl.int64) | |
tmp4 = tmp2 == tmp3 | |
tmp6 = tmp5 + ks0 + ks1 + 2 | |
tmp7 = tmp5 < 0 | |
tmp8 = tl.where(tmp7, tmp6, tmp5) | |
tl.device_assert(((0 <= tmp8) & (tmp8 < 2 + ks0 + ks1)) | ~(xmask & rmask), "index out of bounds: 0 <= tmp8 < 2 + ks0 + ks1") | |
tmp9 = tl.load(in_ptr2 + (tmp8 + (2*x0) + (ks0*x0) + (ks1*x0)), rmask & xmask, eviction_policy='evict_last') | |
tmp10 = tmp9 > tmp3 | |
tmp11 = tmp4 & tmp10 | |
tmp12 = tmp11.to(tl.int8) | |
tmp13 = tmp11.to(tl.int64) | |
tmp14 = tl.broadcast_to(tmp13, [XBLOCK, RBLOCK]) | |
tmp16 = _tmp15 + tmp14 | |
_tmp15 = tl.where(rmask & xmask, tmp16, _tmp15) | |
tl.store(out_ptr0 + (r1 + x0 + (ks0*x0) + (ks1*x0)), tmp12, rmask & xmask) | |
tmp15 = tl.sum(_tmp15, 1)[:, None] | |
tmp17 = 1 + ks0 + ks1 | |
tmp18 = tmp17 - tmp15 | |
tl.store(out_ptr2 + (2*x0), tmp15, xmask) | |
tl.store(out_ptr3 + (2*x0), tmp18, xmask) | |
tl.store(out_ptr4 + (2*x0), tmp15, xmask) | |
tl.store(out_ptr5 + (2*x0), tmp18, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/wx/cwxkhk4qmjok4eeufvw4k4qlxjfqtkmdeh4d6kpathqmv3phoj2y.py | |
# Source Nodes: [count_pairwise_min, count_pairwise_min_1, gather_3, joined_padded_2, padding_mask, rowwise_counts, sum_4], Original ATen: [aten._to_copy, aten.gather, aten.gt, aten.mul, aten.sum] | |
# count_pairwise_min => mul_1 | |
# count_pairwise_min_1 => mul_2 | |
# gather_3 => gather_3 | |
# joined_padded_2 => gather_4 | |
# padding_mask => gt_2 | |
# rowwise_counts => convert_element_type_5 | |
# sum_4 => sum_4 | |
triton_red_fused__to_copy_gather_gt_mul_sum_10 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.reduction( | |
size_hints=[1024, 256], | |
reduction_hint=ReductionHint.DEFAULT, | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: '*i64', 3: '*i64', 4: '*i64', 5: '*i8', 6: '*i64', 7: '*i64', 8: '*i32', 9: 'i32', 10: 'i32', 11: 'i32', 12: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2, 3, 4, 5, 6, 7, 8, 11), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_red_fused__to_copy_gather_gt_mul_sum_10', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True} | |
) | |
@triton.jit | |
def triton_(in_ptr0, in_ptr1, in_ptr2, in_ptr3, in_ptr4, in_ptr5, out_ptr0, out_ptr1, out_ptr3, ks0, ks1, xnumel, rnumel, XBLOCK : tl.constexpr, RBLOCK : tl.constexpr): | |
xnumel = 1024 | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:, None] | |
xmask = xindex < xnumel | |
rbase = tl.arange(0, RBLOCK)[None, :] | |
x0 = xindex | |
_tmp21 = tl.full([XBLOCK, RBLOCK], 0, tl.int64) | |
for roffset in range(0, rnumel, RBLOCK): | |
rindex = roffset + rbase | |
rmask = rindex < rnumel | |
r1 = rindex | |
tmp0 = tl.load(in_ptr0 + (r1 + x0 + (ks0*x0) + (ks1*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp17 = tl.load(in_ptr5 + (r1 + x0 + (ks0*x0) + (ks1*x0)), rmask & xmask, eviction_policy='evict_first', other=0.0) | |
tmp1 = tmp0 + ks0 + ks1 + 2 | |
tmp2 = tmp0 < 0 | |
tmp3 = tl.where(tmp2, tmp1, tmp0) | |
tl.device_assert(((0 <= tmp3) & (tmp3 < 2 + ks0 + ks1)) | ~(xmask & rmask), "index out of bounds: 0 <= tmp3 < 2 + ks0 + ks1") | |
tmp4 = tl.load(in_ptr1 + (tmp3 + (2*x0) + (ks0*x0) + (ks1*x0)), rmask & xmask, eviction_policy='evict_last') | |
tmp5 = tmp0 + ks0 + ks1 + 1 | |
tmp6 = tl.where(tmp2, tmp5, tmp0) | |
tl.device_assert(((0 <= tmp6) & (tmp6 < 1 + ks0 + ks1)) | ~(xmask & rmask), "index out of bounds: 0 <= tmp6 < 1 + ks0 + ks1") | |
tmp7 = tl.load(in_ptr2 + (tmp6 + (2*x0) + (ks0*x0) + (ks1*x0)), rmask & xmask, eviction_policy='evict_last') | |
tmp8 = tl.load(in_ptr3 + (tmp6 + (2*x0) + (ks0*x0) + (ks1*x0)), rmask & xmask, eviction_policy='evict_last') | |
tmp9 = tmp8 + ks0 + ks1 + 2 | |
tmp10 = tmp8 < 0 | |
tmp11 = tl.where(tmp10, tmp9, tmp8) | |
tl.device_assert(((0 <= tmp11) & (tmp11 < 2 + ks0 + ks1)) | ~(xmask & rmask), "index out of bounds: 0 <= tmp11 < 2 + ks0 + ks1") | |
tmp12 = tl.load(in_ptr4 + (tmp11 + (2*x0) + (ks0*x0) + (ks1*x0)), rmask & xmask, eviction_policy='evict_last') | |
tmp13 = tl.full([1, 1], 0, tl.int64) | |
tmp14 = tmp12 > tmp13 | |
tmp15 = tmp14.to(tl.int64) | |
tmp16 = tmp7 * tmp15 | |
tmp18 = tmp17.to(tl.int64) | |
tmp19 = tmp16 * tmp18 | |
tmp20 = tl.broadcast_to(tmp19, [XBLOCK, RBLOCK]) | |
tmp22 = _tmp21 + tmp20 | |
_tmp21 = tl.where(rmask & xmask, tmp22, _tmp21) | |
tl.store(out_ptr0 + (r1 + x0 + (ks0*x0) + (ks1*x0)), tmp4, rmask & xmask) | |
tl.store(out_ptr1 + (r1 + x0 + (ks0*x0) + (ks1*x0)), tmp19, rmask & xmask) | |
tmp21 = tl.sum(_tmp21, 1)[:, None] | |
tmp23 = tmp21.to(tl.int32) | |
tl.store(out_ptr3 + (x0), tmp23, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/3w/c3wcv6zxutzi3no2w2skx7bnsftvelmhnzxpaeviajxxgoohun2r.py | |
# Source Nodes: [unique_val], Original ATen: [aten.index] | |
# unique_val => index_2 | |
triton_poi_fused_index_11 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.pointwise( | |
size_hints=[262144], | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: 'i32', 3: 'i32', 4: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 4), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_index_11', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True}, | |
min_elem_per_thread=0 | |
) | |
@triton.jit | |
def triton_(in_ptr0, out_ptr0, ks0, ks1, xnumel, XBLOCK : tl.constexpr): | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:] | |
xmask = xindex < xnumel | |
x0 = xindex | |
tmp0 = tl.load(in_ptr0 + ((ks0*(x0 // (1 + ks0 + ks1))) + (ks1*(x0 // (1 + ks0 + ks1))) + (x0 // (1 + ks0 + ks1)) + (x0 % (1 + ks0 + ks1))), xmask, eviction_policy='evict_last') | |
tl.store(out_ptr0 + (x0), tmp0, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/qd/cqd4ukqj4klbwsbc2vk5oauespotm4tvxd4yl6toq2ou76rbypvi.py | |
# Source Nodes: [mask_5], Original ATen: [aten.index_select] | |
# mask_5 => index_1 | |
triton_poi_fused_index_select_12 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.pointwise( | |
size_hints=[262144], | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i1', 2: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_index_select_12', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True}, | |
min_elem_per_thread=0 | |
) | |
@triton.jit | |
def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:] | |
xmask = xindex < xnumel | |
x0 = xindex | |
tmp0 = tl.load(in_ptr0 + (x0), xmask) | |
tmp1 = tmp0 + 2048 | |
tmp2 = tmp0 < 0 | |
tmp3 = tl.where(tmp2, tmp1, tmp0) | |
tl.device_assert(((0 <= tmp3) & (tmp3 < 2048)) | ~xmask, "index out of bounds: 0 <= tmp3 < 2048") | |
tmp4 = tmp3 % 2 | |
tmp5 = tl.full([1], 0, tl.int64) | |
tmp6 = tmp4 >= tmp5 | |
tmp7 = tl.full([1], 1, tl.int64) | |
tmp8 = tmp4 < tmp7 | |
tmp9 = tl.full([1], True, tl.int1) | |
tmp10 = tl.full(tmp9.shape, 0.0, tmp9.dtype) | |
tmp11 = tl.where(tmp8, tmp9, tmp10) | |
tmp12 = tmp4 >= tmp7 | |
tmp13 = tl.full([1], 2, tl.int64) | |
tmp14 = tmp4 < tmp13 | |
tmp15 = tl.full([1], False, tl.int1) | |
tmp16 = tl.full(tmp15.shape, 0.0, tmp15.dtype) | |
tmp17 = tl.where(tmp12, tmp15, tmp16) | |
tmp18 = tl.where(tmp8, tmp11, tmp17) | |
tl.store(out_ptr0 + (x0), tmp18, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/4y/c4y4u2zwj2a3ddvkk64qxn7k36hn3pklhua74i4cyqpafigzyyeb.py | |
# Source Nodes: [expanded_ids], Original ATen: [aten.index_select] | |
# expanded_ids => index_5 | |
triton_poi_fused_index_select_13 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.pointwise( | |
size_hints=[8192], | |
filename=__file__, | |
triton_meta={'signature': {0: '*i64', 1: '*i64', 2: 'i64', 3: 'i64'}, 'device': 0, 'device_type': 'cuda', 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1), equal_to_1=())]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_index_select_13', 'mutated_arg_names': ['in_out_ptr0'], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True}, | |
min_elem_per_thread=0 | |
) | |
@triton.jit | |
def triton_(in_out_ptr0, in_ptr0, ks0, xnumel, XBLOCK : tl.constexpr): | |
xoffset = tl.program_id(0).to(tl.int64) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:].to(tl.int64) | |
xmask = xindex < xnumel | |
x0 = xindex | |
tmp0 = tl.load(in_out_ptr0 + (x0), xmask) | |
tmp1 = tmp0 + ks0 | |
tmp2 = tmp0 < 0 | |
tmp3 = tl.where(tmp2, tmp1, tmp0) | |
tl.device_assert(((0 <= tmp3) & (tmp3 < ks0)) | ~xmask, "index out of bounds: 0 <= tmp3 < ks0") | |
tmp4 = tl.load(in_ptr0 + (tmp3), xmask, eviction_policy='evict_last') | |
tl.store(in_out_ptr0 + (x0), tmp4, xmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/v7/cv7i6ynxmyxllr4jgh3p2vywmcplzsxqnecsn24dm5qavbghignv.py | |
# Source Nodes: [cumsum_4], Original ATen: [aten.cumsum] | |
# cumsum_4 => cumsum_4 | |
triton_per_fused_cumsum_14 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton.jit | |
def _triton_helper_fn_add0(arg0_0, arg1_0): | |
tmp0 = arg0_0 + arg1_0 | |
return tmp0 | |
@triton_heuristics.persistent_reduction( | |
size_hints=[1, 1024], | |
reduction_hint=ReductionHint.INNER, | |
filename=__file__, | |
triton_meta={'signature': {0: '*i32', 1: '*i32', 2: 'i32', 3: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {2: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 3), equal_to_1=(2,))]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_per_fused_cumsum_14', 'mutated_arg_names': [], 'no_x_dim': True, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True} | |
) | |
@triton.jit | |
def triton_(in_ptr0, out_ptr0, xnumel, rnumel): | |
xnumel = 1 | |
XBLOCK: tl.constexpr = 1 | |
rnumel = 1024 | |
RBLOCK: tl.constexpr = 1024 | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = tl.full([1], xoffset, tl.int32) | |
xmask = xindex < xnumel | |
rindex = tl.arange(0, RBLOCK)[:] | |
roffset = 0 | |
rmask = rindex < rnumel | |
r0 = rindex | |
tmp0 = tl.load(in_ptr0 + (r0), rmask, other=0.0) | |
tmp1 = tmp0.to(tl.int32) | |
tmp2 = tl.broadcast_to(tmp1, [RBLOCK]) | |
tmp3, = tl.associative_scan((tmp2,), 0, _triton_helper_fn_add0) | |
tl.store(out_ptr0 + (tl.broadcast_to(r0, [RBLOCK])), tmp3, rmask) | |
''', device_str='cuda') | |
# kernel path: /tmp/torchinductor_tarasg/dh/cdh7lau3uqef4ntphscaobz4ageuajvzw3ynqvdzf7gsl7bq7occ.py | |
# Source Nodes: [zeros_2], Original ATen: [aten.zeros] | |
# zeros_2 => full_default_14 | |
triton_poi_fused_zeros_15 = async_compile.triton('triton_', ''' | |
import triton | |
import triton.language as tl | |
from triton.compiler.compiler import AttrsDescriptor | |
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, instance_descriptor | |
@triton_heuristics.pointwise( | |
size_hints=[1], | |
filename=__file__, | |
triton_meta={'signature': {0: '*i32', 1: 'i32'}, 'device': 0, 'device_type': 'cuda', 'constants': {1: 1}, 'configs': [AttrsDescriptor(divisible_by_16=(0,), equal_to_1=(1,))]}, | |
inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_zeros_15', 'mutated_arg_names': [], 'no_x_dim': False, 'backend_hash': '29c521dc5d9e4f14437c83e288b2df511c55eb789833b844b5bcc030e3e984a2', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False, 'is_fbcode': True}, | |
min_elem_per_thread=0 | |
) | |
@triton.jit | |
def triton_(out_ptr0, xnumel, XBLOCK : tl.constexpr): | |
xnumel = 1 | |
xoffset = tl.program_id(0) * XBLOCK | |
xindex = xoffset + tl.arange(0, XBLOCK)[:] | |
xmask = xindex < xnumel | |
tmp0 = tl.full([1], 0, tl.int32) | |
tl.store(out_ptr0 + (tl.full([XBLOCK], 0, tl.int32)), tmp0, None) | |
''', device_str='cuda') | |
async_compile.wait(globals()) | |
del async_compile | |
def call(args): | |
arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1 = args | |
args.clear() | |
s0 = arg0_1 | |
s1 = arg2_1 | |
s2 = arg5_1 | |
s3 = arg7_1 | |
assert_size_stride(arg1_1, (s0, ), (1, )) | |
assert_size_stride(arg3_1, (s1, ), (1, )) | |
assert_size_stride(arg4_1, (1025, ), (1, )) | |
assert_size_stride(arg6_1, (s2, ), (1, )) | |
assert_size_stride(arg8_1, (s3, ), (1, )) | |
assert_size_stride(arg9_1, (1025, ), (1, )) | |
if not s2 >= 2: | |
raise RuntimeError('Runtime assertion failed for s2 >= 2') | |
buf0 = None | |
if not s2 <= 9223372036854775806: | |
raise RuntimeError('Runtime assertion failed for s2 <= 9223372036854775806') | |
buf1 = None | |
if not s1 >= 2: | |
raise RuntimeError('Runtime assertion failed for s1 >= 2') | |
buf2 = None | |
if not s1 <= 9223372036854775806: | |
raise RuntimeError('Runtime assertion failed for s1 <= 9223372036854775806') | |
buf3 = None | |
with torch.cuda._DeviceGuard(0): | |
torch.cuda.set_device(0) | |
ps0 = 1 + s1 | |
buf4 = empty_strided_cuda((1024, 1 + s1), (1 + s1, 1), torch.int64) | |
# Source Nodes: [padded1], Original ATen: [test_inductor_ops.jagged_to_padded_dense] | |
triton_poi_fused_jagged_to_padded_dense_0_xnumel = 1024 + (1024*s1) | |
stream0 = get_raw_stream(0) | |
triton_poi_fused_jagged_to_padded_dense_0.run(arg4_1, arg1_1, buf4, ps0, s0, triton_poi_fused_jagged_to_padded_dense_0_xnumel, grid=grid(triton_poi_fused_jagged_to_padded_dense_0_xnumel), stream=stream0) | |
del arg1_1 | |
del arg4_1 | |
# Source Nodes: [padded1, sort], Original ATen: [aten.sort, test_inductor_ops.jagged_to_padded_dense] | |
buf5 = aten.sort.default(buf4, 1) | |
buf6 = buf5[0] | |
del buf5 | |
buf8 = buf4; del buf4 # reuse | |
# Source Nodes: [unique_0hot], Original ATen: [aten.cat] | |
triton_poi_fused_cat_1_xnumel = 1024 + (1024*s1) | |
triton_poi_fused_cat_1.run(buf6, buf8, ps0, triton_poi_fused_cat_1_xnumel, grid=grid(triton_poi_fused_cat_1_xnumel), stream=stream0) | |
# Source Nodes: [sort_1, unique_0hot], Original ATen: [aten.cat, aten.sort] | |
buf9 = aten.sort.stable(buf8, stable=True, dim=1) | |
buf11 = buf9[1] | |
del buf9 | |
ps1 = 1 + s2 | |
buf12 = empty_strided_cuda((1024, 1 + s2), (1 + s2, 1), torch.int64) | |
# Source Nodes: [padded2], Original ATen: [test_inductor_ops.jagged_to_padded_dense] | |
triton_poi_fused_jagged_to_padded_dense_0_xnumel = 1024 + (1024*s2) | |
triton_poi_fused_jagged_to_padded_dense_0.run(arg9_1, arg8_1, buf12, ps1, s3, triton_poi_fused_jagged_to_padded_dense_0_xnumel, grid=grid(triton_poi_fused_jagged_to_padded_dense_0_xnumel), stream=stream0) | |
del arg8_1 | |
del arg9_1 | |
# Source Nodes: [padded2, sort_2], Original ATen: [aten.sort, test_inductor_ops.jagged_to_padded_dense] | |
buf13 = aten.sort.default(buf12, 1) | |
buf14 = buf13[0] | |
del buf13 | |
buf16 = buf12; del buf12 # reuse | |
# Source Nodes: [unique_0hot_1], Original ATen: [aten.cat] | |
triton_poi_fused_cat_1_xnumel = 1024 + (1024*s2) | |
triton_poi_fused_cat_1.run(buf14, buf16, ps1, triton_poi_fused_cat_1_xnumel, grid=grid(triton_poi_fused_cat_1_xnumel), stream=stream0) | |
# Source Nodes: [sort_3, unique_0hot_1], Original ATen: [aten.cat, aten.sort] | |
buf17 = aten.sort.stable(buf16, stable=True, dim=1) | |
buf19 = buf17[1] | |
del buf17 | |
buf21 = buf8; del buf8 # reuse | |
# Source Nodes: [mask], Original ATen: [aten.zeros_like] | |
triton_poi_fused_zeros_like_2_xnumel = 1024 + (1024*s1) | |
triton_poi_fused_zeros_like_2.run(buf21, triton_poi_fused_zeros_like_2_xnumel, grid=grid(triton_poi_fused_zeros_like_2_xnumel), stream=stream0) | |
# Source Nodes: [adj_diff, diff, gt, mask, n_unique, setitem], Original ATen: [aten._to_copy, aten.gt, aten.index_put, aten.lift_fresh, aten.sub, aten.sum, aten.zeros_like] | |
triton_red_fused__to_copy_gt_index_put_lift_fresh_sub_sum_zeros_like_3.run(buf6, buf21, s1, 1024, s1, grid=grid(1024), stream=stream0) | |
buf23 = empty_strided_cuda((1024, 1 + s1), (1 + s1, 1), torch.int64) | |
buf30 = empty_strided_cuda((1024, 2 + s1 + s2), (2 + s1 + s2, 1), torch.int64) | |
buf24 = reinterpret_tensor(buf30, (1024, 1 + s1), (2 + s1 + s2, 1), 0) # alias | |
# Source Nodes: [cumsum, gather_out, setitem_1], Original ATen: [aten.cumsum, aten.gather, aten.index_put, aten.lift_fresh] | |
triton_red_fused_cumsum_gather_index_put_lift_fresh_4_rnumel = 1 + s1 | |
triton_red_fused_cumsum_gather_index_put_lift_fresh_4.run(buf21, buf11, buf6, buf23, buf24, s1, s2, 1024, triton_red_fused_cumsum_gather_index_put_lift_fresh_4_rnumel, grid=grid(1024), stream=stream0) | |
del buf21 | |
del buf6 | |
buf26 = buf16; del buf16 # reuse | |
# Source Nodes: [mask_2], Original ATen: [aten.zeros_like] | |
triton_poi_fused_zeros_like_2_xnumel = 1024 + (1024*s2) | |
triton_poi_fused_zeros_like_2.run(buf26, triton_poi_fused_zeros_like_2_xnumel, grid=grid(triton_poi_fused_zeros_like_2_xnumel), stream=stream0) | |
# Source Nodes: [adj_diff_1, diff_2, gt_1, mask_2, n_unique_1, setitem_3], Original ATen: [aten._to_copy, aten.gt, aten.index_put, aten.lift_fresh, aten.sub, aten.sum, aten.zeros_like] | |
triton_red_fused__to_copy_gt_index_put_lift_fresh_sub_sum_zeros_like_3.run(buf14, buf26, s2, 1024, s2, grid=grid(1024), stream=stream0) | |
buf28 = empty_strided_cuda((1024, 1 + s2), (1 + s2, 1), torch.int64) | |
buf29 = reinterpret_tensor(buf30, (1024, 1 + s2), (2 + s1 + s2, 1), 1 + s1) # alias | |
# Source Nodes: [cumsum_1, gather_out_1, setitem_4], Original ATen: [aten.cumsum, aten.gather, aten.index_put, aten.lift_fresh] | |
triton_red_fused_cumsum_gather_index_put_lift_fresh_5_rnumel = 1 + s2 | |
triton_red_fused_cumsum_gather_index_put_lift_fresh_5.run(buf26, buf19, buf14, buf28, buf29, s2, s1, 1024, triton_red_fused_cumsum_gather_index_put_lift_fresh_5_rnumel, grid=grid(1024), stream=stream0) | |
del buf14 | |
del buf26 | |
del buf24 | |
del buf29 | |
# Source Nodes: [sort_4], Original ATen: [aten.sort] | |
buf31 = aten.sort.default(buf30, 1) | |
buf32 = buf31[0] | |
buf33 = buf31[1] | |
del buf31 | |
buf34 = empty_strided_cuda((1024, s1), (s1, 1), torch.int64) | |
# Source Nodes: [setitem_2, unique_indices_1], Original ATen: [aten.index_put, aten.lift_fresh, aten.sub] | |
triton_poi_fused_index_put_lift_fresh_sub_6_xnumel = 1024*s1 | |
triton_poi_fused_index_put_lift_fresh_sub_6.run(buf23, buf11, buf34, s1, triton_poi_fused_index_put_lift_fresh_sub_6_xnumel, grid=grid(triton_poi_fused_index_put_lift_fresh_sub_6_xnumel), stream=stream0) | |
del buf11 | |
del buf23 | |
buf35 = empty_strided_cuda((1024, s2), (s2, 1), torch.int64) | |
# Source Nodes: [setitem_5, unique_indices_4], Original ATen: [aten.index_put, aten.lift_fresh, aten.sub] | |
triton_poi_fused_index_put_lift_fresh_sub_6_xnumel = 1024*s2 | |
triton_poi_fused_index_put_lift_fresh_sub_6.run(buf28, buf19, buf35, s2, triton_poi_fused_index_put_lift_fresh_sub_6_xnumel, grid=grid(triton_poi_fused_index_put_lift_fresh_sub_6_xnumel), stream=stream0) | |
del buf19 | |
del buf28 | |
ps2 = 2 + s1 + s2 | |
buf36 = buf30; del buf30 # reuse | |
# Source Nodes: [joined_padded_counts], Original ATen: [aten.cat] | |
triton_poi_fused_cat_7_xnumel = 2048 + (1024*s1) + (1024*s2) | |
triton_poi_fused_cat_7.run(buf34, buf35, buf36, ps2, ps0, s1, s2, ps1, triton_poi_fused_cat_7_xnumel, grid=grid(triton_poi_fused_cat_7_xnumel), stream=stream0) | |
del buf34 | |
del buf35 | |
buf37 = empty_strided_cuda((1024, 2 + s1 + s2), (2 + s1 + s2, 1), torch.int64) | |
# Source Nodes: [min_1], Original ATen: [aten.min] | |
triton_poi_fused_min_8_xnumel = 2048 + (1024*s1) + (1024*s2) | |
triton_poi_fused_min_8.run(buf33, buf36, buf37, ps2, s1, s2, triton_poi_fused_min_8_xnumel, grid=grid(triton_poi_fused_min_8_xnumel), stream=stream0) | |
buf38 = empty_strided_cuda((1024, 1 + s1 + s2), (1 + s1 + s2, 1), torch.int8) | |
buf45 = empty_strided_cuda((1024, 2), (2, 1), torch.int64) | |
buf43 = reinterpret_tensor(buf45, (1024, 1), (2, 1), 0) # alias | |
buf44 = reinterpret_tensor(buf45, (1024, 1), (2, 1), 1) # alias | |
buf56 = empty_strided_cuda((1024, 2), (2, 1), torch.int64) | |
buf54 = reinterpret_tensor(buf56, (1024, 1), (2, 1), 0) # alias | |
buf55 = reinterpret_tensor(buf56, (1024, 1), (2, 1), 1) # alias | |
# Source Nodes: [diff_4, eq, joined_padded_lengths, padded_diff, padding_mask, stack_5, stack_7, to_4], Original ATen: [aten._to_copy, aten.eq, aten.gt, aten.mul, aten.stack, aten.sub, aten.sum] | |
triton_red_fused__to_copy_eq_gt_mul_stack_sub_sum_9_rnumel = 1 + s1 + s2 | |
triton_red_fused__to_copy_eq_gt_mul_stack_sub_sum_9.run(buf32, buf33, buf36, buf38, buf43, buf44, buf54, buf55, s1, s2, 1024, triton_red_fused__to_copy_eq_gt_mul_stack_sub_sum_9_rnumel, grid=grid(1024), stream=stream0) | |
# Source Nodes: [diff_4, eq, padded_diff, padding_mask, sort_5, to_4], Original ATen: [aten._to_copy, aten.eq, aten.gt, aten.mul, aten.sort, aten.sub] | |
buf39 = aten.sort.stable(buf38, stable=True, dim=1, descending=True) | |
del buf38 | |
buf40 = buf39[0] | |
buf41 = buf39[1] | |
del buf39 | |
del buf43 | |
del buf44 | |
# Source Nodes: [mask_5], Original ATen: [aten.repeat_interleave] | |
buf46 = aten.repeat_interleave.Tensor(reinterpret_tensor(buf45, (2048, ), (1, ), 0)) | |
u0 = buf46.size(0) | |
del buf45 | |
buf47 = buf46 | |
del buf46 | |
buf48 = empty_strided_cuda((1024, 1 + s1 + s2), (1 + s1 + s2, 1), torch.int64) | |
buf53 = empty_strided_cuda((1024, 1 + s1 + s2), (1 + s1 + s2, 1), torch.int64) | |
buf66 = empty_strided_cuda((1024, ), (1, ), torch.int32) | |
# Source Nodes: [count_pairwise_min, count_pairwise_min_1, gather_3, joined_padded_2, padding_mask, rowwise_counts, sum_4], Original ATen: [aten._to_copy, aten.gather, aten.gt, aten.mul, aten.sum] | |
triton_red_fused__to_copy_gather_gt_mul_sum_10_rnumel = 1 + s1 + s2 | |
triton_red_fused__to_copy_gather_gt_mul_sum_10.run(buf41, buf32, buf37, buf33, buf36, buf40, buf48, buf53, buf66, s1, s2, 1024, triton_red_fused__to_copy_gather_gt_mul_sum_10_rnumel, grid=grid(1024), stream=stream0) | |
del buf32 | |
del buf33 | |
del buf36 | |
del buf37 | |
del buf40 | |
buf49 = reinterpret_tensor(buf41, (1024 + (1024*s1) + (1024*s2), ), (1, ), 0); del buf41 # reuse | |
# Source Nodes: [unique_val], Original ATen: [aten.index] | |
triton_poi_fused_index_11_xnumel = 1024 + (1024*s1) + (1024*s2) | |
triton_poi_fused_index_11.run(buf48, buf49, s1, s2, triton_poi_fused_index_11_xnumel, grid=grid(triton_poi_fused_index_11_xnumel), stream=stream0) | |
del buf48 | |
buf50 = empty_strided_cuda((1024 + (1024*s1) + (1024*s2), ), (1, ), torch.bool) | |
# Source Nodes: [mask_5], Original ATen: [aten.index_select] | |
triton_poi_fused_index_select_12_xnumel = 1024 + (1024*s1) + (1024*s2) | |
triton_poi_fused_index_select_12.run(buf47, buf50, triton_poi_fused_index_select_12_xnumel, grid=grid(triton_poi_fused_index_select_12_xnumel), stream=stream0) | |
del buf47 | |
# Source Nodes: [unique_val], Original ATen: [aten.index] | |
buf51 = aten.index.Tensor(buf49, [buf50]) | |
u1 = buf51.size(0) | |
del buf49 | |
buf52 = buf51 | |
del buf54 | |
del buf55 | |
# Source Nodes: [mask_7], Original ATen: [aten.repeat_interleave] | |
buf57 = aten.repeat_interleave.Tensor(reinterpret_tensor(buf56, (2048, ), (1, ), 0)) | |
u11 = buf57.size(0) | |
del buf56 | |
buf58 = buf57 | |
del buf57 | |
buf59 = buf50; del buf50 # reuse | |
# Source Nodes: [mask_7], Original ATen: [aten.index_select] | |
triton_poi_fused_index_select_12_xnumel = 1024 + (1024*s1) + (1024*s2) | |
triton_poi_fused_index_select_12.run(buf58, buf59, triton_poi_fused_index_select_12_xnumel, grid=grid(triton_poi_fused_index_select_12_xnumel), stream=stream0) | |
del buf58 | |
# Source Nodes: [counts], Original ATen: [aten.index] | |
buf60 = aten.index.Tensor(reinterpret_tensor(buf53, (1024 + (1024*s1) + (1024*s2), ), (1, ), 0), [buf59]) | |
u2 = buf60.size(0) | |
del buf53 | |
del buf59 | |
buf61 = buf60 | |
del buf60 | |
# Source Nodes: [expanded_ids], Original ATen: [aten.repeat_interleave] | |
buf62 = aten.repeat_interleave.Tensor(buf61) | |
u3 = buf62.size(0) | |
del buf61 | |
buf63 = buf62 | |
buf64 = buf63; del buf63 # reuse | |
# Source Nodes: [expanded_ids], Original ATen: [aten.index_select] | |
triton_poi_fused_index_select_13.run(buf64, buf52, u1, u3, grid=grid(u3), stream=stream0) | |
del buf52 | |
buf69 = empty_strided_cuda((1025, ), (1, ), torch.int32) | |
buf67 = reinterpret_tensor(buf69, (1024, ), (1, ), 1) # alias | |
# Source Nodes: [cumsum_4], Original ATen: [aten.cumsum] | |
triton_per_fused_cumsum_14.run(buf66, buf67, 1, 1024, grid=grid(1), stream=stream0) | |
buf68 = reinterpret_tensor(buf69, (1, ), (1, ), 0) # alias | |
# Source Nodes: [zeros_2], Original ATen: [aten.zeros] | |
triton_poi_fused_zeros_15.run(buf68, 1, grid=grid(1), stream=stream0) | |
return (buf64, buf69, buf66, ) | |
def benchmark_compiled_module(times=10, repeat=10): | |
from torch._dynamo.testing import rand_strided | |
from torch._inductor.utils import print_performance | |
arg0_1 = 96282 | |
arg1_1 = rand_strided((96282, ), (1, ), device='cuda:0', dtype=torch.int64) | |
arg2_1 = 98 | |
arg3_1 = rand_strided((98, ), (1, ), device='cuda:0', dtype=torch.float32) | |
arg4_1 = rand_strided((1025, ), (1, ), device='cuda:0', dtype=torch.int32) | |
arg5_1 = 95 | |
arg6_1 = rand_strided((95, ), (1, ), device='cuda:0', dtype=torch.float32) | |
arg7_1 = 94741 | |
arg8_1 = rand_strided((94741, ), (1, ), device='cuda:0', dtype=torch.int64) | |
arg9_1 = rand_strided((1025, ), (1, ), device='cuda:0', dtype=torch.int32) | |
fn = lambda: call([arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1]) | |
return print_performance(fn, times=times, repeat=repeat) | |
if __name__ == "__main__": | |
from torch._inductor.wrapper_benchmark import compiled_module_main | |
compiled_module_main('None', benchmark_compiled_module) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment