chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,248 @@
|
||||
# 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.
|
||||
|
||||
# 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.
|
||||
# ruff: noqa: F401
|
||||
|
||||
import tempfile
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
import tvm.s_tir.meta_schedule as ms
|
||||
import tvm.testing
|
||||
from tvm import relax
|
||||
from tvm.ir import transform
|
||||
from tvm.ir.module import IRModule
|
||||
from tvm.ir.transform import PassContext
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
|
||||
target = tvm.target.Target({"kind": "llvm", "num-cores": 16})
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class InputModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def tir_matmul(x: T.handle, y: T.handle, z: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "tir_matmul"})
|
||||
k = T.int32()
|
||||
A = T.match_buffer(x, (32, 32))
|
||||
B = T.match_buffer(y, (32, 32))
|
||||
C = T.match_buffer(z, (32, 32))
|
||||
|
||||
for i0, j0, k0 in T.grid(32, 32, 32):
|
||||
with T.sblock():
|
||||
i, j, k = T.axis.remap("SSR", [i0, j0, k0])
|
||||
with T.init():
|
||||
C[i, j] = 0.0
|
||||
C[i, j] += A[i, k] * B[j, k]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def tir_relu(x: T.handle, y: T.handle):
|
||||
T.func_attr({"global_symbol": "tir_relu"})
|
||||
A = T.match_buffer(x, (32, 32))
|
||||
B = T.match_buffer(y, (32, 32))
|
||||
for i, j in T.grid(32, 32):
|
||||
with T.sblock():
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = T.max(A[vi, vj], 0.0)
|
||||
|
||||
@R.function
|
||||
def main(x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")) -> R.Tensor:
|
||||
cls = InputModule
|
||||
with R.dataflow():
|
||||
lv0 = R.call_tir(cls.tir_matmul, (x, w), R.Tensor((32, 32), dtype="float32"))
|
||||
lv1 = R.call_tir(cls.tir_relu, (lv0), R.Tensor((32, 32), dtype="float32"))
|
||||
R.output(lv1)
|
||||
return lv1
|
||||
|
||||
|
||||
# TODO(@sunggg): determine how to pass MS database object across different passes.
|
||||
# PassContext might be an option, but we already have TuningAPI database.
|
||||
# (MS database and TuningAPI database will be unified in the future)
|
||||
# For now, we only support default JSON database config.
|
||||
def test_ms_tuning_irmodule():
|
||||
mod = InputModule
|
||||
assert isinstance(mod, IRModule)
|
||||
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
"""
|
||||
# TODO(@sunggg): revisit when ready
|
||||
with target, PassContext(trace=Trace(mod), opt_level=0):
|
||||
tuning_pass = relax.transform.MetaScheduleTuneIRMod(
|
||||
params={}, work_dir=work_dir, max_trials_global=4
|
||||
)
|
||||
out_mod = tuning_pass(mod)
|
||||
assert PassContext.current().get_trace_stack_size() == 1
|
||||
assert PassContext.current().get_current_trace().size == 1
|
||||
tvm.ir.assert_structural_equal(mod, out_mod)
|
||||
"""
|
||||
|
||||
with target, PassContext(opt_level=0):
|
||||
tuning_pass = relax.transform.MetaScheduleTuneIRMod(
|
||||
params={}, work_dir=work_dir, max_trials_global=4
|
||||
)
|
||||
out_mod = tuning_pass(mod)
|
||||
|
||||
application_pass = relax.transform.MetaScheduleApplyDatabase(work_dir)
|
||||
|
||||
out_mod = application_pass(mod)
|
||||
assert not tvm_ffi.structural_equal(mod, out_mod)
|
||||
|
||||
|
||||
def test_ms_tuning_primfunc():
|
||||
mod = InputModule
|
||||
assert isinstance(mod, IRModule)
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
"""
|
||||
# TODO(@sunggg): revisit when ready
|
||||
with target, PassContext(trace=Trace(mod), opt_level=0):
|
||||
tuning_pass = relax.transform.MetaScheduleTuneTIR(
|
||||
work_dir=work_dir, max_trials_global=4
|
||||
)
|
||||
out_mod = tuning_pass(mod)
|
||||
assert PassContext.current().get_trace_stack_size() == 1
|
||||
# TODO (@sunggg): Need to determine how to track subgraph-level tuning traces.
|
||||
# Currently, we don't track this so the trace size. Revisit this later.
|
||||
tvm.ir.assert_structural_equal(mod, out_mod)
|
||||
"""
|
||||
with target, PassContext(opt_level=0):
|
||||
tuning_pass = relax.transform.MetaScheduleTuneIRMod(
|
||||
params={}, work_dir=work_dir, max_trials_global=4
|
||||
)
|
||||
out_mod = tuning_pass(mod)
|
||||
|
||||
application_pass = relax.transform.MetaScheduleApplyDatabase(work_dir)
|
||||
out_mod = application_pass(mod)
|
||||
assert not tvm_ffi.structural_equal(mod, out_mod)
|
||||
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
with target, PassContext(opt_level=0):
|
||||
tuning_pass = relax.transform.MetaScheduleTuneIRMod(
|
||||
params={},
|
||||
work_dir=work_dir,
|
||||
max_trials_global=4,
|
||||
max_trials_per_task=2,
|
||||
op_names=["matmul"],
|
||||
)
|
||||
tuning_pass(mod)
|
||||
|
||||
db = ms.database.JSONDatabase(
|
||||
work_dir + "/database_workload.json", work_dir + "/database_tuning_record.json"
|
||||
)
|
||||
|
||||
assert len(db.get_all_tuning_records()) == 2
|
||||
|
||||
for rec in db.get_all_tuning_records():
|
||||
assert rec.workload.mod["main"].attrs["global_symbol"] == "tir_matmul"
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class DefaultScheduledModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def tir_matmul(
|
||||
A: T.Buffer((32, 32), "float32"),
|
||||
B: T.Buffer((32, 32), "float32"),
|
||||
C: T.Buffer((32, 32), "float32"),
|
||||
):
|
||||
T.func_attr({"global_symbol": "tir_matmul", "tirx.is_scheduled": True})
|
||||
# with T.sblock("root"):
|
||||
for i0_j0_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
|
||||
for i0_j0_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
||||
for k0 in range(32):
|
||||
with T.sblock(""):
|
||||
i = T.axis.spatial(32, (i0_j0_fused_0 * 1024 + i0_j0_fused_1) // 32)
|
||||
j = T.axis.spatial(32, (i0_j0_fused_0 * 1024 + i0_j0_fused_1) % 32)
|
||||
k = T.axis.reduce(32, k0)
|
||||
T.reads(A[i, k], B[j, k])
|
||||
T.writes(C[i, j])
|
||||
with T.init():
|
||||
C[i, j] = T.float32(0)
|
||||
C[i, j] = C[i, j] + A[i, k] * B[j, k]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def tir_relu(A: T.Buffer((32, 32), "float32"), B: T.Buffer((32, 32), "float32")):
|
||||
T.func_attr({"global_symbol": "tir_relu", "tirx.is_scheduled": True})
|
||||
# with T.sblock("root"):
|
||||
for i_j_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
|
||||
for i_j_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
||||
with T.sblock(""):
|
||||
vi = T.axis.spatial(32, (i_j_fused_0 * 1024 + i_j_fused_1) // 32)
|
||||
vj = T.axis.spatial(32, (i_j_fused_0 * 1024 + i_j_fused_1) % 32)
|
||||
T.reads(A[vi, vj])
|
||||
T.writes(B[vi, vj])
|
||||
B[vi, vj] = T.max(A[vi, vj], T.float32(0))
|
||||
|
||||
@R.function
|
||||
def main(
|
||||
x: R.Tensor((32, 32), dtype="float32"), w: R.Tensor((32, 32), dtype="float32")
|
||||
) -> R.Tensor((32, 32), dtype="float32"):
|
||||
with R.dataflow():
|
||||
lv0 = R.call_tir(
|
||||
DefaultScheduledModule.tir_matmul,
|
||||
(x, w),
|
||||
out_ty=R.Tensor((32, 32), dtype="float32"),
|
||||
)
|
||||
lv1 = R.call_tir(
|
||||
DefaultScheduledModule.tir_relu,
|
||||
(lv0,),
|
||||
out_ty=R.Tensor((32, 32), dtype="float32"),
|
||||
)
|
||||
R.output(lv1)
|
||||
return lv1
|
||||
|
||||
|
||||
def test_ms_database_apply_fallback():
|
||||
mod = InputModule
|
||||
target_cuda = tvm.target.Target("nvidia/geforce-rtx-3090-ti")
|
||||
assert isinstance(mod, IRModule)
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
"""
|
||||
with target_cuda, PassContext(trace=Trace(mod), opt_level=0):
|
||||
tuning_pass = relax.transform.MetaScheduleTuneTIR(
|
||||
work_dir=work_dir, max_trials_global=0
|
||||
)
|
||||
out_mod = tuning_pass(mod)
|
||||
tvm.ir.assert_structural_equal(mod, out_mod)
|
||||
"""
|
||||
with target_cuda, PassContext(opt_level=0):
|
||||
tuning_pass = relax.transform.MetaScheduleTuneTIR(
|
||||
work_dir=work_dir, max_trials_global=0
|
||||
)
|
||||
out_mod = tuning_pass(mod)
|
||||
default_pass = tvm.s_tir.transform.DefaultGPUSchedule()
|
||||
out_mod = default_pass(mod)
|
||||
tvm.ir.assert_structural_equal(out_mod, DefaultScheduledModule)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
Reference in New Issue
Block a user