Skip to content

Bugfix for passing None args to user defined Triton kernel #138260

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
SamGinzburg merged 7 commits into from
Oct 21, 2024
Merged

Bugfix for passing None args to user defined Triton kernel #138260

SamGinzburg merged 7 commits into from
Oct 21, 2024

Conversation

@SamGinzburg
Copy link
Contributor

@SamGinzburg SamGinzburg commented Oct 17, 2024 •
edited
Loading

This is a PR to fix the following issue: #115344

In short, passing None as an arg. to a Triton kernel would cause problems:

Short repro of the specific bug fixed here:

import triton
import triton.language as tl

@triton.autotune( # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
    configs=[
        triton.Config({'BLOCK_SIZE': 32}, num_stages=5, num_warps=2),
        triton.Config({'BLOCK_SIZE': 64}, num_stages=4, num_warps=4),
    ],
    key=['n_elements']
)
@triton.jit # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
def sin_kernel( # E: Function is missing a return type annotation  [no-untyped-def]
    in_ptr0,
    out_ptr,
    n_elements,
    BLOCK_SIZE: "tl.constexpr",
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    if in_ptr0 is not None:
        x = tl.load(in_ptr0 + offsets, mask=mask)
    else:
        x = 0.
    output = tl.sin(x)
    tl.store(out_ptr + offsets, output, mask=mask)

import torch

def sin_triton(x, out):
    n_elements = out.numel()
    sin_kernel[(n_elements,)](x, out, n_elements)

x = torch.randn(65, device="cuda")
out = torch.empty_like(x)
out_compiled = torch.empty_like(x)

sin_triton_compiled = torch.compile(fullgraph=True)(sin_triton)

for first in (x, None):
    sin_triton(first, out)
    sin_triton_compiled(first, out_compiled)
    torch.testing.assert_close(out, out_compiled)

I've added a unit test to catch this issue in the future and the tests in "test/inductor/test_triton_kernels.py"

topic: not user facing

================================================================================
Stack from ghstack (oldest at bottom):

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov

Differential Revision: D64615061

@SamGinzburg
Bugfix for passing None args to user defined Triton kernel
b1086d3 
add test

fewer failing tests

more tests passing

tests passing

lint

[ghstack-poisoned]
@SamGinzburg SamGinzburg requested a review from zou3519 as a code owner October 17, 2024 21:06
@pytorch-bot
Copy link

pytorch-bot bot commented Oct 17, 2024 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/138260

Note: Links to docs will display an error until the docs builds have been completed.

❌ 5 New Failures, 2 Unrelated Failures

As of commit cbc1bc0 with merge base de51ed8 (image):

NEW FAILURES - The following jobs have failed:

FLAKY - The following job failed but was likely due to flakiness present on trunk:

BROKEN TRUNK - The following job failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@linux-foundation-easycla
Copy link

linux-foundation-easycla bot commented Oct 17, 2024 •
edited
Loading

CLA Signed

The committers listed above are authorized under a signed CLA.

@SamGinzburg SamGinzburg requested a review from aakhundov October 17, 2024 21:11
@zou3519 zou3519 removed their request for review October 17, 2024 22:10
@SamGinzburg
Update on "Bugfix for passing None args to user defined Triton kernel"
9c27415 
This is a PR to fix the following issue: #115344

In short, passing None as an arg. to a Triton kernel would cause problems:

Short repro of the specific bug fixed here:
```
import triton
import triton.language as tl

triton.autotune( # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
    configs=[
        triton.Config({'BLOCK_SIZE': 32}, num_stages=5, num_warps=2),
        triton.Config({'BLOCK_SIZE': 64}, num_stages=4, num_warps=4),
    ],
    key=['n_elements']
)
triton.jit # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
def sin_kernel( # E: Function is missing a return type annotation  [no-untyped-def]
    in_ptr0,
    out_ptr,
    n_elements,
    BLOCK_SIZE: "tl.constexpr",
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    if in_ptr0 is not None:
        x = tl.load(in_ptr0 + offsets, mask=mask)
    else:
        x = 0.
    output = tl.sin(x)
    tl.store(out_ptr + offsets, output, mask=mask)

import torch

def sin_triton(x, out):
    n_elements = out.numel()
    sin_kernel[(n_elements,)](x, out, n_elements)

x = torch.randn(65, device="cuda")
out = torch.empty_like(x)
out_compiled = torch.empty_like(x)

sin_triton_compiled = torch.compile(fullgraph=True)(sin_triton)

for first in (x, None):
    sin_triton(first, out)
    sin_triton_compiled(first, out_compiled)
    torch.testing.assert_close(out, out_compiled)
```
I've added a unit test to catch this issue in the future and the tests in "test/inductor/test_triton_kernels.py"


================================================================================


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov 

[ghstack-poisoned]
@SamGinzburg
Copy link
Contributor Author

SamGinzburg commented Oct 17, 2024

@pytorchbot label "topic: not user facing"

@pytorch-bot pytorch-bot bot added the topic: not user facing topic category label Oct 17, 2024
@SamGinzburg
Update on "Bugfix for passing None args to user defined Triton kernel"
aec9202 
This is a PR to fix the following issue: #115344

In short, passing None as an arg. to a Triton kernel would cause problems:

Short repro of the specific bug fixed here:
```
import triton
import triton.language as tl

triton.autotune( # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
    configs=[
        triton.Config({'BLOCK_SIZE': 32}, num_stages=5, num_warps=2),
        triton.Config({'BLOCK_SIZE': 64}, num_stages=4, num_warps=4),
    ],
    key=['n_elements']
)
triton.jit # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
def sin_kernel( # E: Function is missing a return type annotation  [no-untyped-def]
    in_ptr0,
    out_ptr,
    n_elements,
    BLOCK_SIZE: "tl.constexpr",
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    if in_ptr0 is not None:
        x = tl.load(in_ptr0 + offsets, mask=mask)
    else:
        x = 0.
    output = tl.sin(x)
    tl.store(out_ptr + offsets, output, mask=mask)

import torch

def sin_triton(x, out):
    n_elements = out.numel()
    sin_kernel[(n_elements,)](x, out, n_elements)

x = torch.randn(65, device="cuda")
out = torch.empty_like(x)
out_compiled = torch.empty_like(x)

sin_triton_compiled = torch.compile(fullgraph=True)(sin_triton)

for first in (x, None):
    sin_triton(first, out)
    sin_triton_compiled(first, out_compiled)
    torch.testing.assert_close(out, out_compiled)
```
I've added a unit test to catch this issue in the future and the tests in "test/inductor/test_triton_kernels.py"


topic: not user facing

================================================================================


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov 

[ghstack-poisoned]
aakhundov
aakhundov reviewed Oct 17, 2024
View reviewed changes
aakhundov
aakhundov reviewed Oct 17, 2024
View reviewed changes
aakhundov
aakhundov reviewed Oct 17, 2024
View reviewed changes
@SamGinzburg
Update on "Bugfix for passing None args to user defined Triton kernel"
e795dd6 
This is a PR to fix the following issue: #115344

In short, passing None as an arg. to a Triton kernel would cause problems:

Short repro of the specific bug fixed here:
```
import triton
import triton.language as tl

triton.autotune( # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
    configs=[
        triton.Config({'BLOCK_SIZE': 32}, num_stages=5, num_warps=2),
        triton.Config({'BLOCK_SIZE': 64}, num_stages=4, num_warps=4),
    ],
    key=['n_elements']
)
triton.jit # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
def sin_kernel( # E: Function is missing a return type annotation  [no-untyped-def]
    in_ptr0,
    out_ptr,
    n_elements,
    BLOCK_SIZE: "tl.constexpr",
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    if in_ptr0 is not None:
        x = tl.load(in_ptr0 + offsets, mask=mask)
    else:
        x = 0.
    output = tl.sin(x)
    tl.store(out_ptr + offsets, output, mask=mask)

import torch

def sin_triton(x, out):
    n_elements = out.numel()
    sin_kernel[(n_elements,)](x, out, n_elements)

x = torch.randn(65, device="cuda")
out = torch.empty_like(x)
out_compiled = torch.empty_like(x)

sin_triton_compiled = torch.compile(fullgraph=True)(sin_triton)

for first in (x, None):
    sin_triton(first, out)
    sin_triton_compiled(first, out_compiled)
    torch.testing.assert_close(out, out_compiled)
```
I've added a unit test to catch this issue in the future and the tests in "test/inductor/test_triton_kernels.py"


topic: not user facing

================================================================================


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov 

[ghstack-poisoned]
@SamGinzburg
Update on "Bugfix for passing None args to user defined Triton kernel"
6e7d723 
This is a PR to fix the following issue: #115344

In short, passing None as an arg. to a Triton kernel would cause problems:

Short repro of the specific bug fixed here:
```
import triton
import triton.language as tl

triton.autotune( # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
    configs=[
        triton.Config({'BLOCK_SIZE': 32}, num_stages=5, num_warps=2),
        triton.Config({'BLOCK_SIZE': 64}, num_stages=4, num_warps=4),
    ],
    key=['n_elements']
)
triton.jit # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
def sin_kernel( # E: Function is missing a return type annotation  [no-untyped-def]
    in_ptr0,
    out_ptr,
    n_elements,
    BLOCK_SIZE: "tl.constexpr",
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    if in_ptr0 is not None:
        x = tl.load(in_ptr0 + offsets, mask=mask)
    else:
        x = 0.
    output = tl.sin(x)
    tl.store(out_ptr + offsets, output, mask=mask)

import torch

def sin_triton(x, out):
    n_elements = out.numel()
    sin_kernel[(n_elements,)](x, out, n_elements)

x = torch.randn(65, device="cuda")
out = torch.empty_like(x)
out_compiled = torch.empty_like(x)

sin_triton_compiled = torch.compile(fullgraph=True)(sin_triton)

for first in (x, None):
    sin_triton(first, out)
    sin_triton_compiled(first, out_compiled)
    torch.testing.assert_close(out, out_compiled)
```
I've added a unit test to catch this issue in the future and the tests in "test/inductor/test_triton_kernels.py"


topic: not user facing

================================================================================


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov 

[ghstack-poisoned]
SamGinzburg added a commit that referenced this pull request Oct 17, 2024
@SamGinzburg
Bugfix for passing None args to user defined Triton kernel
c1c853c 
add test

fewer failing tests

more tests passing

tests passing

lint

ghstack-source-id: f089650
Pull Request resolved: #138260
aakhundov
aakhundov reviewed Oct 18, 2024
View reviewed changes
Copy link
Contributor

@aakhundov aakhundov left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, @SamGinzburg! Looks good overall. Added some minor comments and questions.

@aakhundov
Copy link
Contributor

aakhundov commented Oct 18, 2024

cc @zou3519 as this fixes our long-standing None arg issue :)

@SamGinzburg
Update on "Bugfix for passing None args to user defined Triton kernel"
75013a8 
This is a PR to fix the following issue: #115344

In short, passing None as an arg. to a Triton kernel would cause problems:

Short repro of the specific bug fixed here:
```
import triton
import triton.language as tl

triton.autotune( # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
    configs=[
        triton.Config({'BLOCK_SIZE': 32}, num_stages=5, num_warps=2),
        triton.Config({'BLOCK_SIZE': 64}, num_stages=4, num_warps=4),
    ],
    key=['n_elements']
)
triton.jit # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
def sin_kernel( # E: Function is missing a return type annotation  [no-untyped-def]
    in_ptr0,
    out_ptr,
    n_elements,
    BLOCK_SIZE: "tl.constexpr",
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    if in_ptr0 is not None:
        x = tl.load(in_ptr0 + offsets, mask=mask)
    else:
        x = 0.
    output = tl.sin(x)
    tl.store(out_ptr + offsets, output, mask=mask)

import torch

def sin_triton(x, out):
    n_elements = out.numel()
    sin_kernel[(n_elements,)](x, out, n_elements)

x = torch.randn(65, device="cuda")
out = torch.empty_like(x)
out_compiled = torch.empty_like(x)

sin_triton_compiled = torch.compile(fullgraph=True)(sin_triton)

for first in (x, None):
    sin_triton(first, out)
    sin_triton_compiled(first, out_compiled)
    torch.testing.assert_close(out, out_compiled)
```
I've added a unit test to catch this issue in the future and the tests in "test/inductor/test_triton_kernels.py"


topic: not user facing

================================================================================


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov 

[ghstack-poisoned]
SamGinzburg added a commit that referenced this pull request Oct 18, 2024
@SamGinzburg
Bugfix for passing None args to user defined Triton kernel
f20812a 
add test

fewer failing tests

more tests passing

tests passing

lint

ghstack-source-id: 749f8ff
Pull Request resolved: #138260
@SamGinzburg
Copy link
Contributor Author

SamGinzburg commented Oct 18, 2024

@SamGinzburg has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@aakhundov aakhundov requested a review from eellison October 18, 2024 19:05
@aakhundov aakhundov requested a review from oulgen October 18, 2024 19:06
@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 18, 2024
@SamGinzburg
Update on "Bugfix for passing None args to user defined Triton kernel"
cbc1bc0 
This is a PR to fix the following issue: #115344

In short, passing None as an arg. to a Triton kernel would cause problems:

Short repro of the specific bug fixed here:
```
import triton
import triton.language as tl

triton.autotune( # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
    configs=[
        triton.Config({'BLOCK_SIZE': 32}, num_stages=5, num_warps=2),
        triton.Config({'BLOCK_SIZE': 64}, num_stages=4, num_warps=4),
    ],
    key=['n_elements']
)
triton.jit # E: Untyped decorator makes function "sin_kernel" untyped  [misc]
def sin_kernel( # E: Function is missing a return type annotation  [no-untyped-def]
    in_ptr0,
    out_ptr,
    n_elements,
    BLOCK_SIZE: "tl.constexpr",
):
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    if in_ptr0 is not None:
        x = tl.load(in_ptr0 + offsets, mask=mask)
    else:
        x = 0.
    output = tl.sin(x)
    tl.store(out_ptr + offsets, output, mask=mask)

import torch

def sin_triton(x, out):
    n_elements = out.numel()
    sin_kernel[(n_elements,)](x, out, n_elements)

x = torch.randn(65, device="cuda")
out = torch.empty_like(x)
out_compiled = torch.empty_like(x)

sin_triton_compiled = torch.compile(fullgraph=True)(sin_triton)

for first in (x, None):
    sin_triton(first, out)
    sin_triton_compiled(first, out_compiled)
    torch.testing.assert_close(out, out_compiled)
```
I've added a unit test to catch this issue in the future and the tests in "test/inductor/test_triton_kernels.py"


topic: not user facing

================================================================================


cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy yf225 chenyang78 kadeng muchulee8 ColinPeppler amjames desertfire chauhang aakhundov

Differential Revision: [D64615061](https://our.internmc.facebook.com/intern/diff/D64615061)

[ghstack-poisoned]
SamGinzburg added a commit that referenced this pull request Oct 18, 2024
@SamGinzburg
Bugfix for passing None args to user defined Triton kernel
bea1b35 
add test

fewer failing tests

more tests passing

tests passing

lint

ghstack-source-id: 4cb2b4f
Pull Request resolved: #138260
@SamGinzburg
Copy link
Contributor Author

SamGinzburg commented Oct 18, 2024

@SamGinzburg has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@SamGinzburg SamGinzburg merged commit d74f5bd into gh/SamGinzburg/5/base Oct 21, 2024
@SamGinzburg
Copy link
Contributor Author

SamGinzburg commented Oct 21, 2024

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

pytorchmergebot commented Oct 21, 2024

Can't merge closed PR #138260

SamGinzburg added a commit that referenced this pull request Oct 21, 2024
@SamGinzburg
Revert "Bugfix for passing None args to user defined Triton kernel (#…
f0652d6 
…138260)"

This reverts commit d74f5bd.
@SamGinzburg SamGinzburg mentioned this pull request Oct 21, 2024
Merged
SamGinzburg added a commit that referenced this pull request Oct 21, 2024
@SamGinzburg
Revert "Bugfix for passing None args to user defined Triton kernel" (#…
b011a30 
…138469)

Revert "Bugfix for passing None args to user defined Triton kernel (#138260)"

This reverts commit d74f5bd.
@aakhundov
Copy link
Contributor

aakhundov commented Oct 21, 2024

This PR was merged manually by mistake. Resubmitted in #138472.

@aakhundov aakhundov mentioned this pull request Oct 21, 2024
Closed
@github-actions github-actions bot deleted the gh/SamGinzburg/5/head branch November 21, 2024 02:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@eellison eellison Awaiting requested review from eellison

@oulgen oulgen Awaiting requested review from oulgen

1 more reviewer

@aakhundov aakhundov aakhundov approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

ciflow/inductor ciflow/trunk Trigger trunk jobs on your pull request module: inductor topic: not user facing topic category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@SamGinzburg @aakhundov @pytorchmergebot