wehub-resource-sync
372 lines
11 KiB
Python
372 lines
11 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
import pytest
|
|
|
|
from tvm.script import tirx as T
|
|
from tvm.script.tirx import tile as Tx
|
|
from tvm.tirx.analysis import verify_tirx_well_formed as verify
|
|
|
|
|
|
def test_root_scope():
|
|
# fmt: off
|
|
@T.prim_func(check_well_formed=False)
|
|
def test1() -> None:
|
|
T.device_entry()
|
|
pass
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test2() -> None:
|
|
pass
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test3() -> None:
|
|
pass
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test4() -> None:
|
|
T.device_entry()
|
|
pass
|
|
|
|
# fmt: on
|
|
|
|
verify(test1)
|
|
verify(test2)
|
|
verify(test3)
|
|
verify(test4)
|
|
|
|
|
|
def test_nested_scope():
|
|
# fmt: off
|
|
@T.prim_func(check_well_formed=False)
|
|
def test1() -> None:
|
|
T.device_entry()
|
|
pass
|
|
pass
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test2() -> None:
|
|
T.device_entry()
|
|
pass
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test3() -> None:
|
|
T.device_entry()
|
|
pass
|
|
@T.prim_func(check_well_formed=False)
|
|
def test4() -> None:
|
|
T.device_entry()
|
|
pass
|
|
pass
|
|
|
|
# fmt: on
|
|
|
|
verify(test1)
|
|
verify(test2)
|
|
verify(test3)
|
|
verify(test4)
|
|
|
|
|
|
def test_scope_id_consistency():
|
|
# fmt: off
|
|
@T.prim_func(check_well_formed=False)
|
|
def test1():
|
|
T.device_entry()
|
|
T.cta_id([32])
|
|
T.warp_id([4])
|
|
T.lane_id([32])
|
|
pass
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test2():
|
|
T.device_entry()
|
|
T.cta_id([32])
|
|
T.warp_id([4])
|
|
T.lane_id([32])
|
|
T.thread_id([128])
|
|
pass
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test3():
|
|
T.device_entry()
|
|
T.cta_id([32])
|
|
T.warp_id([2])
|
|
T.lane_id([32])
|
|
T.thread_id([128])
|
|
pass
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test4():
|
|
T.device_entry()
|
|
bx, by, bz = T.cta_id([8, 10, 12])
|
|
cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
|
|
clx, cly, clz = T.cluster_id([4, 5, 12])
|
|
T.evaluate(bx + by + bz)
|
|
T.evaluate(cbx + cby + cbz)
|
|
T.evaluate(clx + cly + clz)
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test5():
|
|
T.device_entry()
|
|
bx, by, bz = T.cta_id([8, 10, 12])
|
|
cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
|
|
clx, cly, clz = T.cluster_id([3, 5, 12])
|
|
T.evaluate(bx + by + bz)
|
|
T.evaluate(cbx + cby + cbz)
|
|
T.evaluate(clx + cly + clz)
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test6():
|
|
T.device_entry()
|
|
clx, cly, clz = T.cluster_id([4, 5, 12])
|
|
bx, by, bz = T.cta_id([8, 10, 12])
|
|
cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
|
|
T.evaluate(bx + by + bz)
|
|
T.evaluate(cbx + cby + cbz)
|
|
T.evaluate(clx + cly + clz)
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test7():
|
|
T.device_entry()
|
|
clx, cly, clz = T.cluster_id([3, 5, 12])
|
|
bx, by, bz = T.cta_id([8, 10, 12])
|
|
cbx, cby, cbz = T.cta_id_in_cluster([2, 2, 1])
|
|
T.evaluate(bx + by + bz)
|
|
T.evaluate(cbx + cby + cbz)
|
|
T.evaluate(clx + cly + clz)
|
|
|
|
# fmt: on
|
|
|
|
verify(test1)
|
|
verify(test2)
|
|
with pytest.raises(Exception, match="Inconsistent extents for scope"):
|
|
verify(test3)
|
|
verify(test4)
|
|
with pytest.raises(Exception, match="Inconsistent extents|non-divisible extents"):
|
|
verify(test5)
|
|
verify(test6)
|
|
with pytest.raises(Exception, match="Inconsistent extents|non-divisible extents"):
|
|
verify(test7)
|
|
|
|
|
|
def test_layout():
|
|
### TileLayout
|
|
# fmt: off
|
|
@T.prim_func(check_well_formed=False)
|
|
def test1():
|
|
T.device_entry()
|
|
T.cta_id([32])
|
|
T.warp_id([4])
|
|
T.lane_id([32])
|
|
A = T.alloc_buffer((2,), layout=T.TileLayout(T.S[2, 1]))
|
|
|
|
A[0] = 0
|
|
# fmt: on
|
|
verify(test1)
|
|
|
|
### SwizzleLayout
|
|
# fmt: off
|
|
@T.prim_func(check_well_formed=False)
|
|
def test2():
|
|
T.device_entry()
|
|
T.cta_id([32])
|
|
T.warp_id([4])
|
|
T.lane_id([32])
|
|
A = T.alloc_buffer((512,), scope="shared", layout=T.SwizzleLayout(3, 3, 3))
|
|
|
|
A[0] = 0
|
|
# fmt: on
|
|
verify(test2)
|
|
|
|
|
|
def test_host():
|
|
# fmt: off
|
|
@T.prim_func(check_well_formed=False)
|
|
def test1(A_ptr: T.handle):
|
|
A = T.match_buffer(A_ptr, (16, 16), dtype="float32", align=16)
|
|
|
|
A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1)
|
|
T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, "float32", 2, A.data, 16, 16, 64, 16, 16, 1, 1, 0, 0, 0, 0) # noqa: E501
|
|
|
|
T.device_entry()
|
|
for blockIdx in T.thread_binding(1, thread="blockIdx.x"):
|
|
for threadIdx in T.thread_binding(128, thread="threadIdx.x"):
|
|
bar = T.alloc_buffer((1,), "uint64", scope="shared", align=8)
|
|
phase = T.alloc_buffer((1,), "int32", scope="local")
|
|
A_smem = T.alloc_buffer((16, 16), "float32", scope="shared", align=128)
|
|
|
|
phase[0] = 0
|
|
if threadIdx == 0:
|
|
T.ptx.mbarrier.init(bar.data, 1)
|
|
T.ptx.fence.proxy_async("shared::cta")
|
|
T.ptx.cp_async.bulk.tensor.g2c(2, A_smem.data, bar.data, T.address_of(A_map), 0, 1, "", 0, 0) # noqa: E501
|
|
T.ptx.mbarrier.arrive.expect_tx(bar.data, 16*16*4)
|
|
T.ptx.mbarrier.try_wait(bar.data, phase[0])
|
|
phase[0] = phase[0] ^ 1
|
|
T.print_buffer(A_smem.data, "float32", False, False, 2, 16*16)
|
|
# fmt: on
|
|
verify(test1)
|
|
|
|
|
|
def test_device_func():
|
|
# Per-call exec-scope migration: scope is now attached per op via the
|
|
# ``T.op[scope](...)`` subscription surface instead of a ``with T.cta():``
|
|
# region. ``test1`` exercises a per-call-scoped op; ``test2`` the plain
|
|
# (unscoped) op. The old multi-root-scope negative case asserted the removed
|
|
# "only one root scope" verifier rule and no longer has an equivalent, so it
|
|
# is dropped.
|
|
# fmt: off
|
|
@T.prim_func(check_well_formed=False)
|
|
def test1(A: T.Buffer((128,), "float32")):
|
|
T.device_entry()
|
|
T.cta_id([1])
|
|
T.thread_id([128])
|
|
Tx.cta.fill(A, 0.)
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def test2(A: T.Buffer((128,), "float32")):
|
|
T.device_entry()
|
|
T.cta_id([128])
|
|
T.thread_id([128])
|
|
Tx.fill(A, 0.)
|
|
# fmt: on
|
|
verify(test1, device_func=True)
|
|
verify(test2, device_func=True)
|
|
|
|
|
|
def test_preferred_cluster_validation():
|
|
# fmt: off
|
|
# Valid: cluster→cta with preferred_extents matching size
|
|
@T.prim_func(check_well_formed=False)
|
|
def test1() -> None:
|
|
T.device_entry()
|
|
cbx, cby = T.cta_id_in_cluster([2, 1], preferred=[2, 2])
|
|
tx = T.thread_id([128])
|
|
T.evaluate(cbx + cby + tx)
|
|
|
|
# Invalid: preferred size doesn't match extents size (caught at verify time)
|
|
@T.prim_func(check_well_formed=False)
|
|
def test2() -> None:
|
|
T.device_entry()
|
|
cbx, cby = T.cta_id_in_cluster([2, 1], preferred=[2])
|
|
tx = T.thread_id([128])
|
|
T.evaluate(cbx + cby + tx)
|
|
# fmt: on
|
|
|
|
verify(test1)
|
|
with pytest.raises(Exception, match="preferred_extents must have the same size"):
|
|
verify(test2)
|
|
|
|
# Invalid: preferred on a non-cluster→cta scope (caught at IR build time)
|
|
with pytest.raises(Exception):
|
|
# fmt: off
|
|
@T.prim_func(check_well_formed=False)
|
|
def test3() -> None:
|
|
T.device_entry()
|
|
bx = T.cta_id([128], preferred=[256])
|
|
tx = T.thread_id([128])
|
|
T.evaluate(bx + tx)
|
|
# fmt: on
|
|
|
|
|
|
def test_scope_id_deferred_relaxed_at_construction():
|
|
"""Deferred scope_id (no extents) must pass the well-formed check even when
|
|
no sibling provides enough info to resolve it -- strict resolution is
|
|
deferred to LowerTIRx."""
|
|
|
|
# fmt: off
|
|
@T.prim_func(check_well_formed=False)
|
|
def partial_only_cta():
|
|
T.device_entry()
|
|
bx = T.cta_id() # deferred kernel→cta, no closure source
|
|
tx = T.thread_id([128]) # explicit
|
|
T.evaluate(bx + tx)
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def all_deferred():
|
|
T.device_entry()
|
|
bx = T.cta_id()
|
|
wg = T.warpgroup_id()
|
|
warp = T.warp_id_in_wg()
|
|
lane = T.lane_id()
|
|
T.evaluate(bx + wg + warp + lane)
|
|
|
|
@T.prim_func(check_well_formed=False)
|
|
def mixed():
|
|
T.device_entry()
|
|
# kCtaWarp=4, kWarpThread=32 → kCtaThread=128 derivable.
|
|
T.warp_id([4])
|
|
T.lane_id([32])
|
|
T.thread_id() # deferred kCtaThread, resolvable via closure
|
|
pass
|
|
# fmt: on
|
|
|
|
# All three accepted by well-formed: deferred extents are tolerated.
|
|
verify(partial_only_cta)
|
|
verify(all_deferred)
|
|
verify(mixed)
|
|
|
|
|
|
def test_scope_id_deferred_consistency_still_enforced():
|
|
"""Even with deferred defs, known-known consistency between sibling defs
|
|
must still be enforced by the closure check."""
|
|
|
|
# fmt: off
|
|
@T.prim_func(check_well_formed=False)
|
|
def inconsistent():
|
|
# 4 warps * 32 lanes = 128 threads, but explicit thread_id says 64 -> error.
|
|
T.device_entry()
|
|
T.cta_id([32])
|
|
T.warp_id([4])
|
|
T.lane_id([32])
|
|
T.thread_id() # deferred (shouldn't shadow the conflict)
|
|
T.thread_id([64]) # conflicts with derived kCtaThread=128
|
|
pass
|
|
# fmt: on
|
|
|
|
with pytest.raises(Exception, match="Inconsistent extents for scope"):
|
|
verify(inconsistent)
|
|
|
|
|
|
def test_scope_id_deferred_multi_var_rejected():
|
|
"""Deferred form (no extents) requires exactly one Var. Multi-var defers
|
|
have no well-defined recovery from fused closure values."""
|
|
|
|
# The C++ ScopeIdDef ctor enforces this; constructing such a def from the
|
|
# parser path is not currently expressible (parser only emits single-Var
|
|
# deferred), but we exercise the FFI-level guard directly.
|
|
from tvm.tirx.exec_scope import ScopeIdDef
|
|
from tvm.tirx.expr import Var
|
|
|
|
# Single-Var deferred form is fine.
|
|
ScopeIdDef([Var("", "int32")], None, "kernel", "cta")
|
|
|
|
# Two-Var deferred should be rejected.
|
|
with pytest.raises(Exception, match="Deferred ScopeIdDef.*must define exactly one Var"):
|
|
ScopeIdDef([Var("", "int32"), Var("", "int32")], None, "kernel", "cta")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_root_scope()
|
|
test_nested_scope()
|
|
test_scope_id_consistency()
|
|
test_layout()
|
|
test_host()
|
|
test_device_func()
|
|
test_scope_id_deferred_relaxed_at_construction()
|
|
test_scope_id_deferred_consistency_still_enforced()
|
|
test_scope_id_deferred_multi_var_rejected()
|