Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
tilelang
Commits
afa74f4e
Commit
afa74f4e
authored
Apr 23, 2025
by
Lei Wang
Committed by
LeiWang1999
Apr 23, 2025
Browse files
[Bugfix] Fix a bug for simplifier (#425)
* Update submodule 'tvm' to latest commit f4a8f9b * lint fix
parent
2fff0eec
Changes
2
Hide whitespace changes
Inline
Side-by-side
Showing
2 changed files
with
44 additions
and
1 deletion
+44
-1
3rdparty/tvm
3rdparty/tvm
+1
-1
testing/python/language/test_tilelang_language_copy.py
testing/python/language/test_tilelang_language_copy.py
+43
-0
No files found.
tvm
@
b16c9f29
Compare
b2945254
...
b16c9f29
Subproject commit b
2945254932cffa89922ec7f6e868d726aed0f6a
Subproject commit b
16c9f298bc37fa502ffdb2ea809c2793e2a0bd6
testing/python/language/test_tilelang_language_copy.py
0 → 100644
View file @
afa74f4e
import
tilelang
import
tilelang.language
as
T
import
torch
def
tilelang_copy
(
M
,
N
,
block_M
,
block_N
,
dtype
=
"float16"
):
# add decorator @tilelang.jit if you want to return a torch function
@
T
.
prim_func
def
main
(
A
:
T
.
Tensor
((
M
,
N
),
dtype
),
B
:
T
.
Tensor
((
M
,
N
),
dtype
),
):
# Initialize Kernel Context
with
T
.
Kernel
(
T
.
ceildiv
(
N
,
block_N
),
T
.
ceildiv
(
M
,
block_M
),
threads
=
128
)
as
(
bx
,
by
):
for
i
,
j
in
T
.
Parallel
(
block_M
,
block_N
):
B
[
by
*
block_M
+
i
,
bx
*
block_N
+
j
]
=
A
[
by
*
block_M
+
i
,
bx
*
block_N
+
j
]
return
main
def
run_tilelang_copy
(
M
=
1024
,
N
=
1024
,
block_M
=
128
,
block_N
=
128
,
dtype
=
"float16"
):
program
=
tilelang_copy
(
M
,
N
,
block_M
,
block_N
,
dtype
)
kernel
=
tilelang
.
compile
(
program
,
out_idx
=
[
1
],
target
=
"cuda"
,
pass_configs
=
{
"tl.disable_warp_specialized"
:
True
,
"tl.disable_tma_lower"
:
True
})
a
=
torch
.
randn
(
M
,
N
,
device
=
"cuda"
,
dtype
=
getattr
(
torch
,
dtype
))
b
=
kernel
(
a
)
torch
.
testing
.
assert_close
(
b
,
a
,
rtol
=
1e-2
,
atol
=
1e-2
)
def
test_tilelang_copy
():
run_tilelang_copy
(
M
=
1024
,
N
=
1024
,
block_M
=
128
,
block_N
=
128
)
run_tilelang_copy
(
M
=
1024
,
N
=
576
,
block_M
=
32
,
block_N
=
576
)
run_tilelang_copy
(
M
=
1024
,
N
=
576
,
block_M
=
32
,
block_N
=
576
,
dtype
=
"float"
)
if
__name__
==
"__main__"
:
tilelang
.
testing
.
main
()
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment