chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,71 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
|
||||
from tvm.s_tir.meta_schedule.arg_info import ArgInfo, TensorInfo
|
||||
from tvm.script import tirx as T
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
# fmt: off
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def Matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (128, 256), "float32")
|
||||
B = T.match_buffer(b, (256, 512), "float32")
|
||||
C = T.match_buffer(c, (128, 512), "float32")
|
||||
for i, j, k in T.grid(128, 256, 512):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
# fmt: on
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
|
||||
|
||||
def test_meta_schedule_tensor_info_creation():
|
||||
info = TensorInfo("float32", [1, 224, 224, 3])
|
||||
info = str(info)
|
||||
assert info == 'TensorInfo("float32", [1, 224, 224, 3])'
|
||||
|
||||
|
||||
def test_meta_schedule_tensor_info_as_json():
|
||||
info = TensorInfo("float32", [1, 224, 224, 3])
|
||||
info = info.as_json()
|
||||
assert info == ["TENSOR", "float32", [1, 224, 224, 3]]
|
||||
|
||||
|
||||
def test_meta_schedule_tensor_info_from_json():
|
||||
info = ["TENSOR", "float32", [1, 224, 224, 3]]
|
||||
info = TensorInfo.from_json(info)
|
||||
assert str(info) == 'TensorInfo("float32", [1, 224, 224, 3])'
|
||||
|
||||
|
||||
def test_meta_schedule_arg_info_from_prim_func():
|
||||
a_info, b_info, c_info = ArgInfo.from_prim_func(Matmul)
|
||||
assert str(a_info) == 'TensorInfo("float32", [128, 256])'
|
||||
assert str(b_info) == 'TensorInfo("float32", [256, 512])'
|
||||
assert str(c_info) == 'TensorInfo("float32", [128, 512])'
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_meta_schedule_tensor_info_creation()
|
||||
test_meta_schedule_tensor_info_as_json()
|
||||
test_meta_schedule_tensor_info_from_json()
|
||||
test_meta_schedule_arg_info_from_prim_func()
|
||||
@@ -0,0 +1,234 @@
|
||||
# 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.
|
||||
# ruff: noqa: F401
|
||||
"""Test Meta Schedule Builder"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pytest
|
||||
from tvm_ffi import register_global_func
|
||||
|
||||
import tvm.testing
|
||||
from tvm import script
|
||||
from tvm.runtime import Module
|
||||
from tvm.s_tir.meta_schedule.builder import (
|
||||
BuilderInput,
|
||||
BuilderResult,
|
||||
LocalBuilder,
|
||||
PyBuilder,
|
||||
)
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring
|
||||
|
||||
|
||||
@script.ir_module
|
||||
class MatmulModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "matmul", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@script.ir_module
|
||||
class MatmulReluModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul_relu( # pylint: disable=no-self-argument
|
||||
a: T.handle, b: T.handle, d: T.handle
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "matmul_relu", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
D = T.match_buffer(d, (1024, 1024), "float32")
|
||||
C = T.sblock_alloc_buffer((1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j in T.grid(1024, 1024):
|
||||
with T.sblock("relu"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
D[vi, vj] = T.max(C[vi, vj], 0.0)
|
||||
|
||||
|
||||
@script.ir_module
|
||||
class BatchMatmulModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def batch_matmul( # pylint: disable=no-self-argument
|
||||
a: T.handle, b: T.handle, c: T.handle
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "batch_matmul", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, [16, 128, 128])
|
||||
B = T.match_buffer(b, [16, 128, 128])
|
||||
C = T.match_buffer(c, [16, 128, 128])
|
||||
for n, i, j, k in T.grid(16, 128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])
|
||||
with T.init():
|
||||
C[vn, vi, vj] = 0.0
|
||||
C[vn, vi, vj] = C[vn, vi, vj] + A[vn, vi, vk] * B[vn, vj, vk]
|
||||
|
||||
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring
|
||||
|
||||
|
||||
def _check_build_results(builder_results: list[BuilderResult]):
|
||||
"""Simple check whether the build is successful"""
|
||||
for result in builder_results:
|
||||
artifact_path = result.artifact_path
|
||||
error_msg = result.error_msg
|
||||
assert artifact_path is not None
|
||||
assert error_msg is None
|
||||
os.remove(artifact_path)
|
||||
os.rmdir(os.path.dirname(artifact_path))
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_single_build():
|
||||
"""Test meta schedule builder for a single build"""
|
||||
mod = MatmulModule
|
||||
builder = LocalBuilder()
|
||||
builder_inputs = [BuilderInput(mod, Target("llvm"))]
|
||||
builder_results = builder.build(builder_inputs)
|
||||
assert len(builder_results) == len(builder_inputs)
|
||||
_check_build_results(builder_results)
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_multiple_build():
|
||||
"""Test meta schedule builder for multiple builds"""
|
||||
builder = LocalBuilder()
|
||||
builder_inputs = [
|
||||
BuilderInput(MatmulModule, Target("llvm")),
|
||||
BuilderInput(MatmulReluModule, Target("llvm")),
|
||||
BuilderInput(BatchMatmulModule, Target("llvm")),
|
||||
]
|
||||
builder_results = builder.build(builder_inputs)
|
||||
assert len(builder_results) == len(builder_inputs)
|
||||
_check_build_results(builder_results)
|
||||
|
||||
|
||||
def test_meta_schedule_error_handle_test_builder():
|
||||
"""Test the error handing during building"""
|
||||
|
||||
class TestBuilder(PyBuilder):
|
||||
def build( # pylint: disable=no-self-use
|
||||
self,
|
||||
build_inputs: list[BuilderInput],
|
||||
) -> list[BuilderResult]:
|
||||
return [BuilderResult(None, "error") for w in build_inputs]
|
||||
|
||||
builder = TestBuilder()
|
||||
builder_inputs = [
|
||||
BuilderInput(MatmulModule, Target("llvm")),
|
||||
BuilderInput(MatmulReluModule, Target("llvm")),
|
||||
BuilderInput(BatchMatmulModule, Target("llvm")),
|
||||
]
|
||||
builder_results = builder.build(builder_inputs)
|
||||
assert len(builder_results) == len(builder_inputs)
|
||||
for result in builder_results:
|
||||
artifact_path = result.artifact_path
|
||||
error_msg = result.error_msg
|
||||
assert artifact_path is None
|
||||
assert error_msg == "error"
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_error_handle_build_func():
|
||||
"""Test the error handing during building"""
|
||||
|
||||
def initializer():
|
||||
@register_global_func("meta_schedule.builder.test_build")
|
||||
def test_build(mod: Module, target: Target, _) -> None: # pylint: disable=unused-variable
|
||||
raise ValueError("Builder intended Test Error (build func).")
|
||||
|
||||
builder = LocalBuilder(f_build="meta_schedule.builder.test_build", initializer=initializer)
|
||||
builder_inputs = [BuilderInput(MatmulModule, Target("llvm"))]
|
||||
builder_results = builder.build(builder_inputs)
|
||||
assert len(builder_results) == len(builder_inputs)
|
||||
for result in builder_results:
|
||||
artifact_path = result.artifact_path
|
||||
error_msg = result.error_msg
|
||||
assert artifact_path is None
|
||||
assert error_msg.startswith("LocalBuilder: An exception occurred")
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_error_handle_export_func():
|
||||
"""Test the error handing during building"""
|
||||
|
||||
def initializer():
|
||||
@register_global_func("meta_schedule.builder.test_export")
|
||||
def test_build(mod: Module) -> str: # pylint: disable=unused-variable
|
||||
raise ValueError("Builder intended Test Error (export func).")
|
||||
|
||||
builder = LocalBuilder(f_export="meta_schedule.builder.test_export", initializer=initializer)
|
||||
builder_inputs = [BuilderInput(MatmulModule, Target("llvm"))]
|
||||
builder_results = builder.build(builder_inputs)
|
||||
assert len(builder_results) == len(builder_inputs)
|
||||
for result in builder_results:
|
||||
artifact_path = result.artifact_path
|
||||
error_msg = result.error_msg
|
||||
assert artifact_path is None
|
||||
assert error_msg.startswith("LocalBuilder: An exception occurred")
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_error_handle_time_out():
|
||||
"""Test the error handing time out during building"""
|
||||
|
||||
def initializer():
|
||||
@register_global_func("meta_schedule.builder.test_time_out")
|
||||
def timeout_build(mod, target, _): # pylint: disable=unused-argument, unused-variable
|
||||
time.sleep(2)
|
||||
|
||||
builder = LocalBuilder(
|
||||
timeout_sec=1,
|
||||
f_build="meta_schedule.builder.test_time_out",
|
||||
initializer=initializer,
|
||||
)
|
||||
builder_inputs = [BuilderInput(MatmulModule, Target("llvm"))]
|
||||
builder_results = builder.build(builder_inputs)
|
||||
assert len(builder_results) == len(builder_inputs)
|
||||
for result in builder_results:
|
||||
artifact_path = result.artifact_path
|
||||
error_msg = result.error_msg
|
||||
assert artifact_path is None
|
||||
assert error_msg.startswith("LocalBuilder: Timeout")
|
||||
|
||||
|
||||
def test_meta_schedule_missing_build_func():
|
||||
pytest.importorskip("cloudpickle")
|
||||
with pytest.raises(ValueError):
|
||||
LocalBuilder(f_build="wrong-name")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,333 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-docstring
|
||||
# ruff: noqa: F401
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
from functools import partial
|
||||
from importlib.util import find_spec
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.s_tir.meta_schedule.cost_model import PyCostModel, RandomModel, XGBModel
|
||||
from tvm.s_tir.meta_schedule.cost_model.xgb_model import PackSum, _get_custom_call_back
|
||||
from tvm.s_tir.meta_schedule.feature_extractor import RandomFeatureExtractor
|
||||
from tvm.s_tir.meta_schedule.runner import RunnerResult
|
||||
from tvm.s_tir.meta_schedule.search_strategy import MeasureCandidate
|
||||
from tvm.s_tir.meta_schedule.tune_context import TuneContext
|
||||
from tvm.s_tir.schedule.schedule import Schedule
|
||||
from tvm.script import tirx as T
|
||||
|
||||
requires_xgboost = pytest.mark.skipif(
|
||||
find_spec("xgboost") is None, reason="xgboost is not installed"
|
||||
)
|
||||
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class FullModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(T_full: T.Buffer((T.int64(2), T.int64(3)), "float32")):
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
for ax0, ax1 in T.grid(T.int64(2), T.int64(3)):
|
||||
with T.sblock("T_full"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.reads()
|
||||
T.writes(T_full[v_ax0, v_ax1])
|
||||
T_full[v_ax0, v_ax1] = T.float32(1)
|
||||
|
||||
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,disable=unused-argument
|
||||
|
||||
|
||||
def test_meta_schedule_cost_model():
|
||||
@derived_object
|
||||
class FancyCostModel(PyCostModel):
|
||||
def load(self, path: str) -> None:
|
||||
pass
|
||||
|
||||
def save(self, path: str) -> None:
|
||||
pass
|
||||
|
||||
def update(
|
||||
self,
|
||||
context: TuneContext,
|
||||
candidates: list[MeasureCandidate],
|
||||
results: list[RunnerResult],
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
def predict(self, context: TuneContext, candidates: list[MeasureCandidate]) -> np.ndarray:
|
||||
return np.random.rand(10)
|
||||
|
||||
model = FancyCostModel()
|
||||
model.save("fancy_test_location")
|
||||
model.load("fancy_test_location")
|
||||
model.update(TuneContext(), [], [])
|
||||
results = model.predict(
|
||||
TuneContext(), [MeasureCandidate(Schedule(mod=Matmul), []) for _ in range(10)]
|
||||
)
|
||||
assert results.shape == (10,)
|
||||
|
||||
|
||||
def test_meta_schedule_random_model():
|
||||
model = RandomModel()
|
||||
model.update(TuneContext(), [], [])
|
||||
res = model.predict(TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(10)])
|
||||
assert len(res) == 10
|
||||
assert min(res) >= 0 and max(res) <= model.max_range
|
||||
|
||||
|
||||
def test_meta_schedule_random_model_reseed():
|
||||
model = RandomModel(seed=100)
|
||||
res = model.predict(TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(20)])
|
||||
new_model = RandomModel(seed=100)
|
||||
new_res = new_model.predict(
|
||||
TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(20)]
|
||||
)
|
||||
assert (res == new_res).all()
|
||||
|
||||
|
||||
def test_meta_schedule_random_model_reload():
|
||||
model = RandomModel(seed=25973)
|
||||
model.predict(
|
||||
TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(30)]
|
||||
) # change state
|
||||
path = os.path.join(tempfile.mkdtemp(), "test_output_meta_schedule_random_model.npy")
|
||||
model.save(path)
|
||||
res1 = model.predict(TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(70)])
|
||||
model.load(path)
|
||||
res2 = model.predict(TuneContext(), [MeasureCandidate(Schedule(Matmul), []) for i in range(70)])
|
||||
shutil.rmtree(os.path.dirname(path))
|
||||
assert (res1 == res2).all()
|
||||
|
||||
|
||||
def _dummy_candidate():
|
||||
return MeasureCandidate(Schedule(Matmul), [])
|
||||
|
||||
|
||||
def _dummy_result(num_samples: int = 4, max_run_sec: int = 10):
|
||||
return RunnerResult(list(np.random.rand(num_samples) * max_run_sec + 1e-6), None)
|
||||
|
||||
|
||||
@requires_xgboost
|
||||
def test_meta_schedule_xgb_model():
|
||||
extractor = RandomFeatureExtractor()
|
||||
model = XGBModel(extractor=extractor, num_warmup_samples=2, num_tuning_cores=1)
|
||||
update_sample_count = 10
|
||||
predict_sample_count = 100
|
||||
model.update(
|
||||
TuneContext(),
|
||||
[_dummy_candidate() for i in range(update_sample_count)],
|
||||
[_dummy_result() for i in range(update_sample_count)],
|
||||
)
|
||||
model.predict(TuneContext(), [_dummy_candidate() for i in range(predict_sample_count)])
|
||||
|
||||
|
||||
@requires_xgboost
|
||||
def test_meta_schedule_xgb_model_no_feature():
|
||||
model = XGBModel(num_warmup_samples=0, num_tuning_cores=1)
|
||||
tune_ctx = TuneContext(
|
||||
FullModule,
|
||||
num_threads=1,
|
||||
target={"kind": "llvm", "num-cores": 16},
|
||||
space_generator="post-order-apply",
|
||||
search_strategy="evolutionary",
|
||||
)
|
||||
candidate = MeasureCandidate(Schedule(FullModule), [])
|
||||
model.update(tune_ctx, [candidate], [_dummy_result()])
|
||||
model.predict(tune_ctx, [candidate])
|
||||
|
||||
|
||||
@requires_xgboost
|
||||
def test_meta_schedule_xgb_model_reload():
|
||||
extractor = RandomFeatureExtractor()
|
||||
model = XGBModel(extractor=extractor, num_warmup_samples=10, num_tuning_cores=1)
|
||||
update_sample_count = 20
|
||||
predict_sample_count = 30
|
||||
model.update(
|
||||
TuneContext(),
|
||||
[_dummy_candidate() for i in range(update_sample_count)],
|
||||
[_dummy_result() for i in range(update_sample_count)],
|
||||
)
|
||||
model.predict(TuneContext(), [_dummy_candidate() for i in range(predict_sample_count)])
|
||||
with tempfile.NamedTemporaryFile() as path:
|
||||
# Backup
|
||||
random_state = model.extractor.random_state # save feature extractor's random state
|
||||
old_data = model.data
|
||||
old_data_size = model.data_size
|
||||
model.save(path.name)
|
||||
res1 = model.predict(
|
||||
TuneContext(), [_dummy_candidate() for i in range(predict_sample_count)]
|
||||
)
|
||||
# Load
|
||||
model.extractor.random_state = random_state # load feature extractor's random state
|
||||
model.load(path.name)
|
||||
new_data = model.data
|
||||
new_data_size = model.data_size
|
||||
res2 = model.predict(
|
||||
TuneContext(), [_dummy_candidate() for i in range(predict_sample_count)]
|
||||
)
|
||||
assert (res1 == res2).all()
|
||||
assert old_data_size == new_data_size
|
||||
assert len(old_data) == len(new_data)
|
||||
for (k1, g1), (k2, g2) in zip( # pylint: disable=invalid-name
|
||||
old_data.items(), new_data.items()
|
||||
):
|
||||
assert k1 == k2
|
||||
assert k1 == g1.group_hash
|
||||
assert k2 == g2.group_hash
|
||||
assert (g1.costs == g2.costs).all()
|
||||
assert len(g1.features) == len(g2.features)
|
||||
for f1, f2 in zip(g1.features, g2.features): # pylint: disable=invalid-name
|
||||
assert (f1 == f2).all()
|
||||
|
||||
|
||||
@requires_xgboost
|
||||
def test_meta_schedule_xgb_model_reupdate():
|
||||
extractor = RandomFeatureExtractor()
|
||||
model = XGBModel(extractor=extractor, num_warmup_samples=2, num_tuning_cores=1)
|
||||
update_sample_count = 60
|
||||
predict_sample_count = 100
|
||||
model.update(
|
||||
TuneContext(),
|
||||
[_dummy_candidate() for i in range(update_sample_count)],
|
||||
[_dummy_result() for i in range(update_sample_count)],
|
||||
)
|
||||
model.update(
|
||||
TuneContext(),
|
||||
[_dummy_candidate() for i in range(update_sample_count)],
|
||||
[_dummy_result() for i in range(update_sample_count)],
|
||||
)
|
||||
model.update(
|
||||
TuneContext(),
|
||||
[_dummy_candidate() for i in range(update_sample_count)],
|
||||
[_dummy_result() for i in range(update_sample_count)],
|
||||
)
|
||||
model.predict(TuneContext(), [_dummy_candidate() for i in range(predict_sample_count)])
|
||||
|
||||
|
||||
@requires_xgboost
|
||||
def test_meta_schedule_xgb_model_callback_as_function():
|
||||
# pylint: disable=import-outside-toplevel
|
||||
from itertools import chain as itertools_chain
|
||||
|
||||
import xgboost as xgb
|
||||
|
||||
# pylint: enable=import-outside-toplevel
|
||||
|
||||
extractor = RandomFeatureExtractor()
|
||||
model = XGBModel(extractor=extractor, num_warmup_samples=10, num_tuning_cores=1)
|
||||
update_sample_count = 20
|
||||
predict_sample_count = 30
|
||||
|
||||
model.update(
|
||||
TuneContext(),
|
||||
[_dummy_candidate() for i in range(update_sample_count)],
|
||||
[_dummy_result() for i in range(update_sample_count)],
|
||||
)
|
||||
model.predict(TuneContext(), [_dummy_candidate() for i in range(predict_sample_count)])
|
||||
with tempfile.NamedTemporaryFile() as path:
|
||||
# Backup and train on new TrainingCallBack api
|
||||
random_state = model.extractor.random_state # save feature extractor's random state
|
||||
|
||||
model.save(path.name)
|
||||
|
||||
old_booster = model.booster
|
||||
xs = [ # pylint: disable=invalid-name
|
||||
x.numpy().astype("float32")
|
||||
for x in extractor.extract_from(
|
||||
TuneContext(),
|
||||
[_dummy_candidate() for i in range(predict_sample_count)],
|
||||
)
|
||||
]
|
||||
d_test = PackSum(xs=xs, ys=None)
|
||||
pred1 = old_booster.predict(d_test.dmatrix)
|
||||
|
||||
# Load and train on deprecated TrainingCallBack api
|
||||
model.extractor.random_state = random_state # load feature extractor's random state
|
||||
model.load(path.name)
|
||||
d_train = PackSum(
|
||||
xs=list(itertools_chain.from_iterable([g.features for g in model.data.values()])),
|
||||
ys=np.concatenate(
|
||||
[g.min_cost / g.costs for g in model.data.values()],
|
||||
axis=0,
|
||||
),
|
||||
)
|
||||
|
||||
def obj(ys_pred: np.ndarray, d_train1: "xgb.DMatrix"): # type: ignore # pylint: disable = unused-argument
|
||||
return d_train.obj_square_error(ys_pred)
|
||||
|
||||
def rmse(ys_pred: np.ndarray, d_train1: "xgb.DMatrix"): # type: ignore # pylint: disable = unused-argument
|
||||
return d_train.rmse(ys_pred)
|
||||
|
||||
def avg_peak_score(ys_pred: np.ndarray, d_train1: "xgb.DMatrix"): # type: ignore # pylint: disable = unused-argument
|
||||
return d_train.average_peak_score(ys_pred, model.average_peak_n)
|
||||
|
||||
new_booster = xgb.train(
|
||||
model.config.to_dict(),
|
||||
d_train.dmatrix,
|
||||
num_boost_round=10000,
|
||||
obj=obj,
|
||||
callbacks=[
|
||||
_get_custom_call_back(
|
||||
early_stopping_rounds=model.early_stopping_rounds,
|
||||
verbose_eval=model.verbose_eval,
|
||||
fevals=[rmse, avg_peak_score],
|
||||
evals=[(d_train.dmatrix, "tr")],
|
||||
cvfolds=None,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
xs = [ # pylint: disable=invalid-name
|
||||
x.numpy().astype("float32")
|
||||
for x in extractor.extract_from(
|
||||
TuneContext(),
|
||||
[_dummy_candidate() for i in range(predict_sample_count)],
|
||||
)
|
||||
]
|
||||
d_test = PackSum(xs=xs, ys=None)
|
||||
pred2 = new_booster.predict(d_test.dmatrix)
|
||||
|
||||
assert np.allclose(pred1, pred2, rtol=1e-3, atol=1e-3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,613 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: F401
|
||||
"""Test Meta Schedule Database"""
|
||||
|
||||
import os.path as osp
|
||||
import tempfile
|
||||
from collections.abc import Callable
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import tirx
|
||||
from tvm.ir.module import IRModule
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.s_tir import Schedule
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.database import TuningRecord, Workload
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
# fmt: off
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class MatmulRelu:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, d: T.handle) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
B = T.match_buffer(b, (16, 16), "float32")
|
||||
D = T.match_buffer(d, (16, 16), "float32")
|
||||
C = T.sblock_alloc_buffer((16, 16), "float32")
|
||||
for i, j, k in T.grid(16, 16, 16):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock("relu"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
D[vi, vj] = T.max(C[vi, vj], 0.0)
|
||||
|
||||
|
||||
# fmt: on
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
|
||||
|
||||
def _schedule_matmul(sch: Schedule):
|
||||
block = sch.get_sblock("matmul")
|
||||
i, j, k = sch.get_loops(block=block)
|
||||
i_tiles = [1, 1, 2, 512]
|
||||
j_tiles = [1, 512, 1, 2]
|
||||
k_tiles = [256, 4]
|
||||
i_0, i_1, i_2, i_3 = sch.split(loop=i, factors=i_tiles)
|
||||
j_0, j_1, j_2, j_3 = sch.split(loop=j, factors=j_tiles)
|
||||
k_0, k_1 = sch.split(loop=k, factors=k_tiles)
|
||||
sch.reorder(i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3)
|
||||
|
||||
|
||||
def _create_schedule(mod: IRModule, sch_fn: Callable[[Schedule], None]) -> Schedule:
|
||||
sch = tvm.s_tir.Schedule(mod=mod, debug_mask="all")
|
||||
sch_fn(sch)
|
||||
return sch
|
||||
|
||||
|
||||
def _create_tmp_database(tmpdir: str, mod_eq: str = "structural") -> ms.database.JSONDatabase:
|
||||
path_workload = osp.join(tmpdir, "workloads.json")
|
||||
path_tuning_record = osp.join(tmpdir, "tuning_records.json")
|
||||
return ms.database.JSONDatabase(path_workload, path_tuning_record, module_equality=mod_eq)
|
||||
|
||||
|
||||
def _equal_record(a: ms.database.TuningRecord, b: ms.database.TuningRecord):
|
||||
assert str(a.trace) == str(b.trace)
|
||||
assert str(a.run_secs) == str(b.run_secs)
|
||||
# AWAIT(@zxybazh): change to export after fixing "(bool)0"
|
||||
assert str(a.target) == str(b.target)
|
||||
tvm.ir.assert_structural_equal(a.workload.mod, b.workload.mod)
|
||||
for arg0, arg1 in zip(a.args_info, b.args_info):
|
||||
assert str(arg0.as_json()) == str(arg1.as_json())
|
||||
|
||||
|
||||
@derived_object
|
||||
class PyMemoryDatabaseDefault(ms.database.PyDatabase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.tuning_records_: list[TuningRecord] = []
|
||||
self.workloads_: list[Workload] = []
|
||||
|
||||
def has_workload(self, mod: IRModule) -> bool:
|
||||
for workload in self.workloads_:
|
||||
if tvm_ffi.structural_equal(mod, workload.mod):
|
||||
return True
|
||||
|
||||
def commit_workload(self, mod: IRModule) -> ms.database.Workload:
|
||||
if self.has_workload(mod):
|
||||
for workload in self.workloads_:
|
||||
if tvm_ffi.structural_equal(mod, workload.mod):
|
||||
return workload
|
||||
else:
|
||||
workload = ms.database.Workload(mod)
|
||||
self.workloads_.append(workload)
|
||||
return workload
|
||||
|
||||
def commit_tuning_record(self, record: TuningRecord) -> None:
|
||||
self.tuning_records_.append(record)
|
||||
|
||||
def get_all_tuning_records(self) -> list[TuningRecord]:
|
||||
return self.tuning_records_
|
||||
|
||||
def get_top_k(self, workload: ms.database.Workload, top_k: int) -> list[TuningRecord]:
|
||||
return sorted(
|
||||
list(
|
||||
filter(
|
||||
lambda x: tvm_ffi.structural_equal(workload.mod, x.workload.mod),
|
||||
self.tuning_records_,
|
||||
)
|
||||
),
|
||||
key=lambda x: sum(x.run_secs) / len(x.run_secs) if x.run_secs else 1e9,
|
||||
)[:top_k]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.tuning_records_)
|
||||
|
||||
|
||||
@derived_object
|
||||
class PyMemoryDatabaseOverride(ms.database.PyDatabase):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.tuning_records_: list[TuningRecord] = []
|
||||
self.workloads_: list[Workload] = []
|
||||
|
||||
def has_workload(self, mod: IRModule) -> bool:
|
||||
for workload in self.workloads_:
|
||||
if tvm_ffi.structural_equal(mod, workload.mod):
|
||||
return True
|
||||
|
||||
def commit_workload(self, mod: IRModule) -> ms.database.Workload:
|
||||
if self.has_workload(mod):
|
||||
for workload in self.workloads_:
|
||||
if tvm_ffi.structural_equal(mod, workload.mod):
|
||||
return workload
|
||||
else:
|
||||
workload = ms.database.Workload(mod)
|
||||
self.workloads_.append(workload)
|
||||
return workload
|
||||
|
||||
def commit_tuning_record(self, record: TuningRecord) -> None:
|
||||
self.tuning_records_.append(record)
|
||||
|
||||
def get_all_tuning_records(self) -> list[TuningRecord]:
|
||||
return self.tuning_records_
|
||||
|
||||
def get_top_k(self, workload: ms.database.Workload, top_k: int) -> list[TuningRecord]:
|
||||
return sorted(
|
||||
list(
|
||||
filter(
|
||||
lambda x: tvm_ffi.structural_equal(workload.mod, x.workload.mod),
|
||||
self.tuning_records_,
|
||||
)
|
||||
),
|
||||
key=lambda x: sum(x.run_secs) / len(x.run_secs) if x.run_secs else 1e9,
|
||||
)[:top_k]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.tuning_records_)
|
||||
|
||||
def query_tuning_record(
|
||||
self, mod: IRModule, target: Target, workload_name: str | None = None
|
||||
) -> TuningRecord | None:
|
||||
if self.has_workload(mod):
|
||||
records = self.get_top_k(self.commit_workload(mod), 2)
|
||||
if len(records) == 1:
|
||||
return records[0]
|
||||
elif len(records) == 2:
|
||||
return records[1] # return the 2nd best if there are two records
|
||||
return None
|
||||
|
||||
def query_schedule(
|
||||
self, mod: IRModule, target: Target, workload_name: str | None = None
|
||||
) -> Schedule | None:
|
||||
record = self.query_tuning_record(mod, target, workload_name)
|
||||
if record is not None:
|
||||
sch = Schedule(record.workload.mod)
|
||||
record.trace.apply_to_schedule(sch, remove_postproc=False)
|
||||
return sch
|
||||
return None
|
||||
|
||||
def query_ir_module(
|
||||
self, mod: IRModule, target: Target, workload_name: str | None = None
|
||||
) -> IRModule | None:
|
||||
record = self.query_tuning_record(mod, target, workload_name)
|
||||
if record is not None:
|
||||
sch = Schedule(record.workload.mod)
|
||||
record.trace.apply_to_schedule(sch, remove_postproc=False)
|
||||
return sch.mod
|
||||
return None
|
||||
|
||||
|
||||
def test_meta_schedule_tuning_record_round_trip():
|
||||
mod: IRModule = Matmul
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
database = _create_tmp_database(tmpdir)
|
||||
workload = database.commit_workload(mod)
|
||||
record = ms.database.TuningRecord(
|
||||
_create_schedule(mod, _schedule_matmul).trace,
|
||||
workload,
|
||||
[T.float32(1.5), T.float32(2.5), T.float32(1.8)],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
)
|
||||
database.commit_tuning_record(record)
|
||||
new_record = ms.database.TuningRecord.from_json(record.as_json(), workload)
|
||||
_equal_record(record, new_record)
|
||||
|
||||
|
||||
def test_meta_schedule_database_create():
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
database = _create_tmp_database(tmpdir)
|
||||
assert osp.exists(database.path_workload)
|
||||
assert osp.exists(database.path_tuning_record)
|
||||
|
||||
|
||||
def test_meta_schedule_database_has_workload():
|
||||
mod: IRModule = Matmul
|
||||
missing_mod: IRModule = MatmulRelu
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
database = _create_tmp_database(tmpdir)
|
||||
workload = database.commit_workload(mod)
|
||||
record = ms.database.TuningRecord(
|
||||
_create_schedule(mod, _schedule_matmul).trace,
|
||||
workload,
|
||||
[1.5, 2.5, 1.8],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
)
|
||||
database.commit_tuning_record(record)
|
||||
assert len(database) == 1
|
||||
assert database.has_workload(mod)
|
||||
assert not database.has_workload(missing_mod)
|
||||
|
||||
|
||||
def test_meta_schedule_database_add_entry():
|
||||
mod: IRModule = Matmul
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
database = _create_tmp_database(tmpdir)
|
||||
workload = database.commit_workload(mod)
|
||||
record = ms.database.TuningRecord(
|
||||
_create_schedule(mod, _schedule_matmul).trace,
|
||||
workload,
|
||||
[1.5, 2.5, 1.8],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
)
|
||||
database.commit_tuning_record(record)
|
||||
assert len(database) == 1
|
||||
(ret,) = database.get_top_k(workload, 3)
|
||||
_equal_record(ret, record)
|
||||
|
||||
|
||||
def test_meta_schedule_database_missing():
|
||||
mod: IRModule = Matmul
|
||||
mod_2: IRModule = MatmulRelu
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
database = _create_tmp_database(tmpdir)
|
||||
workload = database.commit_workload(mod)
|
||||
workload_2 = database.commit_workload(mod_2)
|
||||
record = ms.database.TuningRecord(
|
||||
_create_schedule(mod, _schedule_matmul).trace,
|
||||
workload,
|
||||
[1.5, 2.5, 1.8],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
)
|
||||
database.commit_tuning_record(record)
|
||||
ret = database.get_top_k(workload_2, 3)
|
||||
assert len(ret) == 0
|
||||
|
||||
|
||||
def test_meta_schedule_database_sorting():
|
||||
mod: IRModule = Matmul
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
database = _create_tmp_database(tmpdir)
|
||||
token = database.commit_workload(mod)
|
||||
trace = _create_schedule(mod, _schedule_matmul).trace
|
||||
records = [
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
token,
|
||||
[7.0, 8.0, 9.0],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
),
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
token,
|
||||
[1.0, 2.0, 3.0],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
),
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
token,
|
||||
[4.0, 5.0, 6.0],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
),
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
token,
|
||||
[1.1, 1.2, 600.0],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
),
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
token,
|
||||
[1.0, 100.0, 6.0],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
),
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
token,
|
||||
[4.0, 9.0, 8.0],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
),
|
||||
]
|
||||
for record in records:
|
||||
database.commit_tuning_record(record)
|
||||
ret = database.get_top_k(token, 2)
|
||||
assert len(ret) == 2
|
||||
try:
|
||||
_equal_record(ret[0], records[2])
|
||||
_equal_record(ret[1], records[1])
|
||||
except AssertionError:
|
||||
_equal_record(ret[0], records[1])
|
||||
_equal_record(ret[1], records[2])
|
||||
|
||||
|
||||
def test_meta_schedule_database_reload():
|
||||
mod: IRModule = Matmul
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
database = _create_tmp_database(tmpdir)
|
||||
token = database.commit_workload(mod)
|
||||
trace = _create_schedule(mod, _schedule_matmul).trace
|
||||
records = [
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
token,
|
||||
[7.0, 8.0, 9.0],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
),
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
token,
|
||||
[1.0, 2.0, 3.0],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
),
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
token,
|
||||
[4.0, 5.0, 6.0],
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
),
|
||||
]
|
||||
for record in records:
|
||||
database.commit_tuning_record(record)
|
||||
new_database = ms.database.JSONDatabase(
|
||||
path_workload=database.path_workload,
|
||||
path_tuning_record=database.path_tuning_record,
|
||||
)
|
||||
token = new_database.commit_workload(mod)
|
||||
ret = new_database.get_top_k(token, 2)
|
||||
assert len(ret) == 2
|
||||
try:
|
||||
_equal_record(ret[0], records[2])
|
||||
_equal_record(ret[1], records[1])
|
||||
except AssertionError:
|
||||
_equal_record(ret[0], records[1])
|
||||
_equal_record(ret[1], records[2])
|
||||
|
||||
|
||||
def test_meta_schedule_database_union():
|
||||
mod: IRModule = Matmul
|
||||
target = tvm.target.Target("llvm")
|
||||
arg_info = ms.arg_info.ArgInfo.from_prim_func(func=mod["main"])
|
||||
db_1 = ms.database.MemoryDatabase()
|
||||
db_2 = ms.database.MemoryDatabase()
|
||||
trace = _create_schedule(mod, _schedule_matmul).trace
|
||||
|
||||
def query(db): # pylint: disable=invalid-name
|
||||
return db.query_tuning_record(mod=mod, target=target, workload_name="main").run_secs
|
||||
|
||||
def commit_record(db, run_sec): # pylint: disable=invalid-name
|
||||
db.commit_tuning_record(
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
workload=db.commit_workload(mod),
|
||||
run_secs=[run_sec],
|
||||
target=target,
|
||||
args_info=arg_info,
|
||||
)
|
||||
)
|
||||
|
||||
commit_record(db_1, 1.0)
|
||||
(run_sec,) = query(db_1)
|
||||
assert run_sec.value == 1.0
|
||||
|
||||
commit_record(db_2, 0.5)
|
||||
(run_sec,) = query(db_2)
|
||||
assert run_sec.value == 0.5
|
||||
|
||||
(run_secs,) = query(ms.database.UnionDatabase(db_1, db_2))
|
||||
assert run_secs.value == 0.5
|
||||
|
||||
(run_secs,) = query(ms.database.OrderedUnionDatabase(db_1, db_2))
|
||||
assert run_secs.value == 1.0
|
||||
|
||||
|
||||
def test_meta_schedule_pydatabase_default_query():
|
||||
mod: IRModule = Matmul
|
||||
target = tvm.target.Target("llvm")
|
||||
arg_info = ms.arg_info.ArgInfo.from_prim_func(func=mod["main"])
|
||||
db = PyMemoryDatabaseDefault() # pylint: disable=invalid-name
|
||||
sch = _create_schedule(mod, _schedule_matmul)
|
||||
trace = sch.trace
|
||||
|
||||
def query(db, mod, target, kind): # pylint: disable=invalid-name
|
||||
return db.query(mod=mod, target=target, workload_name="main", kind=kind)
|
||||
|
||||
def commit_record(trace, db, run_sec): # pylint: disable=invalid-name
|
||||
db.commit_tuning_record(
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
workload=db.commit_workload(mod),
|
||||
run_secs=[run_sec],
|
||||
target=target,
|
||||
args_info=arg_info,
|
||||
)
|
||||
)
|
||||
|
||||
commit_record(trace, db, 1.0)
|
||||
record = query(db, mod, target, "record")
|
||||
assert record is not None and record.run_secs[0].value == 1.0
|
||||
sch_res = query(db, mod, target, "schedule")
|
||||
assert sch_res is not None and tvm_ffi.structural_equal(sch_res.mod, sch.mod)
|
||||
mod_res = query(db, mod, target, "ir_module")
|
||||
assert mod_res is not None and tvm_ffi.structural_equal(mod_res, sch.mod)
|
||||
|
||||
commit_record(Schedule(mod).trace, db, 0.2) # Empty Trace
|
||||
record = query(db, mod, target, "record")
|
||||
assert record is not None and record.run_secs[0].value == 0.2
|
||||
sch_res = query(db, mod, target, "schedule")
|
||||
assert sch_res is not None and tvm_ffi.structural_equal(sch_res.mod, mod)
|
||||
mod_res = query(db, mod, target, "ir_module")
|
||||
assert mod_res is not None and tvm_ffi.structural_equal(mod_res, mod)
|
||||
|
||||
|
||||
def test_meta_schedule_pydatabase_override_query():
|
||||
mod: IRModule = Matmul
|
||||
target = tvm.target.Target("llvm")
|
||||
arg_info = ms.arg_info.ArgInfo.from_prim_func(func=mod["main"])
|
||||
db = PyMemoryDatabaseOverride() # pylint: disable=invalid-name
|
||||
sch = _create_schedule(mod, _schedule_matmul)
|
||||
trace = sch.trace
|
||||
|
||||
def query(db, mod, target, kind): # pylint: disable=invalid-name
|
||||
return db.query(mod=mod, target=target, workload_name="main", kind=kind)
|
||||
|
||||
def commit_record(trace, db, run_sec): # pylint: disable=invalid-name
|
||||
db.commit_tuning_record(
|
||||
ms.database.TuningRecord(
|
||||
trace,
|
||||
workload=db.commit_workload(mod),
|
||||
run_secs=[run_sec],
|
||||
target=target,
|
||||
args_info=arg_info,
|
||||
)
|
||||
)
|
||||
|
||||
commit_record(trace, db, 1.14)
|
||||
record = query(db, mod, target, "record")
|
||||
assert record is not None and record.run_secs[0].value == 1.14
|
||||
sch_res = query(db, mod, target, "schedule")
|
||||
assert sch_res is not None and tvm_ffi.structural_equal(sch_res.mod, sch.mod)
|
||||
mod_res = query(db, mod, target, "ir_module")
|
||||
assert mod_res is not None and tvm_ffi.structural_equal(mod_res, sch.mod)
|
||||
|
||||
commit_record(Schedule(mod).trace, db, 0.514) # Empty Trace
|
||||
record = query(db, mod, target, "record")
|
||||
assert record is not None and record.run_secs[0].value == 1.14 # Override to 2nd best
|
||||
sch_res = query(db, mod, target, "schedule")
|
||||
assert sch_res is not None and tvm_ffi.structural_equal(sch_res.mod, sch.mod)
|
||||
mod_res = query(db, mod, target, "ir_module")
|
||||
assert mod_res is not None and tvm_ffi.structural_equal(mod_res, sch.mod)
|
||||
|
||||
|
||||
def test_meta_schedule_pydatabase_current():
|
||||
db = PyMemoryDatabaseDefault() # pylint: disable=invalid-name
|
||||
with db: # pylint: disable=not-context-manager
|
||||
assert ms.database.Database.current() == db
|
||||
|
||||
|
||||
def call_get_top_k(run_secs_list, database, k):
|
||||
mod: IRModule = Matmul
|
||||
workload = database.commit_workload(mod)
|
||||
for run_secs in run_secs_list:
|
||||
record = ms.database.TuningRecord(
|
||||
_create_schedule(mod, _schedule_matmul).trace,
|
||||
workload,
|
||||
run_secs,
|
||||
tvm.target.Target("llvm"),
|
||||
ms.arg_info.ArgInfo.from_prim_func(func=mod["main"]),
|
||||
)
|
||||
database.commit_tuning_record(record)
|
||||
return [[v.value for v in record.run_secs] for record in database.get_top_k(workload, k)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"k,expected",
|
||||
[
|
||||
(0, []),
|
||||
(1, [[0.0, 2.0]]),
|
||||
(4, [[0.0, 2.0], [2.0], [1.5, 4.5], [3.0, 1e10]]),
|
||||
(5, [[0.0, 2.0], [2.0], [1.5, 4.5], [3.0, 1e10]]),
|
||||
],
|
||||
)
|
||||
def test_memory_database_get_top_k(k, expected):
|
||||
run_secs_list = [[1.5, 4.5], [], [0.0, 2.0], None, [2.0], [3.0, 1e10], [1e10]]
|
||||
database = ms.database.MemoryDatabase()
|
||||
result = call_get_top_k(run_secs_list, database, k)
|
||||
assert result == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"k,expected",
|
||||
[
|
||||
(0, []),
|
||||
(4, [[0.0, 2.0], [2.0], [1.5, 4.5], [3.0, 1e10]]),
|
||||
(5, [[0.0, 2.0], [2.0], [1.5, 4.5], [3.0, 1e10]]),
|
||||
],
|
||||
)
|
||||
def test_json_database_get_top_k(k, expected):
|
||||
run_secs_list = [[1.5, 4.5], [], [0.0, 2.0], None, [2.0], [3.0, 1e10], [1e10]]
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
database = _create_tmp_database(tmpdir)
|
||||
result = call_get_top_k(run_secs_list, database, k)
|
||||
assert result == expected
|
||||
|
||||
|
||||
def MatmulPrimFunc() -> IRModule:
|
||||
return Matmul
|
||||
|
||||
|
||||
@pytest.mark.parametrize("f_mod", [MatmulPrimFunc])
|
||||
@pytest.mark.parametrize("mod_eq", ["structural", "ignore-tensor", "anchor-block"])
|
||||
def test_json_database_commit_workload(f_mod, mod_eq):
|
||||
mod: IRModule = f_mod()
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
database = _create_tmp_database(tmpdir, mod_eq)
|
||||
database.commit_workload(mod)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("f_mod", [MatmulPrimFunc])
|
||||
@pytest.mark.parametrize("mod_eq", ["structural", "ignore-tensor", "anchor-block"])
|
||||
def test_memory_database_commit_workload(f_mod, mod_eq):
|
||||
mod: IRModule = f_mod()
|
||||
database = ms.database.MemoryDatabase(module_equality=mod_eq)
|
||||
database.commit_workload(mod)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,45 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
|
||||
import numpy as np
|
||||
|
||||
import tvm.runtime
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.s_tir.meta_schedule import TuneContext
|
||||
from tvm.s_tir.meta_schedule.feature_extractor import PyFeatureExtractor
|
||||
from tvm.s_tir.meta_schedule.search_strategy import MeasureCandidate
|
||||
|
||||
|
||||
def test_meta_schedule_feature_extractor():
|
||||
@derived_object
|
||||
class FancyFeatureExtractor(PyFeatureExtractor):
|
||||
def extract_from(
|
||||
self,
|
||||
context: TuneContext, # pylint: disable = unused-argument
|
||||
candidates: list[MeasureCandidate], # pylint: disable = unused-argument
|
||||
) -> list[np.ndarray]:
|
||||
return [tvm.runtime.tensor(np.random.rand(4, 5))]
|
||||
|
||||
extractor = FancyFeatureExtractor()
|
||||
features = extractor.extract_from(TuneContext(), [])
|
||||
assert len(features) == 1
|
||||
assert features[0].shape == (4, 5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_meta_schedule_feature_extractor()
|
||||
+1645
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,164 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.schedule import Schedule
|
||||
from tvm.script import tirx as T
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,
|
||||
# fmt: off
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
# fmt: on
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
|
||||
|
||||
def test_meta_schedule_measure_callback():
|
||||
@derived_object
|
||||
class FancyMeasureCallback(ms.measure_callback.PyMeasureCallback):
|
||||
def apply(
|
||||
self,
|
||||
task_scheduler: ms.task_scheduler.TaskScheduler,
|
||||
task_id: int,
|
||||
measure_candidates: list[ms.MeasureCandidate],
|
||||
builder_results: list[ms.builder.BuilderResult],
|
||||
runner_results: list[ms.runner.RunnerResult],
|
||||
) -> None:
|
||||
assert len(measure_candidates) == 1
|
||||
tvm.ir.assert_structural_equal(measure_candidates[0].sch.mod, Matmul)
|
||||
assert (
|
||||
len(builder_results) == 1
|
||||
and builder_results[0].error_msg is None
|
||||
and builder_results[0].artifact_path == "test_build"
|
||||
)
|
||||
assert (
|
||||
len(runner_results) == 1
|
||||
and runner_results[0].error_msg is None
|
||||
and len(runner_results[0].run_secs) == 2
|
||||
)
|
||||
|
||||
measure_callback = FancyMeasureCallback()
|
||||
measure_callback.apply(
|
||||
ms.task_scheduler.RoundRobin(),
|
||||
0,
|
||||
[ms.MeasureCandidate(Schedule(Matmul), None)],
|
||||
[ms.builder.BuilderResult("test_build", None)],
|
||||
[ms.runner.RunnerResult([1.0, 2.1], None)],
|
||||
)
|
||||
|
||||
|
||||
def test_meta_schedule_measure_callback_fail():
|
||||
@derived_object
|
||||
class FailingMeasureCallback(ms.measure_callback.PyMeasureCallback):
|
||||
def apply(
|
||||
self,
|
||||
task_scheduler: ms.task_scheduler.TaskScheduler,
|
||||
task_id: int,
|
||||
measure_candidates: list[ms.MeasureCandidate],
|
||||
builder_results: list[ms.builder.BuilderResult],
|
||||
runner_results: list[ms.runner.RunnerResult],
|
||||
) -> None:
|
||||
raise ValueError("test")
|
||||
|
||||
measure_callback = FailingMeasureCallback()
|
||||
with pytest.raises(ValueError, match="test"):
|
||||
measure_callback.apply(
|
||||
ms.task_scheduler.RoundRobin(),
|
||||
0,
|
||||
[ms.MeasureCandidate(Schedule(Matmul), None)],
|
||||
[ms.builder.BuilderResult("test_build", None)],
|
||||
[ms.runner.RunnerResult([1.0, 2.1], None)],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_measure_callback_update_cost_model_with_zero():
|
||||
@derived_object
|
||||
class AllZeroRunnerFuture(ms.runner.PyRunnerFuture):
|
||||
def done(self) -> bool:
|
||||
return True
|
||||
|
||||
def result(self) -> ms.runner.RunnerResult:
|
||||
return ms.runner.RunnerResult([0.0, 0.0], None)
|
||||
|
||||
@derived_object
|
||||
class AllZeroRunner(ms.runner.PyRunner):
|
||||
def run(self, runner_inputs: list[ms.runner.RunnerInput]) -> list[ms.runner.RunnerResult]:
|
||||
return [AllZeroRunnerFuture() for _ in runner_inputs]
|
||||
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
ms.tune_tir(
|
||||
mod=Matmul,
|
||||
target={"kind": "llvm", "num-cores": 1},
|
||||
work_dir=work_dir,
|
||||
max_trials_global=10,
|
||||
runner=AllZeroRunner(),
|
||||
measure_callbacks=[ms.measure_callback.UpdateCostModel()],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_measure_callback_update_cost_model_with_runtime_error():
|
||||
@derived_object
|
||||
class EmptyRunnerFuture(ms.runner.PyRunnerFuture):
|
||||
def done(self) -> bool:
|
||||
return True
|
||||
|
||||
def result(self) -> ms.runner.RunnerResult:
|
||||
return ms.runner.RunnerResult(None, "error")
|
||||
|
||||
@derived_object
|
||||
class EmptyRunner(ms.runner.PyRunner):
|
||||
def run(self, runner_inputs: list[ms.runner.RunnerInput]) -> list[ms.runner.RunnerResult]:
|
||||
return [EmptyRunnerFuture() for _ in runner_inputs]
|
||||
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
ms.tune_tir(
|
||||
mod=Matmul,
|
||||
target={"kind": "llvm", "num-cores": 1},
|
||||
work_dir=work_dir,
|
||||
max_trials_global=10,
|
||||
runner=EmptyRunner(),
|
||||
measure_callbacks=[ms.measure_callback.UpdateCostModel()],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_meta_schedule_measure_callback()
|
||||
test_meta_schedule_measure_callback_fail()
|
||||
test_meta_schedule_measure_callback_update_cost_model_with_zero()
|
||||
test_meta_schedule_measure_callback_update_cost_model_with_runtime_error()
|
||||
@@ -0,0 +1,342 @@
|
||||
# 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.
|
||||
# ruff: noqa: E501, F841
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.s_tir.tensor_intrin # pylint: disable=unused-import
|
||||
import tvm.testing
|
||||
from tvm.s_tir.schedule import Schedule
|
||||
from tvm.script import tirx as T
|
||||
from tvm.testing import env
|
||||
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
M, N, K = 4096, 4096, 4096
|
||||
np.random.seed(0)
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Gemm_F16F16F16:
|
||||
# fmt: off
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
A: T.Buffer((M, K), "float16"), # type: ignore
|
||||
B: T.Buffer((K, N), "float16"), # type: ignore
|
||||
C: T.Buffer((M, N), "float16"), # type: ignore
|
||||
):
|
||||
for i, j, k in T.grid(M, N, K):
|
||||
with T.sblock("C"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Gemm_F16F16F32:
|
||||
# fmt: off
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
A: T.Buffer((M, K), "float16"), # type: ignore
|
||||
B: T.Buffer((K, N), "float16"), # type: ignore
|
||||
C: T.Buffer((M, N), "float32"), # type: ignore
|
||||
):
|
||||
for i, j, k in T.grid(M, N, K):
|
||||
with T.sblock("C"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + T.cast(A[vi, vk], "float32") * T.cast(B[vk, vj], "float32")
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_run_target(mod=None, tgt_str=None, in_dtype="float16", out_dtype="float16"):
|
||||
if mod is None:
|
||||
return
|
||||
tgt_str = tgt_str or "cuda"
|
||||
target = tvm.target.Target(target=tgt_str)
|
||||
with tvm.transform.PassContext(opt_level=3):
|
||||
lib: tvm.runtime.Module = tvm.compile(mod, target=target)
|
||||
|
||||
a_np = np.random.rand(M, K).astype(in_dtype)
|
||||
b_np = np.random.rand(K, N).astype(in_dtype)
|
||||
c_np = np.ones((M, N), dtype=out_dtype)
|
||||
f = lib["main"]
|
||||
|
||||
def run_and_check():
|
||||
dev = tvm.device(tgt_str, 0)
|
||||
a = tvm.runtime.tensor(a_np, dev)
|
||||
b = tvm.runtime.tensor(b_np, dev)
|
||||
c = tvm.runtime.tensor(c_np, dev)
|
||||
f(a, b, c)
|
||||
c_th = torch.matmul(torch.tensor(a_np).to(tgt_str), torch.tensor(b_np).to(tgt_str)).to(
|
||||
torch.float32 if out_dtype == "float32" else torch.float16
|
||||
)
|
||||
c_f = torch.tensor(c.numpy()).to(tgt_str)
|
||||
assert torch.allclose(c_th, c_f, rtol=0.05, atol=0.05)
|
||||
|
||||
tvm.testing.run_with_gpu_lock(run_and_check)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_f16f16f16_mma_gemm():
|
||||
# fmt: off
|
||||
mod = Gemm_F16F16F16
|
||||
sch = Schedule(mod)
|
||||
b0 = sch.get_sblock(name="C", func_name="main")
|
||||
b1 = sch.get_sblock(name="root", func_name="main")
|
||||
sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")
|
||||
b2 = sch.reindex(block=b0, buffer=("write", 0))
|
||||
b3 = sch.reindex(block=b0, buffer=("read", 0))
|
||||
b4 = sch.reindex(block=b0, buffer=("read", 1))
|
||||
sch.transform_layout(block=b0, buffer=("read", 0), index_map=lambda vi, vk: (vi, vk,), pad_value=None, assume_injective_transform=True)
|
||||
sch.transform_layout(block=b0, buffer=("read", 1), index_map=lambda vj, vk: (vk, vj,), pad_value=None, assume_injective_transform=True)
|
||||
sch.transform_layout(block=b0, buffer=("write", 0), index_map=lambda vi, vj: (vi, vj,), pad_value=None, assume_injective_transform=True)
|
||||
sch.transform_block_layout(block=b2, index_map=lambda vi, vj: (vi, vj,))
|
||||
sch.transform_block_layout(block=b3, index_map=lambda vi, vk: (vi, vk,))
|
||||
sch.transform_block_layout(block=b4, index_map=lambda vj, vk: (vk, vj,))
|
||||
sch.transform_block_layout(block=b0, index_map=lambda vi, vj, vk: (vi, vj, vk,))
|
||||
l5, l6, l7 = sch.get_loops(block=b0)
|
||||
l8, l9 = sch.split(loop=l7, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
|
||||
l10, l11 = sch.split(loop=l6, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
|
||||
l12, l13 = sch.split(loop=l5, factors=[None, 16], preserve_unit_iters=True, disable_predication=False)
|
||||
l14, l15, l16, l17, l18, l19 = sch.get_loops(block=b0)
|
||||
sch.reorder(l16, l18, l13, l11, l9)
|
||||
b20 = sch.blockize(target=l13, preserve_unit_iters=True)
|
||||
sch.annotate(block_or_loop=b20, ann_key="meta_schedule.auto_tensorize", ann_val="mma_sync_m16n8k8_f16f16f16")
|
||||
sch.annotate(block_or_loop=b20, ann_key="meta_schedule.auto_tensorize_init", ann_val="mma_init_m16n8k8_f16")
|
||||
sch.annotate(block_or_loop=b20, ann_key="warp_execution", ann_val=1)
|
||||
l21, l22, l23 = sch.get_loops(block=b20)
|
||||
v24, v25, v26, v27, v28 = sch.sample_partitioned_tile(loop=l21, n=5, partition_pos=3, innerpart_factor=2, decision=[2, 16, 4, 1, 2])
|
||||
l29, l30, l31, l32, l33 = sch.split(loop=l21, factors=[v24, v25, v26, v27, v28], preserve_unit_iters=True, disable_predication=False)
|
||||
v34, v35, v36, v37, v38 = sch.sample_partitioned_tile(loop=l22, n=5, partition_pos=3, innerpart_factor=4, decision=[2, 16, 4, 1, 4])
|
||||
l39, l40, l41, l42, l43 = sch.split(loop=l22, factors=[v34, v35, v36, v37, v38], preserve_unit_iters=True, disable_predication=False)
|
||||
v44, v45, v46 = sch.sample_perfect_tile(loop=l23, n=3, max_innermost_factor=4, decision=[128, 1, 4])
|
||||
l47, l48, l49 = sch.split(loop=l23, factors=[v44, v45, v46], preserve_unit_iters=True, disable_predication=False)
|
||||
sch.reorder(l29, l39, l30, l40, l31, l41, l47, l48, l32, l42, l49, l33, l43)
|
||||
l50 = sch.fuse(l29, l39, preserve_unit_iters=True)
|
||||
sch.bind(loop=l50, thread_axis="blockIdx.y")
|
||||
l51 = sch.fuse(l30, l40, preserve_unit_iters=True)
|
||||
sch.bind(loop=l51, thread_axis="blockIdx.x")
|
||||
l52 = sch.fuse(l31, l41, preserve_unit_iters=True)
|
||||
sch.bind(loop=l52, thread_axis="threadIdx.y")
|
||||
sch.annotate(block_or_loop=b20, ann_key="meta_schedule.thread_extent_low_inclusive", ann_val=32)
|
||||
sch.annotate(block_or_loop=b20, ann_key="meta_schedule.thread_extent_high_inclusive", ann_val=1024)
|
||||
b53 = sch.write_at(loop=l52, block=b20, write_buffer_index=0, storage_scope="m16n8k8.matrixC")
|
||||
sch.reverse_compute_inline(block=b2)
|
||||
b54 = sch.read_at(loop=l47, block=b20, read_buffer_index=0, storage_scope="shared.dyn")
|
||||
sch.annotate(block_or_loop=b54, ann_key="permuted_layout", ann_val="g2s_A")
|
||||
b55 = sch.read_at(loop=l47, block=b20, read_buffer_index=1, storage_scope="shared.dyn")
|
||||
sch.annotate(block_or_loop=b55, ann_key="permuted_layout", ann_val="g2s_B")
|
||||
b56 = sch.cache_read(block=b20, read_buffer_index=0, storage_scope="m16n8k8.matrixA")
|
||||
sch.compute_at(block=b56, loop=l48, preserve_unit_loops=True, index=-1)
|
||||
l57, l58, l59, l60, l61, l62, l63 = sch.get_loops(block=b56)
|
||||
l64, l65 = sch.split(loop=l63, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
|
||||
l66, l67 = sch.split(loop=l62, factors=[None, 32], preserve_unit_iters=True, disable_predication=False)
|
||||
l68, l69, l70, l71, l72, l73, l74, l75, l76 = sch.get_loops(block=b56)
|
||||
sch.reorder(l75, l67, l65)
|
||||
b77 = sch.blockize(target=l67, preserve_unit_iters=True)
|
||||
sch.annotate(block_or_loop=b77, ann_key="meta_schedule.auto_tensorize", ann_val="mma_load_m16n8k8_f16_A_shared_dyn")
|
||||
sch.annotate(block_or_loop=b77, ann_key="permuted_layout", ann_val="s2l_A")
|
||||
b78 = sch.cache_read(block=b20, read_buffer_index=1, storage_scope="m16n8k8.matrixB")
|
||||
sch.compute_at(block=b78, loop=l48, preserve_unit_loops=True, index=-1)
|
||||
l79, l80, l81, l82, l83, l84, l85 = sch.get_loops(block=b78)
|
||||
l86, l87 = sch.split(loop=l85, factors=[None, 32], preserve_unit_iters=True, disable_predication=False)
|
||||
l88, l89 = sch.split(loop=l84, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
|
||||
l90, l91, l92, l93, l94, l95, l96, l97, l98 = sch.get_loops(block=b78)
|
||||
sch.reorder(l97, l89, l87)
|
||||
b99 = sch.blockize(target=l89, preserve_unit_iters=True)
|
||||
sch.annotate(block_or_loop=b99, ann_key="meta_schedule.auto_tensorize", ann_val="mma_load_m16n8k8_f16_B_shared_dyn")
|
||||
sch.annotate(block_or_loop=b99, ann_key="permuted_layout", ann_val="s2l_B")
|
||||
b100, = sch.get_producers(block=b54)
|
||||
sch.compute_inline(block=b100)
|
||||
sch.storage_align(block=b54, buffer_index=0, axis=-2, factor=32, offset=8)
|
||||
b101, = sch.get_producers(block=b55)
|
||||
sch.compute_inline(block=b101)
|
||||
sch.storage_align(block=b55, buffer_index=0, axis=-2, factor=32, offset=8)
|
||||
sch.annotate(block_or_loop=b54, ann_key="vector_bytes", ann_val=16)
|
||||
sch.annotate(block_or_loop=b55, ann_key="vector_bytes", ann_val=16)
|
||||
sch.annotate(block_or_loop=l48, ann_key="software_pipeline_stage", ann_val=[0, 0, 1])
|
||||
sch.annotate(block_or_loop=l48, ann_key="software_pipeline_order", ann_val=[0, 1, 2])
|
||||
sch.annotate(block_or_loop=l47, ann_key="software_pipeline_async_stages", ann_val=[0])
|
||||
sch.annotate(block_or_loop=l47, ann_key="software_pipeline_stage", ann_val=[0, 0, 1, 2, 2])
|
||||
sch.annotate(block_or_loop=l47, ann_key="software_pipeline_order", ann_val=[0, 1, 3, 2, 4])
|
||||
v102 = sch.sample_categorical(candidates=[0, 16, 64, 512, 1024], probs=[0.20000000000000001, 0.20000000000000001, 0.20000000000000001, 0.20000000000000001, 0.20000000000000001], decision=0)
|
||||
sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v102)
|
||||
sch.enter_postproc()
|
||||
b103 = sch.get_sblock(name="root", func_name="main")
|
||||
sch.unannotate(block_or_loop=b103, ann_key="meta_schedule.unroll_explicit")
|
||||
b104, b105, b106, b107, b108, b109 = sch.get_child_blocks(b103)
|
||||
l110, l111, l112, l113 = sch.get_loops(block=b104)
|
||||
l114, l115, l116, l117 = sch.get_loops(block=b105)
|
||||
l118, l119, l120, l121, l122, l123, l124 = sch.get_loops(block=b106)
|
||||
l125, l126, l127, l128, l129, l130, l131 = sch.get_loops(block=b107)
|
||||
l132, l133, l134, l135, l136, l137, l138, l139, l140, l141 = sch.get_loops(block=b108)
|
||||
l142, l143, l144 = sch.get_loops(block=b109)
|
||||
b145 = sch.get_sblock(name="C_o", func_name="main")
|
||||
l146, l147, l148, l149, l150, l151, l152, l153, l154, l155 = sch.get_loops(block=b145)
|
||||
b156 = sch.decompose_reduction(block=b145, loop=l149)
|
||||
sch.unannotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize")
|
||||
sch.annotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize", ann_val="mma_init_m16n8k8_f16")
|
||||
sch.unannotate(block_or_loop=b145, ann_key="meta_schedule.auto_tensorize_init")
|
||||
sch.unannotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize_init")
|
||||
b157 = sch.get_sblock(name="C_o_init", func_name="main")
|
||||
sch.unannotate(block_or_loop=b157, ann_key="meta_schedule.auto_tensorize")
|
||||
sch.tensorize(block_or_loop=b157, tensor_intrin="mma_init_m16n8k8_f16", preserve_unit_iters=True)
|
||||
b158 = sch.get_sblock(name="A_reindex_shared.dyn_m16n8k8.matrixA_o", func_name="main")
|
||||
sch.unannotate(block_or_loop=b158, ann_key="meta_schedule.auto_tensorize")
|
||||
sch.tensorize(block_or_loop=b158, tensor_intrin="mma_load_m16n8k8_f16_A_shared_dyn", preserve_unit_iters=True)
|
||||
b159 = sch.get_sblock(name="B_reindex_shared.dyn_m16n8k8.matrixB_o", func_name="main")
|
||||
sch.unannotate(block_or_loop=b159, ann_key="meta_schedule.auto_tensorize")
|
||||
sch.tensorize(block_or_loop=b159, tensor_intrin="mma_load_m16n8k8_f16_B_shared_dyn", preserve_unit_iters=True)
|
||||
b160 = sch.get_sblock(name="C_o_update", func_name="main")
|
||||
sch.unannotate(block_or_loop=b160, ann_key="meta_schedule.auto_tensorize")
|
||||
sch.tensorize(block_or_loop=b160, tensor_intrin="mma_sync_m16n8k8_f16f16f16", preserve_unit_iters=True)
|
||||
mod = sch.mod
|
||||
test_run_target(mod)
|
||||
|
||||
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_f16f16f32_mma_gemm():
|
||||
mod = Gemm_F16F16F32
|
||||
sch = Schedule(mod)
|
||||
# fmt: off
|
||||
sch = Schedule(mod)
|
||||
b0 = sch.get_sblock(name="C", func_name="main")
|
||||
b1 = sch.get_sblock(name="root", func_name="main")
|
||||
sch.annotate(block_or_loop=b0, ann_key="meta_schedule.tiling_structure", ann_val="SSSRRSRS")
|
||||
b2 = sch.reindex(block=b0, buffer=("write", 0))
|
||||
b3 = sch.reindex(block=b0, buffer=("read", 0))
|
||||
b4 = sch.reindex(block=b0, buffer=("read", 1))
|
||||
sch.transform_layout(block=b0, buffer=("read", 0), index_map=lambda vi, vk: (vi, vk,), pad_value=None, assume_injective_transform=True)
|
||||
sch.transform_layout(block=b0, buffer=("read", 1), index_map=lambda vj, vk: (vk, vj,), pad_value=None, assume_injective_transform=True)
|
||||
sch.transform_layout(block=b0, buffer=("write", 0), index_map=lambda vi, vj: (vi, vj,), pad_value=None, assume_injective_transform=True)
|
||||
sch.transform_block_layout(block=b2, index_map=lambda vi, vj: (vi, vj,))
|
||||
sch.transform_block_layout(block=b3, index_map=lambda vi, vk: (vi, vk,))
|
||||
sch.transform_block_layout(block=b4, index_map=lambda vj, vk: (vk, vj,))
|
||||
sch.transform_block_layout(block=b0, index_map=lambda vi, vj, vk: (vi, vj, vk,))
|
||||
l5, l6, l7 = sch.get_loops(block=b0)
|
||||
l8, l9 = sch.split(loop=l7, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
|
||||
l10, l11 = sch.split(loop=l6, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
|
||||
l12, l13 = sch.split(loop=l5, factors=[None, 16], preserve_unit_iters=True, disable_predication=False)
|
||||
l14, l15, l16, l17, l18, l19 = sch.get_loops(block=b0)
|
||||
sch.reorder(l16, l18, l13, l11, l9)
|
||||
b20 = sch.blockize(target=l13, preserve_unit_iters=True)
|
||||
sch.annotate(block_or_loop=b20, ann_key="meta_schedule.auto_tensorize", ann_val="mma_sync_m16n8k8_f16f16f32")
|
||||
sch.annotate(block_or_loop=b20, ann_key="meta_schedule.auto_tensorize_init", ann_val="mma_init_m16n8k8_f32")
|
||||
sch.annotate(block_or_loop=b20, ann_key="warp_execution", ann_val=1)
|
||||
l21, l22, l23 = sch.get_loops(block=b20)
|
||||
v24, v25, v26, v27, v28 = sch.sample_partitioned_tile(loop=l21, n=5, partition_pos=3, innerpart_factor=2, decision=[1, 16, 2, 2, 4])
|
||||
l29, l30, l31, l32, l33 = sch.split(loop=l21, factors=[v24, v25, v26, v27, v28], preserve_unit_iters=True, disable_predication=False)
|
||||
v34, v35, v36, v37, v38 = sch.sample_partitioned_tile(loop=l22, n=5, partition_pos=3, innerpart_factor=4, decision=[2, 16, 2, 4, 2])
|
||||
l39, l40, l41, l42, l43 = sch.split(loop=l22, factors=[v34, v35, v36, v37, v38], preserve_unit_iters=True, disable_predication=False)
|
||||
v44, v45, v46 = sch.sample_perfect_tile(loop=l23, n=3, max_innermost_factor=4, decision=[128, 1, 4])
|
||||
l47, l48, l49 = sch.split(loop=l23, factors=[v44, v45, v46], preserve_unit_iters=True, disable_predication=False)
|
||||
sch.reorder(l29, l39, l30, l40, l31, l41, l47, l48, l32, l42, l49, l33, l43)
|
||||
l50 = sch.fuse(l29, l39, preserve_unit_iters=True)
|
||||
sch.bind(loop=l50, thread_axis="blockIdx.y")
|
||||
l51 = sch.fuse(l30, l40, preserve_unit_iters=True)
|
||||
sch.bind(loop=l51, thread_axis="blockIdx.x")
|
||||
l52 = sch.fuse(l31, l41, preserve_unit_iters=True)
|
||||
sch.bind(loop=l52, thread_axis="threadIdx.y")
|
||||
sch.annotate(block_or_loop=b20, ann_key="meta_schedule.thread_extent_low_inclusive", ann_val=32)
|
||||
sch.annotate(block_or_loop=b20, ann_key="meta_schedule.thread_extent_high_inclusive", ann_val=1024)
|
||||
b53 = sch.write_at(loop=l52, block=b20, write_buffer_index=0, storage_scope="m16n8k8.matrixC")
|
||||
sch.reverse_compute_inline(block=b2)
|
||||
b54 = sch.read_at(loop=l47, block=b20, read_buffer_index=0, storage_scope="shared.dyn")
|
||||
sch.annotate(block_or_loop=b54, ann_key="permuted_layout", ann_val="g2s_A")
|
||||
b55 = sch.read_at(loop=l47, block=b20, read_buffer_index=1, storage_scope="shared.dyn")
|
||||
sch.annotate(block_or_loop=b55, ann_key="permuted_layout", ann_val="g2s_B")
|
||||
b56 = sch.cache_read(block=b20, read_buffer_index=0, storage_scope="m16n8k8.matrixA")
|
||||
sch.compute_at(block=b56, loop=l48, preserve_unit_loops=True, index=-1)
|
||||
l57, l58, l59, l60, l61, l62, l63 = sch.get_loops(block=b56)
|
||||
l64, l65 = sch.split(loop=l63, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
|
||||
l66, l67 = sch.split(loop=l62, factors=[None, 32], preserve_unit_iters=True, disable_predication=False)
|
||||
l68, l69, l70, l71, l72, l73, l74, l75, l76 = sch.get_loops(block=b56)
|
||||
sch.reorder(l75, l67, l65)
|
||||
b77 = sch.blockize(target=l67, preserve_unit_iters=True)
|
||||
sch.annotate(block_or_loop=b77, ann_key="meta_schedule.auto_tensorize", ann_val="mma_load_m16n8k8_f16_A_shared_dyn")
|
||||
sch.annotate(block_or_loop=b77, ann_key="permuted_layout", ann_val="s2l_A")
|
||||
b78 = sch.cache_read(block=b20, read_buffer_index=1, storage_scope="m16n8k8.matrixB")
|
||||
sch.compute_at(block=b78, loop=l48, preserve_unit_loops=True, index=-1)
|
||||
l79, l80, l81, l82, l83, l84, l85 = sch.get_loops(block=b78)
|
||||
l86, l87 = sch.split(loop=l85, factors=[None, 32], preserve_unit_iters=True, disable_predication=False)
|
||||
l88, l89 = sch.split(loop=l84, factors=[None, 8], preserve_unit_iters=True, disable_predication=False)
|
||||
l90, l91, l92, l93, l94, l95, l96, l97, l98 = sch.get_loops(block=b78)
|
||||
sch.reorder(l97, l89, l87)
|
||||
b99 = sch.blockize(target=l89, preserve_unit_iters=True)
|
||||
sch.annotate(block_or_loop=b99, ann_key="meta_schedule.auto_tensorize", ann_val="mma_load_m16n8k8_f16_B_shared_dyn")
|
||||
sch.annotate(block_or_loop=b99, ann_key="permuted_layout", ann_val="s2l_B")
|
||||
b100, = sch.get_producers(block=b54)
|
||||
sch.compute_inline(block=b100)
|
||||
sch.storage_align(block=b54, buffer_index=0, axis=-2, factor=32, offset=8)
|
||||
b101, = sch.get_producers(block=b55)
|
||||
sch.compute_inline(block=b101)
|
||||
sch.storage_align(block=b55, buffer_index=0, axis=-2, factor=32, offset=8)
|
||||
sch.annotate(block_or_loop=b54, ann_key="vector_bytes", ann_val=16)
|
||||
sch.annotate(block_or_loop=b55, ann_key="vector_bytes", ann_val=16)
|
||||
sch.annotate(block_or_loop=l48, ann_key="software_pipeline_stage", ann_val=[0, 0, 1])
|
||||
sch.annotate(block_or_loop=l48, ann_key="software_pipeline_order", ann_val=[0, 1, 2])
|
||||
sch.annotate(block_or_loop=l47, ann_key="software_pipeline_async_stages", ann_val=[0])
|
||||
sch.annotate(block_or_loop=l47, ann_key="software_pipeline_stage", ann_val=[0, 0, 1, 2, 2])
|
||||
sch.annotate(block_or_loop=l47, ann_key="software_pipeline_order", ann_val=[0, 1, 3, 2, 4])
|
||||
v102 = sch.sample_categorical(candidates=[0, 16, 64, 512, 1024], probs=[0.20000000000000001, 0.20000000000000001, 0.20000000000000001, 0.20000000000000001, 0.20000000000000001], decision=0)
|
||||
sch.annotate(block_or_loop=b1, ann_key="meta_schedule.unroll_explicit", ann_val=v102)
|
||||
sch.enter_postproc()
|
||||
b103 = sch.get_sblock(name="root", func_name="main")
|
||||
sch.unannotate(block_or_loop=b103, ann_key="meta_schedule.unroll_explicit")
|
||||
b104, b105, b106, b107, b108, b109 = sch.get_child_blocks(b103)
|
||||
l110, l111, l112, l113 = sch.get_loops(block=b104)
|
||||
l114, l115, l116, l117 = sch.get_loops(block=b105)
|
||||
l118, l119, l120, l121, l122, l123, l124 = sch.get_loops(block=b106)
|
||||
l125, l126, l127, l128, l129, l130, l131 = sch.get_loops(block=b107)
|
||||
l132, l133, l134, l135, l136, l137, l138, l139, l140, l141 = sch.get_loops(block=b108)
|
||||
sch.annotate(block_or_loop=l132, ann_key="pragma_auto_unroll_max_step", ann_val=0)
|
||||
sch.annotate(block_or_loop=l132, ann_key="pragma_unroll_explicit", ann_val=1)
|
||||
l142, l143, l144 = sch.get_loops(block=b109)
|
||||
b145 = sch.get_sblock(name="C_o", func_name="main")
|
||||
l146, l147, l148, l149, l150, l151, l152, l153, l154, l155 = sch.get_loops(block=b145)
|
||||
b156 = sch.decompose_reduction(block=b145, loop=l149)
|
||||
sch.unannotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize")
|
||||
sch.annotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize", ann_val="mma_init_m16n8k8_f32")
|
||||
sch.unannotate(block_or_loop=b145, ann_key="meta_schedule.auto_tensorize_init")
|
||||
sch.unannotate(block_or_loop=b156, ann_key="meta_schedule.auto_tensorize_init")
|
||||
b157 = sch.get_sblock(name="C_o_init", func_name="main")
|
||||
sch.unannotate(block_or_loop=b157, ann_key="meta_schedule.auto_tensorize")
|
||||
sch.tensorize(block_or_loop=b157, tensor_intrin="mma_init_m16n8k8_f32", preserve_unit_iters=True)
|
||||
b158 = sch.get_sblock(name="A_reindex_shared.dyn_m16n8k8.matrixA_o", func_name="main")
|
||||
sch.unannotate(block_or_loop=b158, ann_key="meta_schedule.auto_tensorize")
|
||||
sch.tensorize(block_or_loop=b158, tensor_intrin="mma_load_m16n8k8_f16_A_shared_dyn", preserve_unit_iters=True)
|
||||
b159 = sch.get_sblock(name="B_reindex_shared.dyn_m16n8k8.matrixB_o", func_name="main")
|
||||
sch.unannotate(block_or_loop=b159, ann_key="meta_schedule.auto_tensorize")
|
||||
sch.tensorize(block_or_loop=b159, tensor_intrin="mma_load_m16n8k8_f16_B_shared_dyn", preserve_unit_iters=True)
|
||||
b160 = sch.get_sblock(name="C_o_update", func_name="main")
|
||||
sch.unannotate(block_or_loop=b160, ann_key="meta_schedule.auto_tensorize")
|
||||
sch.tensorize(block_or_loop=b160, tensor_intrin="mma_sync_m16n8k8_f16f16f32", preserve_unit_iters=True)
|
||||
mod = sch.mod
|
||||
test_run_target(mod, out_dtype="float32")
|
||||
|
||||
|
||||
if __name__ == """__main__""":
|
||||
test_f16f16f16_mma_gemm()
|
||||
test_f16f16f32_mma_gemm()
|
||||
+91
@@ -0,0 +1,91 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
from tvm.s_tir import Schedule
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# pylint: disable=invalid-name, no-member
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def add(a: T.handle, b: T.handle) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, [2048, 2048, 2048], dtype="float32")
|
||||
B = T.match_buffer(b, [2048, 2048, 2048], dtype="float32")
|
||||
A_cached = T.sblock_alloc_buffer([2048, 2048, 2048], dtype="float32")
|
||||
# body
|
||||
for i, j, k in T.grid(2048, 2048, 2048):
|
||||
with T.sblock("move"):
|
||||
vi, vj, vk = T.axis.remap("SSS", [i, j, k])
|
||||
T.reads([A[vi, vj, vk]])
|
||||
T.writes([A_cached[vi, vj, vk]])
|
||||
A_cached[vi, vj, vk] = A[vi, vj, vk]
|
||||
for i0, j0, i1, j1, k0, i2, j2, k1 in T.grid(128, 64, 4, 4, 64, 4, 8, 32):
|
||||
with T.sblock("add"):
|
||||
vi = T.axis.spatial(2048, i0 * 16 + i1 * 4 + i2)
|
||||
vj = T.axis.spatial(2048, j0 * 32 + j1 * 8 + j2)
|
||||
vk = T.axis.spatial(2048, k0 * 32 + k1)
|
||||
T.reads([A_cached[vi, vj, vk]])
|
||||
T.writes([B[vi, vj, vk]])
|
||||
B[vi, vj, vk] = A_cached[vi, vj, vk] + T.float32(1)
|
||||
|
||||
|
||||
# pylint: enable=invalid-name, no-member
|
||||
|
||||
|
||||
def _sch(decision: int) -> Schedule:
|
||||
sch = Schedule(add, debug_mask="all")
|
||||
# pylint: disable=invalid-name
|
||||
b0 = sch.get_sblock(name="move", func_name="main")
|
||||
l1 = sch.sample_compute_location(block=b0, decision=decision)
|
||||
sch.compute_at(block=b0, loop=l1, preserve_unit_loops=True)
|
||||
# pylint: enable=invalid-name
|
||||
return sch
|
||||
|
||||
|
||||
def _make_mutator(target: Target) -> ms.Mutator:
|
||||
ctx = ms.TuneContext(
|
||||
mod=add,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[],
|
||||
mutator_probs={ms.mutator.MutateComputeLocation(): 1.0},
|
||||
),
|
||||
)
|
||||
return next(iter(ctx.space_generator.mutator_probs.keys()))
|
||||
|
||||
|
||||
def test_mutate_compute_location_add():
|
||||
mutator = _make_mutator(
|
||||
target=Target("llvm"),
|
||||
)
|
||||
sch = _sch(decision=4)
|
||||
results = set()
|
||||
for _ in range(100):
|
||||
trace = mutator.apply(sch.trace)
|
||||
decision = trace.decisions[trace.insts[-2]]
|
||||
assert not decision == 4
|
||||
results.add(decision)
|
||||
assert len(results) == 9
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_mutate_compute_location_add()
|
||||
@@ -0,0 +1,118 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
|
||||
from tvm.s_tir import Schedule
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# pylint: disable=invalid-name, no-member
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [512, 512])
|
||||
B = T.match_buffer(b, [512, 512])
|
||||
C = T.match_buffer(c, [512, 512])
|
||||
for i, j, k in T.grid(512, 512, 512): # type: ignore
|
||||
with T.sblock("C"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k]) # type: ignore
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0 # type: ignore
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
|
||||
# pylint: enable=invalid-name, no-member
|
||||
|
||||
|
||||
def _sch(decisions: list[list[int]], ann_val: int) -> Schedule:
|
||||
sch = Schedule(matmul, debug_mask="all")
|
||||
# pylint: disable=invalid-name
|
||||
d0, d1, d2 = decisions
|
||||
b0 = sch.get_sblock(name="C", func_name="main")
|
||||
root = sch.get_sblock(name="root", func_name="main")
|
||||
sch.get_consumers(block=b0)
|
||||
b1 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")
|
||||
l2, l3, l4 = sch.get_loops(block=b0)
|
||||
v5, v6, v7, v8 = sch.sample_perfect_tile(
|
||||
loop=l2,
|
||||
n=4,
|
||||
max_innermost_factor=64,
|
||||
decision=d0,
|
||||
)
|
||||
l9, l10, l11, l12 = sch.split(loop=l2, factors=[v5, v6, v7, v8])
|
||||
v13, v14, v15, v16 = sch.sample_perfect_tile(
|
||||
loop=l3,
|
||||
n=4,
|
||||
max_innermost_factor=64,
|
||||
decision=d1,
|
||||
)
|
||||
l17, l18, l19, l20 = sch.split(loop=l3, factors=[v13, v14, v15, v16])
|
||||
v21, v22 = sch.sample_perfect_tile(
|
||||
loop=l4,
|
||||
n=2,
|
||||
max_innermost_factor=64,
|
||||
decision=d2,
|
||||
)
|
||||
l23, l24 = sch.split(loop=l4, factors=[v21, v22])
|
||||
sch.reorder(l9, l17, l10, l18, l23, l11, l19, l24, l12, l20)
|
||||
sch.reverse_compute_at(block=b1, loop=l18, preserve_unit_loops=True)
|
||||
sch.annotate(block_or_loop=root, ann_key="meta_schedule.parallel", ann_val=ann_val)
|
||||
# pylint: enable=invalid-name
|
||||
return sch
|
||||
|
||||
|
||||
def _make_mutator(target: Target, max_jobs_per_core: int) -> ms.Mutator:
|
||||
ctx = ms.TuneContext(
|
||||
mod=matmul,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[],
|
||||
mutator_probs={ms.mutator.MutateParallel(max_jobs_per_core): 1.0},
|
||||
),
|
||||
)
|
||||
return next(iter(ctx.space_generator.mutator_probs.keys()))
|
||||
|
||||
|
||||
def test_mutate_parallel_matmul():
|
||||
mutator = _make_mutator(
|
||||
target=Target({"kind": "llvm", "num-cores": 16}),
|
||||
max_jobs_per_core=256,
|
||||
)
|
||||
sch = _sch(
|
||||
decisions=[
|
||||
[4, 32, 4, 1],
|
||||
[8, 4, 8, 2],
|
||||
[512, 1],
|
||||
],
|
||||
ann_val=64,
|
||||
)
|
||||
results = set()
|
||||
for _ in range(100):
|
||||
trace = mutator.apply(sch.trace)
|
||||
ann_val = int(trace.insts[-1].inputs[1])
|
||||
results.add(ann_val)
|
||||
if len(results) == 3:
|
||||
break
|
||||
assert len(results) == 3
|
||||
assert results == {4, 32, 4096}
|
||||
|
||||
|
||||
if __name__ == """__main__""":
|
||||
test_mutate_parallel_matmul()
|
||||
@@ -0,0 +1,92 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
from tvm.s_tir import Schedule
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# pylint: disable=invalid-name, no-member
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def element_wise(var_A: T.handle, var_B: T.handle) -> None:
|
||||
A = T.match_buffer(var_A, [512, 512], dtype="float32")
|
||||
B = T.match_buffer(var_B, [512, 512], dtype="float32")
|
||||
for i, j in T.grid(512, 512):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = A[vi, vj] + 1.0
|
||||
|
||||
|
||||
# pylint: enable=invalid-name, no-member
|
||||
|
||||
|
||||
def _sch() -> Schedule:
|
||||
sch = Schedule(element_wise, debug_mask="all")
|
||||
# pylint: disable=invalid-name
|
||||
b0 = sch.get_sblock(name="C", func_name="main")
|
||||
l1, l2 = sch.get_loops(block=b0)
|
||||
l3 = sch.fuse(l1, l2)
|
||||
v4 = sch.sample_categorical(
|
||||
candidates=[32, 64, 128, 256, 512, 1024],
|
||||
probs=[
|
||||
0.16666666666666666,
|
||||
0.16666666666666666,
|
||||
0.16666666666666666,
|
||||
0.16666666666666666,
|
||||
0.16666666666666666,
|
||||
0.16666666666666666,
|
||||
],
|
||||
decision=3,
|
||||
)
|
||||
l5, l6 = sch.split(loop=l3, factors=[None, v4])
|
||||
sch.bind(loop=l5, thread_axis="blockIdx.x")
|
||||
sch.bind(loop=l6, thread_axis="threadIdx.x")
|
||||
# pylint: enable=invalid-name
|
||||
return sch
|
||||
|
||||
|
||||
def _make_mutator(target: Target) -> ms.Mutator:
|
||||
ctx = ms.TuneContext(
|
||||
mod=element_wise,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[],
|
||||
mutator_probs={ms.mutator.MutateThreadBinding(): 1.0},
|
||||
),
|
||||
)
|
||||
return next(iter(ctx.space_generator.mutator_probs.keys()))
|
||||
|
||||
|
||||
def test_mutate_thread_binding():
|
||||
mutator = _make_mutator(target=Target("cuda"))
|
||||
sch = _sch()
|
||||
results = set()
|
||||
for _ in range(100):
|
||||
trace = mutator.apply(sch.trace)
|
||||
decision = trace.decisions[trace.insts[-4]]
|
||||
results.add(decision)
|
||||
if len(results) == 5:
|
||||
break
|
||||
assert len(results) == 5
|
||||
assert results == {0, 1, 2, 4, 5}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_mutate_thread_binding()
|
||||
@@ -0,0 +1,111 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
import operator
|
||||
from functools import reduce
|
||||
|
||||
from tvm.s_tir import Schedule
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# pylint: disable=invalid-name, no-member
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [512, 512])
|
||||
B = T.match_buffer(b, [512, 512])
|
||||
C = T.match_buffer(c, [512, 512])
|
||||
for i, j, k in T.grid(512, 512, 512): # type: ignore
|
||||
with T.sblock("C"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k]) # type: ignore
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0 # type: ignore
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
|
||||
# pylint: enable=invalid-name, no-member
|
||||
|
||||
|
||||
def _sch(decisions: list[list[int]]) -> Schedule:
|
||||
sch = Schedule(matmul, debug_mask="all")
|
||||
# pylint: disable=invalid-name
|
||||
(d0,) = decisions
|
||||
b0 = sch.get_sblock(name="C", func_name="main")
|
||||
sch.get_consumers(block=b0)
|
||||
b1 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")
|
||||
l2, l3, l4 = sch.get_loops(block=b0)
|
||||
v5, v6, v7, v8 = sch.sample_perfect_tile(
|
||||
loop=l2,
|
||||
n=4,
|
||||
max_innermost_factor=64,
|
||||
decision=d0,
|
||||
)
|
||||
l9, l10, l11, l12 = sch.split(loop=l2, factors=[v5, v6, v7, v8])
|
||||
l17, l18, l19, l20 = sch.split(loop=l3, factors=[8, 4, 8, 2])
|
||||
l23, l24 = sch.split(loop=l4, factors=[512, 1])
|
||||
sch.reorder(l9, l17, l10, l18, l23, l11, l19, l24, l12, l20)
|
||||
sch.reverse_compute_at(block=b1, loop=l18, preserve_unit_loops=True)
|
||||
# pylint: enable=invalid-name
|
||||
return sch
|
||||
|
||||
|
||||
def _make_mutator(target: Target) -> ms.Mutator:
|
||||
ctx = ms.TuneContext(
|
||||
mod=matmul,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[],
|
||||
mutator_probs={ms.mutator.MutateTileSize(): 1.0},
|
||||
),
|
||||
)
|
||||
return next(iter(ctx.space_generator.mutator_probs.keys()))
|
||||
|
||||
|
||||
def test_mutate_tile_size_matmul():
|
||||
mutator = _make_mutator(
|
||||
target=Target({"kind": "llvm", "num-cores": 16}),
|
||||
)
|
||||
results = {}
|
||||
sch = _sch(decisions=[[4, 32, 4, 1]])
|
||||
for _ in range(1000):
|
||||
trace = mutator.apply(sch.trace)
|
||||
assert trace.insts[4].kind.name == "SamplePerfectTile"
|
||||
decision = trace.decisions[trace.insts[4]]
|
||||
decision = [int(x) for x in decision]
|
||||
results[str(decision)] = decision
|
||||
assert reduce(operator.mul, decision, 1) == 512
|
||||
assert len(results) > 15
|
||||
|
||||
|
||||
def test_mutate_sample_categorical_single_candidate():
|
||||
mutator = _make_mutator(
|
||||
target=Target({"kind": "llvm", "num-cores": 16}),
|
||||
)
|
||||
sch = Schedule(matmul, debug_mask="all")
|
||||
sch.sample_categorical(candidates=[1], probs=[1.0], decision=0)
|
||||
|
||||
# The mutator finds the SampleCategorical has only one candidate, and thus skips it.
|
||||
trace = mutator.apply(sch.trace)
|
||||
assert trace is None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_mutate_tile_size_matmul()
|
||||
test_mutate_sample_categorical_single_candidate()
|
||||
@@ -0,0 +1,119 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
|
||||
from tvm.s_tir import Schedule
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# pylint: disable=invalid-name, no-member
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [512, 512])
|
||||
B = T.match_buffer(b, [512, 512])
|
||||
C = T.match_buffer(c, [512, 512])
|
||||
for i, j, k in T.grid(512, 512, 512): # type: ignore
|
||||
with T.sblock("C"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k]) # type: ignore
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0 # type: ignore
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
|
||||
# pylint: enable=invalid-name, no-member
|
||||
|
||||
|
||||
def _sch(decisions: list[list[int]]) -> Schedule:
|
||||
sch = Schedule(matmul, debug_mask="all")
|
||||
# pylint: disable=invalid-name
|
||||
d0, d1, d2 = decisions
|
||||
b0 = sch.get_sblock(name="C", func_name="main")
|
||||
root = sch.get_sblock(name="root", func_name="main")
|
||||
sch.get_consumers(block=b0)
|
||||
b1 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="global")
|
||||
l2, l3, l4 = sch.get_loops(block=b0)
|
||||
v5, v6, v7, v8 = sch.sample_perfect_tile(
|
||||
loop=l2,
|
||||
n=4,
|
||||
max_innermost_factor=64,
|
||||
decision=d0,
|
||||
)
|
||||
l9, l10, l11, l12 = sch.split(loop=l2, factors=[v5, v6, v7, v8])
|
||||
v13, v14, v15, v16 = sch.sample_perfect_tile(
|
||||
loop=l3,
|
||||
n=4,
|
||||
max_innermost_factor=64,
|
||||
decision=d1,
|
||||
)
|
||||
l17, l18, l19, l20 = sch.split(loop=l3, factors=[v13, v14, v15, v16])
|
||||
v21, v22 = sch.sample_perfect_tile(
|
||||
loop=l4,
|
||||
n=2,
|
||||
max_innermost_factor=64,
|
||||
decision=d2,
|
||||
)
|
||||
l23, l24 = sch.split(loop=l4, factors=[v21, v22])
|
||||
sch.reorder(l9, l17, l10, l18, l23, l11, l19, l24, l12, l20)
|
||||
sch.reverse_compute_at(block=b1, loop=l18, preserve_unit_loops=True)
|
||||
v57 = sch.sample_categorical(
|
||||
candidates=[0, 16, 64, 512],
|
||||
probs=[0.25, 0.25, 0.25, 0.25],
|
||||
decision=0,
|
||||
)
|
||||
sch.annotate(block_or_loop=root, ann_key="meta_schedule.unroll_explicit", ann_val=v57)
|
||||
# pylint: enable=invalid-name
|
||||
return sch
|
||||
|
||||
|
||||
def _make_mutator(target: Target) -> ms.Mutator:
|
||||
ctx = ms.TuneContext(
|
||||
mod=matmul,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[],
|
||||
mutator_probs={ms.mutator.MutateUnroll(): 1.0},
|
||||
),
|
||||
)
|
||||
return next(iter(ctx.space_generator.mutator_probs.keys()))
|
||||
|
||||
|
||||
def test_mutate_unroll_matmul():
|
||||
mutator = _make_mutator(target=Target({"kind": "llvm", "num-cores": 16}))
|
||||
sch = _sch(
|
||||
decisions=[
|
||||
[4, 32, 4, 1],
|
||||
[8, 4, 8, 2],
|
||||
[512, 1],
|
||||
],
|
||||
)
|
||||
results = set()
|
||||
for _ in range(100):
|
||||
trace = mutator.apply(sch.trace)
|
||||
decision = trace.decisions[trace.insts[-2]]
|
||||
results.add(decision)
|
||||
if len(results) == 3:
|
||||
break
|
||||
assert len(results) == 3
|
||||
assert results == {1, 2, 3}
|
||||
|
||||
|
||||
if __name__ == """__main__""":
|
||||
test_mutate_unroll_matmul()
|
||||
@@ -0,0 +1,534 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: F401
|
||||
|
||||
import math
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import tvm_ffi
|
||||
from tvm_ffi import register_global_func
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import te
|
||||
from tvm.ir.module import IRModule
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.s_tir.meta_schedule import TuneContext
|
||||
from tvm.s_tir.meta_schedule.schedule_rule import PyScheduleRule
|
||||
from tvm.s_tir.meta_schedule.space_generator import PostOrderApply
|
||||
from tvm.s_tir.schedule import SBlockRV, Schedule
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,
|
||||
# fmt: off
|
||||
|
||||
|
||||
def get_matmul_packed(m, n, k, lhs_type="int8", rhs_dtype="int8", acc_dtype="int32"):
|
||||
X = te.placeholder((m, k), name="X", dtype=lhs_type)
|
||||
W = te.placeholder((n, k), name="W", dtype=rhs_dtype)
|
||||
|
||||
ak = te.reduce_axis((0, k), name="k")
|
||||
matmul = te.compute(
|
||||
(m, n),
|
||||
lambda i, j: te.sum(
|
||||
X[i, ak].astype(acc_dtype) * W[j, ak].astype(acc_dtype),
|
||||
axis=ak,
|
||||
),
|
||||
name="compute",
|
||||
)
|
||||
return te.create_prim_func([X, W, matmul])
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class DuplicateMatmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class TrinityMatmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, d: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.sblock_alloc_buffer((1024, 1024), "float32")
|
||||
C = T.sblock_alloc_buffer((1024, 1024), "float32")
|
||||
D = T.match_buffer(d, (1024, 1024), "float32")
|
||||
for i, j in T.grid(1024, 1024):
|
||||
with T.sblock("A"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = A[vi, vj] * 2.0
|
||||
for i, j in T.grid(1024, 1024):
|
||||
with T.sblock("B"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
C[vi, vj] = B[vi, vj] + 3.0
|
||||
for i, j in T.grid(1024, 1024):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
D[vi, vj] = C[vi, vj] * 5.0
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class TrinityMatmulProcessedForReference:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, d: T.handle) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, [1024, 1024], dtype="float32")
|
||||
D = T.match_buffer(d, [1024, 1024], dtype="float32")
|
||||
# body
|
||||
# with tirx.block("root")
|
||||
B = T.sblock_alloc_buffer([1024, 1024], dtype="float32")
|
||||
for i0_0, i1_0, i0_1, i1_1 in T.grid(16, 64, 64, 16):
|
||||
with T.sblock("A"):
|
||||
vi = T.axis.S(1024, i0_0 * 64 + i0_1)
|
||||
vj = T.axis.S(1024, i1_0 * 16 + i1_1)
|
||||
T.reads([A[vi, vj]])
|
||||
T.writes([B[vi, vj]])
|
||||
B[vi, vj] = A[vi, vj] * T.float32(2)
|
||||
for i0_0, i1_0, i0_1, i1_1 in T.grid(16, 64, 64, 16):
|
||||
with T.sblock("C"):
|
||||
vi = T.axis.S(1024, i0_0 * 64 + i0_1)
|
||||
vj = T.axis.S(1024, i1_0 * 16 + i1_1)
|
||||
T.reads([B[vi, vj]])
|
||||
T.writes([D[vi, vj]])
|
||||
D[vi, vj] = (B[vi, vj] + T.float32(3)) * T.float32(5)
|
||||
|
||||
|
||||
# fmt: on
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
|
||||
|
||||
def _is_root(sch: Schedule, block: SBlockRV) -> bool:
|
||||
return sch.get_sref(block).parent is None
|
||||
|
||||
|
||||
def _check_correct(schedule: Schedule):
|
||||
trace = schedule.trace
|
||||
for inst in trace.decisions:
|
||||
assert math.prod(trace.decisions[inst]) == 1024
|
||||
|
||||
|
||||
@derived_object
|
||||
class WowSoFancyScheduleRule(PyScheduleRule):
|
||||
def _initialize_with_tune_context(self, context: "TuneContext") -> None:
|
||||
pass
|
||||
|
||||
def apply(self, sch: Schedule, block: SBlockRV) -> list[Schedule]:
|
||||
if _is_root(sch, block):
|
||||
return [sch]
|
||||
new_sch = sch.copy()
|
||||
i, j, k = new_sch.get_loops(block=block)
|
||||
i_0, i_1, i_2, i_3 = new_sch.split(loop=i, factors=[2, 4, 64, 2])
|
||||
j_0, j_1, j_2, j_3 = new_sch.split(loop=j, factors=[4, 64, 2, 2])
|
||||
k_0, k_1 = new_sch.split(loop=k, factors=[32, 32])
|
||||
new_sch.reorder(i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3)
|
||||
return [new_sch]
|
||||
|
||||
|
||||
@derived_object
|
||||
class DoubleScheduleRule(PyScheduleRule):
|
||||
def _initialize_with_tune_context(self, context: "TuneContext") -> None:
|
||||
pass
|
||||
|
||||
def apply(self, sch: Schedule, block: SBlockRV) -> list[Schedule]:
|
||||
if _is_root(sch, block):
|
||||
return [sch]
|
||||
new_sch = sch.copy()
|
||||
i, j, k = new_sch.get_loops(block=block)
|
||||
i_0, i_1, i_2, i_3 = new_sch.split(loop=i, factors=[4, 64, 2, 2])
|
||||
j_0, j_1, j_2, j_3 = new_sch.split(loop=j, factors=[2, 4, 64, 2])
|
||||
k_0, k_1 = new_sch.split(loop=k, factors=[32, 32])
|
||||
new_sch.reorder(i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3)
|
||||
result = [new_sch]
|
||||
new_sch = sch.copy()
|
||||
i, j, k = new_sch.get_loops(block=block)
|
||||
i_0, i_1, i_2, i_3 = new_sch.split(loop=i, factors=[4, 64, 2, 2])
|
||||
j_0, j_1, j_2, j_3 = new_sch.split(loop=j, factors=[2, 4, 64, 2])
|
||||
k_0, k_1 = new_sch.split(loop=k, factors=[32, 32])
|
||||
new_sch.reorder(i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3)
|
||||
result.append(new_sch)
|
||||
return result
|
||||
|
||||
|
||||
@derived_object
|
||||
class TrinityDoubleRule(PyScheduleRule):
|
||||
def _initialize_with_tune_context(self, context: "TuneContext") -> None:
|
||||
pass
|
||||
|
||||
def apply(self, sch: Schedule, block: SBlockRV) -> list[Schedule]:
|
||||
if _is_root(sch, block):
|
||||
return [sch]
|
||||
new_sch = sch.copy()
|
||||
i, j = new_sch.get_loops(block=block)
|
||||
i_0, i_1 = new_sch.split(loop=i, factors=[16, 64])
|
||||
j_0, j_1 = new_sch.split(loop=j, factors=[64, 16])
|
||||
new_sch.reorder(i_0, j_0, i_1, j_1)
|
||||
result = [new_sch]
|
||||
new_sch = sch.copy()
|
||||
i, j = new_sch.get_loops(block=block)
|
||||
i_0, i_1 = new_sch.split(loop=i, factors=[2, 512])
|
||||
j_0, j_1 = new_sch.split(loop=j, factors=[2, 512])
|
||||
new_sch.reorder(i_0, j_0, i_1, j_1)
|
||||
result.append(new_sch)
|
||||
return result
|
||||
|
||||
|
||||
@derived_object
|
||||
class ReorderScheduleRule(PyScheduleRule):
|
||||
def _initialize_with_tune_context(self, context: "TuneContext") -> None:
|
||||
pass
|
||||
|
||||
def apply(self, sch: Schedule, block: SBlockRV) -> list[Schedule]:
|
||||
if _is_root(sch, block):
|
||||
return [sch]
|
||||
new_sch = sch.copy()
|
||||
i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3 = new_sch.get_loops(block=block)
|
||||
new_sch.reorder(i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3, i_0, j_0)
|
||||
result = [new_sch]
|
||||
new_sch = sch.copy()
|
||||
i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3 = new_sch.get_loops(block=block)
|
||||
new_sch.reorder(i_1, j_3, i_0, j_0, j_1, k_0, i_2, j_2, k_1, i_3)
|
||||
result.append(new_sch)
|
||||
return result
|
||||
|
||||
|
||||
def test_meta_schedule_post_order_apply():
|
||||
mod = Matmul
|
||||
context = TuneContext(
|
||||
mod=mod,
|
||||
target=Target("llvm"),
|
||||
task_name="Test Task",
|
||||
space_generator=PostOrderApply(
|
||||
sch_rules=[WowSoFancyScheduleRule()],
|
||||
postprocs=[],
|
||||
mutator_probs={},
|
||||
),
|
||||
)
|
||||
post_order_apply = context.space_generator
|
||||
schs = post_order_apply.generate_design_space(mod)
|
||||
assert len(schs) == 1
|
||||
assert not tvm_ffi.structural_equal(schs[0].mod, mod)
|
||||
_check_correct(schs[0])
|
||||
|
||||
|
||||
def test_meta_schedule_post_order_apply_double():
|
||||
mod = Matmul
|
||||
context = TuneContext(
|
||||
mod=mod,
|
||||
target=Target("llvm"),
|
||||
task_name="Double Rules Task",
|
||||
space_generator=PostOrderApply(
|
||||
sch_rules=[DoubleScheduleRule()],
|
||||
postprocs=[],
|
||||
mutator_probs={},
|
||||
),
|
||||
)
|
||||
post_order_apply = context.space_generator
|
||||
schs = post_order_apply.generate_design_space(mod)
|
||||
assert len(schs) == 2
|
||||
for sch in schs:
|
||||
assert not tvm_ffi.structural_equal(sch.mod, mod)
|
||||
_check_correct(sch)
|
||||
|
||||
|
||||
def test_meta_schedule_post_order_apply_multiple():
|
||||
mod = Matmul
|
||||
context = TuneContext(
|
||||
mod=mod,
|
||||
target=Target("llvm"),
|
||||
task_name="Double Rules Task",
|
||||
space_generator=PostOrderApply(
|
||||
sch_rules=[DoubleScheduleRule(), ReorderScheduleRule()],
|
||||
postprocs=[],
|
||||
mutator_probs={},
|
||||
),
|
||||
)
|
||||
post_order_apply = context.space_generator
|
||||
schs = post_order_apply.generate_design_space(mod)
|
||||
assert len(schs) == 4
|
||||
for sch in schs:
|
||||
assert not tvm_ffi.structural_equal(sch.mod, mod)
|
||||
_check_correct(sch)
|
||||
|
||||
|
||||
def test_meta_schedule_post_order_apply_duplicate_matmul():
|
||||
mod = DuplicateMatmul
|
||||
context = TuneContext(
|
||||
mod=mod,
|
||||
target=Target("llvm"),
|
||||
task_name="Duplicate Matmul Task",
|
||||
space_generator=PostOrderApply(
|
||||
sch_rules=[WowSoFancyScheduleRule()],
|
||||
postprocs=[],
|
||||
mutator_probs={},
|
||||
),
|
||||
)
|
||||
post_order_apply = context.space_generator
|
||||
with pytest.raises(
|
||||
RuntimeError,
|
||||
match=r".*Duplicated block name matmul in function main not supported!",
|
||||
):
|
||||
post_order_apply.generate_design_space(mod)
|
||||
|
||||
|
||||
def test_meta_schedule_post_order_apply_remove_block():
|
||||
@derived_object
|
||||
class RemoveBlock(PyScheduleRule):
|
||||
def _initialize_with_tune_context(self, context: "TuneContext") -> None:
|
||||
pass
|
||||
|
||||
def apply(self, sch: Schedule, block: SBlockRV) -> list[Schedule]:
|
||||
if _is_root(sch, block):
|
||||
return [sch]
|
||||
sch = sch.copy()
|
||||
if sch.get(block).name_hint == "B":
|
||||
sch.compute_inline(block)
|
||||
return [sch]
|
||||
|
||||
def correct_trace(a, b, c, d):
|
||||
return "\n".join(
|
||||
[
|
||||
"# from tvm import s_tir",
|
||||
"def apply_trace(sch: s_tir.Schedule) -> None:",
|
||||
' b0 = sch.get_sblock(name="A", func_name="main")',
|
||||
' b1 = sch.get_sblock(name="B", func_name="main")',
|
||||
' b2 = sch.get_sblock(name="C", func_name="main")',
|
||||
" sch.compute_inline(block=b1)",
|
||||
" l3, l4 = sch.get_loops(block=b2)",
|
||||
" l5, l6 = sch.split(loop=l3, factors="
|
||||
+ str(a)
|
||||
+ ", preserve_unit_iters=True, disable_predication=False)",
|
||||
" l7, l8 = sch.split(loop=l4, factors="
|
||||
+ str(b)
|
||||
+ ", preserve_unit_iters=True, disable_predication=False)",
|
||||
" sch.reorder(l5, l7, l6, l8)",
|
||||
" l9, l10 = sch.get_loops(block=b0)",
|
||||
" l11, l12 = sch.split(loop=l9, factors="
|
||||
+ str(c)
|
||||
+ ", preserve_unit_iters=True, disable_predication=False)",
|
||||
" l13, l14 = sch.split(loop=l10, factors="
|
||||
+ str(d)
|
||||
+ ", preserve_unit_iters=True, disable_predication=False)",
|
||||
" sch.reorder(l11, l13, l12, l14)",
|
||||
]
|
||||
)
|
||||
|
||||
mod = TrinityMatmul
|
||||
context = TuneContext(
|
||||
mod=mod,
|
||||
target=Target("llvm"),
|
||||
task_name="Remove Block Task",
|
||||
space_generator=PostOrderApply(
|
||||
sch_rules=[RemoveBlock(), TrinityDoubleRule()],
|
||||
postprocs=[],
|
||||
mutator_probs={},
|
||||
),
|
||||
)
|
||||
post_order_apply = context.space_generator
|
||||
schs = post_order_apply.generate_design_space(mod)
|
||||
assert len(schs) == 4
|
||||
for sch in schs:
|
||||
with pytest.raises(
|
||||
tvm.s_tir.schedule.schedule.ScheduleError,
|
||||
match="ScheduleError: An error occurred in the schedule primitive 'get-block'.",
|
||||
):
|
||||
sch.get_sblock("B", "main")
|
||||
sch_trace = sch.trace.simplified(True)
|
||||
assert (
|
||||
str(sch_trace) == correct_trace([16, 64], [64, 16], [2, 512], [2, 512])
|
||||
or str(sch_trace) == correct_trace([2, 512], [2, 512], [2, 512], [2, 512])
|
||||
or str(sch_trace) == correct_trace([16, 64], [64, 16], [16, 64], [64, 16])
|
||||
or str(sch_trace) == correct_trace([2, 512], [2, 512], [16, 64], [64, 16])
|
||||
)
|
||||
|
||||
|
||||
def test_target_sblocks_search_space():
|
||||
# Test that specific blocks of trinity matmul can be targeted.
|
||||
def filter_fn(block, target_names) -> bool:
|
||||
return block.name_hint in target_names
|
||||
|
||||
def _get_sch(filter_fn):
|
||||
mod = TrinityMatmul
|
||||
context = TuneContext(
|
||||
mod=mod,
|
||||
target=Target("llvm"),
|
||||
task_name="Custom Search Space Task",
|
||||
space_generator=PostOrderApply(
|
||||
f_block_filter=filter_fn,
|
||||
sch_rules=[TrinityDoubleRule()],
|
||||
postprocs=[],
|
||||
mutator_probs={},
|
||||
),
|
||||
)
|
||||
post_order_apply = context.space_generator
|
||||
schs = post_order_apply.generate_design_space(mod)
|
||||
return schs
|
||||
|
||||
# Start by checking that by default each block has a space generated.
|
||||
schs = _get_sch(None)
|
||||
assert len(schs) == 8
|
||||
|
||||
# Next check that we can target a specific block and only get its' revelant schedules.
|
||||
schs = _get_sch(lambda block: filter_fn(block, ["B"]))
|
||||
assert len(schs) == 2
|
||||
|
||||
## Check that extracting two blocks works.
|
||||
schs = _get_sch(lambda block: filter_fn(block, ["A", "C"]))
|
||||
assert len(schs) == 4
|
||||
|
||||
## Finally check that all blocks can be extracted by name.
|
||||
schs = _get_sch(lambda block: filter_fn(block, ["A", "B", "C"]))
|
||||
assert len(schs) == 8
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"target,mod,expected_intr",
|
||||
[
|
||||
(
|
||||
Target(
|
||||
{
|
||||
"kind": "llvm",
|
||||
"device": "arm_cpu",
|
||||
"mtriple": "aarch64-linux-gnu",
|
||||
"mattr": ["+neon"],
|
||||
"num-cores": 2,
|
||||
}
|
||||
),
|
||||
IRModule({"main": get_matmul_packed(128, 128, 128, "int8", "int8", "int32")}),
|
||||
"dot_4x4_i8i8s32_neon",
|
||||
),
|
||||
(
|
||||
Target(
|
||||
{
|
||||
"kind": "llvm",
|
||||
"device": "arm_cpu",
|
||||
"mtriple": "aarch64-linux-gnu",
|
||||
"mattr": ["+neon", "+v8.2a", "+dotprod"],
|
||||
"num-cores": 2,
|
||||
}
|
||||
),
|
||||
IRModule({"main": get_matmul_packed(128, 128, 128, "int8", "int8", "int32")}),
|
||||
"dot_4x4_i8i8s32_sdot",
|
||||
),
|
||||
(
|
||||
Target(
|
||||
{
|
||||
"kind": "llvm",
|
||||
"device": "arm_cpu",
|
||||
"mtriple": "aarch64-linux-gnu",
|
||||
"mattr": ["+neon", "+v8.2a", "+dotprod"],
|
||||
"num-cores": 2,
|
||||
}
|
||||
),
|
||||
IRModule({"main": get_matmul_packed(128, 128, 128, "uint8", "uint8", "uint32")}),
|
||||
"dot_4x4_u8u8u32_udot",
|
||||
),
|
||||
(
|
||||
Target(
|
||||
{
|
||||
"kind": "llvm",
|
||||
"device": "arm_cpu",
|
||||
"mtriple": "aarch64-linux-gnu",
|
||||
"mattr": ["+neon", "+v8.2a", "+dotprod"],
|
||||
"num-cores": 2,
|
||||
}
|
||||
),
|
||||
IRModule({"main": get_matmul_packed(128, 128, 128, "uint8", "uint8", "int32")}),
|
||||
"dot_4x4_u8u8i32_hdot",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_meta_schedule_post_order_apply_arm_intrin(target, mod, expected_intr):
|
||||
context = TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
task_name="Arm Intrinsic Task",
|
||||
space_generator=PostOrderApply(), # Triggers default generator
|
||||
rand_state=1, # Change it while all tests are not passing
|
||||
)
|
||||
post_order_apply = context.space_generator
|
||||
schs = post_order_apply.generate_design_space(mod)
|
||||
|
||||
assert len(schs) != 0
|
||||
|
||||
for sch in schs:
|
||||
sch.enter_postproc()
|
||||
|
||||
for proc in context.space_generator.postprocs:
|
||||
proc.apply(sch)
|
||||
|
||||
assert any(["call_llvm_pure_intrin" in sch.mod.script() for sch in schs])
|
||||
assert any([expected_intr in str(sch.trace) for sch in schs])
|
||||
|
||||
|
||||
def test_meta_schedule_derived_object():
|
||||
@derived_object
|
||||
class RemoveBlock(PyScheduleRule):
|
||||
@classmethod
|
||||
def class_construct(cls):
|
||||
return cls()
|
||||
|
||||
@staticmethod
|
||||
def static_construct():
|
||||
return RemoveBlock()
|
||||
|
||||
inst_by_init = RemoveBlock()
|
||||
assert isinstance(inst_by_init, RemoveBlock)
|
||||
|
||||
inst_by_classmethod = RemoveBlock.class_construct()
|
||||
assert isinstance(inst_by_classmethod, RemoveBlock)
|
||||
|
||||
inst_by_staticmethod = RemoveBlock.static_construct()
|
||||
assert isinstance(inst_by_staticmethod, RemoveBlock)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
+111
@@ -0,0 +1,111 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: F401
|
||||
|
||||
import tvm
|
||||
from tvm import tirx
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def _target() -> Target:
|
||||
return Target("qcom/hexagon-v68", host="llvm")
|
||||
|
||||
|
||||
def _create_context(mod, target) -> ms.TuneContext:
|
||||
ctx = ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[
|
||||
ms.postproc.DisallowAsyncStridedMemCopy(),
|
||||
],
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
return ctx
|
||||
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
# fmt: off
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
# fmt: on
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
|
||||
|
||||
def test_postproc_disallow_async_strided_mem_copy_allows():
|
||||
mod = Matmul
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
|
||||
matmul_block = sch.get_sblock("matmul")
|
||||
|
||||
loops = sch.get_loops(matmul_block)
|
||||
cache_read = sch.cache_read(matmul_block, 0, "global.vtcm")
|
||||
|
||||
sch.compute_at(cache_read, loops[1])
|
||||
|
||||
sch.annotate(loops[1], "software_pipeline_stage", [0, 1])
|
||||
sch.annotate(loops[1], "software_pipeline_order", [0, 1])
|
||||
sch.annotate(loops[1], "software_pipeline_async_stages", [0])
|
||||
|
||||
ctx = _create_context(sch.mod, target=_target())
|
||||
sch.mod.show()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
|
||||
|
||||
def test_postproc_disallow_async_strided_mem_copy_disallows():
|
||||
mod = Matmul
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
|
||||
matmul_block = sch.get_sblock("matmul")
|
||||
|
||||
loops = sch.get_loops(matmul_block)
|
||||
# Make it a strided mem copy.
|
||||
cache_read = sch.cache_read(matmul_block, 1, "global.vtcm")
|
||||
|
||||
sch.compute_at(cache_read, loops[1])
|
||||
sch.annotate(loops[1], "software_pipeline_stage", [0, 1])
|
||||
sch.annotate(loops[1], "software_pipeline_order", [0, 1])
|
||||
sch.annotate(loops[1], "software_pipeline_async_stages", [0])
|
||||
|
||||
sch.mod.show()
|
||||
ctx = _create_context(sch.mod, target=_target())
|
||||
assert not ctx.space_generator.postprocs[0].apply(sch)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_postproc_disallow_async_strided_mem_copy_allows()
|
||||
test_postproc_disallow_async_strided_mem_copy_disallows()
|
||||
+102
@@ -0,0 +1,102 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: F401
|
||||
|
||||
import tvm
|
||||
from tvm import tirx
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def _target() -> Target:
|
||||
return Target("cuda", host="llvm")
|
||||
|
||||
|
||||
def _create_context(mod, target) -> ms.TuneContext:
|
||||
ctx = ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[
|
||||
ms.postproc.DisallowDynamicLoop(),
|
||||
],
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
return ctx
|
||||
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
# fmt: off
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class DynamicLoop:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j in T.grid(1024, 1024):
|
||||
for k in T.serial(0, i):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
# fmt: on
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
|
||||
|
||||
def test_postproc_disallow_dynamic_loops():
|
||||
mod = Matmul
|
||||
ctx = _create_context(mod, target=_target())
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
|
||||
|
||||
def test_postproc_disallow_dynamic_loops_fail():
|
||||
mod = DynamicLoop
|
||||
ctx = _create_context(mod, target=_target())
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
assert not ctx.space_generator.postprocs[0].apply(sch)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_postproc_disallow_dynamic_loops()
|
||||
test_postproc_disallow_dynamic_loops_fail()
|
||||
+303
@@ -0,0 +1,303 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: E501, F401
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import tirx
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing import te_workload
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
from tvm.te import create_prim_func
|
||||
|
||||
|
||||
def _target() -> Target:
|
||||
return Target("cuda", host="llvm")
|
||||
|
||||
|
||||
def _create_context(mod, target) -> ms.TuneContext:
|
||||
ctx = ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[
|
||||
ms.postproc.RewriteCooperativeFetch(),
|
||||
],
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
return ctx
|
||||
|
||||
|
||||
# fmt: off
|
||||
# pylint: disable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
|
||||
@tvm.script.ir_module
|
||||
class AfterRewrite0:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_A: T.handle, var_B: T.handle, var_C: T.handle) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(var_A, [512, 512], dtype="float32")
|
||||
B = T.match_buffer(var_B, [512, 512], dtype="float32")
|
||||
C = T.match_buffer(var_C, [512, 512], dtype="float32")
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C_local = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="local")
|
||||
A_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
B_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
for i0_0_i1_0_fused in T.thread_binding(0, 16, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_fused in T.thread_binding(0, 16, thread="vthread.x"):
|
||||
for i0_2_i1_2_fused in T.thread_binding(0, 8, thread="threadIdx.x"):
|
||||
for i2_0 in T.serial(0, 1):
|
||||
for ax0_ax1_fused_0 in T.serial(0, 32768):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(0, 8, thread="threadIdx.x"):
|
||||
with T.sblock("A_shared"):
|
||||
v0 = T.axis.spatial(512, (ax0_ax1_fused_0 * 8 + ax0_ax1_fused_1) // 512)
|
||||
v1 = T.axis.spatial(512, (ax0_ax1_fused_0 * 8 + ax0_ax1_fused_1) % 512)
|
||||
T.reads([A[v0, v1]])
|
||||
T.writes([A_shared[v0, v1]])
|
||||
A_shared[v0, v1] = A[v0, v1]
|
||||
for ax0_ax1_fused_0 in T.serial(0, 1024):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(0, 8, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_2 in T.vectorized(0, 2):
|
||||
with T.sblock("B_shared"):
|
||||
v0 = T.axis.spatial(512, (ax0_ax1_fused_0 * 16 + ax0_ax1_fused_1 * 2 + ax0_ax1_fused_2) // 32)
|
||||
v1 = T.axis.spatial(512, i0_0_i1_0_fused * 32 + (ax0_ax1_fused_0 * 16 + ax0_ax1_fused_1 * 2 + ax0_ax1_fused_2) % 32)
|
||||
T.reads([B[v0, v1]])
|
||||
T.writes([B_shared[v0, v1]])
|
||||
B_shared[v0, v1] = B[v0, v1]
|
||||
for i2_1, i0_3, i1_3, i2_2, i0_4, i1_4 in T.grid(16, 2, 2, 32, 16, 2):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(512, i0_1_i1_1_fused * 32 + i0_3 * 16 + i0_4)
|
||||
j = T.axis.spatial(512, i0_0_i1_0_fused * 32 + i0_2_i1_2_fused * 4 + i1_3 * 2 + i1_4)
|
||||
k = T.axis.reduce(512, i2_0 * 512 + i2_1 * 32 + i2_2)
|
||||
T.reads([A_shared[i, k], B_shared[k, j]])
|
||||
T.writes([C_local[i, j]])
|
||||
with T.init():
|
||||
C_local[i, j] = T.float32(0)
|
||||
C_local[i, j] = C_local[i, j] + A_shared[i, k] * B_shared[k, j]
|
||||
for ax0, ax1 in T.grid(32, 4):
|
||||
with T.sblock("C_local"):
|
||||
v0 = T.axis.spatial(512, i0_1_i1_1_fused * 32 + ax0)
|
||||
v1 = T.axis.spatial(512, i0_0_i1_0_fused * 32 + i0_2_i1_2_fused * 4 + ax1)
|
||||
T.reads([C_local[v0, v1]])
|
||||
T.writes([C[v0, v1]])
|
||||
C[v0, v1] = C_local[v0, v1]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class WarpExecutionAfterRewrite:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
C: T.Buffer((512, 512), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C_local = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="local")
|
||||
A_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
B_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
for i0_0_i1_0_fused in T.thread_binding(0, 16, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_fused in T.thread_binding(0, 16, thread="vthread.x"):
|
||||
for i0_2_i1_2_fused in T.thread_binding(0, 8, thread="threadIdx.y"):
|
||||
for i2_0 in T.serial(0, 1):
|
||||
for ax0_ax1_fused_0 in T.serial(0, 1024):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(0, 8, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(
|
||||
0, 32, thread="threadIdx.x"
|
||||
):
|
||||
with T.sblock("A_shared"):
|
||||
v0 = T.axis.spatial(
|
||||
512,
|
||||
(
|
||||
ax0_ax1_fused_0 * 256
|
||||
+ ax0_ax1_fused_1 * 32
|
||||
+ ax0_ax1_fused_2
|
||||
)
|
||||
// 512,
|
||||
)
|
||||
v1 = T.axis.spatial(
|
||||
512,
|
||||
(
|
||||
ax0_ax1_fused_0 * 256
|
||||
+ ax0_ax1_fused_1 * 32
|
||||
+ ax0_ax1_fused_2
|
||||
)
|
||||
% 512,
|
||||
)
|
||||
T.reads([A[v0, v1]])
|
||||
T.writes([A_shared[v0, v1]])
|
||||
A_shared[v0, v1] = A[v0, v1]
|
||||
for ax0_ax1_fused_0 in T.serial(0, 32):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(0, 8, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(
|
||||
0, 32, thread="threadIdx.x"
|
||||
):
|
||||
for ax0_ax1_fused_3 in T.vectorized(0, 2):
|
||||
with T.sblock("B_shared"):
|
||||
v0 = T.axis.spatial(
|
||||
512,
|
||||
(
|
||||
ax0_ax1_fused_0 * 512
|
||||
+ ax0_ax1_fused_1 * 64
|
||||
+ ax0_ax1_fused_2 * 2
|
||||
+ ax0_ax1_fused_3
|
||||
)
|
||||
// 32,
|
||||
)
|
||||
v1 = T.axis.spatial(
|
||||
512,
|
||||
i0_0_i1_0_fused * 32
|
||||
+ (
|
||||
ax0_ax1_fused_0 * 512
|
||||
+ ax0_ax1_fused_1 * 64
|
||||
+ ax0_ax1_fused_2 * 2
|
||||
+ ax0_ax1_fused_3
|
||||
)
|
||||
% 32,
|
||||
)
|
||||
T.reads([B[v0, v1]])
|
||||
T.writes([B_shared[v0, v1]])
|
||||
B_shared[v0, v1] = B[v0, v1]
|
||||
for i2_1, i0_3, i1_3, i2_2, i0_4, i1_4 in T.grid(16, 2, 2, 32, 16, 2):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(512, i0_1_i1_1_fused * 32 + i0_3 * 16 + i0_4)
|
||||
j = T.axis.spatial(
|
||||
512,
|
||||
i0_0_i1_0_fused * 32 + i0_2_i1_2_fused * 4 + i1_3 * 2 + i1_4,
|
||||
)
|
||||
k = T.axis.reduce(512, i2_0 * 512 + i2_1 * 32 + i2_2)
|
||||
T.reads([A_shared[i, k], B_shared[k, j]])
|
||||
T.writes([C_local[i, j]])
|
||||
T.sblock_attr({"warp_execution": 1})
|
||||
with T.init():
|
||||
C_local[i, j] = T.float32(0)
|
||||
C_local[i, j] = C_local[i, j] + A_shared[i, k] * B_shared[k, j]
|
||||
for ax0, ax1 in T.grid(32, 4):
|
||||
with T.sblock("C_local"):
|
||||
v0 = T.axis.spatial(512, i0_1_i1_1_fused * 32 + ax0)
|
||||
v1 = T.axis.spatial(
|
||||
512, i0_0_i1_0_fused * 32 + i0_2_i1_2_fused * 4 + ax1
|
||||
)
|
||||
T.reads([C_local[v0, v1]])
|
||||
T.writes([C[v0, v1]])
|
||||
C[v0, v1] = C_local[v0, v1]
|
||||
|
||||
|
||||
# pylint: enable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
# fmt: on
|
||||
|
||||
|
||||
def test_rewrite_cooperative_fetch():
|
||||
mod = create_prim_func(te_workload.matmul(n=512, m=512, k=512))
|
||||
target = _target()
|
||||
ctx = _create_context(mod, target)
|
||||
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
# fmt: off
|
||||
# pylint: disable=line-too-long,invalid-name
|
||||
b0 = sch.get_sblock(name="C", func_name="main")
|
||||
b1 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="local")
|
||||
l2, l3, l4 = sch.get_loops(block=b0)
|
||||
v5, v6, v7, v8, v9 = sch.sample_perfect_tile(loop=l2, n=5, max_innermost_factor=64, decision=[1, 16, 1, 2, 16])
|
||||
l10, l11, l12, l13, l14 = sch.split(loop=l2, factors=[v5, v6, v7, v8, v9])
|
||||
v15, v16, v17, v18, v19 = sch.sample_perfect_tile(loop=l3, n=5, max_innermost_factor=64, decision=[16, 1, 8, 2, 2])
|
||||
l20, l21, l22, l23, l24 = sch.split(loop=l3, factors=[v15, v16, v17, v18, v19])
|
||||
v25, v26, v27 = sch.sample_perfect_tile(loop=l4, n=3, max_innermost_factor=64, decision=[1, 16, 32])
|
||||
l28, l29, l30 = sch.split(loop=l4, factors=[v25, v26, v27])
|
||||
sch.reorder(l10, l20, l11, l21, l12, l22, l28, l29, l13, l23, l30, l14, l24)
|
||||
l31 = sch.fuse(l10, l20)
|
||||
sch.bind(loop=l31, thread_axis="blockIdx.x")
|
||||
l32 = sch.fuse(l11, l21)
|
||||
sch.bind(loop=l32, thread_axis="vthread.x")
|
||||
l33 = sch.fuse(l12, l22)
|
||||
sch.bind(loop=l33, thread_axis="threadIdx.x")
|
||||
b34 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")
|
||||
sch.compute_at(block=b34, loop=l28, preserve_unit_loops=True)
|
||||
_, _, _, _, l39, l40 = sch.get_loops(block=b34)
|
||||
l41 = sch.fuse(l39, l40)
|
||||
_, v43 = sch.sample_perfect_tile(loop=l41, n=2, max_innermost_factor=4, decision=[262144, 1])
|
||||
sch.annotate(block_or_loop=b34, ann_key="meta_schedule.cooperative_fetch", ann_val=v43)
|
||||
b44 = sch.cache_read(block=b0, read_buffer_index=1, storage_scope="shared")
|
||||
sch.compute_at(block=b44, loop=l28, preserve_unit_loops=True)
|
||||
_, _, _, _, l49, l50 = sch.get_loops(block=b44)
|
||||
l51 = sch.fuse(l49, l50)
|
||||
_, v53 = sch.sample_perfect_tile(loop=l51, n=2, max_innermost_factor=4, decision=[8192, 2])
|
||||
sch.annotate(block_or_loop=b44, ann_key="meta_schedule.cooperative_fetch", ann_val=v53)
|
||||
sch.reverse_compute_at(block=b1, loop=l33, preserve_unit_loops=True)
|
||||
# pylint: enable=line-too-long,invalid-name
|
||||
# fmt: on
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
tvm.ir.assert_structural_equal(sch.mod, AfterRewrite0)
|
||||
|
||||
|
||||
def test_rewrite_warp_execution():
|
||||
mod = create_prim_func(te_workload.matmul(n=512, m=512, k=512))
|
||||
target = _target()
|
||||
ctx = _create_context(mod, target)
|
||||
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
# fmt: off
|
||||
# pylint: disable=line-too-long,invalid-name
|
||||
b0 = sch.get_sblock(name="C", func_name="main")
|
||||
b1 = sch.cache_write(block=b0, write_buffer_index=0, storage_scope="local")
|
||||
l2, l3, l4 = sch.get_loops(block=b0)
|
||||
sch.annotate(b0, "warp_execution", 1)
|
||||
v5, v6, v7, v8, v9 = sch.sample_perfect_tile(loop=l2, n=5, max_innermost_factor=64, decision=[1, 16, 1, 2, 16])
|
||||
l10, l11, l12, l13, l14 = sch.split(loop=l2, factors=[v5, v6, v7, v8, v9])
|
||||
v15, v16, v17, v18, v19 = sch.sample_perfect_tile(loop=l3, n=5, max_innermost_factor=64, decision=[16, 1, 8, 2, 2])
|
||||
l20, l21, l22, l23, l24 = sch.split(loop=l3, factors=[v15, v16, v17, v18, v19])
|
||||
v25, v26, v27 = sch.sample_perfect_tile(loop=l4, n=3, max_innermost_factor=64, decision=[1, 16, 32])
|
||||
l28, l29, l30 = sch.split(loop=l4, factors=[v25, v26, v27])
|
||||
sch.reorder(l10, l20, l11, l21, l12, l22, l28, l29, l13, l23, l30, l14, l24)
|
||||
l31 = sch.fuse(l10, l20)
|
||||
sch.bind(loop=l31, thread_axis="blockIdx.x")
|
||||
l32 = sch.fuse(l11, l21)
|
||||
sch.bind(loop=l32, thread_axis="vthread.x")
|
||||
l33 = sch.fuse(l12, l22)
|
||||
sch.bind(loop=l33, thread_axis="threadIdx.y")
|
||||
b34 = sch.cache_read(block=b0, read_buffer_index=0, storage_scope="shared")
|
||||
sch.compute_at(block=b34, loop=l28, preserve_unit_loops=True)
|
||||
_, _, _, _, l39, l40 = sch.get_loops(block=b34)
|
||||
l41 = sch.fuse(l39, l40)
|
||||
_, v43 = sch.sample_perfect_tile(loop=l41, n=2, max_innermost_factor=4, decision=[262144, 1])
|
||||
sch.annotate(block_or_loop=b34, ann_key="meta_schedule.cooperative_fetch", ann_val=v43)
|
||||
b44 = sch.cache_read(block=b0, read_buffer_index=1, storage_scope="shared")
|
||||
sch.compute_at(block=b44, loop=l28, preserve_unit_loops=True)
|
||||
_, _, _, _, l49, l50 = sch.get_loops(block=b44)
|
||||
l51 = sch.fuse(l49, l50)
|
||||
_, v53 = sch.sample_perfect_tile(loop=l51, n=2, max_innermost_factor=4, decision=[8192, 2])
|
||||
sch.annotate(block_or_loop=b44, ann_key="meta_schedule.cooperative_fetch", ann_val=v53)
|
||||
sch.reverse_compute_at(block=b1, loop=l33, preserve_unit_loops=True)
|
||||
# pylint: enable=line-too-long,invalid-name
|
||||
# fmt: on
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
tvm.ir.assert_structural_equal(sch.mod, WarpExecutionAfterRewrite)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,637 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: E501
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def _target() -> Target:
|
||||
return Target("cuda", host="llvm")
|
||||
|
||||
|
||||
def _create_context(mod, target) -> ms.TuneContext:
|
||||
ctx = ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[
|
||||
ms.postproc.RewriteLayout(),
|
||||
],
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
return ctx
|
||||
|
||||
|
||||
def _apply_rewrite_layout(mod):
|
||||
"""Apply the RewriteLayout postproc transformation."""
|
||||
target = Target("cuda", host="llvm")
|
||||
ctx = ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[
|
||||
ms.postproc.RewriteLayout(),
|
||||
],
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
if not ctx.space_generator.postprocs[0].apply(sch):
|
||||
raise RuntimeError("RewriteLayout postproc failed")
|
||||
return sch.mod
|
||||
|
||||
|
||||
def test_tir_matmul():
|
||||
"""Main functionality test
|
||||
|
||||
A new block should be inserted to transform the layout, with the
|
||||
compute block operating on the temporary transformed buffer.
|
||||
"""
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(
|
||||
A: T.Buffer((16, 16), "float32"),
|
||||
B: T.Buffer((16, 16), "float32"),
|
||||
C: T.Buffer((16, 16), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"layout_free_buffers": [1]})
|
||||
for i0, j, k0, i1, k1 in T.grid(4, 16, 4, 4, 4):
|
||||
with T.sblock("matmul"):
|
||||
vi = T.axis.S(16, i0 * 4 + i1)
|
||||
vj = T.axis.S(16, j)
|
||||
vk = T.axis.R(16, k0 * 4 + k1)
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(
|
||||
A: T.Buffer((16, 16), "float32"),
|
||||
B: T.Buffer((16, 16), "float32"),
|
||||
C: T.Buffer((16, 16), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"layout_free_buffers": [1]})
|
||||
B_reindex = T.sblock_alloc_buffer([16, 4, 4], dtype="float32")
|
||||
for ax0, ax1 in T.grid(16, 16):
|
||||
with T.sblock("layout_rewrite"):
|
||||
i0, i1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.sblock_attr({"meta_schedule.layout_rewrite_preproc": True})
|
||||
B_reindex[i1, i0 // 4, i0 % 4] = B[i0, i1]
|
||||
for i0, j, k0, i1, k1 in T.grid(4, 16, 4, 4, 4):
|
||||
with T.sblock("matmul"):
|
||||
vi = T.axis.spatial(16, i0 * 4 + i1)
|
||||
vj = T.axis.spatial(16, j)
|
||||
vk = T.axis.reduce(16, k0 * 4 + k1)
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B_reindex[vj, vk // 4, vk % 4]
|
||||
|
||||
mod = tvm.IRModule.from_expr(before)
|
||||
mod = _apply_rewrite_layout(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
def test_rewritten_buffers_must_occur_within_block():
|
||||
"""Buffers must occur within a Block"""
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(
|
||||
A: T.Buffer((16, 16), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"layout_free_buffers": [0]})
|
||||
for i, j in T.grid(16, 16):
|
||||
T.evaluate(A[i, j])
|
||||
|
||||
mod = tvm.IRModule.from_expr(before)
|
||||
with pytest.raises(RuntimeError):
|
||||
_apply_rewrite_layout(mod)
|
||||
|
||||
|
||||
def test_extent_one():
|
||||
"""Buffers with dimensions of extent 1 can be transformed
|
||||
|
||||
Regression test for a previous bug, in which the removal of
|
||||
trivial variables resulted in an error in `IndexMap::Inverse`.
|
||||
"""
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(
|
||||
A: T.Buffer((16, 1), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"layout_free_buffers": [0]})
|
||||
for i, j in T.grid(16, 1):
|
||||
with T.sblock("block"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
T.evaluate(A[vi, vj])
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(A: T.Buffer((16, 1), "float32")):
|
||||
T.func_attr({"layout_free_buffers": [0]})
|
||||
|
||||
A_global = T.sblock_alloc_buffer([16], dtype="float32")
|
||||
for ax0, ax1 in T.grid(16, 1):
|
||||
with T.sblock("A_global"):
|
||||
v0, v1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.sblock_attr({"meta_schedule.layout_rewrite_preproc": True})
|
||||
A_global[v0] = A[v0, v1]
|
||||
|
||||
for i, j in T.grid(16, 1):
|
||||
with T.sblock("block"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
T.evaluate(A_global[vi])
|
||||
|
||||
mod = tvm.IRModule.from_expr(before)
|
||||
mod = _apply_rewrite_layout(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def tir_matmul(
|
||||
A: T.Buffer((16, 16), "float32"),
|
||||
B: T.Buffer((16, 16), "float32"),
|
||||
C: T.Buffer((16, 16), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"layout_free_buffers": [1]})
|
||||
for i0, j, k0, i1, k1 in T.grid(4, 16, 4, 4, 4):
|
||||
with T.sblock("matmul"):
|
||||
vi = T.axis.S(16, i0 * 4 + i1)
|
||||
vj = T.axis.S(16, j)
|
||||
vk = T.axis.R(16, k0 * 4 + k1)
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def rewritten_tir_matmul(
|
||||
A: T.Buffer((16, 16), "float32"),
|
||||
B: T.Buffer((16, 16), "float32"),
|
||||
C: T.Buffer((16, 16), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"layout_free_buffers": [1]})
|
||||
B_reindex = T.sblock_alloc_buffer([16, 4, 4], dtype="float32")
|
||||
for ax0, ax1 in T.grid(16, 16):
|
||||
with T.sblock("layout_rewrite"):
|
||||
i0, i1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.sblock_attr({"meta_schedule.layout_rewrite_preproc": True})
|
||||
B_reindex[i1, i0 // 4, i0 % 4] = B[i0, i1]
|
||||
for i0, j, k0, i1, k1 in T.grid(4, 16, 4, 4, 4):
|
||||
with T.sblock("matmul"):
|
||||
vi = T.axis.spatial(16, i0 * 4 + i1)
|
||||
vj = T.axis.spatial(16, j)
|
||||
vk = T.axis.reduce(16, k0 * 4 + k1)
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B_reindex[vj, vk // 4, vk % 4]
|
||||
|
||||
|
||||
def test_layout_rewrite():
|
||||
target = _target()
|
||||
ctx = _create_context(tir_matmul, target)
|
||||
sch = tvm.s_tir.Schedule(tir_matmul, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
assert_structural_equal_ignore_global_symbol(sch.mod["main"], rewritten_tir_matmul)
|
||||
|
||||
|
||||
# fmt: off
|
||||
@tvm.script.ir_module
|
||||
class Conv2dCacheRead:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(p0: T.Buffer((1, 56, 56, 64), "float32"), p1: T.Buffer((3, 3, 64, 64), "float32"), conv2d_nhwc: T.Buffer((1, 56, 56, 64), "float32")):
|
||||
T.func_attr({"layout_free_buffers": [1], "tirx.noalias": True, "global_symbol": "main"})
|
||||
pad_temp = T.sblock_alloc_buffer([1, 58, 58, 64], dtype="float32")
|
||||
conv2d_nhwc_global = T.sblock_alloc_buffer([1, 56, 56, 64], dtype="float32")
|
||||
pad_temp_global = T.sblock_alloc_buffer([1, 58, 58, 64], dtype="float32")
|
||||
p1_global = T.sblock_alloc_buffer([3, 3, 64, 64], dtype="float32")
|
||||
for i0_0_i1_0_i2_0_fused in T.parallel(4, annotations={"pragma_auto_unroll_max_step":16, "pragma_unroll_explicit":1}):
|
||||
for ax0, ax1, ax2 in T.grid(1, 30, 30):
|
||||
for ax3_fused in T.vectorized(64):
|
||||
with T.sblock("pad_temp"):
|
||||
i0 = T.axis.spatial(1, ax0)
|
||||
i1 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused // 2 * 28 + ax1)
|
||||
i2 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused % 2 * 28 + ax2)
|
||||
i3 = T.axis.spatial(64, ax3_fused)
|
||||
T.reads(p0[i0, i1 - 1, i2 - 1, i3])
|
||||
T.writes(pad_temp[i0, i1, i2, i3])
|
||||
pad_temp[i0, i1, i2, i3] = T.if_then_else(1 <= i1 and i1 < 57 and 1 <= i2 and i2 < 57, p0[i0, i1 - 1, i2 - 1, i3], T.float32(0), dtype="float32")
|
||||
for i3_0 in T.serial(16):
|
||||
for ax0_ax1_ax2_ax3_fused in T.serial(57600):
|
||||
with T.sblock("pad_temp_global"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused // 2 * 28 + ax0_ax1_ax2_ax3_fused // 1920)
|
||||
v2 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused % 2 * 28 + ax0_ax1_ax2_ax3_fused % 1920 // 64)
|
||||
v3 = T.axis.spatial(64, ax0_ax1_ax2_ax3_fused % 64)
|
||||
T.reads(pad_temp[v0, v1, v2, v3])
|
||||
T.writes(pad_temp_global[v0, v1, v2, v3])
|
||||
pad_temp_global[v0, v1, v2, v3] = pad_temp[v0, v1, v2, v3]
|
||||
for ax0_ax1_ax2_ax3_fused in T.serial(2304):
|
||||
with T.sblock("p1_global"):
|
||||
v0 = T.axis.spatial(3, ax0_ax1_ax2_ax3_fused // 768)
|
||||
v1 = T.axis.spatial(3, ax0_ax1_ax2_ax3_fused % 768 // 256)
|
||||
v2 = T.axis.spatial(64, ax0_ax1_ax2_ax3_fused % 256 // 4)
|
||||
v3 = T.axis.spatial(64, i3_0 * 4 + ax0_ax1_ax2_ax3_fused % 4)
|
||||
T.reads(p1[v0, v1, v2, v3])
|
||||
T.writes(p1_global[v0, v1, v2, v3])
|
||||
p1_global[v0, v1, v2, v3] = p1[v0, v1, v2, v3]
|
||||
for i0_1, i1_1, i2_1, i3_1 in T.grid(1, 7, 2, 1):
|
||||
for i0_2_init, i1_2_init, i2_2_init, i3_2_init, i0_3_init, i1_3_init, i2_3_init in T.grid(1, 1, 14, 2, 1, 4, 1):
|
||||
for i3_3_fused_init in T.vectorized(2):
|
||||
with T.sblock("conv2d_nhwc_init"):
|
||||
nn = T.axis.spatial(1, i0_1 + i0_2_init + i0_3_init)
|
||||
yy = T.axis.spatial(56, i0_0_i1_0_i2_0_fused // 2 * 28 + i1_1 * 4 + i1_2_init * 4 + i1_3_init)
|
||||
xx = T.axis.spatial(56, i2_3_init + i0_0_i1_0_i2_0_fused % 2 * 28 + i2_1 * 14 + i2_2_init)
|
||||
ff = T.axis.spatial(64, i3_0 * 4 + i3_1 * 4 + i3_2_init * 2 + i3_3_fused_init)
|
||||
T.reads()
|
||||
T.writes(conv2d_nhwc_global[nn, yy, xx, ff])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure":"SSRSRS"})
|
||||
conv2d_nhwc_global[nn, yy, xx, ff] = T.float32(0)
|
||||
for i4_0, i5_0, i6_0, i0_2, i1_2, i2_2, i3_2, i4_1, i5_1, i6_1, i0_3, i1_3, i2_3 in T.grid(1, 1, 2, 1, 1, 14, 2, 3, 3, 32, 1, 4, 1):
|
||||
for i3_3_fused in T.vectorized(2):
|
||||
with T.sblock("conv2d_nhwc_update"):
|
||||
nn = T.axis.spatial(1, i0_1 + i0_2 + i0_3)
|
||||
yy = T.axis.spatial(56, i0_0_i1_0_i2_0_fused // 2 * 28 + i1_1 * 4 + i1_2 * 4 + i1_3)
|
||||
xx = T.axis.spatial(56, i2_3 + i0_0_i1_0_i2_0_fused % 2 * 28 + i2_1 * 14 + i2_2)
|
||||
ff = T.axis.spatial(64, i3_0 * 4 + i3_1 * 4 + i3_2 * 2 + i3_3_fused)
|
||||
ry = T.axis.reduce(3, i4_0 * 3 + i4_1)
|
||||
rx = T.axis.reduce(3, i5_0 * 3 + i5_1)
|
||||
rc = T.axis.reduce(64, i6_0 * 32 + i6_1)
|
||||
T.reads(conv2d_nhwc_global[nn, yy, xx, ff], pad_temp_global[nn, yy + ry, xx + rx, rc], p1_global[ry, rx, rc, ff])
|
||||
T.writes(conv2d_nhwc_global[nn, yy, xx, ff])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure":"SSRSRS"})
|
||||
conv2d_nhwc_global[nn, yy, xx, ff] = conv2d_nhwc_global[nn, yy, xx, ff] + pad_temp_global[nn, yy + ry, xx + rx, rc] * p1_global[ry, rx, rc, ff]
|
||||
for ax0, ax1, ax2 in T.grid(1, 4, 14):
|
||||
for ax3_fused in T.vectorized(4):
|
||||
with T.sblock("conv2d_nhwc_global"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(56, i0_0_i1_0_i2_0_fused // 2 * 28 + i1_1 * 4 + ax1)
|
||||
v2 = T.axis.spatial(56, i0_0_i1_0_i2_0_fused % 2 * 28 + i2_1 * 14 + ax2)
|
||||
v3 = T.axis.spatial(64, i3_0 * 4 + ax3_fused)
|
||||
T.reads(conv2d_nhwc_global[v0, v1, v2, v3])
|
||||
T.writes(conv2d_nhwc[v0, v1, v2, v3])
|
||||
conv2d_nhwc[v0, v1, v2, v3] = conv2d_nhwc_global[v0, v1, v2, v3]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2dCacheReadRewritten:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(p0: T.Buffer((1, 56, 56, 64), "float32"), p1: T.Buffer((3, 3, 64, 64), "float32"), conv2d_nhwc: T.Buffer((1, 56, 56, 64), "float32")):
|
||||
T.func_attr({"layout_free_buffers": [1], "tirx.noalias": True, "global_symbol": "main"})
|
||||
pad_temp = T.sblock_alloc_buffer([1, 58, 58, 64], dtype="float32")
|
||||
conv2d_nhwc_global = T.sblock_alloc_buffer([1, 56, 56, 64], dtype="float32")
|
||||
pad_temp_global = T.sblock_alloc_buffer([1, 58, 58, 64], dtype="float32")
|
||||
p1_global = T.sblock_alloc_buffer([16, 2, 2, 3, 3, 32, 2], dtype="float32")
|
||||
p1_global_1 = T.sblock_alloc_buffer([16, 2, 2, 3, 3, 32, 2], dtype="float32")
|
||||
for ax0, ax1, ax2, ax3 in T.grid(3, 3, 64, 64):
|
||||
with T.sblock("p1_global"):
|
||||
v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
||||
T.reads(p1[v0, v1, v2, v3])
|
||||
T.writes(p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2])
|
||||
T.sblock_attr({"meta_schedule.layout_rewrite_preproc":True})
|
||||
p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2] = p1[v0, v1, v2, v3]
|
||||
for i0_0_i1_0_i2_0_fused in T.parallel(4, annotations={"pragma_auto_unroll_max_step":16, "pragma_unroll_explicit":1}):
|
||||
for ax0, ax1, ax2 in T.grid(1, 30, 30):
|
||||
for ax3_fused in T.vectorized(64):
|
||||
with T.sblock("pad_temp"):
|
||||
i0 = T.axis.spatial(1, ax0)
|
||||
i1 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused // 2 * 28 + ax1)
|
||||
i2 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused % 2 * 28 + ax2)
|
||||
i3 = T.axis.spatial(64, ax3_fused)
|
||||
T.reads(p0[i0, i1 - 1, i2 - 1, i3])
|
||||
T.writes(pad_temp[i0, i1, i2, i3])
|
||||
pad_temp[i0, i1, i2, i3] = T.if_then_else(1 <= i1 and i1 < 57 and 1 <= i2 and i2 < 57, p0[i0, i1 - 1, i2 - 1, i3], T.float32(0), dtype="float32")
|
||||
for i3_0 in T.serial(16):
|
||||
for ax0_ax1_ax2_ax3_fused in T.serial(57600):
|
||||
with T.sblock("pad_temp_global"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused // 2 * 28 + ax0_ax1_ax2_ax3_fused // 1920)
|
||||
v2 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused % 2 * 28 + ax0_ax1_ax2_ax3_fused % 1920 // 64)
|
||||
v3 = T.axis.spatial(64, ax0_ax1_ax2_ax3_fused % 64)
|
||||
T.reads(pad_temp[v0, v1, v2, v3])
|
||||
T.writes(pad_temp_global[v0, v1, v2, v3])
|
||||
pad_temp_global[v0, v1, v2, v3] = pad_temp[v0, v1, v2, v3]
|
||||
for ax0_ax1_ax2_ax3_fused in T.serial(2304):
|
||||
with T.sblock("p1_global"):
|
||||
v0 = T.axis.spatial(3, ax0_ax1_ax2_ax3_fused // 768)
|
||||
v1 = T.axis.spatial(3, ax0_ax1_ax2_ax3_fused % 768 // 256)
|
||||
v2 = T.axis.spatial(64, ax0_ax1_ax2_ax3_fused % 256 // 4)
|
||||
v3 = T.axis.spatial(64, i3_0 * 4 + ax0_ax1_ax2_ax3_fused % 4)
|
||||
T.reads(p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2])
|
||||
T.writes(p1_global[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2])
|
||||
p1_global[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2] = p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2]
|
||||
for i0_1, i1_1, i2_1, i3_1 in T.grid(1, 7, 2, 1):
|
||||
for i0_2_init, i1_2_init, i2_2_init, i3_2_init, i0_3_init, i1_3_init, i2_3_init in T.grid(1, 1, 14, 2, 1, 4, 1):
|
||||
for i3_3_fused_init in T.vectorized(2):
|
||||
with T.sblock("conv2d_nhwc_init"):
|
||||
nn = T.axis.spatial(1, i0_1 + i0_2_init + i0_3_init)
|
||||
yy = T.axis.spatial(56, i0_0_i1_0_i2_0_fused // 2 * 28 + i1_1 * 4 + i1_2_init * 4 + i1_3_init)
|
||||
xx = T.axis.spatial(56, i2_3_init + i0_0_i1_0_i2_0_fused % 2 * 28 + i2_1 * 14 + i2_2_init)
|
||||
ff = T.axis.spatial(64, i3_0 * 4 + i3_1 * 4 + i3_2_init * 2 + i3_3_fused_init)
|
||||
T.reads()
|
||||
T.writes(conv2d_nhwc_global[nn, yy, xx, ff])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure":"SSRSRS"})
|
||||
conv2d_nhwc_global[nn, yy, xx, ff] = T.float32(0)
|
||||
for i4_0, i5_0, i6_0, i0_2, i1_2, i2_2, i3_2, i4_1, i5_1, i6_1, i0_3, i1_3, i2_3 in T.grid(1, 1, 2, 1, 1, 14, 2, 3, 3, 32, 1, 4, 1):
|
||||
for i3_3_fused in T.vectorized(2):
|
||||
with T.sblock("conv2d_nhwc_update"):
|
||||
nn = T.axis.spatial(1, i0_1 + i0_2 + i0_3)
|
||||
yy = T.axis.spatial(56, i0_0_i1_0_i2_0_fused // 2 * 28 + i1_1 * 4 + i1_2 * 4 + i1_3)
|
||||
xx = T.axis.spatial(56, i2_3 + i0_0_i1_0_i2_0_fused % 2 * 28 + i2_1 * 14 + i2_2)
|
||||
ff = T.axis.spatial(64, i3_0 * 4 + i3_1 * 4 + i3_2 * 2 + i3_3_fused)
|
||||
ry = T.axis.reduce(3, i4_0 * 3 + i4_1)
|
||||
rx = T.axis.reduce(3, i5_0 * 3 + i5_1)
|
||||
rc = T.axis.reduce(64, i6_0 * 32 + i6_1)
|
||||
T.reads(conv2d_nhwc_global[nn, yy, xx, ff], pad_temp_global[nn, yy + ry, xx + rx, rc], p1_global[ff // 4, rc // 32, ff % 4 // 2, ry, rx, rc % 32, ff % 2])
|
||||
T.writes(conv2d_nhwc_global[nn, yy, xx, ff])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure":"SSRSRS"})
|
||||
conv2d_nhwc_global[nn, yy, xx, ff] = conv2d_nhwc_global[nn, yy, xx, ff] + pad_temp_global[nn, yy + ry, xx + rx, rc] * p1_global[ff // 4, rc // 32, ff % 4 // 2, ry, rx, rc % 32, ff % 2]
|
||||
for ax0, ax1, ax2 in T.grid(1, 4, 14):
|
||||
for ax3_fused in T.vectorized(4):
|
||||
with T.sblock("conv2d_nhwc_global"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(56, i0_0_i1_0_i2_0_fused // 2 * 28 + i1_1 * 4 + ax1)
|
||||
v2 = T.axis.spatial(56, i0_0_i1_0_i2_0_fused % 2 * 28 + i2_1 * 14 + ax2)
|
||||
v3 = T.axis.spatial(64, i3_0 * 4 + ax3_fused)
|
||||
T.reads(conv2d_nhwc_global[v0, v1, v2, v3])
|
||||
T.writes(conv2d_nhwc[v0, v1, v2, v3])
|
||||
conv2d_nhwc[v0, v1, v2, v3] = conv2d_nhwc_global[v0, v1, v2, v3]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2dCacheReadMultipleRewritten:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(p0: T.Buffer((1, 56, 56, 64), "float32"), p1: T.Buffer((3, 3, 64, 64), "float32"), conv2d_nhwc: T.Buffer((1, 56, 56, 64), "float32")):
|
||||
T.func_attr({"layout_free_buffers": [1], "tirx.noalias": True, "global_symbol": "main"})
|
||||
pad_temp = T.sblock_alloc_buffer([1, 58, 58, 64], dtype="float32")
|
||||
conv2d_nhwc_global = T.sblock_alloc_buffer([1, 56, 56, 64], dtype="float32")
|
||||
pad_temp_global = T.sblock_alloc_buffer([1, 58, 58, 64], dtype="float32")
|
||||
p1_global = T.sblock_alloc_buffer([16, 2, 2, 3, 3, 32, 2], dtype="float32")
|
||||
p1_global2 = T.sblock_alloc_buffer([16, 2, 2, 3, 3, 32, 2], dtype="float32", scope="global2")
|
||||
p1_global_1 = T.sblock_alloc_buffer([16, 2, 2, 3, 3, 32, 2], dtype="float32")
|
||||
for ax0, ax1, ax2, ax3 in T.grid(3, 3, 64, 64):
|
||||
with T.sblock("p1_global"):
|
||||
v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
||||
T.reads(p1[v0, v1, v2, v3])
|
||||
T.writes(p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2])
|
||||
T.sblock_attr({"meta_schedule.layout_rewrite_preproc":True})
|
||||
p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2] = p1[v0, v1, v2, v3]
|
||||
for ax0, ax1, ax2, ax3 in T.grid(3, 3, 64, 64):
|
||||
with T.sblock("p1_global2"):
|
||||
v0, v1, v2, v3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
||||
T.reads(p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2])
|
||||
T.writes(p1_global2[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2])
|
||||
p1_global2[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2] = p1_global_1[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2]
|
||||
for i0_0_i1_0_i2_0_fused in T.parallel(4, annotations={"pragma_auto_unroll_max_step":16, "pragma_unroll_explicit":1}):
|
||||
for ax0, ax1, ax2 in T.grid(1, 30, 30):
|
||||
for ax3_fused in T.vectorized(64):
|
||||
with T.sblock("pad_temp"):
|
||||
i0 = T.axis.spatial(1, ax0)
|
||||
i1 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused // 2 * 28 + ax1)
|
||||
i2 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused % 2 * 28 + ax2)
|
||||
i3 = T.axis.spatial(64, ax3_fused)
|
||||
T.reads(p0[i0, i1 - 1, i2 - 1, i3])
|
||||
T.writes(pad_temp[i0, i1, i2, i3])
|
||||
pad_temp[i0, i1, i2, i3] = T.if_then_else(1 <= i1 and i1 < 57 and 1 <= i2 and i2 < 57, p0[i0, i1 - 1, i2 - 1, i3], T.float32(0), dtype="float32")
|
||||
for i3_0 in T.serial(16):
|
||||
for ax0_ax1_ax2_ax3_fused in T.serial(57600):
|
||||
with T.sblock("pad_temp_global"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused // 2 * 28 + ax0_ax1_ax2_ax3_fused // 1920)
|
||||
v2 = T.axis.spatial(58, i0_0_i1_0_i2_0_fused % 2 * 28 + ax0_ax1_ax2_ax3_fused % 1920 // 64)
|
||||
v3 = T.axis.spatial(64, ax0_ax1_ax2_ax3_fused % 64)
|
||||
T.reads(pad_temp[v0, v1, v2, v3])
|
||||
T.writes(pad_temp_global[v0, v1, v2, v3])
|
||||
pad_temp_global[v0, v1, v2, v3] = pad_temp[v0, v1, v2, v3]
|
||||
for ax0_ax1_ax2_ax3_fused in T.serial(2304):
|
||||
with T.sblock("p1_global"):
|
||||
v0 = T.axis.spatial(3, ax0_ax1_ax2_ax3_fused // 768)
|
||||
v1 = T.axis.spatial(3, ax0_ax1_ax2_ax3_fused % 768 // 256)
|
||||
v2 = T.axis.spatial(64, ax0_ax1_ax2_ax3_fused % 256 // 4)
|
||||
v3 = T.axis.spatial(64, i3_0 * 4 + ax0_ax1_ax2_ax3_fused % 4)
|
||||
T.reads(p1_global2[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2])
|
||||
T.writes(p1_global[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2])
|
||||
p1_global[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2] = p1_global2[v3 // 4, v2 // 32, v3 % 4 // 2, v0, v1, v2 % 32, v3 % 2]
|
||||
for i0_1, i1_1, i2_1, i3_1 in T.grid(1, 7, 2, 1):
|
||||
for i0_2_init, i1_2_init, i2_2_init, i3_2_init, i0_3_init, i1_3_init, i2_3_init in T.grid(1, 1, 14, 2, 1, 4, 1):
|
||||
for i3_3_fused_init in T.vectorized(2):
|
||||
with T.sblock("conv2d_nhwc_init"):
|
||||
nn = T.axis.spatial(1, i0_1 + i0_2_init + i0_3_init)
|
||||
yy = T.axis.spatial(56, i0_0_i1_0_i2_0_fused // 2 * 28 + i1_1 * 4 + i1_2_init * 4 + i1_3_init)
|
||||
xx = T.axis.spatial(56, i2_3_init + i0_0_i1_0_i2_0_fused % 2 * 28 + i2_1 * 14 + i2_2_init)
|
||||
ff = T.axis.spatial(64, i3_0 * 4 + i3_1 * 4 + i3_2_init * 2 + i3_3_fused_init)
|
||||
T.reads()
|
||||
T.writes(conv2d_nhwc_global[nn, yy, xx, ff])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure":"SSRSRS"})
|
||||
conv2d_nhwc_global[nn, yy, xx, ff] = T.float32(0)
|
||||
for i4_0, i5_0, i6_0, i0_2, i1_2, i2_2, i3_2, i4_1, i5_1, i6_1, i0_3, i1_3, i2_3 in T.grid(1, 1, 2, 1, 1, 14, 2, 3, 3, 32, 1, 4, 1):
|
||||
for i3_3_fused in T.vectorized(2):
|
||||
with T.sblock("conv2d_nhwc_update"):
|
||||
nn = T.axis.spatial(1, i0_1 + i0_2 + i0_3)
|
||||
yy = T.axis.spatial(56, i0_0_i1_0_i2_0_fused // 2 * 28 + i1_1 * 4 + i1_2 * 4 + i1_3)
|
||||
xx = T.axis.spatial(56, i2_3 + i0_0_i1_0_i2_0_fused % 2 * 28 + i2_1 * 14 + i2_2)
|
||||
ff = T.axis.spatial(64, i3_0 * 4 + i3_1 * 4 + i3_2 * 2 + i3_3_fused)
|
||||
ry = T.axis.reduce(3, i4_0 * 3 + i4_1)
|
||||
rx = T.axis.reduce(3, i5_0 * 3 + i5_1)
|
||||
rc = T.axis.reduce(64, i6_0 * 32 + i6_1)
|
||||
T.reads(conv2d_nhwc_global[nn, yy, xx, ff], pad_temp_global[nn, yy + ry, xx + rx, rc], p1_global[ff // 4, rc // 32, ff % 4 // 2, ry, rx, rc % 32, ff % 2])
|
||||
T.writes(conv2d_nhwc_global[nn, yy, xx, ff])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure":"SSRSRS"})
|
||||
conv2d_nhwc_global[nn, yy, xx, ff] = conv2d_nhwc_global[nn, yy, xx, ff] + pad_temp_global[nn, yy + ry, xx + rx, rc] * p1_global[ff // 4, rc // 32, ff % 4 // 2, ry, rx, rc % 32, ff % 2]
|
||||
for ax0, ax1, ax2 in T.grid(1, 4, 14):
|
||||
for ax3_fused in T.vectorized(4):
|
||||
with T.sblock("conv2d_nhwc_global"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(56, i0_0_i1_0_i2_0_fused // 2 * 28 + i1_1 * 4 + ax1)
|
||||
v2 = T.axis.spatial(56, i0_0_i1_0_i2_0_fused % 2 * 28 + i2_1 * 14 + ax2)
|
||||
v3 = T.axis.spatial(64, i3_0 * 4 + ax3_fused)
|
||||
T.reads(conv2d_nhwc_global[v0, v1, v2, v3])
|
||||
T.writes(conv2d_nhwc[v0, v1, v2, v3])
|
||||
conv2d_nhwc[v0, v1, v2, v3] = conv2d_nhwc_global[v0, v1, v2, v3]
|
||||
|
||||
# fmt: on
|
||||
|
||||
|
||||
def test_layout_rewrite_cache_read():
|
||||
target = Target("llvm")
|
||||
ctx = _create_context(Conv2dCacheRead, target)
|
||||
sch = tvm.s_tir.Schedule(Conv2dCacheRead, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
tvm.ir.assert_structural_equal(sch.mod, Conv2dCacheReadRewritten)
|
||||
|
||||
|
||||
def test_layout_rewrite_cache_read_multiple():
|
||||
target = Target("llvm")
|
||||
ctx = _create_context(Conv2dCacheRead, target)
|
||||
sch = tvm.s_tir.Schedule(Conv2dCacheRead, debug_mask="all")
|
||||
sch.cache_read(sch.get_sblock("p1_global"), 0, "global2")
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
tvm.ir.assert_structural_equal(sch.mod, Conv2dCacheReadMultipleRewritten)
|
||||
|
||||
|
||||
def test_layout_rewrite_int64_index():
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def before(
|
||||
p0: T.Buffer((T.int64(12), T.int64(197), T.int64(64)), "int8"),
|
||||
p1: T.Buffer((T.int64(12), T.int64(197), T.int64(64)), "int8"),
|
||||
T_batch_matmul_NT: T.Buffer((T.int64(12), T.int64(197), T.int64(197)), "int32"),
|
||||
):
|
||||
T.func_attr({"layout_free_buffers": [1], "tirx.noalias": True})
|
||||
for b_0_i_0_fused in T.parallel(T.int64(394)):
|
||||
for j_0 in T.serial(T.int64(1)):
|
||||
for b_1, i_1, j_1 in T.grid(T.int64(1), T.int64(1), T.int64(1)):
|
||||
for b_2_init, i_2_init, j_2_init, b_3_init, i_3_init, j_3_init in T.grid(
|
||||
T.int64(6), T.int64(1), T.int64(197), T.int64(1), T.int64(1), T.int64(1)
|
||||
):
|
||||
with T.sblock("T_batch_matmul_NT_init"):
|
||||
v_b = T.axis.spatial(
|
||||
T.int64(12),
|
||||
b_3_init
|
||||
+ b_0_i_0_fused // T.int64(197) * T.int64(6)
|
||||
+ b_1 * T.int64(6)
|
||||
+ b_2_init,
|
||||
)
|
||||
v_i = T.axis.spatial(
|
||||
T.int64(197),
|
||||
b_0_i_0_fused % T.int64(197) + i_1 + i_2_init + i_3_init,
|
||||
)
|
||||
v_j = T.axis.spatial(
|
||||
T.int64(197),
|
||||
j_3_init + j_0 * T.int64(197) + j_1 * T.int64(197) + j_2_init,
|
||||
)
|
||||
T_batch_matmul_NT[v_b, v_i, v_j] = 0
|
||||
for k_0, b_2, i_2, j_2, k_1, b_3, i_3, j_3 in T.grid(
|
||||
T.int64(64),
|
||||
T.int64(6),
|
||||
T.int64(1),
|
||||
T.int64(197),
|
||||
T.int64(1),
|
||||
T.int64(1),
|
||||
T.int64(1),
|
||||
T.int64(1),
|
||||
):
|
||||
with T.sblock("T_batch_matmul_NT_update"):
|
||||
v_b = T.axis.spatial(
|
||||
T.int64(12),
|
||||
b_3
|
||||
+ b_0_i_0_fused // T.int64(197) * T.int64(6)
|
||||
+ b_1 * T.int64(6)
|
||||
+ b_2,
|
||||
)
|
||||
v_i = T.axis.spatial(
|
||||
T.int64(197), b_0_i_0_fused % T.int64(197) + i_1 + i_2 + i_3
|
||||
)
|
||||
v_j = T.axis.spatial(
|
||||
T.int64(197), j_3 + j_0 * T.int64(197) + j_1 * T.int64(197) + j_2
|
||||
)
|
||||
v_k = T.axis.reduce(T.int64(64), k_0 + k_1)
|
||||
T_batch_matmul_NT[v_b, v_i, v_j] = T_batch_matmul_NT[
|
||||
v_b, v_i, v_j
|
||||
] + T.Cast("int32", p0[v_b, v_i, v_k]) * T.Cast(
|
||||
"int32", p1[v_b, v_j, v_k]
|
||||
)
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def expected(
|
||||
p0: T.Buffer((T.int64(12), T.int64(197), T.int64(64)), "int8"),
|
||||
p1: T.Buffer((T.int64(12), T.int64(197), T.int64(64)), "int8"),
|
||||
T_batch_matmul_NT: T.Buffer((T.int64(12), T.int64(197), T.int64(197)), "int32"),
|
||||
):
|
||||
T.func_attr({"tirx.noalias": True, "layout_free_buffers": [1]})
|
||||
p1_global = T.sblock_alloc_buffer(
|
||||
[T.int64(2), T.int64(64), T.int64(6), T.int64(197)], dtype="int8"
|
||||
)
|
||||
for ax0, ax1, ax2 in T.grid(T.int64(12), T.int64(197), T.int64(64)):
|
||||
with T.sblock("p1_global"):
|
||||
v0, v1, v2 = T.axis.remap("SSS", [ax0, ax1, ax2])
|
||||
T.reads(p1[v0, v1, v2])
|
||||
T.writes(p1_global[v0 // T.int64(6), v2, v0 % T.int64(6), v1])
|
||||
T.sblock_attr({"meta_schedule.layout_rewrite_preproc": True})
|
||||
p1_global[v0 // T.int64(6), v2, v0 % T.int64(6), v1] = p1[v0, v1, v2]
|
||||
for b_0_i_0_fused in T.parallel(T.int64(394)):
|
||||
for j_0, b_1, i_1, j_1 in T.grid(T.int64(1), T.int64(1), T.int64(1), T.int64(1)):
|
||||
for b_2_init, i_2_init, j_2_init, b_3_init, i_3_init, j_3_init in T.grid(
|
||||
T.int64(6), T.int64(1), T.int64(197), T.int64(1), T.int64(1), T.int64(1)
|
||||
):
|
||||
with T.sblock("T_batch_matmul_NT_init"):
|
||||
v_b = T.axis.spatial(
|
||||
T.int64(12),
|
||||
b_3_init
|
||||
+ b_0_i_0_fused // T.int64(197) * T.int64(6)
|
||||
+ b_1 * T.int64(6)
|
||||
+ b_2_init,
|
||||
)
|
||||
v_i = T.axis.spatial(
|
||||
T.int64(197), b_0_i_0_fused % T.int64(197) + i_1 + i_2_init + i_3_init
|
||||
)
|
||||
v_j = T.axis.spatial(
|
||||
T.int64(197),
|
||||
j_3_init + j_0 * T.int64(197) + j_1 * T.int64(197) + j_2_init,
|
||||
)
|
||||
T_batch_matmul_NT[v_b, v_i, v_j] = 0
|
||||
for k_0, b_2, i_2, j_2, k_1, b_3, i_3, j_3 in T.grid(
|
||||
T.int64(64),
|
||||
T.int64(6),
|
||||
T.int64(1),
|
||||
T.int64(197),
|
||||
T.int64(1),
|
||||
T.int64(1),
|
||||
T.int64(1),
|
||||
T.int64(1),
|
||||
):
|
||||
with T.sblock("T_batch_matmul_NT_update"):
|
||||
v_b = T.axis.spatial(
|
||||
T.int64(12),
|
||||
b_3
|
||||
+ b_0_i_0_fused // T.int64(197) * T.int64(6)
|
||||
+ b_1 * T.int64(6)
|
||||
+ b_2,
|
||||
)
|
||||
v_i = T.axis.spatial(
|
||||
T.int64(197), b_0_i_0_fused % T.int64(197) + i_1 + i_2 + i_3
|
||||
)
|
||||
v_j = T.axis.spatial(
|
||||
T.int64(197), j_3 + j_0 * T.int64(197) + j_1 * T.int64(197) + j_2
|
||||
)
|
||||
v_k = T.axis.reduce(T.int64(64), k_0 + k_1)
|
||||
T_batch_matmul_NT[v_b, v_i, v_j] = T_batch_matmul_NT[
|
||||
v_b, v_i, v_j
|
||||
] + T.Cast("int32", p0[v_b, v_i, v_k]) * T.Cast(
|
||||
"int32", p1_global[v_b // T.int64(6), v_k, v_b % T.int64(6), v_j]
|
||||
)
|
||||
|
||||
mod = tvm.IRModule.from_expr(before)
|
||||
mod = _apply_rewrite_layout(mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
+297
@@ -0,0 +1,297 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: E501, F841
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.s_tir.meta_schedule.postproc import RewriteParallelVectorizeUnroll
|
||||
from tvm.s_tir.schedule import Schedule
|
||||
from tvm.s_tir.schedule.testing import assert_structural_equal_ignore_global_symbol
|
||||
from tvm.script import tirx as T
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,not-callable,misplaced-comparison-constant
|
||||
# fmt: off
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Move_PUV:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, [1024, 1024, 1024], dtype="float32")
|
||||
B = T.match_buffer(b, [1024, 1024, 1024], dtype="float32")
|
||||
# body
|
||||
with T.sblock("root"):
|
||||
T.sblock_attr({"meta_schedule.parallel":128, "meta_schedule.vectorize":32})
|
||||
for i0, j0, i1, j1, k0, i2, j2, k1 in T.grid(128, 64, 4, 4, 64, 4, 8, 32):
|
||||
with T.sblock("move"):
|
||||
vi = T.axis.spatial(1024, i0 * 16 + i1 * 4 + i2)
|
||||
vj = T.axis.spatial(1024, j0 * 32 + j1 * 8 + j2)
|
||||
vk = T.axis.spatial(1024, k0 * 32 + k1)
|
||||
T.where((i0 * 4 + i1) * 4 + i2 < 1024 and (j0 * 4 + j1) * 8 + j2 < 1024 and k0 * 32 + k1 < 1024)
|
||||
T.reads([A[vi, vj, vk]])
|
||||
T.writes([B[vi, vj, vk]])
|
||||
B[vi, vj, vk] = A[vi, vj, vk]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def Move_PUV0(a: T.handle, b: T.handle) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, [1024, 1024, 1024], dtype="float32")
|
||||
B = T.match_buffer(b, [1024, 1024, 1024], dtype="float32")
|
||||
# body
|
||||
with T.sblock("root"):
|
||||
for i0_j0_fused in T.parallel(0, 8192):
|
||||
for i1, j1, k0, i2, j2 in T.grid(4, 4, 64, 4, 8):
|
||||
for k1_fused in T.vectorized(0, 32):
|
||||
with T.sblock("move"):
|
||||
vi = T.axis.spatial(1024, i0_j0_fused // 64 * 16 + i1 * 4 + i2)
|
||||
vj = T.axis.spatial(1024, i0_j0_fused % 64 * 32 + j1 * 8 + j2)
|
||||
vk = T.axis.spatial(1024, k0 * 32 + k1_fused)
|
||||
T.where(
|
||||
i0_j0_fused < 4064
|
||||
and i0_j0_fused % 64 < 32
|
||||
and k0 < 32
|
||||
)
|
||||
T.reads([A[vi, vj, vk]])
|
||||
T.writes([B[vi, vj, vk]])
|
||||
B[vi, vj, vk] = A[vi, vj, vk]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Fused_NN_Dense:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(placeholder: T.Buffer((64, 768), "float32"), placeholder_1: T.Buffer((768, 768), "float32"), T_matmul_NT: T.Buffer((64, 768), "float32")) -> None:
|
||||
for i0, i1, i2 in T.grid(64, 768, 768):
|
||||
with T.sblock("T_matmul_NT"):
|
||||
i, j, k = T.axis.remap("SSR", [i0, i1, i2])
|
||||
T.reads(placeholder[i, k], placeholder_1[j, k])
|
||||
T.writes(T_matmul_NT[i, j])
|
||||
with T.init():
|
||||
T_matmul_NT[i, j] = T.float32(0)
|
||||
T_matmul_NT[i, j] = T_matmul_NT[i, j] + placeholder[i, k] * placeholder_1[j, k]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def before_matmul_vectorize(
|
||||
placeholder: T.Buffer((64, 768), "float32"),
|
||||
placeholder_1: T.Buffer((768, 768), "float32"),
|
||||
T_matmul_NT: T.Buffer((64, 768), "float32"),
|
||||
) -> None:
|
||||
with T.sblock("root"):
|
||||
T.reads()
|
||||
T.writes()
|
||||
T.sblock_attr({"meta_schedule.vectorize":64})
|
||||
T_matmul_NT_global = T.sblock_alloc_buffer([64, 768], dtype="float32")
|
||||
for i0_0, i1_0, i0_1, i1_1 in T.grid(1, 16, 1, 3):
|
||||
for i2_0, i0_2, i1_2, i2_1, i0_3, i1_3 in T.grid(48, 8, 1, 16, 8, 16):
|
||||
with T.sblock("T_matmul_NT"):
|
||||
i = T.axis.spatial(64, i0_2 * 8 + i0_3)
|
||||
j = T.axis.spatial(768, i1_0 * 48 + i1_1 * 16 + i1_3)
|
||||
k = T.axis.reduce(768, i2_0 * 16 + i2_1)
|
||||
T.reads(placeholder[i, k], placeholder_1[j, k])
|
||||
T.writes(T_matmul_NT_global[i, j])
|
||||
with T.init():
|
||||
T_matmul_NT_global[i, j] = T.float32(0)
|
||||
T_matmul_NT_global[i, j] = T_matmul_NT_global[i, j] + placeholder[i, k] * placeholder_1[j, k]
|
||||
for ax0, ax1 in T.grid(64, 16):
|
||||
with T.sblock("T_matmul_NT_global"):
|
||||
v0 = T.axis.spatial(64, ax0)
|
||||
v1 = T.axis.spatial(768, i1_0 * 48 + i1_1 * 16 + ax1)
|
||||
T.reads(T_matmul_NT_global[v0, v1])
|
||||
T.writes(T_matmul_NT[v0, v1])
|
||||
T_matmul_NT[v0, v1] = T_matmul_NT_global[v0, v1]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def after_matmul_vectorize(
|
||||
placeholder: T.Buffer((64, 768), "float32"),
|
||||
placeholder_1: T.Buffer((768, 768), "float32"),
|
||||
T_matmul_NT: T.Buffer((64, 768), "float32"),
|
||||
) -> None:
|
||||
T_matmul_NT_global = T.sblock_alloc_buffer([64, 768], dtype="float32")
|
||||
for i0_0, i1_0, i0_1, i1_1 in T.grid(1, 16, 1, 3):
|
||||
for i2_0, i0_2, i1_2, i2_1, i0_3 in T.grid(48, 8, 1, 16, 8):
|
||||
for i1_3_fused in T.vectorized(16):
|
||||
with T.sblock("T_matmul_NT"):
|
||||
i = T.axis.spatial(64, i0_2 * 8 + i0_3)
|
||||
j = T.axis.spatial(768, i1_0 * 48 + i1_1 * 16 + i1_3_fused)
|
||||
k = T.axis.reduce(768, i2_0 * 16 + i2_1)
|
||||
T.reads(placeholder[i, k], placeholder_1[j, k])
|
||||
T.writes(T_matmul_NT_global[i, j])
|
||||
with T.init():
|
||||
T_matmul_NT_global[i, j] = T.float32(0)
|
||||
T_matmul_NT_global[i, j] = T_matmul_NT_global[i, j] + placeholder[i, k] * placeholder_1[j, k]
|
||||
for ax0 in T.serial(64):
|
||||
for ax1_fused in T.vectorized(16):
|
||||
with T.sblock("T_matmul_NT_global"):
|
||||
v0 = T.axis.spatial(64, ax0)
|
||||
v1 = T.axis.spatial(768, i1_0 * 48 + i1_1 * 16 + ax1_fused)
|
||||
T.reads(T_matmul_NT_global[v0, v1])
|
||||
T.writes(T_matmul_NT[v0, v1])
|
||||
T_matmul_NT[v0, v1] = T_matmul_NT_global[v0, v1]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def before_postproc_add(
|
||||
lhs: T.Buffer((1, 8, 56, 56, 32), "uint8"),
|
||||
rhs: T.Buffer((1, 8, 56, 56, 32), "uint8"),
|
||||
add_compute: T.Buffer((1, 8, 56, 56, 32), "uint8"),
|
||||
) -> None:
|
||||
with T.sblock("root"):
|
||||
T.sblock_attr({"meta_schedule.parallel":64, "meta_schedule.vectorize":128})
|
||||
for n, c0, h, w, c1 in T.grid(1, 8, 56, 56, 32):
|
||||
with T.sblock("add_compute"):
|
||||
v0, v1, v2, v3, v4 = T.axis.remap("SSSSS", [n, c0, h, w, c1])
|
||||
T.reads(lhs[v0, v1, v2, v3, v4], rhs[v0, v1, v2, v3, v4])
|
||||
T.writes(add_compute[v0, v1, v2, v3, v4])
|
||||
add_compute[v0, v1, v2, v3, v4] = lhs[v0, v1, v2, v3, v4] + rhs[v0, v1, v2, v3, v4]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def after_postproc_add(
|
||||
lhs: T.Buffer((1, 8, 56, 56, 32), "uint8"),
|
||||
rhs: T.Buffer((1, 8, 56, 56, 32), "uint8"),
|
||||
add_compute: T.Buffer((1, 8, 56, 56, 32), "uint8"),
|
||||
) -> None:
|
||||
with T.sblock("root"):
|
||||
for n_c0_h_w_c1_fused_0 in T.parallel(0, 6272):
|
||||
for n_c0_h_w_c1_fused_1 in T.vectorized(0, 128):
|
||||
with T.sblock("add_compute"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(8, (n_c0_h_w_c1_fused_0 * 128 + n_c0_h_w_c1_fused_1) // 100352)
|
||||
v2 = T.axis.spatial(56, (n_c0_h_w_c1_fused_0 * 128 + n_c0_h_w_c1_fused_1) % 100352 // 1792)
|
||||
v3 = T.axis.spatial(56, (n_c0_h_w_c1_fused_0 * 128 + n_c0_h_w_c1_fused_1) % 1792 // 32)
|
||||
v4 = T.axis.spatial(32, (n_c0_h_w_c1_fused_0 * 128 + n_c0_h_w_c1_fused_1) % 32)
|
||||
T.reads(lhs[v0, v1, v2, v3, v4], rhs[v0, v1, v2, v3, v4])
|
||||
T.writes(add_compute[v0, v1, v2, v3, v4])
|
||||
add_compute[v0, v1, v2, v3, v4] = lhs[v0, v1, v2, v3, v4] + rhs[v0, v1, v2, v3, v4]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def before_postproc_dynamic_shape_vectorize(
|
||||
a: T.handle,
|
||||
b: T.handle,
|
||||
) -> None:
|
||||
n = T.int64()
|
||||
A = T.match_buffer(a, (n,), dtype="float32")
|
||||
B = T.match_buffer(b, (n,), dtype="float32")
|
||||
with T.sblock("root"):
|
||||
T.sblock_attr({"meta_schedule.vectorize": 64})
|
||||
for i in T.serial(0, n):
|
||||
with T.sblock("copy"):
|
||||
vi = T.axis.spatial(n, i)
|
||||
T.reads(A[vi])
|
||||
T.writes(B[vi])
|
||||
B[vi] = A[vi]
|
||||
|
||||
|
||||
# fmt: on
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,not-callable
|
||||
|
||||
|
||||
def test_meta_schedule_postproc_rewrite_parallel_unroll_vectorize():
|
||||
postproc = RewriteParallelVectorizeUnroll()
|
||||
sch = Schedule(Move_PUV)
|
||||
assert postproc.apply(sch)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(sch.mod)
|
||||
tvm.ir.assert_structural_equal(mod["main"], Move_PUV0)
|
||||
|
||||
|
||||
def test_vectorize_inner_loop():
|
||||
sch = Schedule(before_matmul_vectorize)
|
||||
rule = RewriteParallelVectorizeUnroll()
|
||||
assert rule.apply(sch)
|
||||
assert_structural_equal_ignore_global_symbol(sch.mod["main"], after_matmul_vectorize)
|
||||
|
||||
|
||||
def test_parallel_vectorize_add():
|
||||
sch = Schedule(before_postproc_add)
|
||||
rule = RewriteParallelVectorizeUnroll()
|
||||
assert rule.apply(sch)
|
||||
assert_structural_equal_ignore_global_symbol(sch.mod["main"], after_postproc_add)
|
||||
|
||||
|
||||
def test_no_unroll_for_spatial_block():
|
||||
# fmt: off
|
||||
@T.prim_func(s_tir=True)
|
||||
def layer_norm(A: T.Buffer((1, 4, 4, 32), "float32"), B: T.Buffer((4, 4, 32), "float32"), C: T.Buffer((4, 4, 32), "float32"), T_layer_norm: T.Buffer((1, 4, 4, 32), "float32")):
|
||||
with T.sblock("root"):
|
||||
T.sblock_attr({"meta_schedule.unroll_explicit": 512})
|
||||
A_red_temp_v0 = T.sblock_alloc_buffer((1,))
|
||||
A_red_temp_v1 = T.sblock_alloc_buffer((1,))
|
||||
for ax0, k1, k2, k3 in T.grid(1, 4, 4, 32):
|
||||
with T.sblock("A_red_temp"):
|
||||
v_ax0, v_k1, v_k2, v_k3 = T.axis.remap("SRRR", [ax0, k1, k2, k3])
|
||||
T.reads(A[v_ax0, v_k1, v_k2, v_k3])
|
||||
T.writes(A_red_temp_v0[v_ax0], A_red_temp_v1[v_ax0])
|
||||
with T.init():
|
||||
A_red_temp_v0[v_ax0] = T.float32(0)
|
||||
A_red_temp_v1[v_ax0] = T.float32(0)
|
||||
v_A_red_temp_v0: T.let[T.float32] = A_red_temp_v0[v_ax0] + A[v_ax0, v_k1, v_k2, v_k3]
|
||||
v_A_red_temp_v1: T.let[T.float32] = A_red_temp_v1[v_ax0] + A[v_ax0, v_k1, v_k2, v_k3] * A[v_ax0, v_k1, v_k2, v_k3]
|
||||
A_red_temp_v0[v_ax0] = v_A_red_temp_v0
|
||||
A_red_temp_v1[v_ax0] = v_A_red_temp_v1
|
||||
for ax0, ax1, ax2, ax3 in T.grid(1, 4, 4, 32):
|
||||
with T.sblock("T_layer_norm"):
|
||||
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
|
||||
T.reads(A[v_ax0, v_ax1, v_ax2, v_ax3], A_red_temp_v0[v_ax0], A_red_temp_v1[v_ax0], B[v_ax1, v_ax2, v_ax3], C[v_ax1, v_ax2, v_ax3])
|
||||
T.writes(T_layer_norm[v_ax0, v_ax1, v_ax2, v_ax3])
|
||||
T_layer_norm[v_ax0, v_ax1, v_ax2, v_ax3] = (A[v_ax0, v_ax1, v_ax2, v_ax3] - A_red_temp_v0[v_ax0] * T.float32(0.001953125)) * T.rsqrt(A_red_temp_v1[v_ax0] * T.float32(0.001953125) - A_red_temp_v0[v_ax0] * T.float32(0.001953125) * (A_red_temp_v0[v_ax0] * T.float32(0.001953125)) + T.float32(1.0000000000000001e-05)) * B[v_ax1, v_ax2, v_ax3] + C[v_ax1, v_ax2, v_ax3]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def expected(A: T.Buffer((1, 4, 4, 32), "float32"), B: T.Buffer((4, 4, 32), "float32"), C: T.Buffer((4, 4, 32), "float32"), T_layer_norm: T.Buffer((1, 4, 4, 32), "float32")):
|
||||
with T.sblock("root"):
|
||||
A_red_temp_v0 = T.sblock_alloc_buffer((1,))
|
||||
A_red_temp_v1 = T.sblock_alloc_buffer((1,))
|
||||
for ax0 in T.serial(1, annotations={"pragma_auto_unroll_max_step": 512, "pragma_unroll_explicit": 1}):
|
||||
for k1, k2, k3 in T.grid(4, 4, 32):
|
||||
with T.sblock("A_red_temp"):
|
||||
v_ax0 = T.axis.spatial(1, 0)
|
||||
v_k1, v_k2, v_k3 = T.axis.remap("RRR", [k1, k2, k3])
|
||||
T.reads(A[0, v_k1, v_k2, v_k3])
|
||||
T.writes(A_red_temp_v0[0], A_red_temp_v1[0])
|
||||
with T.init():
|
||||
A_red_temp_v0[0] = T.float32(0)
|
||||
A_red_temp_v1[0] = T.float32(0)
|
||||
v_A_red_temp_v0: T.let[T.float32] = A_red_temp_v0[0] + A[0, v_k1, v_k2, v_k3]
|
||||
v_A_red_temp_v1: T.let[T.float32] = A_red_temp_v1[0] + A[0, v_k1, v_k2, v_k3] * A[0, v_k1, v_k2, v_k3]
|
||||
A_red_temp_v0[0] = v_A_red_temp_v0
|
||||
A_red_temp_v1[0] = v_A_red_temp_v1
|
||||
for ax0, ax1, ax2, ax3 in T.grid(1, 4, 4, 32):
|
||||
with T.sblock("T_layer_norm"):
|
||||
v_ax0 = T.axis.spatial(1, 0)
|
||||
v_ax1, v_ax2, v_ax3 = T.axis.remap("SSS", [ax1, ax2, ax3])
|
||||
T.reads(A[0, v_ax1, v_ax2, v_ax3], A_red_temp_v0[0], A_red_temp_v1[0], B[v_ax1, v_ax2, v_ax3], C[v_ax1, v_ax2, v_ax3])
|
||||
T.writes(T_layer_norm[0, v_ax1, v_ax2, v_ax3])
|
||||
T_layer_norm[0, v_ax1, v_ax2, v_ax3] = (A[0, v_ax1, v_ax2, v_ax3] - A_red_temp_v0[0] * T.float32(0.001953125)) * T.rsqrt(A_red_temp_v1[0] * T.float32(0.001953125) - A_red_temp_v0[0] * T.float32(0.001953125) * (A_red_temp_v0[0] * T.float32(0.001953125)) + T.float32(1.0000000000000001e-05)) * B[v_ax1, v_ax2, v_ax3] + C[v_ax1, v_ax2, v_ax3]
|
||||
# fmt: on
|
||||
|
||||
postproc = RewriteParallelVectorizeUnroll()
|
||||
sch = Schedule(layer_norm)
|
||||
assert postproc.apply(sch)
|
||||
mod = tvm.tirx.transform.StmtSimplify()(sch.mod)
|
||||
assert_structural_equal_ignore_global_symbol(mod["main"], expected)
|
||||
|
||||
|
||||
def test_rewrite_parallel_vectorize_unroll_dynamic_shape_no_crash():
|
||||
sch = Schedule(before_postproc_dynamic_shape_vectorize)
|
||||
rule = RewriteParallelVectorizeUnroll()
|
||||
assert rule.apply(sch)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
+224
@@ -0,0 +1,224 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: E501, F401
|
||||
|
||||
import tvm
|
||||
from tvm import tirx
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def _target() -> Target:
|
||||
return Target("cuda", host="llvm")
|
||||
|
||||
|
||||
def _create_context(mod, target) -> ms.TuneContext:
|
||||
ctx = ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[
|
||||
ms.postproc.RewriteReductionBlock(),
|
||||
],
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
return ctx
|
||||
|
||||
|
||||
# fmt: off
|
||||
# pylint: disable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul_before_rewrite:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_A: T.handle, var_B: T.handle, var_C: T.handle) -> None:
|
||||
A = T.match_buffer(var_A, [512, 512], dtype="float32")
|
||||
B = T.match_buffer(var_B, [512, 512], dtype="float32")
|
||||
C = T.match_buffer(var_C, [512, 512], dtype="float32")
|
||||
C_local = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="local")
|
||||
A_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
B_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
for i0_0_i1_0_fused in T.thread_binding(0, 16, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_fused in T.thread_binding(0, 16, thread="vthread.x"):
|
||||
for i0_2_i1_2_fused in T.thread_binding(0, 8, thread="threadIdx.x"):
|
||||
for i2_0 in T.serial(0, 1):
|
||||
for ax0_ax1_fused_0 in T.serial(0, 32768):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(0, 8, thread="threadIdx.x"):
|
||||
with T.sblock("A_shared"):
|
||||
v0 = T.axis.spatial(512, (ax0_ax1_fused_0 * 8 + ax0_ax1_fused_1) // 512)
|
||||
v1 = T.axis.spatial(512, (ax0_ax1_fused_0 * 8 + ax0_ax1_fused_1) % 512)
|
||||
T.reads([A[v0, v1]])
|
||||
T.writes([A_shared[v0, v1]])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch":1})
|
||||
A_shared[v0, v1] = A[v0, v1]
|
||||
for ax0_ax1_fused_0 in T.serial(0, 1024):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(0, 8, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_2 in T.vectorized(0, 2):
|
||||
with T.sblock("B_shared"):
|
||||
v0 = T.axis.spatial(512, (ax0_ax1_fused_0 * 16 + ax0_ax1_fused_1 * 2 + ax0_ax1_fused_2) // 32)
|
||||
v1 = T.axis.spatial(512, i0_0_i1_0_fused * 32 + (ax0_ax1_fused_0 * 16 + ax0_ax1_fused_1 * 2 + ax0_ax1_fused_2) % 32)
|
||||
T.reads([B[v0, v1]])
|
||||
T.writes([B_shared[v0, v1]])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch":2})
|
||||
B_shared[v0, v1] = B[v0, v1]
|
||||
for i2_1, i0_3, i1_3, i2_2, i0_4, i1_4 in T.grid(16, 2, 2, 32, 16, 2):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(512, i0_1_i1_1_fused * 32 + i0_3 * 16 + i0_4)
|
||||
j = T.axis.spatial(512, i0_0_i1_0_fused * 32 + i0_2_i1_2_fused * 4 + i1_3 * 2 + i1_4)
|
||||
k = T.axis.reduce(512, i2_1 * 32 + i2_2)
|
||||
T.reads([C_local[i, j], A_shared[i, k], B_shared[k, j]])
|
||||
T.writes([C_local[i, j]])
|
||||
with T.init():
|
||||
C_local[i, j] = T.float32(0)
|
||||
C_local[i, j] = C_local[i, j] + A_shared[i, k] * B_shared[k, j]
|
||||
for ax0, ax1 in T.grid(32, 4):
|
||||
with T.sblock("C_local"):
|
||||
v0 = T.axis.spatial(512, i0_1_i1_1_fused * 32 + ax0)
|
||||
v1 = T.axis.spatial(512, i0_0_i1_0_fused * 32 + i0_2_i1_2_fused * 4 + ax1)
|
||||
T.reads([C_local[v0, v1]])
|
||||
T.writes([C[v0, v1]])
|
||||
C[v0, v1] = C_local[v0, v1]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul_after_rewrite:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_A: T.handle, var_B: T.handle, var_C: T.handle) -> None:
|
||||
A = T.match_buffer(var_A, [512, 512], dtype="float32")
|
||||
B = T.match_buffer(var_B, [512, 512], dtype="float32")
|
||||
C = T.match_buffer(var_C, [512, 512], dtype="float32")
|
||||
C_local = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="local")
|
||||
A_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
B_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
for i0_0_i1_0_fused in T.thread_binding(0, 16, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_fused in T.thread_binding(0, 16, thread="vthread.x"):
|
||||
for i0_2_i1_2_fused in T.thread_binding(0, 8, thread="threadIdx.x"):
|
||||
for i2_0 in T.serial(0, 1):
|
||||
for ax0_ax1_fused_0 in T.serial(0, 32768):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(0, 8, thread="threadIdx.x"):
|
||||
with T.sblock("A_shared"):
|
||||
v0 = T.axis.spatial(512, (ax0_ax1_fused_0 * 8 + ax0_ax1_fused_1) // 512)
|
||||
v1 = T.axis.spatial(512, (ax0_ax1_fused_0 * 8 + ax0_ax1_fused_1) % 512)
|
||||
T.reads([A[v0, v1]])
|
||||
T.writes([A_shared[v0, v1]])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch":1})
|
||||
A_shared[v0, v1] = A[v0, v1]
|
||||
for ax0_ax1_fused_0 in T.serial(0, 1024):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(0, 8, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_2 in T.vectorized(0, 2):
|
||||
with T.sblock("B_shared"):
|
||||
v0 = T.axis.spatial(512, (ax0_ax1_fused_0 * 16 + ax0_ax1_fused_1 * 2 + ax0_ax1_fused_2) // 32)
|
||||
v1 = T.axis.spatial(512, i0_0_i1_0_fused * 32 + (ax0_ax1_fused_0 * 16 + ax0_ax1_fused_1 * 2 + ax0_ax1_fused_2) % 32)
|
||||
T.reads([B[v0, v1]])
|
||||
T.writes([B_shared[v0, v1]])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch":2})
|
||||
B_shared[v0, v1] = B[v0, v1]
|
||||
for i0_3_init, i1_3_init, i0_4_init, i1_4_init in T.grid(2, 2, 16, 2):
|
||||
with T.sblock("C_init"):
|
||||
i = T.axis.spatial(512, i0_1_i1_1_fused * 32 + i0_3_init * 16 + i0_4_init)
|
||||
j = T.axis.spatial(512, i0_0_i1_0_fused * 32 + i0_2_i1_2_fused * 4 + i1_3_init * 2 + i1_4_init)
|
||||
T.reads([])
|
||||
T.writes([C_local[i, j]])
|
||||
C_local[i, j] = T.float32(0)
|
||||
for i2_1, i0_3, i1_3, i2_2, i0_4, i1_4 in T.grid(16, 2, 2, 32, 16, 2):
|
||||
with T.sblock("C_update"):
|
||||
i = T.axis.spatial(512, i0_1_i1_1_fused * 32 + i0_3 * 16 + i0_4)
|
||||
j = T.axis.spatial(512, i0_0_i1_0_fused * 32 + i0_2_i1_2_fused * 4 + i1_3 * 2 + i1_4)
|
||||
k = T.axis.reduce(512, i2_1 * 32 + i2_2)
|
||||
T.reads([C_local[i, j], A_shared[i, k], B_shared[k, j]])
|
||||
T.writes([C_local[i, j]])
|
||||
C_local[i, j] = C_local[i, j] + A_shared[i, k] * B_shared[k, j]
|
||||
for ax0, ax1 in T.grid(32, 4):
|
||||
with T.sblock("C_local"):
|
||||
v0 = T.axis.spatial(512, i0_1_i1_1_fused * 32 + ax0)
|
||||
v1 = T.axis.spatial(512, i0_0_i1_0_fused * 32 + i0_2_i1_2_fused * 4 + ax1)
|
||||
T.reads([C_local[v0, v1]])
|
||||
T.writes([C[v0, v1]])
|
||||
C[v0, v1] = C_local[v0, v1]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Softmax_cross_thread_reduction:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")) -> None:
|
||||
T_softmax_maxelem_shared = T.sblock_alloc_buffer([256], dtype="float32", scope="shared")
|
||||
T_softmax_expsum_shared = T.sblock_alloc_buffer([256], dtype="float32", scope="shared")
|
||||
for i0 in T.serial(256):
|
||||
for ax0, ax1_0 in T.grid(1, 8):
|
||||
for ax1_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
i0_1 = T.axis.spatial(256, i0)
|
||||
k = T.axis.reduce(256, ax1_0 * 32 + ax1_1)
|
||||
T.reads(T_softmax_maxelem_shared[i0_1], A[i0_1, k])
|
||||
T.writes(T_softmax_maxelem_shared[i0_1])
|
||||
with T.init():
|
||||
T_softmax_maxelem_shared[i0_1] = T.float32(-3.4028234663852886e+38)
|
||||
T_softmax_maxelem_shared[i0_1] = T.max(T_softmax_maxelem_shared[i0_1], A[i0_1, k])
|
||||
for ax0, ax1_0 in T.grid(1, 8):
|
||||
for ax1_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
i0_2 = T.axis.spatial(256, i0)
|
||||
k = T.axis.reduce(256, ax1_0 * 32 + ax1_1)
|
||||
T.reads(T_softmax_expsum_shared[i0_2], A[i0_2, k], T_softmax_maxelem_shared[i0_2])
|
||||
T.writes(T_softmax_expsum_shared[i0_2])
|
||||
with T.init():
|
||||
T_softmax_expsum_shared[i0_2] = T.float32(0)
|
||||
T_softmax_expsum_shared[i0_2] = T_softmax_expsum_shared[i0_2] + T.exp(A[i0_2, k] - T_softmax_maxelem_shared[i0_2], dtype="float32")
|
||||
for i1_0 in T.serial(8):
|
||||
for i1_1 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
i0_3 = T.axis.spatial(256, i0)
|
||||
i1 = T.axis.spatial(256, i1_0 * 32 + i1_1)
|
||||
T.reads(A[i0_3, i1], T_softmax_maxelem_shared[i0_3], T_softmax_expsum_shared[i0_3])
|
||||
T.writes(T_softmax_norm[i0_3, i1])
|
||||
T.sblock_attr({"axis":1})
|
||||
T_softmax_norm[i0_3, i1] = T.exp(A[i0_3, i1] - T_softmax_maxelem_shared[i0_3], dtype="float32") / T_softmax_expsum_shared[i0_3]
|
||||
|
||||
|
||||
# pylint: enable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
# fmt: on
|
||||
|
||||
|
||||
def test_rewrite_tiled_matmul():
|
||||
mod = Matmul_before_rewrite
|
||||
target = _target()
|
||||
ctx = _create_context(mod, target)
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
tvm.ir.assert_structural_equal(sch.mod, Matmul_after_rewrite)
|
||||
|
||||
|
||||
def test_rewrite_softmax():
|
||||
mod = Softmax_cross_thread_reduction
|
||||
target = _target()
|
||||
ctx = _create_context(mod, target)
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
# The module should not be rewritten
|
||||
tvm.ir.assert_structural_equal(sch.mod, Softmax_cross_thread_reduction)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_rewrite_tiled_matmul()
|
||||
test_rewrite_softmax()
|
||||
@@ -0,0 +1,513 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: F401, F841
|
||||
import tvm
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.tensor_intrin import cuda, rocm, x86
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2dNCHWcVNNIModuleTiled:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"),
|
||||
placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"),
|
||||
conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32"),
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
for (
|
||||
i0_0,
|
||||
i1_0,
|
||||
i2_0,
|
||||
i3_0,
|
||||
i4_0_0,
|
||||
i0_1,
|
||||
i1_1,
|
||||
i2_1,
|
||||
i3_1,
|
||||
i4_0_1,
|
||||
i5_0,
|
||||
i6_0,
|
||||
i7_0,
|
||||
i8_0,
|
||||
i9_0_0,
|
||||
i0_2,
|
||||
i1_2,
|
||||
i2_2,
|
||||
i3_2,
|
||||
i4_0_2,
|
||||
i5_1,
|
||||
i6_1,
|
||||
i7_1,
|
||||
i8_1,
|
||||
i9_0_1,
|
||||
i0_3,
|
||||
i1_3,
|
||||
i2_3,
|
||||
i3_3,
|
||||
i4_0_3,
|
||||
) in T.grid(
|
||||
1,
|
||||
1,
|
||||
2,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
4,
|
||||
1,
|
||||
14,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
4,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
4,
|
||||
7,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
4,
|
||||
1,
|
||||
1,
|
||||
1,
|
||||
4,
|
||||
4,
|
||||
1,
|
||||
):
|
||||
with T.sblock("conv2d_NCHWc_int8_o"):
|
||||
n = T.axis.spatial(1, 0)
|
||||
oc_chunk = T.axis.spatial(16, i1_1 * 4 + i1_2)
|
||||
oh = T.axis.spatial(56, i2_0 * 28 + i2_2 * 4 + i2_3)
|
||||
ow = T.axis.spatial(56, i3_1 * 4 + i3_3)
|
||||
oc_block_o = T.axis.spatial(1, 0)
|
||||
kh = T.axis.reduce(1, 0)
|
||||
kw = T.axis.reduce(1, 0)
|
||||
ic_outer, ic_f_inner = T.axis.remap("RR", [i7_0, i8_1])
|
||||
ic_s_inner_o = T.axis.reduce(1, 0)
|
||||
T.reads(
|
||||
placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 : ic_f_inner * 4 + 4],
|
||||
placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, 0:16, 0:4],
|
||||
)
|
||||
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, 0:16])
|
||||
T.sblock_attr({"meta_schedule.auto_tensorize": "dot_16x4_vnni"})
|
||||
with T.init():
|
||||
for i4_1 in T.serial(16):
|
||||
with T.sblock("conv2d_NCHWc_int8_init"):
|
||||
oc_block_init = T.axis.spatial(16, i4_1)
|
||||
T.reads()
|
||||
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block_init])
|
||||
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block_init] = 0
|
||||
for i4_1, i9_1 in T.grid(16, 4):
|
||||
with T.sblock("conv2d_NCHWc_int8"):
|
||||
oc_block, ic_s_inner = T.axis.remap("SR", [i4_1, i9_1])
|
||||
T.reads(
|
||||
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block],
|
||||
placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
|
||||
placeholder_1[
|
||||
oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner
|
||||
],
|
||||
)
|
||||
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SSRSRS"})
|
||||
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = conv2d_NCHWc_int8[
|
||||
n, oc_chunk, oh, ow, oc_block
|
||||
] + T.cast(
|
||||
placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
|
||||
"int32",
|
||||
) * T.cast(
|
||||
placeholder_1[
|
||||
oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner
|
||||
],
|
||||
"int32",
|
||||
)
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2dNCHWcVNNIModuleTensorized:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"),
|
||||
placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"),
|
||||
conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
for i0_0, i1_0, i2_0, i3_0, i4_0_0, i0_1, i1_1, i2_1, i3_1, i4_0_1, i5_0, i6_0 in T.grid(
|
||||
1, 1, 2, 1, 1, 1, 4, 1, 14, 1, 1, 1
|
||||
):
|
||||
for i1_2_init, i2_2_init, i2_3_init, i3_3_init in T.grid(4, 7, 4, 4):
|
||||
with T.sblock("conv2d_NCHWc_int8_o_init"):
|
||||
n = T.axis.spatial(1, 0)
|
||||
oc_chunk = T.axis.spatial(16, i1_1 * 4 + i1_2_init)
|
||||
oh = T.axis.spatial(56, i2_0 * 28 + i2_2_init * 4 + i2_3_init)
|
||||
ow = T.axis.spatial(56, i3_1 * 4 + i3_3_init)
|
||||
oc_block_o = T.axis.spatial(1, 0)
|
||||
T.reads()
|
||||
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, 0:16])
|
||||
for i4_1 in T.vectorized(16):
|
||||
with T.sblock("conv2d_NCHWc_int8_init"):
|
||||
oc_block_init = T.axis.spatial(16, i4_1)
|
||||
T.reads()
|
||||
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block_init])
|
||||
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block_init] = 0
|
||||
for (
|
||||
i7_0,
|
||||
i8_0,
|
||||
i9_0_0,
|
||||
i0_2,
|
||||
i1_2,
|
||||
i2_2,
|
||||
i3_2,
|
||||
i4_0_2,
|
||||
i5_1,
|
||||
i6_1,
|
||||
i7_1,
|
||||
i8_1,
|
||||
i9_0_1,
|
||||
i0_3,
|
||||
i1_3,
|
||||
i2_3,
|
||||
i3_3,
|
||||
i4_0_3,
|
||||
) in T.grid(4, 1, 1, 1, 4, 7, 1, 1, 1, 1, 1, 4, 1, 1, 1, 4, 4, 1):
|
||||
with T.sblock("conv2d_NCHWc_int8_o_update"):
|
||||
n = T.axis.spatial(1, 0)
|
||||
oc_chunk = T.axis.spatial(16, i1_1 * 4 + i1_2)
|
||||
oh = T.axis.spatial(56, i2_0 * 28 + i2_2 * 4 + i2_3)
|
||||
ow = T.axis.spatial(56, i3_1 * 4 + i3_3)
|
||||
oc_block_o = T.axis.spatial(1, 0)
|
||||
kh = T.axis.reduce(1, 0)
|
||||
kw = T.axis.reduce(1, 0)
|
||||
ic_outer, ic_f_inner = T.axis.remap("RR", [i7_0, i8_1])
|
||||
ic_s_inner_o = T.axis.reduce(1, 0)
|
||||
T.reads(
|
||||
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, 0:16],
|
||||
placeholder[
|
||||
n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 : ic_f_inner * 4 + 4
|
||||
],
|
||||
placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, 0:16, 0:4],
|
||||
)
|
||||
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, 0:16])
|
||||
A = T.match_buffer(
|
||||
placeholder[
|
||||
n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 : ic_f_inner * 4 + 4
|
||||
],
|
||||
[4],
|
||||
dtype="uint8",
|
||||
offset_factor=1,
|
||||
)
|
||||
B = T.match_buffer(
|
||||
placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, 0:16, 0:4],
|
||||
[16, 4],
|
||||
dtype="int8",
|
||||
offset_factor=1,
|
||||
)
|
||||
C = T.match_buffer(
|
||||
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, 0:16],
|
||||
[16],
|
||||
dtype="int32",
|
||||
offset_factor=1,
|
||||
)
|
||||
A_u8x4 = A.vload([0], "uint8x4")
|
||||
A_i32 = T.reinterpret(A_u8x4, dtype="int32")
|
||||
B_i8x64 = B.vload([0, 0], dtype="int8x64")
|
||||
B_i32x16 = T.reinterpret(B_i8x64, dtype="int32x16")
|
||||
C_i32x16 = C.vload([0], dtype="int32x16")
|
||||
C[T.ramp(0, 1, 16)] = T.call_llvm_pure_intrin(
|
||||
T.llvm_lookup_intrinsic_id("llvm.x86.avx512.vpdpbusd.512"),
|
||||
C_i32x16,
|
||||
T.broadcast(A_i32, 16),
|
||||
B_i32x16,
|
||||
dtype="int32x16",
|
||||
)
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class DenseDP4ATiled:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
X: T.Buffer((128, 128), "int8"),
|
||||
W: T.Buffer((128, 128), "int8"),
|
||||
compute: T.Buffer((128, 128), "int32"),
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
compute_local = T.sblock_alloc_buffer([128, 128], dtype="int32", scope="local")
|
||||
X_shared = T.sblock_alloc_buffer([128, 128], dtype="int8", scope="shared")
|
||||
W_shared = T.sblock_alloc_buffer([128, 128], dtype="int8", scope="shared")
|
||||
for i0_0_i1_0_fused in T.thread_binding(16, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_fused in T.thread_binding(2, thread="vthread.x"):
|
||||
for i0_2_i1_2_fused in T.thread_binding(2, thread="threadIdx.x"):
|
||||
for i2_0_0 in T.serial(2):
|
||||
for ax0_ax1_fused in T.serial(1024):
|
||||
with T.sblock("X_shared"):
|
||||
v0 = T.axis.spatial(
|
||||
128, i0_0_i1_0_fused // 2 * 16 + ax0_ax1_fused // 64
|
||||
)
|
||||
v1 = T.axis.spatial(128, i2_0_0 * 64 + ax0_ax1_fused % 64)
|
||||
T.reads(X[v0, v1])
|
||||
T.writes(X_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 4})
|
||||
X_shared[v0, v1] = X[v0, v1]
|
||||
for ax0_ax1_fused in T.serial(4096):
|
||||
with T.sblock("W_shared"):
|
||||
v0 = T.axis.spatial(
|
||||
128, i0_0_i1_0_fused % 2 * 64 + ax0_ax1_fused // 64
|
||||
)
|
||||
v1 = T.axis.spatial(128, i2_0_0 * 64 + ax0_ax1_fused % 64)
|
||||
T.reads(W[v0, v1])
|
||||
T.writes(W_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 1})
|
||||
W_shared[v0, v1] = W[v0, v1]
|
||||
for i2_0_1, i0_3, i1_3, i2_0_2, i0_4, i1_4 in T.grid(2, 4, 16, 8, 4, 1):
|
||||
with T.sblock("compute_o"):
|
||||
i = T.axis.spatial(128, i0_0_i1_0_fused // 2 * 16 + i0_3 * 4 + i0_4)
|
||||
j = T.axis.spatial(
|
||||
128,
|
||||
i0_0_i1_0_fused % 2 * 64
|
||||
+ i0_1_i1_1_fused * 32
|
||||
+ i0_2_i1_2_fused * 16
|
||||
+ i1_3,
|
||||
)
|
||||
k_o = T.axis.reduce(32, i2_0_0 * 16 + i2_0_1 * 8 + i2_0_2)
|
||||
T.reads(
|
||||
X_shared[i, k_o * 4 : k_o * 4 + 4],
|
||||
W_shared[j, k_o * 4 : k_o * 4 + 4],
|
||||
)
|
||||
T.writes(compute_local[i, j])
|
||||
T.sblock_attr({"meta_schedule.auto_tensorize": "dp4a_s8s8s32"})
|
||||
with T.init():
|
||||
with T.sblock("compute_init"):
|
||||
T.reads()
|
||||
T.writes(compute_local[i, j])
|
||||
compute_local[i, j] = 0
|
||||
for i2_1 in T.serial(4):
|
||||
with T.sblock("compute"):
|
||||
k = T.axis.reduce(4, i2_1)
|
||||
T.reads(
|
||||
compute_local[i, j],
|
||||
X_shared[i, k_o * 4 + k],
|
||||
W_shared[j, k_o * 4 + k],
|
||||
)
|
||||
T.writes(compute_local[i, j])
|
||||
T.sblock_attr(
|
||||
{"meta_schedule.tiling_structure": "SSSRRSRS"}
|
||||
)
|
||||
compute_local[i, j] = compute_local[i, j] + T.cast(
|
||||
X_shared[i, k_o * 4 + k], "int32"
|
||||
) * T.cast(W_shared[j, k_o * 4 + k], "int32")
|
||||
for ax0, ax1 in T.grid(16, 16):
|
||||
with T.sblock("compute_local"):
|
||||
v0 = T.axis.spatial(128, i0_0_i1_0_fused // 2 * 16 + ax0)
|
||||
v1 = T.axis.spatial(
|
||||
128,
|
||||
i0_0_i1_0_fused % 2 * 64
|
||||
+ i0_1_i1_1_fused * 32
|
||||
+ i0_2_i1_2_fused * 16
|
||||
+ ax1,
|
||||
)
|
||||
T.reads(compute_local[v0, v1])
|
||||
T.writes(compute[v0, v1])
|
||||
compute[v0, v1] = compute_local[v0, v1]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class DenseDP4ATensorized:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
X: T.Buffer((128, 128), "int8"),
|
||||
W: T.Buffer((128, 128), "int8"),
|
||||
compute: T.Buffer((128, 128), "int32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
compute_local = T.sblock_alloc_buffer([128, 128], dtype="int32", scope="local")
|
||||
X_shared = T.sblock_alloc_buffer([128, 128], dtype="int8", scope="shared")
|
||||
W_shared = T.sblock_alloc_buffer([128, 128], dtype="int8", scope="shared")
|
||||
for i0_0_i1_0_fused in T.thread_binding(16, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_fused in T.thread_binding(2, thread="vthread.x"):
|
||||
for i0_2_i1_2_fused in T.thread_binding(2, thread="threadIdx.x"):
|
||||
for i0_3_init, i1_3_init, i0_4_init in T.grid(4, 16, 4):
|
||||
with T.sblock("compute_o_init"):
|
||||
i = T.axis.spatial(
|
||||
128, i0_0_i1_0_fused // 2 * 16 + i0_3_init * 4 + i0_4_init
|
||||
)
|
||||
j = T.axis.spatial(
|
||||
128,
|
||||
i0_0_i1_0_fused % 2 * 64
|
||||
+ i0_1_i1_1_fused * 32
|
||||
+ i0_2_i1_2_fused * 16
|
||||
+ i1_3_init,
|
||||
)
|
||||
T.reads()
|
||||
T.writes(compute_local[i, j])
|
||||
T.sblock_attr({"meta_schedule.auto_tensorize": ""})
|
||||
with T.sblock("compute_init"):
|
||||
T.reads()
|
||||
T.writes(compute_local[i, j])
|
||||
compute_local[i, j] = 0
|
||||
for i2_0_0 in T.serial(2):
|
||||
for ax0_ax1_fused in T.serial(1024):
|
||||
with T.sblock("X_shared"):
|
||||
v0 = T.axis.spatial(
|
||||
128, i0_0_i1_0_fused // 2 * 16 + ax0_ax1_fused // 64
|
||||
)
|
||||
v1 = T.axis.spatial(128, i2_0_0 * 64 + ax0_ax1_fused % 64)
|
||||
T.reads(X[v0, v1])
|
||||
T.writes(X_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 4})
|
||||
X_shared[v0, v1] = X[v0, v1]
|
||||
for ax0_ax1_fused in T.serial(4096):
|
||||
with T.sblock("W_shared"):
|
||||
v0 = T.axis.spatial(
|
||||
128, i0_0_i1_0_fused % 2 * 64 + ax0_ax1_fused // 64
|
||||
)
|
||||
v1 = T.axis.spatial(128, i2_0_0 * 64 + ax0_ax1_fused % 64)
|
||||
T.reads(W[v0, v1])
|
||||
T.writes(W_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 1})
|
||||
W_shared[v0, v1] = W[v0, v1]
|
||||
for i2_0_1, i0_3, i1_3, i2_0_2, i0_4, i1_4 in T.grid(2, 4, 16, 8, 4, 1):
|
||||
with T.sblock("compute_o_update"):
|
||||
i = T.axis.spatial(128, i0_0_i1_0_fused // 2 * 16 + i0_3 * 4 + i0_4)
|
||||
j = T.axis.spatial(
|
||||
128,
|
||||
i0_0_i1_0_fused % 2 * 64
|
||||
+ i0_1_i1_1_fused * 32
|
||||
+ i0_2_i1_2_fused * 16
|
||||
+ i1_3,
|
||||
)
|
||||
k_o = T.axis.reduce(32, i2_0_0 * 16 + i2_0_1 * 8 + i2_0_2)
|
||||
T.reads(
|
||||
compute_local[i, j],
|
||||
X_shared[i, k_o * 4 : k_o * 4 + 4],
|
||||
W_shared[j, k_o * 4 : k_o * 4 + 4],
|
||||
)
|
||||
T.writes(compute_local[i, j])
|
||||
A = T.match_buffer(
|
||||
X_shared[i, k_o * 4 : k_o * 4 + 4],
|
||||
[4],
|
||||
dtype="int8",
|
||||
scope="shared",
|
||||
align=4,
|
||||
offset_factor=1,
|
||||
)
|
||||
B = T.match_buffer(
|
||||
W_shared[j, k_o * 4 : k_o * 4 + 4],
|
||||
[4],
|
||||
dtype="int8",
|
||||
scope="shared",
|
||||
align=4,
|
||||
offset_factor=1,
|
||||
)
|
||||
C = T.match_buffer(
|
||||
compute_local[i, j],
|
||||
[1],
|
||||
dtype="int32",
|
||||
scope="local",
|
||||
align=4,
|
||||
offset_factor=1,
|
||||
)
|
||||
C[0] = C[0] + T.call_pure_extern(
|
||||
"__dp4a",
|
||||
A[T.ramp(0, 1, 4)],
|
||||
B[T.ramp(0, 1, 4)],
|
||||
0,
|
||||
dtype="int32",
|
||||
)
|
||||
for ax0, ax1 in T.grid(16, 16):
|
||||
with T.sblock("compute_local"):
|
||||
v0 = T.axis.spatial(128, i0_0_i1_0_fused // 2 * 16 + ax0)
|
||||
v1 = T.axis.spatial(
|
||||
128,
|
||||
i0_0_i1_0_fused % 2 * 64
|
||||
+ i0_1_i1_1_fused * 32
|
||||
+ i0_2_i1_2_fused * 16
|
||||
+ ax1,
|
||||
)
|
||||
T.reads(compute_local[v0, v1])
|
||||
T.writes(compute[v0, v1])
|
||||
compute[v0, v1] = compute_local[v0, v1]
|
||||
|
||||
|
||||
def _create_context(mod, target, postprocs) -> ms.TuneContext:
|
||||
ctx = ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=postprocs,
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
return ctx
|
||||
|
||||
|
||||
def test_rewrite_tensorize_conv2d_nchwc_vnni():
|
||||
mod = Conv2dNCHWcVNNIModuleTiled
|
||||
target = tvm.target.Target({"kind": "llvm", "mcpu": "cascadelake", "num-cores": 4})
|
||||
ctx = _create_context(
|
||||
mod,
|
||||
target,
|
||||
[
|
||||
ms.postproc.RewriteReductionBlock(),
|
||||
ms.postproc.RewriteTensorize(True),
|
||||
],
|
||||
)
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
|
||||
for proc in ctx.space_generator.postprocs:
|
||||
proc.apply(sch)
|
||||
|
||||
tvm.ir.assert_structural_equal(sch.mod, Conv2dNCHWcVNNIModuleTensorized)
|
||||
|
||||
|
||||
def test_rewrite_tensorize_dense_dp4a():
|
||||
mod = DenseDP4ATiled
|
||||
target = tvm.target.Target("nvidia/geforce-rtx-3070")
|
||||
ctx = _create_context(
|
||||
mod,
|
||||
target,
|
||||
[
|
||||
ms.postproc.RewriteCooperativeFetch(),
|
||||
ms.postproc.RewriteReductionBlock(),
|
||||
ms.postproc.RewriteTensorize(),
|
||||
],
|
||||
)
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
|
||||
for proc in ctx.space_generator.postprocs:
|
||||
proc.apply(sch)
|
||||
|
||||
tvm.ir.assert_structural_equal(sch.mod, DenseDP4ATensorized)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_rewrite_tensorize_conv2d_nchwc_vnni()
|
||||
test_rewrite_tensorize_dense_dp4a()
|
||||
+378
@@ -0,0 +1,378 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: F401
|
||||
|
||||
import tvm
|
||||
from tvm import tirx
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def _target() -> Target:
|
||||
return Target({"kind": "cuda", "max_threads_per_block": 1024}, host="llvm")
|
||||
|
||||
|
||||
def _create_context(mod, target) -> ms.TuneContext:
|
||||
ctx = ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[ms.postproc.RewriteUnboundBlock()],
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
return ctx
|
||||
|
||||
|
||||
# pylint: disable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Before_cooperative_fetch:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_A: T.handle, var_B: T.handle) -> None:
|
||||
A = T.match_buffer(var_A, [512, 512], dtype="float32")
|
||||
B = T.match_buffer(var_B, [512, 512], dtype="float32")
|
||||
for i, j in T.grid(512, 512):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = A[vi, vj] + 1.0
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class After_cooperative_fetch:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_A: T.handle, var_B: T.handle) -> None:
|
||||
A = T.match_buffer(var_A, [512, 512], dtype="float32")
|
||||
B = T.match_buffer(var_B, [512, 512], dtype="float32")
|
||||
for i_j_fused_0 in T.thread_binding(256, thread="blockIdx.x"):
|
||||
for i_j_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
||||
with T.sblock("C"):
|
||||
vi = T.axis.spatial(512, (i_j_fused_0 * 1024 + i_j_fused_1) // 512)
|
||||
vj = T.axis.spatial(512, (i_j_fused_0 * 1024 + i_j_fused_1) % 512)
|
||||
B[vi, vj] = A[vi, vj] + 1.0
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Before_norm_bmn:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((1, 256, 256), "float32"), D: T.Buffer((1,), "float32")) -> None:
|
||||
C = T.sblock_alloc_buffer([1], dtype="float32")
|
||||
for i0, i1, i2 in T.grid(1, 256, 256):
|
||||
with T.sblock("C"):
|
||||
b, i, j = T.axis.remap("SRR", [i0, i1, i2])
|
||||
with T.init():
|
||||
C[b] = T.float32(0)
|
||||
C[b] = C[b] + A[b, i, j] * A[b, i, j]
|
||||
for i0 in T.serial(1):
|
||||
with T.sblock("D"):
|
||||
b = T.axis.S(1, i0)
|
||||
D[b] = T.sqrt(C[b], dtype="float32")
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class After_norm_bmn:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((1, 256, 256), "float32"), D: T.Buffer((1,), "float32")) -> None:
|
||||
C = T.sblock_alloc_buffer([1], dtype="float32")
|
||||
for i0_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
|
||||
for i0_fused_1 in T.thread_binding(1, thread="threadIdx.x"):
|
||||
for i1, i2 in T.grid(256, 256):
|
||||
with T.sblock("C"):
|
||||
b = T.axis.S(1, 0)
|
||||
i, j = T.axis.remap("RR", [i1, i2])
|
||||
with T.init():
|
||||
C[b] = T.float32(0)
|
||||
C[b] = C[b] + A[b, i, j] * A[b, i, j]
|
||||
for i0_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
|
||||
for i0_fused_1 in T.thread_binding(1, thread="threadIdx.x"):
|
||||
with T.sblock("D"):
|
||||
b = T.axis.S(1, 0)
|
||||
D[b] = T.sqrt(C[b], dtype="float32")
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Bert_fused_reshape_transpose_reshape:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
placeholder: T.Buffer((12, 64, 64), "float32"), T_reshape: T.Buffer((64, 768), "float32")
|
||||
) -> None:
|
||||
for i0_i1_fused_0, i0_i1_fused_1 in T.grid(1536, 32):
|
||||
with T.sblock("T_reshape_1"):
|
||||
ax0 = T.axis.spatial(64, (i0_i1_fused_0 * 32 + i0_i1_fused_1) // 768)
|
||||
ax1 = T.axis.spatial(768, (i0_i1_fused_0 * 32 + i0_i1_fused_1) % 768)
|
||||
T.reads(placeholder[ax1 % 768 // 64, (ax1 // 768 + ax0) % 64, ax1 % 64])
|
||||
T.writes(T_reshape[ax0, ax1])
|
||||
T_reshape[ax0, ax1] = placeholder[
|
||||
((ax1 % 64 // 64 + (ax1 // 768 + ax0) % 64) // 64 + ax1 % 768 // 64) % 12,
|
||||
(ax1 % 64 // 64 + (ax1 // 768 + ax0) % 64) % 64,
|
||||
ax1 % 64 % 64,
|
||||
]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Bert_fused_reshape_transpose_reshape_large:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
placeholder: T.Buffer((12, 64, 64), "float32"), T_reshape: T.Buffer((64, 768), "float32")
|
||||
) -> None:
|
||||
for i0_i1_fused_0, i0_i1_fused_1 in T.grid(1536000, 32):
|
||||
with T.sblock("T_reshape_1"):
|
||||
ax0 = T.axis.spatial(64, (i0_i1_fused_0 * 32 + i0_i1_fused_1) // 768)
|
||||
ax1 = T.axis.spatial(768, (i0_i1_fused_0 * 32 + i0_i1_fused_1) % 768)
|
||||
T.reads(placeholder[ax1 % 768 // 64, (ax1 // 768 + ax0) % 64, ax1 % 64])
|
||||
T.writes(T_reshape[ax0, ax1])
|
||||
T_reshape[ax0, ax1] = placeholder[
|
||||
((ax1 % 64 // 64 + (ax1 // 768 + ax0) % 64) // 64 + ax1 % 768 // 64) % 12,
|
||||
(ax1 % 64 // 64 + (ax1 // 768 + ax0) % 64) % 64,
|
||||
ax1 % 64 % 64,
|
||||
]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Bert_fused_reshape_transpose_reshape_after_rub:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
placeholder: T.Buffer((12, 64, 64), "float32"), T_reshape: T.Buffer((64, 768), "float32")
|
||||
) -> None:
|
||||
for i0_i1_fused_0_i0_i1_fused_1_fused_0 in T.thread_binding(48, thread="blockIdx.x"):
|
||||
for i0_i1_fused_0_i0_i1_fused_1_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
||||
with T.sblock("T_reshape_1"):
|
||||
ax0 = T.axis.spatial(
|
||||
64,
|
||||
(
|
||||
i0_i1_fused_0_i0_i1_fused_1_fused_0 * 1024
|
||||
+ i0_i1_fused_0_i0_i1_fused_1_fused_1
|
||||
)
|
||||
// 768,
|
||||
)
|
||||
ax1 = T.axis.spatial(
|
||||
768,
|
||||
(
|
||||
i0_i1_fused_0_i0_i1_fused_1_fused_0 * 1024
|
||||
+ i0_i1_fused_0_i0_i1_fused_1_fused_1
|
||||
)
|
||||
% 768,
|
||||
)
|
||||
T.reads(placeholder[ax1 % 768 // 64, (ax1 // 768 + ax0) % 64, ax1 % 64])
|
||||
T.writes(T_reshape[ax0, ax1])
|
||||
T_reshape[ax0, ax1] = placeholder[
|
||||
((ax1 % 64 // 64 + (ax1 // 768 + ax0) % 64) // 64 + ax1 % 768 // 64) % 12,
|
||||
(ax1 % 64 // 64 + (ax1 // 768 + ax0) % 64) % 64,
|
||||
ax1 % 64 % 64,
|
||||
]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Bert_fused_reshape_transpose_reshape_after_rub_large:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
placeholder: T.Buffer((12, 64, 64), "float32"), T_reshape: T.Buffer((64, 768), "float32")
|
||||
) -> None:
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
for i0_i1_fused_0_i0_i1_fused_1_fused_1 in T.thread_binding(256, thread="blockIdx.x"):
|
||||
for i0_i1_fused_0_i0_i1_fused_1_fused_2 in T.thread_binding(1024, thread="threadIdx.x"):
|
||||
for i0_i1_fused_0_i0_i1_fused_1_fused_0 in T.serial(188):
|
||||
with T.sblock("T_reshape_1"):
|
||||
ax0 = T.axis.spatial(
|
||||
64,
|
||||
(
|
||||
i0_i1_fused_0_i0_i1_fused_1_fused_0 * 262144
|
||||
+ i0_i1_fused_0_i0_i1_fused_1_fused_1 * 1024
|
||||
+ i0_i1_fused_0_i0_i1_fused_1_fused_2
|
||||
)
|
||||
// 768,
|
||||
)
|
||||
ax1 = T.axis.spatial(
|
||||
768,
|
||||
(
|
||||
i0_i1_fused_0_i0_i1_fused_1_fused_0 * 262144
|
||||
+ i0_i1_fused_0_i0_i1_fused_1_fused_1 * 1024
|
||||
+ i0_i1_fused_0_i0_i1_fused_1_fused_2
|
||||
)
|
||||
% 768,
|
||||
)
|
||||
T.where(
|
||||
(
|
||||
i0_i1_fused_0_i0_i1_fused_1_fused_0 * 256
|
||||
+ i0_i1_fused_0_i0_i1_fused_1_fused_1
|
||||
)
|
||||
* 1024
|
||||
+ i0_i1_fused_0_i0_i1_fused_1_fused_2
|
||||
< 49152000
|
||||
)
|
||||
T.reads(placeholder[ax1 % 768 // 64, (ax1 // 768 + ax0) % 64, ax1 % 64])
|
||||
T.writes(T_reshape[ax0, ax1])
|
||||
T_reshape[ax0, ax1] = placeholder[
|
||||
((ax1 % 64 // 64 + (ax1 // 768 + ax0) % 64) // 64 + ax1 % 768 // 64)
|
||||
% 12,
|
||||
(ax1 % 64 // 64 + (ax1 // 768 + ax0) % 64) % 64,
|
||||
ax1 % 64 % 64,
|
||||
]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def before_unrolled_loop(
|
||||
placeholder: T.Buffer((1, 56, 56, 64), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
bgemm = T.sblock_alloc_buffer([6, 6, 196, 64], dtype="float32")
|
||||
inverse = T.sblock_alloc_buffer([4, 4, 196, 64], dtype="float32")
|
||||
for i2_0, i3_0, i2_1, i3_1 in T.grid(98, 4, 2, 16):
|
||||
for i0 in T.unroll(4):
|
||||
for i1 in T.unroll(4):
|
||||
for i4 in T.unroll(6):
|
||||
for i5 in T.unroll(6):
|
||||
with T.sblock("inverse"):
|
||||
vh, vw = T.axis.remap("SS", [i0, i1])
|
||||
p = T.axis.spatial(196, i2_0 * 2 + i2_1)
|
||||
co = T.axis.spatial(64, i3_0 * 16 + i3_1)
|
||||
r_a, r_b = T.axis.remap("RR", [i4, i5])
|
||||
T.reads(bgemm[r_a, r_b, p, co])
|
||||
T.writes(inverse[vh, vw, p, co])
|
||||
with T.init():
|
||||
inverse[vh, vw, p, co] = T.float32(0)
|
||||
inverse[vh, vw, p, co] = inverse[vh, vw, p, co] + bgemm[r_a, r_b, p, co]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def after_unrolled_loop(
|
||||
placeholder: T.Buffer((1, 56, 56, 64), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
bgemm = T.sblock_alloc_buffer([6, 6, 196, 64], dtype="float32")
|
||||
inverse = T.sblock_alloc_buffer([4, 4, 196, 64], dtype="float32")
|
||||
for i2_0_i3_0_i2_1_i3_1_fused_0 in T.thread_binding(13, thread="blockIdx.x"):
|
||||
for i2_0_i3_0_i2_1_i3_1_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
||||
for i0 in T.unroll(4):
|
||||
for i1 in T.unroll(4):
|
||||
for i4 in T.unroll(6):
|
||||
for i5 in T.unroll(6):
|
||||
with T.sblock("inverse"):
|
||||
vh, vw = T.axis.remap("SS", [i0, i1])
|
||||
p = T.axis.spatial(
|
||||
196,
|
||||
(
|
||||
i2_0_i3_0_i2_1_i3_1_fused_0 * 1024
|
||||
+ i2_0_i3_0_i2_1_i3_1_fused_1
|
||||
)
|
||||
// 128
|
||||
* 2
|
||||
+ (
|
||||
i2_0_i3_0_i2_1_i3_1_fused_0 * 1024
|
||||
+ i2_0_i3_0_i2_1_i3_1_fused_1
|
||||
)
|
||||
% 32
|
||||
// 16,
|
||||
)
|
||||
co = T.axis.spatial(
|
||||
64,
|
||||
(
|
||||
i2_0_i3_0_i2_1_i3_1_fused_0 * 1024
|
||||
+ i2_0_i3_0_i2_1_i3_1_fused_1
|
||||
)
|
||||
% 128
|
||||
// 32
|
||||
* 16
|
||||
+ (
|
||||
i2_0_i3_0_i2_1_i3_1_fused_0 * 1024
|
||||
+ i2_0_i3_0_i2_1_i3_1_fused_1
|
||||
)
|
||||
% 16,
|
||||
)
|
||||
r_a, r_b = T.axis.remap("RR", [i4, i5])
|
||||
T.where(
|
||||
i2_0_i3_0_i2_1_i3_1_fused_0 * 1024 + i2_0_i3_0_i2_1_i3_1_fused_1
|
||||
< 12544
|
||||
)
|
||||
T.reads(bgemm[r_a, r_b, p, co])
|
||||
T.writes(inverse[vh, vw, p, co])
|
||||
with T.init():
|
||||
inverse[vh, vw, p, co] = T.float32(0)
|
||||
inverse[vh, vw, p, co] = (
|
||||
inverse[vh, vw, p, co] + bgemm[r_a, r_b, p, co]
|
||||
)
|
||||
|
||||
|
||||
# pylint: enable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
# fmt: on
|
||||
|
||||
|
||||
def test_rewrite_cooperative_fetch():
|
||||
mod = Before_cooperative_fetch
|
||||
target = _target()
|
||||
ctx = _create_context(mod, target)
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
tvm.ir.assert_structural_equal(sch.mod, After_cooperative_fetch)
|
||||
|
||||
|
||||
def test_rewrite_norm_bmn():
|
||||
mod = Before_norm_bmn
|
||||
target = _target()
|
||||
ctx = _create_context(mod, target)
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
tvm.ir.assert_structural_equal(sch.mod, After_norm_bmn)
|
||||
|
||||
|
||||
def test_rewrite_cuda_loop_split_no_reduction():
|
||||
mod = Bert_fused_reshape_transpose_reshape
|
||||
target = Target("nvidia/nvidia-v100", host="llvm")
|
||||
ctx = _create_context(mod, target)
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
tvm.ir.assert_structural_equal(sch.mod, Bert_fused_reshape_transpose_reshape_after_rub)
|
||||
|
||||
|
||||
def test_rewrite_cuda_loop_split_no_reduction_large():
|
||||
mod = Bert_fused_reshape_transpose_reshape_large
|
||||
target = Target("nvidia/nvidia-v100", host="llvm")
|
||||
ctx = _create_context(mod, target)
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
tvm.ir.assert_structural_equal(sch.mod, Bert_fused_reshape_transpose_reshape_after_rub_large)
|
||||
|
||||
|
||||
def test_rewrite_cuda_loop_split_for_kind():
|
||||
mod = before_unrolled_loop
|
||||
target = Target("nvidia/nvidia-v100", host="llvm")
|
||||
ctx = _create_context(mod, target)
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
sch.enter_postproc()
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
tvm.ir.assert_structural_equal(sch.mod["main"], after_unrolled_loop)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_rewrite_cooperative_fetch()
|
||||
test_rewrite_norm_bmn()
|
||||
test_rewrite_cuda_loop_split_no_reduction()
|
||||
test_rewrite_cuda_loop_split_no_reduction_large()
|
||||
test_rewrite_cuda_loop_split_for_kind()
|
||||
@@ -0,0 +1,807 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: E501, F401
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import tirx
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def _target() -> Target:
|
||||
return Target("nvidia/geforce-rtx-3080")
|
||||
|
||||
|
||||
def _create_context(mod, target) -> ms.TuneContext:
|
||||
return ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[ms.postproc.VerifyGPUCode()],
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,not-callable,misplaced-comparison-constant
|
||||
# fmt: off
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2dCuda0:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "T.noalias": True})
|
||||
# var definition
|
||||
threadIdx_x = T.env_thread("threadIdx.x")
|
||||
threadIdx_y = T.env_thread("threadIdx.y")
|
||||
blockIdx_x = T.env_thread("blockIdx.x")
|
||||
blockIdx_y = T.env_thread("blockIdx.y")
|
||||
blockIdx_z = T.env_thread("blockIdx.z")
|
||||
A = T.match_buffer(a, [14*14*256*256], dtype="float32")
|
||||
B = T.match_buffer(b, [14*14*512*256], dtype="float32")
|
||||
# body
|
||||
T.launch_thread(blockIdx_z, 196)
|
||||
B_local = T.decl_buffer([64], "float32", scope="local")
|
||||
Apad_shared = T.decl_buffer([512], "float32", scope="shared")
|
||||
Apad_shared_local = T.decl_buffer([8], "float32", scope="local")
|
||||
T.launch_thread(blockIdx_y, 8)
|
||||
T.launch_thread(blockIdx_x, 4)
|
||||
T.launch_thread(threadIdx_y, 8)
|
||||
T.launch_thread(threadIdx_x, 8)
|
||||
for ff_c_init, nn_c_init in T.grid(8, 8):
|
||||
B_local[ff_c_init * 8 + nn_c_init] = T.float32(0)
|
||||
for rc_outer, ry, rx in T.grid(32, 3, 3):
|
||||
for ax3_inner_outer in T.serial(0, 2):
|
||||
Apad_shared[T.ramp(threadIdx_y * 64 + threadIdx_x * 8 + ax3_inner_outer * 4, 1, 4)] = T.if_then_else(
|
||||
1 <= blockIdx_z // 14 + ry and blockIdx_z // 14 + ry < 15 and 1 <= rx + blockIdx_z % 14 and rx + blockIdx_z % 14 < 15,
|
||||
A[T.ramp(ry * 917504 + blockIdx_z * 65536 + rx * 65536 + rc_outer * 2048 + threadIdx_y * 256 + blockIdx_x * 64 + threadIdx_x * 8 + ax3_inner_outer * 4 - 983040, 1, 4)],
|
||||
T.broadcast(T.float32(0), 4),
|
||||
dtype="float32x4",
|
||||
)
|
||||
for rc_inner in T.serial(0, 8):
|
||||
for ax3 in T.serial(0, 8):
|
||||
Apad_shared_local[ax3] = Apad_shared[rc_inner * 64 + threadIdx_x * 8 + ax3]
|
||||
for ff_c, nn_c in T.grid(8, 8):
|
||||
B_local[ff_c * 8 + nn_c] = B_local[ff_c * 8 + nn_c] + Apad_shared_local[nn_c]
|
||||
for ff_inner_inner_inner, nn_inner_inner_inner in T.grid(8, 8):
|
||||
B[blockIdx_z * 131072 + blockIdx_y * 16384 + threadIdx_y * 2048 + ff_inner_inner_inner * 256 + blockIdx_x * 64 + threadIdx_x * 8 + nn_inner_inner_inner] = B_local[ff_inner_inner_inner * 8 + nn_inner_inner_inner]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2dCuda1:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "T.noalias": True})
|
||||
# var definition
|
||||
threadIdx_x = T.env_thread("threadIdx.x")
|
||||
threadIdx_y = T.env_thread("threadIdx.y")
|
||||
blockIdx_x = T.env_thread("blockIdx.x")
|
||||
blockIdx_y = T.env_thread("blockIdx.y")
|
||||
blockIdx_z = T.env_thread("blockIdx.z")
|
||||
A = T.match_buffer(a, [14*14*256*256], dtype="float32")
|
||||
B = T.match_buffer(b, [14*14*512*256], dtype="float32")
|
||||
# body
|
||||
T.launch_thread(blockIdx_z, 196)
|
||||
B_local = T.decl_buffer([6400000], "float32", scope="local")
|
||||
Apad_shared = T.decl_buffer([512], "float32", scope="shared")
|
||||
Apad_shared_local = T.decl_buffer([8], "float32", scope="local")
|
||||
T.launch_thread(blockIdx_y, 8)
|
||||
T.launch_thread(blockIdx_x, 4)
|
||||
T.launch_thread(threadIdx_y, 8)
|
||||
T.launch_thread(threadIdx_x, 8)
|
||||
for ff_c_init, nn_c_init in T.grid(8, 8):
|
||||
B_local[ff_c_init * 8 + nn_c_init] = T.float32(0)
|
||||
# Access of the last element of B_local prevents buffer
|
||||
# compacting from reducing the amount of shared memory
|
||||
# used.
|
||||
B_local[6400000-1 + ff_c_init*8] = 0.0
|
||||
for rc_outer, ry, rx in T.grid(32, 3, 3):
|
||||
for ax3_inner_outer in T.serial(0, 2):
|
||||
Apad_shared[T.ramp(threadIdx_y * 64 + threadIdx_x * 8 + ax3_inner_outer * 4, 1, 4)] = T.if_then_else(
|
||||
1 <= blockIdx_z // 14 + ry and blockIdx_z // 14 + ry < 15 and 1 <= rx + blockIdx_z % 14 and rx + blockIdx_z % 14 < 15,
|
||||
A[T.ramp(ry * 917504 + blockIdx_z * 65536 + rx * 65536 + rc_outer * 2048 + threadIdx_y * 256 + blockIdx_x * 64 + threadIdx_x * 8 + ax3_inner_outer * 4 - 983040, 1, 4)],
|
||||
T.broadcast(T.float32(0), 4),
|
||||
dtype="float32x4",
|
||||
)
|
||||
for rc_inner in T.serial(0, 8):
|
||||
for ax3 in T.serial(0, 8):
|
||||
Apad_shared_local[ax3] = Apad_shared[rc_inner * 64 + threadIdx_x * 8 + ax3]
|
||||
for ff_c, nn_c in T.grid(8, 8):
|
||||
B_local[ff_c * 8 + nn_c] = B_local[ff_c * 8 + nn_c] + Apad_shared_local[nn_c]
|
||||
for ff_inner_inner_inner, nn_inner_inner_inner in T.grid(8, 8):
|
||||
B[blockIdx_z * 131072 + blockIdx_y * 16384 + threadIdx_y * 2048 + ff_inner_inner_inner * 256 + blockIdx_x * 64 + threadIdx_x * 8 + nn_inner_inner_inner] = B_local[ff_inner_inner_inner * 8 + nn_inner_inner_inner]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2dCuda2:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "T.noalias": True})
|
||||
# var definition
|
||||
threadIdx_x = T.env_thread("threadIdx.x")
|
||||
threadIdx_y = T.env_thread("threadIdx.y")
|
||||
blockIdx_x = T.env_thread("blockIdx.x")
|
||||
blockIdx_y = T.env_thread("blockIdx.y")
|
||||
blockIdx_z = T.env_thread("blockIdx.z")
|
||||
A = T.match_buffer(a, [14*14*256*256], dtype="float32")
|
||||
B = T.match_buffer(b, [14*14*512*256], dtype="float32")
|
||||
# body
|
||||
T.launch_thread(blockIdx_z, 196)
|
||||
B_local = T.decl_buffer([64], "float32", scope="local")
|
||||
Apad_shared = T.decl_buffer([512000], "float32", scope="shared")
|
||||
Apad_shared_local = T.decl_buffer([8], "float32", scope="local")
|
||||
T.launch_thread(blockIdx_y, 8)
|
||||
T.launch_thread(blockIdx_x, 4)
|
||||
T.launch_thread(threadIdx_y, 8)
|
||||
T.launch_thread(threadIdx_x, 8)
|
||||
for ff_c_init, nn_c_init in T.grid(8, 8):
|
||||
B_local[ff_c_init * 8 + nn_c_init] = T.float32(0)
|
||||
for rc_outer, ry, rx in T.grid(32, 3, 3):
|
||||
for ax3_inner_outer in T.serial(0, 2):
|
||||
Apad_shared[T.ramp(threadIdx_y * 64 + threadIdx_x * 8 + ax3_inner_outer * 4, 1, 4)] = T.if_then_else(
|
||||
1 <= blockIdx_z // 14 + ry and blockIdx_z // 14 + ry < 15 and 1 <= rx + blockIdx_z % 14 and rx + blockIdx_z % 14 < 15,
|
||||
A[T.ramp(ry * 917504 + blockIdx_z * 65536 + rx * 65536 + rc_outer * 2048 + threadIdx_y * 256 + blockIdx_x * 64 + threadIdx_x * 8 + ax3_inner_outer * 4 - 983040, 1, 4)],
|
||||
T.broadcast(T.float32(0), 4),
|
||||
dtype="float32x4",
|
||||
)
|
||||
# Access of the last element of Apad_shared prevents
|
||||
# buffer compacting from reducing the amount of shared
|
||||
# memory used.
|
||||
Apad_shared[512000-1] = 0.0
|
||||
for rc_inner in T.serial(0, 8):
|
||||
for ax3 in T.serial(0, 8):
|
||||
Apad_shared_local[ax3] = Apad_shared[rc_inner * 64 + threadIdx_x * 8 + ax3]
|
||||
for ff_c, nn_c in T.grid(8, 8):
|
||||
B_local[ff_c * 8 + nn_c] = B_local[ff_c * 8 + nn_c] + Apad_shared_local[nn_c]
|
||||
for ff_inner_inner_inner, nn_inner_inner_inner in T.grid(8, 8):
|
||||
B[blockIdx_z * 131072 + blockIdx_y * 16384 + threadIdx_y * 2048 + ff_inner_inner_inner * 256 + blockIdx_x * 64 + threadIdx_x * 8 + nn_inner_inner_inner] = B_local[ff_inner_inner_inner * 8 + nn_inner_inner_inner]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2dCuda3:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "T.noalias": True})
|
||||
# var definition
|
||||
threadIdx_x = T.env_thread("threadIdx.x")
|
||||
threadIdx_y = T.env_thread("threadIdx.y")
|
||||
blockIdx_x = T.env_thread("blockIdx.x")
|
||||
blockIdx_y = T.env_thread("blockIdx.y")
|
||||
blockIdx_z = T.env_thread("blockIdx.z")
|
||||
A = T.match_buffer(a, [14*14*256*256], dtype="float32")
|
||||
B = T.match_buffer(b, [14*14*512*256], dtype="float32")
|
||||
# body
|
||||
T.launch_thread(blockIdx_z, 196)
|
||||
B_local = T.decl_buffer([64], "float32", scope="local")
|
||||
Apad_shared = T.decl_buffer([512], "float32", scope="shared")
|
||||
Apad_shared_local = T.decl_buffer([8], "float32", scope="local")
|
||||
T.launch_thread(blockIdx_y, 8)
|
||||
T.launch_thread(blockIdx_x, 4)
|
||||
T.launch_thread(threadIdx_y, 8)
|
||||
T.launch_thread(threadIdx_x, 800000)
|
||||
for ff_c_init, nn_c_init in T.grid(8, 8):
|
||||
B_local[ff_c_init * 8 + nn_c_init] = T.float32(0)
|
||||
for rc_outer, ry, rx in T.grid(32, 3, 3):
|
||||
for ax3_inner_outer in T.serial(0, 2):
|
||||
Apad_shared[T.ramp(threadIdx_y * 64 + threadIdx_x * 8 + ax3_inner_outer * 4, 1, 4)] = T.if_then_else(
|
||||
1 <= blockIdx_z // 14 + ry and blockIdx_z // 14 + ry < 15 and 1 <= rx + blockIdx_z % 14 and rx + blockIdx_z % 14 < 15,
|
||||
A[T.ramp(ry * 917504 + blockIdx_z * 65536 + rx * 65536 + rc_outer * 2048 + threadIdx_y * 256 + blockIdx_x * 64 + threadIdx_x * 8 + ax3_inner_outer * 4 - 983040, 1, 4)],
|
||||
T.broadcast(T.float32(0), 4),
|
||||
dtype="float32x4",
|
||||
)
|
||||
for rc_inner in T.serial(0, 8):
|
||||
for ax3 in T.serial(0, 8):
|
||||
Apad_shared_local[ax3] = Apad_shared[rc_inner * 64 + threadIdx_x * 8 + ax3]
|
||||
for ff_c, nn_c in T.grid(8, 8):
|
||||
B_local[ff_c * 8 + nn_c] = B_local[ff_c * 8 + nn_c] + Apad_shared_local[nn_c]
|
||||
for ff_inner_inner_inner, nn_inner_inner_inner in T.grid(8, 8):
|
||||
B[blockIdx_z * 131072 + blockIdx_y * 16384 + threadIdx_y * 2048 + ff_inner_inner_inner * 256 + blockIdx_x * 64 + threadIdx_x * 8 + nn_inner_inner_inner] = B_local[ff_inner_inner_inner * 8 + nn_inner_inner_inner]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def GmmCuda0(X: T.Buffer((1, 128, 128), "float32"), Y: T.Buffer((1, 128, 128), "float32"), Z: T.Buffer((1, 128, 128), "float32")) -> None:
|
||||
Z_local = T.sblock_alloc_buffer([1, 128, 128], dtype="float32", scope="local")
|
||||
X_shared = T.sblock_alloc_buffer([1, 128, 128], dtype="float32", scope="shared")
|
||||
Y_shared = T.sblock_alloc_buffer([1, 128, 128], dtype="float32", scope="shared")
|
||||
for i0_0_i1_0_i2_0_fused in T.thread_binding(16, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_i2_1_fused in T.thread_binding(1, thread="vthread.x"):
|
||||
for i0_2_i1_2_i2_2_fused in T.thread_binding(128, thread="threadIdx.x"):
|
||||
for i1_3_init, i2_4_init in T.grid(4, 2):
|
||||
with T.sblock("Z_init"):
|
||||
b = T.axis.spatial(1, 0)
|
||||
i = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + i0_2_i1_2_i2_2_fused // 16 * 4 + i1_3_init)
|
||||
j = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + i0_2_i1_2_i2_2_fused % 16 * 2 + i2_4_init)
|
||||
T.reads()
|
||||
T.writes(Z_local[b, i, j])
|
||||
Z_local[b, i, j] = T.float32(0)
|
||||
for i3_0 in T.serial(4):
|
||||
for ax0_ax1_ax2_fused_0 in T.serial(4):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(128, thread="threadIdx.x"):
|
||||
for ax0_ax1_ax2_fused_2 in T.vectorized(2):
|
||||
with T.sblock("X_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + (ax0_ax1_ax2_fused_0 * 256 + ax0_ax1_ax2_fused_1 * 2 + ax0_ax1_ax2_fused_2) // 32)
|
||||
v2 = T.axis.spatial(128, i3_0 * 32 + (ax0_ax1_ax2_fused_0 * 256 + ax0_ax1_ax2_fused_1 * 2 + ax0_ax1_ax2_fused_2) % 32)
|
||||
T.reads(X[v0, v1, v2])
|
||||
T.writes(X_shared[v0, v1, v2])
|
||||
X_shared[v0, v1, v2] = X[v0, v1, v2]
|
||||
for ax0_ax1_ax2_fused_0 in T.serial(8):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(128, thread="threadIdx.x"):
|
||||
with T.sblock("Y_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(128, i3_0 * 32 + (ax0_ax1_ax2_fused_0 * 128 + ax0_ax1_ax2_fused_1) // 32)
|
||||
v2 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + (ax0_ax1_ax2_fused_0 * 128 + ax0_ax1_ax2_fused_1) % 32)
|
||||
T.reads(Y[v0, v1, v2])
|
||||
T.writes(Y_shared[v0, v1, v2])
|
||||
Y_shared[v0, v1, v2] = Y[v0, v1, v2]
|
||||
for i3_1, i0_3, i1_3, i2_3, i3_2, i0_4, i1_4, i2_4 in T.grid(1, 1, 4, 1, 32, 1, 1, 2):
|
||||
with T.sblock("Z_update"):
|
||||
b = T.axis.spatial(1, 0)
|
||||
i = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + i0_2_i1_2_i2_2_fused // 16 * 4 + i1_3)
|
||||
j = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + i0_2_i1_2_i2_2_fused % 16 * 2 + i2_4)
|
||||
k = T.axis.reduce(128, i3_0 * 32 + i3_2)
|
||||
T.reads(Z_local[b, i, j], X_shared[b, i, k], Y_shared[b, k, j])
|
||||
T.writes(Z_local[b, i, j])
|
||||
Z_local[b, i, j] = Z_local[b, i, j] + X_shared[b, i, k] * Y_shared[b, k, j]
|
||||
for ax0, ax1, ax2 in T.grid(1, 4, 2):
|
||||
with T.sblock("Z_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + i0_2_i1_2_i2_2_fused // 16 * 4 + ax1)
|
||||
v2 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + i0_2_i1_2_i2_2_fused % 16 * 2 + ax2)
|
||||
T.reads(Z_local[v0, v1, v2])
|
||||
T.writes(Z[v0, v1, v2])
|
||||
Z[v0, v1, v2] = Z_local[v0, v1, v2]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def GmmCuda1(X: T.Buffer((1, 128, 128), "float32"), Y: T.Buffer((1, 128, 128), "float32"), Z: T.Buffer((1, 128, 128), "float32")) -> None:
|
||||
Z_local = T.sblock_alloc_buffer([1, 128, 128], dtype="float32", scope="local")
|
||||
X_shared = T.sblock_alloc_buffer([1, 128, 128], dtype="float32", scope="shared")
|
||||
Y_shared = T.sblock_alloc_buffer([1, 128, 128], dtype="float32", scope="shared")
|
||||
for i0_0_i1_0_i2_0_fused in T.thread_binding(16, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_i2_1_fused in T.thread_binding(1, thread="vthread.x"):
|
||||
for i0_2_i1_2_i2_2_fused in T.thread_binding(128, thread="threadIdx.x"):
|
||||
for i1_3_init, i2_4_init in T.grid(4, 2):
|
||||
with T.sblock("Z_init"):
|
||||
b = T.axis.spatial(1, 0)
|
||||
i = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + i0_2_i1_2_i2_2_fused // 16 * 4 + i1_3_init)
|
||||
j = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + i0_2_i1_2_i2_2_fused % 16 * 2 + i2_4_init)
|
||||
T.reads()
|
||||
T.writes(Z_local[b, i, j])
|
||||
Z_local[b, i, j] = T.float32(0)
|
||||
for i3_0 in T.serial(4):
|
||||
for ax0_ax1_ax2_fused_0 in T.serial(4):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(128, thread="threadIdx.x"):
|
||||
for ax0_ax1_ax2_fused_2 in T.vectorized(2):
|
||||
with T.sblock("X_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + (ax0_ax1_ax2_fused_0 * 256 + ax0_ax1_ax2_fused_1 * 2 + ax0_ax1_ax2_fused_2) // 32)
|
||||
v2 = T.axis.spatial(128, i3_0 * 32 + (ax0_ax1_ax2_fused_0 * 256 + ax0_ax1_ax2_fused_1 * 2 + ax0_ax1_ax2_fused_2) % 32)
|
||||
T.reads(X[v0, v1, v2])
|
||||
T.writes(X_shared[v0, v1, v2])
|
||||
X_shared[v0, v1, v2] = X[v0, v1, v2]
|
||||
for ax0_ax1_ax2_fused_0 in T.serial(8):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(128, thread="threadIdx.x"):
|
||||
with T.sblock("Y_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(128, i3_0 * 32 + (ax0_ax1_ax2_fused_0 * 128 + ax0_ax1_ax2_fused_1) // 32)
|
||||
v2 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + (ax0_ax1_ax2_fused_0 * 128 + ax0_ax1_ax2_fused_1) % 32)
|
||||
T.reads(Y[v0, v1, v2])
|
||||
T.writes(Y_shared[v0, v1, v2])
|
||||
Y_shared[v0, v1, v2] = Y[v0, v1, v2]
|
||||
for i3_1, i0_3, i1_3, i2_3, i3_2, i0_4, i1_4, i2_4 in T.grid(1, 1, 4, 1, 32, 1, 1, 2):
|
||||
with T.sblock("Z_update"):
|
||||
b = T.axis.spatial(1, 0)
|
||||
i = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + i0_2_i1_2_i2_2_fused // 16 * 4 + i1_3)
|
||||
j = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + i0_2_i1_2_i2_2_fused % 16 * 2 + i2_4)
|
||||
k = T.axis.reduce(128, i3_0 * 32 + i3_2)
|
||||
T.sblock_attr({
|
||||
"meta_schedule.thread_extent_low_inclusive": 0,
|
||||
"meta_schedule.thread_extent_high_inclusive": 32,
|
||||
})
|
||||
T.reads(Z_local[b, i, j], X_shared[b, i, k], Y_shared[b, k, j])
|
||||
T.writes(Z_local[b, i, j])
|
||||
Z_local[b, i, j] = Z_local[b, i, j] + X_shared[b, i, k] * Y_shared[b, k, j]
|
||||
for ax0, ax1, ax2 in T.grid(1, 4, 2):
|
||||
with T.sblock("Z_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + i0_2_i1_2_i2_2_fused // 16 * 4 + ax1)
|
||||
v2 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + i0_2_i1_2_i2_2_fused % 16 * 2 + ax2)
|
||||
T.reads(Z_local[v0, v1, v2])
|
||||
T.writes(Z[v0, v1, v2])
|
||||
Z[v0, v1, v2] = Z_local[v0, v1, v2]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def GmmCuda2(X: T.Buffer((1, 128, 128), "float32"), Y: T.Buffer((1, 128, 128), "float32"), Z: T.Buffer((1, 128, 128), "float32")) -> None:
|
||||
Z_local = T.sblock_alloc_buffer([1, 128, 128], dtype="float32", scope="local")
|
||||
X_shared = T.sblock_alloc_buffer([1, 128, 128], dtype="float32", scope="shared")
|
||||
Y_shared = T.sblock_alloc_buffer([1, 128, 128], dtype="float32", scope="shared")
|
||||
for i0_0_i1_0_i2_0_fused in T.thread_binding(16, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_i2_1_fused in T.thread_binding(1, thread="vthread.x"):
|
||||
for i0_2_i1_2_i2_2_fused in T.thread_binding(128, thread="threadIdx.x"):
|
||||
for i1_3_init, i2_4_init in T.grid(4, 2):
|
||||
with T.sblock("Z_init"):
|
||||
b = T.axis.spatial(1, 0)
|
||||
i = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + i0_2_i1_2_i2_2_fused // 16 * 4 + i1_3_init)
|
||||
j = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + i0_2_i1_2_i2_2_fused % 16 * 2 + i2_4_init)
|
||||
T.reads()
|
||||
T.writes(Z_local[b, i, j])
|
||||
Z_local[b, i, j] = T.float32(0)
|
||||
for i3_0 in T.serial(4):
|
||||
for ax0_ax1_ax2_fused_0 in T.serial(4):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(128, thread="threadIdx.x"):
|
||||
for ax0_ax1_ax2_fused_2 in T.vectorized(2):
|
||||
with T.sblock("X_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + (ax0_ax1_ax2_fused_0 * 256 + ax0_ax1_ax2_fused_1 * 2 + ax0_ax1_ax2_fused_2) // 32)
|
||||
v2 = T.axis.spatial(128, i3_0 * 32 + (ax0_ax1_ax2_fused_0 * 256 + ax0_ax1_ax2_fused_1 * 2 + ax0_ax1_ax2_fused_2) % 32)
|
||||
T.reads(X[v0, v1, v2])
|
||||
T.writes(X_shared[v0, v1, v2])
|
||||
X_shared[v0, v1, v2] = X[v0, v1, v2]
|
||||
for ax0_ax1_ax2_fused_0 in T.serial(8):
|
||||
for ax0_ax1_ax2_fused_1 in T.thread_binding(128, thread="threadIdx.x"):
|
||||
with T.sblock("Y_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(128, i3_0 * 32 + (ax0_ax1_ax2_fused_0 * 128 + ax0_ax1_ax2_fused_1) // 32)
|
||||
v2 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + (ax0_ax1_ax2_fused_0 * 128 + ax0_ax1_ax2_fused_1) % 32)
|
||||
T.reads(Y[v0, v1, v2])
|
||||
T.writes(Y_shared[v0, v1, v2])
|
||||
Y_shared[v0, v1, v2] = Y[v0, v1, v2]
|
||||
for i3_1, i0_3, i1_3, i2_3, i3_2, i0_4, i1_4, i2_4 in T.grid(1, 1, 4, 1, 32, 1, 1, 2):
|
||||
with T.sblock("Z_update"):
|
||||
b = T.axis.spatial(1, 0)
|
||||
i = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + i0_2_i1_2_i2_2_fused // 16 * 4 + i1_3)
|
||||
j = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + i0_2_i1_2_i2_2_fused % 16 * 2 + i2_4)
|
||||
k = T.axis.reduce(128, i3_0 * 32 + i3_2)
|
||||
T.sblock_attr({
|
||||
"meta_schedule.thread_extent_low_inclusive": 1024,
|
||||
"meta_schedule.thread_extent_high_inclusive": 1024,
|
||||
})
|
||||
T.reads(Z_local[b, i, j], X_shared[b, i, k], Y_shared[b, k, j])
|
||||
T.writes(Z_local[b, i, j])
|
||||
Z_local[b, i, j] = Z_local[b, i, j] + X_shared[b, i, k] * Y_shared[b, k, j]
|
||||
for ax0, ax1, ax2 in T.grid(1, 4, 2):
|
||||
with T.sblock("Z_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused // 4 * 32 + i0_2_i1_2_i2_2_fused // 16 * 4 + ax1)
|
||||
v2 = T.axis.spatial(128, i0_0_i1_0_i2_0_fused % 4 * 32 + i0_2_i1_2_i2_2_fused % 16 * 2 + ax2)
|
||||
T.reads(Z_local[v0, v1, v2])
|
||||
T.writes(Z[v0, v1, v2])
|
||||
Z[v0, v1, v2] = Z_local[v0, v1, v2]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def GMMCUDATensorCore(
|
||||
X: T.Buffer((1024, 1024), "float16"),
|
||||
Y: T.Buffer((1024, 1024), "float16"),
|
||||
Z: T.Buffer((1024, 1024), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
s0 = T.int32()
|
||||
s0_1 = T.int32()
|
||||
s0_2 = T.int32()
|
||||
s1 = T.int32()
|
||||
s1_1 = T.int32()
|
||||
s1_2 = T.int32()
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
Z_wmma_accumulator = T.sblock_alloc_buffer([1024, 1024], dtype="float32", scope="wmma.accumulator")
|
||||
X_shared = T.sblock_alloc_buffer([1024, 1024], dtype="float16", scope="shared")
|
||||
Y_shared = T.sblock_alloc_buffer([1024, 1024], dtype="float16", scope="shared")
|
||||
X_shared_wmma_matrix_a = T.sblock_alloc_buffer([1024, 1024], dtype="float16", scope="wmma.matrix_a")
|
||||
Y_shared_wmma_matrix_b = T.sblock_alloc_buffer([1024, 1024], dtype="float16", scope="wmma.matrix_b")
|
||||
for ax0_0_ax1_0_0_ax2_0_0_fused in T.thread_binding(64, thread="blockIdx.x"):
|
||||
for ax0_1_ax1_0_1_ax2_0_1_fused in T.thread_binding(2, thread="blockIdx.y"):
|
||||
for ax0_2_ax1_0_2_ax2_0_2_fused in T.thread_binding(2, thread="threadIdx.y"):
|
||||
for ax1_0_3_init, ax2_0_3_init, ax1_0_4_init, ax2_0_4_init in T.grid(2, 1, 2, 4):
|
||||
with T.sblock("Z_o_init"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(
|
||||
64,
|
||||
ax0_0_ax1_0_0_ax2_0_0_fused % 64 // 16 * 16
|
||||
+ ax0_1_ax1_0_1_ax2_0_1_fused % 2 * 8
|
||||
+ ax0_2_ax1_0_2_ax2_0_2_fused % 2 * 4
|
||||
+ ax1_0_3_init * 2
|
||||
+ ax1_0_4_init,
|
||||
)
|
||||
v2_o = T.axis.spatial(
|
||||
64,
|
||||
(ax0_0_ax1_0_0_ax2_0_0_fused % 16 + 0 + 0 + ax2_0_3_init) * 4
|
||||
+ ax2_0_4_init,
|
||||
)
|
||||
T.reads()
|
||||
T.writes(
|
||||
Z_wmma_accumulator[
|
||||
v1_o * 16 : v1_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
|
||||
]
|
||||
)
|
||||
T.sblock_attr(
|
||||
{
|
||||
"meta_schedule.thread_extent_high_inclusive": 1024,
|
||||
"meta_schedule.thread_extent_low_inclusive": 32,
|
||||
"warp_execution": 1,
|
||||
}
|
||||
)
|
||||
C = T.match_buffer(
|
||||
Z_wmma_accumulator[
|
||||
v1_o * 16 : v1_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
|
||||
],
|
||||
[16, 16],
|
||||
dtype="float32",
|
||||
scope="wmma.accumulator",
|
||||
offset_factor=16,
|
||||
)
|
||||
T.evaluate(
|
||||
T.tvm_fill_fragment(
|
||||
C.data,
|
||||
16,
|
||||
16,
|
||||
16,
|
||||
C.elem_offset // 256 + C.elem_offset % 256 // 16,
|
||||
T.float32(0),
|
||||
dtype="handle",
|
||||
)
|
||||
)
|
||||
for ax3_0_0 in T.serial(32):
|
||||
for ax0_ax1_fused_0 in T.serial(16):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(2, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.vectorized(4):
|
||||
with T.sblock("X_shared"):
|
||||
v0 = T.axis.spatial(
|
||||
1024,
|
||||
ax0_0_ax1_0_0_ax2_0_0_fused // 16 * 256
|
||||
+ ax0_1_ax1_0_1_ax2_0_1_fused * 128
|
||||
+ (
|
||||
ax0_ax1_fused_0 * 256
|
||||
+ ax0_ax1_fused_1 * 128
|
||||
+ ax0_ax1_fused_2 * 4
|
||||
+ ax0_ax1_fused_3
|
||||
)
|
||||
// 32,
|
||||
)
|
||||
v1 = T.axis.spatial(
|
||||
1024,
|
||||
ax3_0_0 * 32
|
||||
+ (
|
||||
ax0_ax1_fused_0 * 256
|
||||
+ ax0_ax1_fused_1 * 128
|
||||
+ ax0_ax1_fused_2 * 4
|
||||
+ ax0_ax1_fused_3
|
||||
)
|
||||
% 32,
|
||||
)
|
||||
T.reads(X[v0, v1])
|
||||
T.writes(X_shared[v0, v1])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]]})
|
||||
X_shared[v0, v1] = X[v0, v1]
|
||||
for ax0_ax1_fused_0 in T.serial(8):
|
||||
for ax0_ax1_fused_1 in T.thread_binding(2, thread="threadIdx.y"):
|
||||
for ax0_ax1_fused_2 in T.thread_binding(32, thread="threadIdx.x"):
|
||||
for ax0_ax1_fused_3 in T.vectorized(4):
|
||||
with T.sblock("Y_shared"):
|
||||
v0 = T.axis.spatial(
|
||||
1024,
|
||||
ax3_0_0 * 32
|
||||
+ (
|
||||
ax0_ax1_fused_0 * 256
|
||||
+ ax0_ax1_fused_1 * 128
|
||||
+ ax0_ax1_fused_2 * 4
|
||||
+ ax0_ax1_fused_3
|
||||
)
|
||||
// 64,
|
||||
)
|
||||
v1 = T.axis.spatial(
|
||||
1024,
|
||||
ax0_0_ax1_0_0_ax2_0_0_fused % 16 * 64
|
||||
+ (
|
||||
ax0_ax1_fused_0 * 256
|
||||
+ ax0_ax1_fused_1 * 128
|
||||
+ ax0_ax1_fused_2 * 4
|
||||
+ ax0_ax1_fused_3
|
||||
)
|
||||
% 64,
|
||||
)
|
||||
T.reads(Y[v0, v1])
|
||||
T.writes(Y_shared[v0, v1])
|
||||
T.sblock_attr({"buffer_dim_align": [[0, 0, 32, 8]]})
|
||||
Y_shared[v0, v1] = Y[v0, v1]
|
||||
for ax3_0_1 in T.serial(2):
|
||||
for ax0_0, ax1_0 in T.grid(4, 1):
|
||||
with T.sblock("X_shared_wmma.matrix_a_o"):
|
||||
v0_o = T.axis.spatial(
|
||||
64,
|
||||
ax0_0_ax1_0_0_ax2_0_0_fused // 16 * 16
|
||||
+ ax0_1_ax1_0_1_ax2_0_1_fused * 8
|
||||
+ ax0_2_ax1_0_2_ax2_0_2_fused * 4
|
||||
+ ax0_0,
|
||||
)
|
||||
v1_o = T.axis.spatial(64, ax3_0_0 * 2 + ax3_0_1)
|
||||
T.reads(
|
||||
X_shared[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16]
|
||||
)
|
||||
T.writes(
|
||||
X_shared_wmma_matrix_a[
|
||||
v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
|
||||
]
|
||||
)
|
||||
A = T.match_buffer(
|
||||
X_shared[
|
||||
v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
|
||||
],
|
||||
[16, 16],
|
||||
dtype="float16",
|
||||
strides=[s1, s0],
|
||||
scope="shared",
|
||||
offset_factor=16,
|
||||
)
|
||||
C_1 = T.match_buffer(
|
||||
X_shared_wmma_matrix_a[
|
||||
v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
|
||||
],
|
||||
[16, 16],
|
||||
dtype="float16",
|
||||
scope="wmma.matrix_a",
|
||||
offset_factor=16,
|
||||
)
|
||||
T.evaluate(
|
||||
T.tvm_load_matrix_sync(
|
||||
C_1.data,
|
||||
16,
|
||||
16,
|
||||
16,
|
||||
C_1.elem_offset // 256 + C_1.elem_offset % 256 // 16,
|
||||
T.tvm_access_ptr(
|
||||
T.type_annotation(dtype="float16"),
|
||||
A.data,
|
||||
A.elem_offset,
|
||||
s1 * 16,
|
||||
1,
|
||||
dtype="handle",
|
||||
),
|
||||
s1,
|
||||
"row_major",
|
||||
dtype="handle",
|
||||
)
|
||||
)
|
||||
for ax0_0, ax1_0 in T.grid(1, 4):
|
||||
with T.sblock("Y_shared_wmma.matrix_b_o"):
|
||||
v0_o = T.axis.spatial(64, ax3_0_0 * 2 + ax3_0_1)
|
||||
v1_o = T.axis.spatial(
|
||||
64, ax0_0_ax1_0_0_ax2_0_0_fused % 16 * 4 + ax1_0
|
||||
)
|
||||
T.reads(
|
||||
Y_shared[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16]
|
||||
)
|
||||
T.writes(
|
||||
Y_shared_wmma_matrix_b[
|
||||
v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
|
||||
]
|
||||
)
|
||||
A_1 = T.match_buffer(
|
||||
Y_shared[
|
||||
v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
|
||||
],
|
||||
[16, 16],
|
||||
dtype="float16",
|
||||
strides=[s1_1, s0_1],
|
||||
scope="shared",
|
||||
offset_factor=16,
|
||||
)
|
||||
C_2 = T.match_buffer(
|
||||
Y_shared_wmma_matrix_b[
|
||||
v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
|
||||
],
|
||||
[16, 16],
|
||||
dtype="float16",
|
||||
scope="wmma.matrix_b",
|
||||
offset_factor=16,
|
||||
)
|
||||
T.evaluate(
|
||||
T.tvm_load_matrix_sync(
|
||||
C_2.data,
|
||||
16,
|
||||
16,
|
||||
16,
|
||||
C_2.elem_offset // 256 + C_2.elem_offset % 256 // 16,
|
||||
T.tvm_access_ptr(
|
||||
T.type_annotation(dtype="float16"),
|
||||
A_1.data,
|
||||
A_1.elem_offset,
|
||||
s1_1 * 16,
|
||||
1,
|
||||
dtype="handle",
|
||||
),
|
||||
s1_1,
|
||||
"row_major",
|
||||
dtype="handle",
|
||||
)
|
||||
)
|
||||
for ax0_3, ax1_0_3, ax2_0_3, ax3_0_2, ax0_4, ax1_0_4, ax2_0_4 in T.grid(
|
||||
1, 2, 1, 1, 1, 2, 4
|
||||
):
|
||||
with T.sblock("Z_o_update"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1_o = T.axis.spatial(
|
||||
64,
|
||||
ax0_0_ax1_0_0_ax2_0_0_fused % 64 // 16 * 16
|
||||
+ ax0_1_ax1_0_1_ax2_0_1_fused % 2 * 8
|
||||
+ ax0_2_ax1_0_2_ax2_0_2_fused % 2 * 4
|
||||
+ ax1_0_3 * 2
|
||||
+ ax1_0_4,
|
||||
)
|
||||
v2_o = T.axis.spatial(
|
||||
64,
|
||||
(ax0_0_ax1_0_0_ax2_0_0_fused % 16 + 0 + 0 + ax2_0_3) * 4
|
||||
+ ax2_0_4,
|
||||
)
|
||||
v3_o = T.axis.reduce(64, ax3_0_0 * 2 + ax3_0_1 + ax3_0_2)
|
||||
T.reads(
|
||||
Z_wmma_accumulator[
|
||||
v1_o * 16 : v1_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
|
||||
],
|
||||
X_shared_wmma_matrix_a[
|
||||
v1_o * 16 : v1_o * 16 + 16, v3_o * 16 : v3_o * 16 + 16
|
||||
],
|
||||
Y_shared_wmma_matrix_b[
|
||||
v3_o * 16 : v3_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
|
||||
],
|
||||
)
|
||||
T.writes(
|
||||
Z_wmma_accumulator[
|
||||
v1_o * 16 : v1_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
|
||||
]
|
||||
)
|
||||
T.sblock_attr(
|
||||
{
|
||||
"meta_schedule.thread_extent_high_inclusive": 1024,
|
||||
"meta_schedule.thread_extent_low_inclusive": 32,
|
||||
"warp_execution": 1,
|
||||
}
|
||||
)
|
||||
A_2 = T.match_buffer(
|
||||
X_shared_wmma_matrix_a[
|
||||
v1_o * 16 : v1_o * 16 + 16, v3_o * 16 : v3_o * 16 + 16
|
||||
],
|
||||
[16, 16],
|
||||
dtype="float16",
|
||||
scope="wmma.matrix_a",
|
||||
offset_factor=16,
|
||||
)
|
||||
B = T.match_buffer(
|
||||
Y_shared_wmma_matrix_b[
|
||||
v3_o * 16 : v3_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
|
||||
],
|
||||
[16, 16],
|
||||
dtype="float16",
|
||||
scope="wmma.matrix_b",
|
||||
offset_factor=16,
|
||||
)
|
||||
C_3 = T.match_buffer(
|
||||
Z_wmma_accumulator[
|
||||
v1_o * 16 : v1_o * 16 + 16, v2_o * 16 : v2_o * 16 + 16
|
||||
],
|
||||
[16, 16],
|
||||
dtype="float32",
|
||||
scope="wmma.accumulator",
|
||||
offset_factor=16,
|
||||
)
|
||||
T.evaluate(
|
||||
T.tvm_mma_sync(
|
||||
C_3.data,
|
||||
C_3.elem_offset // 256 + C_3.elem_offset % 256 // 16,
|
||||
A_2.data,
|
||||
A_2.elem_offset // 256,
|
||||
B.data,
|
||||
B.elem_offset // 256,
|
||||
C_3.data,
|
||||
C_3.elem_offset // 256 + C_3.elem_offset % 256 // 16,
|
||||
dtype="handle",
|
||||
)
|
||||
)
|
||||
for ax0_0, ax1_0 in T.grid(4, 4):
|
||||
with T.sblock("Z_wmma.accumulator_o"):
|
||||
v0_o = T.axis.spatial(
|
||||
64,
|
||||
ax0_0_ax1_0_0_ax2_0_0_fused // 16 * 16
|
||||
+ ax0_1_ax1_0_1_ax2_0_1_fused * 8
|
||||
+ ax0_2_ax1_0_2_ax2_0_2_fused * 4
|
||||
+ ax0_0,
|
||||
)
|
||||
v1_o = T.axis.spatial(64, ax0_0_ax1_0_0_ax2_0_0_fused % 16 * 4 + ax1_0)
|
||||
T.reads(
|
||||
Z_wmma_accumulator[
|
||||
v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
|
||||
]
|
||||
)
|
||||
T.writes(Z[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16])
|
||||
A_3 = T.match_buffer(
|
||||
Z_wmma_accumulator[
|
||||
v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16
|
||||
],
|
||||
[16, 16],
|
||||
dtype="float32",
|
||||
scope="wmma.accumulator",
|
||||
offset_factor=16,
|
||||
)
|
||||
C_4 = T.match_buffer(
|
||||
Z[v0_o * 16 : v0_o * 16 + 16, v1_o * 16 : v1_o * 16 + 16],
|
||||
[16, 16],
|
||||
dtype="float32",
|
||||
strides=[s1_2, s0_2],
|
||||
offset_factor=16,
|
||||
)
|
||||
T.evaluate(
|
||||
T.tvm_store_matrix_sync(
|
||||
A_3.data,
|
||||
16,
|
||||
16,
|
||||
16,
|
||||
A_3.elem_offset // 256 + A_3.elem_offset % 256 // 16,
|
||||
T.tvm_access_ptr(
|
||||
T.type_annotation(dtype="float32"),
|
||||
C_4.data,
|
||||
C_4.elem_offset,
|
||||
s1_2 * 16,
|
||||
2,
|
||||
dtype="handle",
|
||||
),
|
||||
s1_2,
|
||||
"row_major",
|
||||
dtype="handle",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# fmt: on
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,not-callable,misplaced-comparison-constant
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mod", [Conv2dCuda0, Conv2dCuda1, GmmCuda0, GMMCUDATensorCore])
|
||||
def test_postproc_check_pass(mod):
|
||||
ctx = _create_context(mod, target=_target())
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mod",
|
||||
[
|
||||
Conv2dCuda2, # Should fail due to too much local memory per block (large Apad_shared allocation)
|
||||
Conv2dCuda3, # Should fail due to too many threads per block (large threadIdx.x extent)
|
||||
GmmCuda1,
|
||||
GmmCuda2,
|
||||
],
|
||||
)
|
||||
def test_postproc_check_fail(mod):
|
||||
ctx = _create_context(mod, target=_target())
|
||||
sch = tvm.s_tir.Schedule(mod, debug_mask="all")
|
||||
assert not ctx.space_generator.postprocs[0].apply(sch)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,133 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: E501, F401, F841
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm import tirx
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
def _create_context(mod, target) -> ms.TuneContext:
|
||||
return ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=[],
|
||||
postprocs=[ms.postproc.VerifyVTCMLimit()],
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument,not-callable,misplaced-comparison-constant
|
||||
# fmt: off
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2dNCHWcVTCM:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(p0: T.Buffer((T.int64(1), T.int64(2), T.int64(56), T.int64(56), T.int64(32)), "uint8"), p1: T.Buffer((T.int64(2), T.int64(2), T.int64(3), T.int64(3), T.int64(8), T.int64(32), T.int64(4)), "uint8"), conv2d_NCHWc_int8: T.Buffer((T.int64(1), T.int64(2), T.int64(54), T.int64(54), T.int64(32)), "int32")):
|
||||
T.func_attr({"tirx.noalias": True, "global_symbol": "main"})
|
||||
p0_global_vtcm = T.sblock_alloc_buffer([T.int64(1), T.int64(2), T.int64(56), T.int64(56), T.int64(32)], dtype="uint8", scope="global.vtcm")
|
||||
p1_global_vtcm = T.sblock_alloc_buffer([T.int64(2), T.int64(2), T.int64(3), T.int64(3), T.int64(8), T.int64(32), T.int64(4)], dtype="uint8", scope="global.vtcm")
|
||||
for n_0 in T.serial(T.int64(1), annotations={"pragma_auto_unroll_max_step":16, "pragma_unroll_explicit":1}):
|
||||
for oc_chunk_0, oh_0, ow_0, oc_block_0_0 in T.grid(T.int64(2), T.int64(2), T.int64(2), T.int64(1)):
|
||||
for oc_chunk_1_init, oh_1_init, ow_1_init, oc_chunk_2_init, oh_2_init, ow_2_init in T.grid(T.int64(1), T.int64(27), T.int64(3), T.int64(1), T.int64(1), T.int64(9)):
|
||||
with T.sblock("conv2d_NCHWc_int8_o_init"):
|
||||
v_n = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v_oc_chunk = T.axis.spatial(T.int64(2), oc_chunk_1_init + oc_chunk_2_init + oc_chunk_0)
|
||||
v_oh = T.axis.spatial(T.int64(54), oh_2_init + oh_0 * T.int64(27) + oh_1_init)
|
||||
v_ow = T.axis.spatial(T.int64(54), ow_0 * T.int64(27) + ow_1_init * T.int64(9) + ow_2_init)
|
||||
v_oc_block_o = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
T.reads()
|
||||
T.writes(conv2d_NCHWc_int8[v_n, v_oc_chunk, v_oh, v_ow, T.int64(0) : T.int64(32)])
|
||||
for oc_block_1 in T.vectorized(T.int64(32)):
|
||||
with T.sblock("conv2d_NCHWc_int8_init"):
|
||||
v_oc_block_i_init = T.axis.spatial(T.int64(32), oc_block_1)
|
||||
T.reads()
|
||||
T.writes(conv2d_NCHWc_int8[v_n, v_oc_chunk, v_oh, v_ow, v_oc_block_i_init])
|
||||
conv2d_NCHWc_int8[v_n, v_oc_chunk, v_oh, v_ow, v_oc_block_i_init] = 0
|
||||
for kh_0_kw_0_ic_outer_0_ic_f_inner_0_ic_s_inner_0_0_fused in T.serial(T.int64(2), annotations={"software_pipeline_async_stages":[0], "software_pipeline_order":[0, 1, 2], "software_pipeline_stage":[0, 0, 1]}):
|
||||
for ax0_ax1_ax2_ax3_ax4_fused in T.serial(T.int64(26912)):
|
||||
with T.sblock("p0_global.vtcm"):
|
||||
v0 = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v1 = T.axis.spatial(T.int64(2), ax0_ax1_ax2_ax3_ax4_fused // T.int64(13456))
|
||||
v2 = T.axis.spatial(T.int64(56), oh_0 * T.int64(27) + ax0_ax1_ax2_ax3_ax4_fused % T.int64(13456) // T.int64(464))
|
||||
v3 = T.axis.spatial(T.int64(56), ow_0 * T.int64(27) + ax0_ax1_ax2_ax3_ax4_fused % T.int64(464) // T.int64(16))
|
||||
v4 = T.axis.spatial(T.int64(32), kh_0_kw_0_ic_outer_0_ic_f_inner_0_ic_s_inner_0_0_fused * T.int64(16) + ax0_ax1_ax2_ax3_ax4_fused % T.int64(16))
|
||||
T.reads(p0[v0, v1, v2, v3, v4])
|
||||
T.writes(p0_global_vtcm[v0, v1, v2, v3, v4])
|
||||
p0_global_vtcm[v0, v1, v2, v3, v4] = p0[v0, v1, v2, v3, v4]
|
||||
for ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused in T.serial(T.int64(9216)):
|
||||
with T.sblock("p1_global.vtcm"):
|
||||
v0 = T.axis.spatial(T.int64(2), oc_chunk_0)
|
||||
v1 = T.axis.spatial(T.int64(2), ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused // T.int64(4608))
|
||||
v2 = T.axis.spatial(T.int64(3), ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused % T.int64(4608) // T.int64(1536))
|
||||
v3 = T.axis.spatial(T.int64(3), ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused % T.int64(1536) // T.int64(512))
|
||||
v4 = T.axis.spatial(T.int64(8), kh_0_kw_0_ic_outer_0_ic_f_inner_0_ic_s_inner_0_0_fused * T.int64(4) + ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused % T.int64(512) // T.int64(128))
|
||||
v5 = T.axis.spatial(T.int64(32), ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused % T.int64(128) // T.int64(4))
|
||||
v6 = T.axis.spatial(T.int64(4), ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused % T.int64(4))
|
||||
T.reads(p1[v0, v1, v2, v3, v4, v5, v6])
|
||||
T.writes(p1_global_vtcm[v0, v1, v2, v3, v4, v5, v6])
|
||||
p1_global_vtcm[v0, v1, v2, v3, v4, v5, v6] = p1[v0, v1, v2, v3, v4, v5, v6]
|
||||
for n_1, oc_chunk_1, oh_1, ow_1, oc_block_0_1, kh_1, kw_1, ic_outer_1, ic_f_inner_1, ic_s_inner_0_1, n_2, oc_chunk_2, oh_2, ow_2, oc_block_0_2 in T.grid(T.int64(1), T.int64(1), T.int64(27), T.int64(3), T.int64(1), T.int64(3), T.int64(3), T.int64(2), T.int64(4), T.int64(1), T.int64(1), T.int64(1), T.int64(1), T.int64(9), T.int64(1)):
|
||||
with T.sblock("conv2d_NCHWc_int8_o_update"):
|
||||
v_n = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v_oc_chunk = T.axis.spatial(T.int64(2), oc_chunk_1 + oc_chunk_2 + oc_chunk_0)
|
||||
v_oh = T.axis.spatial(T.int64(54), oh_2 + oh_0 * T.int64(27) + oh_1)
|
||||
v_ow = T.axis.spatial(T.int64(54), ow_0 * T.int64(27) + ow_1 * T.int64(9) + ow_2)
|
||||
v_oc_block_o = T.axis.spatial(T.int64(1), T.int64(0))
|
||||
v_kh, v_kw, v_ic_outer = T.axis.remap("RRR", [kh_1, kw_1, ic_outer_1])
|
||||
v_ic_f_inner = T.axis.reduce(T.int64(8), kh_0_kw_0_ic_outer_0_ic_f_inner_0_ic_s_inner_0_0_fused * T.int64(4) + ic_f_inner_1)
|
||||
v_ic_s_inner_o = T.axis.reduce(T.int64(1), T.int64(0))
|
||||
T.reads(conv2d_NCHWc_int8[v_n, v_oc_chunk, v_oh, v_ow, T.int64(0) : T.int64(32)], p0_global_vtcm[v_n, v_ic_outer, v_oh + v_kh, v_ow + v_kw, v_ic_f_inner * T.int64(4) : v_ic_f_inner * T.int64(4) + T.int64(4)], p1_global_vtcm[v_oc_chunk, v_ic_outer, v_kh, v_kw, v_ic_f_inner, T.int64(0) : T.int64(32), T.int64(0) : T.int64(4)])
|
||||
T.writes(conv2d_NCHWc_int8[v_n, v_oc_chunk, v_oh, v_ow, T.int64(0) : T.int64(32)])
|
||||
for oc_block_1, ic_s_inner_1 in T.grid(T.int64(32), T.int64(4)):
|
||||
with T.sblock("conv2d_NCHWc_int8"):
|
||||
v_oc_block_i, v_ic_s_inner_i = T.axis.remap("SR", [oc_block_1, ic_s_inner_1])
|
||||
T.reads(conv2d_NCHWc_int8[v_n, v_oc_chunk, v_oh, v_ow, v_oc_block_i], p0_global_vtcm[v_n, v_ic_outer, v_oh + v_kh, v_ow + v_kw, v_ic_f_inner * T.int64(4) + v_ic_s_inner_i], p1_global_vtcm[v_oc_chunk, v_ic_outer, v_kh, v_kw, v_ic_f_inner, v_oc_block_i, v_ic_s_inner_i])
|
||||
T.writes(conv2d_NCHWc_int8[v_n, v_oc_chunk, v_oh, v_ow, v_oc_block_i])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure":"SRSRS"})
|
||||
conv2d_NCHWc_int8[v_n, v_oc_chunk, v_oh, v_ow, v_oc_block_i] = conv2d_NCHWc_int8[v_n, v_oc_chunk, v_oh, v_ow, v_oc_block_i] + T.Cast("int32", p0_global_vtcm[v_n, v_ic_outer, v_oh + v_kh, v_ow + v_kw, v_ic_f_inner * T.int64(4) + v_ic_s_inner_i]) * T.Cast("int32", p1_global_vtcm[v_oc_chunk, v_ic_outer, v_kh, v_kw, v_ic_f_inner, v_oc_block_i, v_ic_s_inner_i])
|
||||
|
||||
#fmt on
|
||||
|
||||
|
||||
def test_conv2d_vtcm():
|
||||
def get_target(vtcm_cap):
|
||||
target = tvm.target.Target({
|
||||
"kind": "hexagon", "mtriple": "hexagon", "mcpu": "hexagonv68",
|
||||
"mattr": ["+hvxv68", "+hvx-length128b", "+hvx-qfloat", "-hvx-ieee-fp"],
|
||||
"num-cores": 4, "vtcm-capacity": vtcm_cap,
|
||||
"llvm-options": ["-force-hvx-float"],
|
||||
})
|
||||
return tvm.target.Target(target, host=target)
|
||||
|
||||
sch = tvm.s_tir.Schedule(Conv2dNCHWcVTCM, debug_mask="all")
|
||||
|
||||
ctx = _create_context(Conv2dNCHWcVTCM, target=get_target(70000))
|
||||
assert not ctx.space_generator.postprocs[0].apply(sch)
|
||||
|
||||
ctx = _create_context(Conv2dNCHWcVTCM, target=get_target(75000))
|
||||
assert ctx.space_generator.postprocs[0].apply(sch)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,47 @@
|
||||
# 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.
|
||||
"""Test Meta Schedule Profiler"""
|
||||
|
||||
import time
|
||||
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
|
||||
|
||||
def test_meta_schedule_profiler_context_manager():
|
||||
with ms.Profiler() as profiler:
|
||||
time.sleep(1)
|
||||
with ms.Profiler.timeit("Level0"):
|
||||
time.sleep(1)
|
||||
with ms.Profiler.timeit("Level1"):
|
||||
time.sleep(2)
|
||||
# Note that the results are in seconds
|
||||
|
||||
result = profiler.get()
|
||||
assert len(result) == 3
|
||||
assert 3.9 <= result["Total"] <= 4.1
|
||||
assert 2.9 <= result["Level0"] <= 3.1
|
||||
assert 1.9 <= result["Level1"] <= 2.1
|
||||
|
||||
|
||||
def test_meta_schedule_no_context():
|
||||
with ms.Profiler.timeit("Level0"):
|
||||
assert ms.Profiler.current() is None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_meta_schedule_profiler_context_manager()
|
||||
test_meta_schedule_no_context()
|
||||
@@ -0,0 +1,925 @@
|
||||
# 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.
|
||||
# ruff: noqa: F401
|
||||
"""Test Meta Schedule Runner"""
|
||||
|
||||
import itertools
|
||||
import sys
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from tvm_ffi import register_global_func
|
||||
|
||||
pytest.importorskip("tornado") # tvm.rpc.tracker (LocalRPC) requires tornado
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.rpc import RPCSession
|
||||
from tvm.runtime import Device, Module
|
||||
from tvm.s_tir.meta_schedule.arg_info import TensorInfo
|
||||
from tvm.s_tir.meta_schedule.builder import BuilderInput, LocalBuilder
|
||||
from tvm.s_tir.meta_schedule.runner import (
|
||||
EvaluatorConfig,
|
||||
LocalRunner,
|
||||
PyRunner,
|
||||
RPCConfig,
|
||||
RPCRunner,
|
||||
RunnerFuture,
|
||||
RunnerInput,
|
||||
)
|
||||
from tvm.s_tir.meta_schedule.runner.local_runner import (
|
||||
default_alloc_argument as local_default_alloc_argument,
|
||||
)
|
||||
from tvm.s_tir.meta_schedule.runner.rpc_runner import (
|
||||
T_ARG_INFO_JSON_OBJ_LIST,
|
||||
T_ARGUMENT_LIST,
|
||||
)
|
||||
from tvm.s_tir.meta_schedule.runner.rpc_runner import (
|
||||
default_alloc_argument as rpc_default_alloc_argument,
|
||||
)
|
||||
from tvm.s_tir.meta_schedule.testing.local_rpc import LocalRPC
|
||||
from tvm.s_tir.meta_schedule.utils import (
|
||||
get_global_func_with_default_on_worker,
|
||||
)
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
from tvm.tirx import FloatImm
|
||||
|
||||
MATMUL_N = 16
|
||||
MATMUL_M = 32
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring,unbalanced-tuple-unpacking
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class MatmulModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
B = T.match_buffer(b, (16, 16), "float32")
|
||||
C = T.match_buffer(c, (16, 16), "float32")
|
||||
for i, j, k in T.grid(16, 16, 16):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class MatmulReluModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, d: T.handle) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, (16, 16), "float32")
|
||||
B = T.match_buffer(b, (16, 16), "float32")
|
||||
D = T.match_buffer(d, (16, 16), "float32")
|
||||
C = T.sblock_alloc_buffer((16, 16), "float32")
|
||||
for i, j, k in T.grid(16, 16, 16):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j in T.grid(16, 16):
|
||||
with T.sblock("relu"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
D[vi, vj] = T.max(C[vi, vj], 0.0)
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class BatchMatmulModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, [16, 32, 32])
|
||||
B = T.match_buffer(b, [16, 32, 32])
|
||||
C = T.match_buffer(c, [16, 32, 32])
|
||||
for n, i, j, k in T.grid(16, 32, 32, 32):
|
||||
with T.sblock("update"):
|
||||
vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])
|
||||
with T.init():
|
||||
C[vn, vi, vj] = 0.0
|
||||
C[vn, vi, vj] = C[vn, vi, vj] + A[vn, vi, vk] * B[vn, vj, vk]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class AddModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, [32], "float32")
|
||||
B = T.match_buffer(b, [32], "float32")
|
||||
C = T.match_buffer(c, [32], "float32")
|
||||
for i in range(32):
|
||||
with T.sblock("add"):
|
||||
vi = T.axis.S(32, i)
|
||||
C[vi] = A[vi] + B[vi]
|
||||
|
||||
|
||||
# A huge matmul that must cause timeout in the timeout test below.
|
||||
@tvm.script.ir_module
|
||||
class MatmulHugeModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, (4096, 4096), "float32")
|
||||
B = T.match_buffer(b, (4096, 4096), "float32")
|
||||
C = T.match_buffer(c, (4096, 4096), "float32")
|
||||
for i, j, k in T.grid(4096, 4096, 4096):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring
|
||||
|
||||
|
||||
def _clean_build(artifact_path: str) -> None:
|
||||
f_clean_build = get_global_func_with_default_on_worker(
|
||||
"s_tir.meta_schedule.remove_build_dir", None
|
||||
)
|
||||
if f_clean_build is not None:
|
||||
f_clean_build(artifact_path)
|
||||
else:
|
||||
raise RuntimeError("Unable to find remove_build_dir function.")
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_rpc_single_run():
|
||||
"""Test meta schedule rpc runner for a single run"""
|
||||
# Build the module
|
||||
mod = MatmulModule
|
||||
builder = LocalBuilder()
|
||||
(builder_result,) = builder.build([BuilderInput(mod, Target("llvm"))])
|
||||
assert builder_result.artifact_path is not None
|
||||
assert builder_result.error_msg is None
|
||||
|
||||
runner_input = RunnerInput(
|
||||
builder_result.artifact_path,
|
||||
"llvm",
|
||||
[
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
],
|
||||
)
|
||||
|
||||
with LocalRPC() as rpc:
|
||||
rpc_config = RPCConfig(
|
||||
tracker_host=rpc.tracker_host,
|
||||
tracker_port=rpc.tracker_port,
|
||||
tracker_key=rpc.tracker_key,
|
||||
session_priority=1,
|
||||
session_timeout_sec=100,
|
||||
)
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
runner = RPCRunner(rpc_config, evaluator_config)
|
||||
# Run the module
|
||||
(runner_future,) = runner.run([runner_input])
|
||||
runner_result = runner_future.result()
|
||||
assert runner_result.error_msg is None
|
||||
for result in runner_result.run_secs:
|
||||
if isinstance(result, FloatImm):
|
||||
result = result.value
|
||||
assert isinstance(result, float)
|
||||
assert result >= 0.0
|
||||
_clean_build(builder_result.artifact_path)
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_local_single_run():
|
||||
"""Test meta schedule local runner for a single run"""
|
||||
# Build the module
|
||||
mod = MatmulModule
|
||||
builder = LocalBuilder()
|
||||
(builder_result,) = builder.build([BuilderInput(mod, Target("llvm"))])
|
||||
assert builder_result.artifact_path is not None
|
||||
assert builder_result.error_msg is None
|
||||
|
||||
runner_input = RunnerInput(
|
||||
builder_result.artifact_path,
|
||||
"llvm",
|
||||
[
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
],
|
||||
)
|
||||
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
runner = LocalRunner(timeout_sec=100, evaluator_config=evaluator_config)
|
||||
# Run the module
|
||||
(runner_future,) = runner.run([runner_input])
|
||||
runner_result = runner_future.result()
|
||||
assert runner_result.error_msg is None
|
||||
for result in runner_result.run_secs:
|
||||
if isinstance(result, FloatImm):
|
||||
result = result.value
|
||||
assert isinstance(result, float)
|
||||
assert result >= 0.0
|
||||
_clean_build(builder_result.artifact_path)
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_rpc_multiple_runs():
|
||||
"""Test meta schedule rpc runner for multiple runs"""
|
||||
# Build the module
|
||||
mods = [
|
||||
MatmulModule,
|
||||
MatmulReluModule,
|
||||
BatchMatmulModule,
|
||||
]
|
||||
builder = LocalBuilder()
|
||||
builder_inputs = [BuilderInput(mod, Target("llvm")) for mod in mods]
|
||||
builder_results = builder.build(builder_inputs)
|
||||
for builder_result in builder_results:
|
||||
assert builder_result.artifact_path is not None
|
||||
assert builder_result.error_msg is None
|
||||
|
||||
args_infos = [
|
||||
[
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
],
|
||||
[
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
],
|
||||
[
|
||||
TensorInfo("float32", [16, MATMUL_M, MATMUL_M]),
|
||||
TensorInfo("float32", [16, MATMUL_M, MATMUL_M]),
|
||||
TensorInfo("float32", [16, MATMUL_M, MATMUL_M]),
|
||||
],
|
||||
]
|
||||
|
||||
runner_inputs = [
|
||||
RunnerInput(builder_results[i].artifact_path, "llvm", args_infos[i])
|
||||
for i in range(len(mods))
|
||||
]
|
||||
|
||||
with LocalRPC() as rpc:
|
||||
rpc_config = RPCConfig(
|
||||
tracker_host=rpc.tracker_host,
|
||||
tracker_port=rpc.tracker_port,
|
||||
tracker_key=rpc.tracker_key,
|
||||
session_priority=1,
|
||||
session_timeout_sec=100,
|
||||
)
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
runner = RPCRunner(rpc_config, evaluator_config)
|
||||
# Run the module
|
||||
runner_futures = runner.run(runner_inputs)
|
||||
runner_results = [runner_future.result() for runner_future in runner_futures]
|
||||
|
||||
for runner_result in runner_results:
|
||||
assert runner_result.error_msg is None
|
||||
for result in runner_result.run_secs:
|
||||
if isinstance(result, FloatImm):
|
||||
result = result.value
|
||||
assert isinstance(result, float)
|
||||
assert result >= 0.0
|
||||
|
||||
for builder_result in builder_results:
|
||||
_clean_build(builder_result.artifact_path)
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_local_multiple_runs():
|
||||
"""Test meta schedule local runner for multiple runs"""
|
||||
# Build the module
|
||||
mods = [
|
||||
MatmulModule,
|
||||
MatmulReluModule,
|
||||
BatchMatmulModule,
|
||||
]
|
||||
builder = LocalBuilder()
|
||||
builder_inputs = [BuilderInput(mod, Target("llvm")) for mod in mods]
|
||||
builder_results = builder.build(builder_inputs)
|
||||
for builder_result in builder_results:
|
||||
assert builder_result.artifact_path is not None
|
||||
assert builder_result.error_msg is None
|
||||
|
||||
args_infos = [
|
||||
[
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
],
|
||||
[
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
],
|
||||
[
|
||||
TensorInfo("float32", [16, MATMUL_M, MATMUL_M]),
|
||||
TensorInfo("float32", [16, MATMUL_M, MATMUL_M]),
|
||||
TensorInfo("float32", [16, MATMUL_M, MATMUL_M]),
|
||||
],
|
||||
]
|
||||
|
||||
runner_inputs = [
|
||||
RunnerInput(builder_results[i].artifact_path, "llvm", args_infos[i])
|
||||
for i in range(len(mods))
|
||||
]
|
||||
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
|
||||
runner = LocalRunner(timeout_sec=100, evaluator_config=evaluator_config)
|
||||
|
||||
# Run the module
|
||||
runner_futures = runner.run(runner_inputs)
|
||||
runner_results = [runner_future.result() for runner_future in runner_futures]
|
||||
|
||||
for runner_result in runner_results:
|
||||
assert runner_result.error_msg is None
|
||||
for result in runner_result.run_secs:
|
||||
if isinstance(result, FloatImm):
|
||||
result = result.value
|
||||
assert isinstance(result, float)
|
||||
assert result >= 0.0
|
||||
|
||||
for builder_result in builder_results:
|
||||
_clean_build(builder_result.artifact_path)
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_py_runner():
|
||||
"""Test meta schedule PyRunner"""
|
||||
|
||||
@derived_object
|
||||
class TestRunner(PyRunner):
|
||||
def run(self, runner_inputs: list[RunnerInput]) -> list[RunnerFuture]:
|
||||
raise ValueError("TestRunner")
|
||||
|
||||
runner = TestRunner()
|
||||
with pytest.raises(ValueError, match="TestRunner"):
|
||||
runner.run([])
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
@tvm.testing.skip_if_32bit(reason="skipping test for i386.")
|
||||
def test_meta_schedule_rpc_runner_time_out():
|
||||
"""Test meta schedule RPC Runner time out by using a super large workload"""
|
||||
|
||||
builder = LocalBuilder()
|
||||
builder_inputs = [BuilderInput(MatmulHugeModule, Target("llvm"))]
|
||||
builder_results = builder.build(builder_inputs)
|
||||
builder_results[0].artifact_path
|
||||
|
||||
runner_input = RunnerInput(
|
||||
builder_results[0].artifact_path,
|
||||
"llvm",
|
||||
[
|
||||
TensorInfo("float32", (4096, 4096)),
|
||||
TensorInfo("float32", (4096, 4096)),
|
||||
TensorInfo("float32", (4096, 4096)),
|
||||
],
|
||||
)
|
||||
|
||||
with LocalRPC() as rpc:
|
||||
rpc_config = RPCConfig(
|
||||
tracker_host=rpc.tracker_host,
|
||||
tracker_port=rpc.tracker_port,
|
||||
tracker_key=rpc.tracker_key,
|
||||
session_priority=1,
|
||||
session_timeout_sec=1,
|
||||
)
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
runner = RPCRunner(
|
||||
rpc_config,
|
||||
evaluator_config,
|
||||
)
|
||||
# Run the module
|
||||
(runner_future,) = runner.run([runner_input])
|
||||
runner_result = runner_future.result()
|
||||
|
||||
assert runner_result.error_msg is not None and runner_result.error_msg.startswith(
|
||||
"RPCRunner: An exception occurred"
|
||||
)
|
||||
assert runner_result.run_secs is None
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_local_runner_time_out():
|
||||
"""Test meta schedule Local Runner time out"""
|
||||
mod = MatmulModule
|
||||
builder = LocalBuilder()
|
||||
(builder_result,) = builder.build([BuilderInput(mod, Target("llvm"))])
|
||||
assert builder_result.artifact_path is not None
|
||||
assert builder_result.error_msg is None
|
||||
|
||||
runner_input = RunnerInput(
|
||||
builder_result.artifact_path,
|
||||
"llvm",
|
||||
[
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
],
|
||||
)
|
||||
|
||||
def initializer():
|
||||
@register_global_func("s_tir.meta_schedule.runner.test_time_out")
|
||||
def timeout_session_creator( # pylint: disable=unused-variable
|
||||
device: Device, # pylint: disable=unused-argument
|
||||
args_info: T_ARG_INFO_JSON_OBJ_LIST, # pylint: disable=unused-argument
|
||||
alloc_repeat: int, # pylint: disable=unused-argument
|
||||
) -> RPCSession:
|
||||
time.sleep(2)
|
||||
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
|
||||
runner = LocalRunner(
|
||||
timeout_sec=1,
|
||||
evaluator_config=evaluator_config,
|
||||
initializer=initializer,
|
||||
f_alloc_argument="s_tir.meta_schedule.runner.test_time_out",
|
||||
)
|
||||
|
||||
# Run the module
|
||||
(runner_future,) = runner.run([runner_input])
|
||||
runner_result = runner_future.result()
|
||||
|
||||
assert runner_result.error_msg is not None and runner_result.error_msg.startswith(
|
||||
"LocalRunner: Timeout, killed after"
|
||||
)
|
||||
assert runner_result.run_secs is None
|
||||
_clean_build(builder_result.artifact_path)
|
||||
|
||||
|
||||
@pytest.mark.skip("Disable this test to unblock CI.")
|
||||
def test_meta_schedule_rpc_runner_exception():
|
||||
"""Test meta schedule RPC Runner exception"""
|
||||
|
||||
def initializer():
|
||||
@register_global_func("s_tir.meta_schedule.runner.test_exception")
|
||||
def exception_session_creator( # pylint: disable=unused-variable
|
||||
rpc_config: RPCConfig, # pylint: disable=unused-argument
|
||||
) -> RPCSession:
|
||||
raise Exception("Test")
|
||||
|
||||
runner_input = RunnerInput(
|
||||
"test",
|
||||
"llvm",
|
||||
[
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
],
|
||||
)
|
||||
|
||||
with LocalRPC() as rpc:
|
||||
rpc_config = RPCConfig(
|
||||
tracker_host=rpc.tracker_host,
|
||||
tracker_port=rpc.tracker_port,
|
||||
tracker_key=rpc.tracker_key,
|
||||
session_priority=1,
|
||||
session_timeout_sec=100,
|
||||
)
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
runner = RPCRunner(
|
||||
rpc_config,
|
||||
evaluator_config,
|
||||
initializer=initializer,
|
||||
f_create_session="s_tir.meta_schedule.runner.test_exception",
|
||||
)
|
||||
(runner_future,) = runner.run([runner_input])
|
||||
runner_result = runner_future.result()
|
||||
|
||||
assert runner_result.error_msg is not None and runner_result.error_msg.startswith(
|
||||
"RPCRunner: An exception occurred\n"
|
||||
)
|
||||
assert runner_result.run_secs is None
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_local_runner_exception():
|
||||
"""Test meta schedule Local Runner exception"""
|
||||
mod = MatmulModule
|
||||
builder = LocalBuilder()
|
||||
(builder_result,) = builder.build([BuilderInput(mod, Target("llvm"))])
|
||||
assert builder_result.artifact_path is not None
|
||||
assert builder_result.error_msg is None
|
||||
|
||||
runner_input = RunnerInput(
|
||||
builder_result.artifact_path,
|
||||
"llvm",
|
||||
[
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
],
|
||||
)
|
||||
|
||||
def initializer():
|
||||
@register_global_func("s_tir.meta_schedule.runner.test_exception")
|
||||
def timeout_session_creator( # pylint: disable=unused-variable
|
||||
device: Device, # pylint: disable=unused-argument
|
||||
args_info: T_ARG_INFO_JSON_OBJ_LIST, # pylint: disable=unused-argument
|
||||
alloc_repeat: int, # pylint: disable=unused-argument
|
||||
) -> RPCSession:
|
||||
raise Exception("Test")
|
||||
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
|
||||
runner = LocalRunner(
|
||||
evaluator_config=evaluator_config,
|
||||
initializer=initializer,
|
||||
f_alloc_argument="s_tir.meta_schedule.runner.test_exception",
|
||||
)
|
||||
|
||||
# Run the module
|
||||
(runner_future,) = runner.run([runner_input])
|
||||
runner_result = runner_future.result()
|
||||
|
||||
assert runner_result.error_msg is not None and runner_result.error_msg.startswith(
|
||||
"LocalRunner: An exception occurred\n"
|
||||
)
|
||||
assert runner_result.run_secs is None
|
||||
_clean_build(builder_result.artifact_path)
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_runner_matmul_test():
|
||||
"""Test meta schedule runner with add module"""
|
||||
|
||||
def _check_correct_matmul(
|
||||
args_before: list[np.ndarray],
|
||||
args_after: list[np.ndarray],
|
||||
) -> None:
|
||||
a_before, b_before, c_before = args_before
|
||||
a_after, b_after, c_after = args_after
|
||||
c_before = np.matmul(a_before, b_before)
|
||||
assert (a_before == a_after).all()
|
||||
assert (b_before == b_after).all()
|
||||
tvm.testing.assert_allclose(c_before, c_after, rtol=1e-5)
|
||||
|
||||
def test_alloc_argument(
|
||||
session: RPCSession,
|
||||
device: Device,
|
||||
args_info: Any,
|
||||
alloc_repeat: int,
|
||||
) -> list[Any]:
|
||||
global repeated_args_before # pylint: disable=global-variable-undefined, invalid-name
|
||||
repeated_args_before = [] # type: ignore
|
||||
repeated_args = rpc_default_alloc_argument(session, device, args_info, alloc_repeat)
|
||||
for args in repeated_args:
|
||||
repeated_args_before.append([arg.numpy() for arg in args]) # type: ignore
|
||||
return repeated_args
|
||||
|
||||
def test_run_evaluator(
|
||||
session: RPCSession, # pylint: disable=unused-argument
|
||||
rt_mod: Module,
|
||||
device: Device,
|
||||
evaluator_config: EvaluatorConfig,
|
||||
repeated_args: list[Any],
|
||||
) -> list[float]:
|
||||
global repeated_args_before # pylint: disable=global-variable-undefined, invalid-name
|
||||
repeated_args_after = []
|
||||
evaluator = rt_mod.time_evaluator(
|
||||
func_name=rt_mod.entry_name,
|
||||
dev=device,
|
||||
number=evaluator_config.number,
|
||||
repeat=evaluator_config.repeat,
|
||||
min_repeat_ms=evaluator_config.min_repeat_ms,
|
||||
f_preproc=(
|
||||
"cache_flush_cpu_non_first_arg" if evaluator_config.enable_cpu_cache_flush else ""
|
||||
),
|
||||
)
|
||||
repeated_costs: list[list[float]] = []
|
||||
for args in repeated_args:
|
||||
device.sync()
|
||||
profile_result = evaluator(*args)
|
||||
repeated_costs.append(profile_result.results)
|
||||
repeated_args_after.append([arg.numpy() for arg in args])
|
||||
costs = [float(cost) for cost in itertools.chain.from_iterable(repeated_costs)]
|
||||
for args_before, args_after in zip(
|
||||
repeated_args_before, # type: ignore
|
||||
repeated_args_after,
|
||||
):
|
||||
_check_correct_matmul(args_before, args_after)
|
||||
del repeated_args_before # type: ignore
|
||||
return costs
|
||||
|
||||
# Build the module
|
||||
mod = MatmulModule
|
||||
builder = LocalBuilder()
|
||||
(builder_result,) = builder.build([BuilderInput(mod, Target("llvm"))])
|
||||
assert builder_result.artifact_path is not None
|
||||
assert builder_result.error_msg is None
|
||||
|
||||
runner_input = RunnerInput(
|
||||
builder_result.artifact_path,
|
||||
"llvm",
|
||||
[
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
|
||||
],
|
||||
)
|
||||
|
||||
with LocalRPC() as rpc:
|
||||
rpc_config = RPCConfig(
|
||||
tracker_host=rpc.tracker_host,
|
||||
tracker_port=rpc.tracker_port,
|
||||
tracker_key=rpc.tracker_key,
|
||||
session_priority=1,
|
||||
session_timeout_sec=100,
|
||||
)
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
runner = RPCRunner(
|
||||
rpc_config,
|
||||
evaluator_config,
|
||||
f_alloc_argument=test_alloc_argument,
|
||||
f_run_evaluator=test_run_evaluator,
|
||||
)
|
||||
# Run the module
|
||||
(runner_future,) = runner.run([runner_input])
|
||||
runner_result = runner_future.result()
|
||||
assert runner_result.error_msg is None
|
||||
for result in runner_result.run_secs:
|
||||
if isinstance(result, FloatImm):
|
||||
result = result.value
|
||||
assert isinstance(result, float)
|
||||
assert result >= 0.0
|
||||
_clean_build(builder_result.artifact_path)
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_runner_add_test():
|
||||
"""Test meta schedule runner with add module"""
|
||||
|
||||
def _check_correct_add(args_before: list[np.ndarray], args_after: list[np.ndarray]) -> None:
|
||||
a_before, b_before, c_before = args_before
|
||||
a_after, b_after, c_after = args_after
|
||||
c_before = a_before + b_before
|
||||
assert (a_before == a_after).all()
|
||||
assert (b_before == b_after).all()
|
||||
assert (c_before == c_after).all()
|
||||
|
||||
def test_alloc_argument(
|
||||
session: RPCSession,
|
||||
device: Device,
|
||||
args_info: Any,
|
||||
alloc_repeat: int,
|
||||
) -> list[Any]:
|
||||
global repeated_args_before # pylint: disable=global-variable-undefined, invalid-name
|
||||
repeated_args_before = [] # type: ignore
|
||||
repeated_args = rpc_default_alloc_argument(
|
||||
session,
|
||||
device,
|
||||
args_info,
|
||||
alloc_repeat,
|
||||
)
|
||||
for args in repeated_args:
|
||||
repeated_args_before.append([arg.numpy() for arg in args]) # type: ignore
|
||||
return repeated_args
|
||||
|
||||
def test_run_evaluator(
|
||||
session: RPCSession, # pylint: disable=unused-argument
|
||||
rt_mod: Module,
|
||||
device: Device,
|
||||
evaluator_config: EvaluatorConfig,
|
||||
repeated_args: list[Any],
|
||||
) -> list[float]:
|
||||
global repeated_args_before # pylint: disable=global-variable-undefined, invalid-name
|
||||
repeated_args_after = []
|
||||
evaluator = rt_mod.time_evaluator(
|
||||
func_name=rt_mod.entry_name,
|
||||
dev=device,
|
||||
number=evaluator_config.number,
|
||||
repeat=evaluator_config.repeat,
|
||||
min_repeat_ms=evaluator_config.min_repeat_ms,
|
||||
f_preproc=(
|
||||
"cache_flush_cpu_non_first_arg" if evaluator_config.enable_cpu_cache_flush else ""
|
||||
),
|
||||
)
|
||||
repeated_costs: list[list[float]] = []
|
||||
for args in repeated_args:
|
||||
device.sync()
|
||||
profile_result = evaluator(*args)
|
||||
repeated_costs.append(profile_result.results)
|
||||
repeated_args_after.append([arg.numpy() for arg in args])
|
||||
costs = [float(cost) for cost in itertools.chain.from_iterable(repeated_costs)]
|
||||
for args_before, args_after in zip(
|
||||
repeated_args_before, # type: ignore
|
||||
repeated_args_after,
|
||||
):
|
||||
_check_correct_add(args_before, args_after)
|
||||
del repeated_args_before # type: ignore
|
||||
return costs
|
||||
|
||||
# Build the module
|
||||
mod = AddModule
|
||||
builder = LocalBuilder()
|
||||
(builder_result,) = builder.build([BuilderInput(mod, Target("llvm"))])
|
||||
assert builder_result.artifact_path is not None
|
||||
assert builder_result.error_msg is None
|
||||
|
||||
runner_input = RunnerInput(
|
||||
builder_result.artifact_path,
|
||||
"llvm",
|
||||
[
|
||||
TensorInfo("float32", [MATMUL_M]),
|
||||
TensorInfo("float32", [MATMUL_M]),
|
||||
TensorInfo("float32", [MATMUL_M]),
|
||||
],
|
||||
)
|
||||
|
||||
with LocalRPC() as rpc:
|
||||
rpc_config = RPCConfig(
|
||||
tracker_host=rpc.tracker_host,
|
||||
tracker_port=rpc.tracker_port,
|
||||
tracker_key=rpc.tracker_key,
|
||||
session_priority=1,
|
||||
session_timeout_sec=100,
|
||||
)
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
runner = RPCRunner(
|
||||
rpc_config,
|
||||
evaluator_config,
|
||||
f_alloc_argument=test_alloc_argument,
|
||||
f_run_evaluator=test_run_evaluator,
|
||||
)
|
||||
# Run the module
|
||||
(runner_future,) = runner.run([runner_input])
|
||||
runner_result = runner_future.result()
|
||||
assert runner_result.error_msg is None
|
||||
for result in runner_result.run_secs:
|
||||
if isinstance(result, FloatImm):
|
||||
result = result.value
|
||||
assert isinstance(result, float)
|
||||
assert result >= 0.0
|
||||
_clean_build(builder_result.artifact_path)
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_meta_schedule_local_runner_add_test():
|
||||
"""Test meta schedule local runner with add module"""
|
||||
|
||||
def _check_correct_add(args_before: list[np.array], args_after: list[np.array]) -> None:
|
||||
a_before, b_before, c_before = args_before
|
||||
a_after, b_after, c_after = args_after
|
||||
c_before = a_before + b_before
|
||||
assert (a_before == a_after).all()
|
||||
assert (b_before == b_after).all()
|
||||
assert (c_before == c_after).all()
|
||||
|
||||
def test_alloc_argument(
|
||||
device: Device,
|
||||
args_info: T_ARG_INFO_JSON_OBJ_LIST, # pylint: disable=unused-argument
|
||||
alloc_repeat: int,
|
||||
) -> list[T_ARGUMENT_LIST]:
|
||||
global repeated_args_before # pylint: disable=global-variable-undefined, invalid-name
|
||||
repeated_args_before = []
|
||||
repeated_args = local_default_alloc_argument(device, args_info, alloc_repeat)
|
||||
for args in repeated_args:
|
||||
repeated_args_before.append([arg.numpy() for arg in args])
|
||||
return repeated_args
|
||||
|
||||
def test_run_evaluator(
|
||||
rt_mod: Module,
|
||||
device: Device,
|
||||
evaluator_config: EvaluatorConfig,
|
||||
repeated_args: list[Any],
|
||||
) -> list[float]:
|
||||
global repeated_args_before # pylint: disable=global-variable-undefined, invalid-name
|
||||
repeated_args_after = []
|
||||
evaluator = rt_mod.time_evaluator(
|
||||
func_name=rt_mod.entry_name,
|
||||
dev=device,
|
||||
number=evaluator_config.number,
|
||||
repeat=evaluator_config.repeat,
|
||||
min_repeat_ms=evaluator_config.min_repeat_ms,
|
||||
f_preproc=(
|
||||
"cache_flush_cpu_non_first_arg" if evaluator_config.enable_cpu_cache_flush else ""
|
||||
),
|
||||
)
|
||||
repeated_costs: list[list[float]] = []
|
||||
for args in repeated_args:
|
||||
device.sync()
|
||||
profile_result = evaluator(*args)
|
||||
repeated_costs.append(profile_result.results)
|
||||
repeated_args_after.append([arg.numpy() for arg in args])
|
||||
costs = [float(cost) for cost in itertools.chain.from_iterable(repeated_costs)]
|
||||
for args_before, args_after in zip(repeated_args_before, repeated_args_after):
|
||||
_check_correct_add(args_before, args_after)
|
||||
del repeated_args_before
|
||||
return costs
|
||||
|
||||
# Build the module
|
||||
mod = AddModule
|
||||
builder = LocalBuilder()
|
||||
(builder_result,) = builder.build([BuilderInput(mod, Target("llvm"))])
|
||||
assert builder_result.artifact_path is not None
|
||||
assert builder_result.error_msg is None
|
||||
|
||||
runner_input = RunnerInput(
|
||||
builder_result.artifact_path,
|
||||
"llvm",
|
||||
[
|
||||
TensorInfo("float32", [MATMUL_M]),
|
||||
TensorInfo("float32", [MATMUL_M]),
|
||||
TensorInfo("float32", [MATMUL_M]),
|
||||
],
|
||||
)
|
||||
|
||||
evaluator_config = EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=0,
|
||||
enable_cpu_cache_flush=False,
|
||||
)
|
||||
runner = LocalRunner(
|
||||
timeout_sec=100,
|
||||
evaluator_config=evaluator_config,
|
||||
f_alloc_argument=test_alloc_argument,
|
||||
f_run_evaluator=test_run_evaluator,
|
||||
)
|
||||
# Run the module
|
||||
(runner_future,) = runner.run([runner_input])
|
||||
runner_result = runner_future.result()
|
||||
assert runner_result.error_msg is None
|
||||
for result in runner_result.run_secs:
|
||||
if isinstance(result, FloatImm):
|
||||
result = result.value
|
||||
assert isinstance(result, float)
|
||||
assert result >= 0.0
|
||||
_clean_build(builder_result.artifact_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,294 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing import te_workload
|
||||
from tvm.s_tir.meta_schedule.testing.space_generation import (
|
||||
check_sketches,
|
||||
generate_design_space,
|
||||
)
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
from tvm.te import create_prim_func
|
||||
|
||||
|
||||
def test_cpu_matmul():
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_matmul_0(
|
||||
A: T.Buffer((4, 512), "float32"),
|
||||
B: T.Buffer((512, 4), "float32"),
|
||||
C: T.Buffer((4, 4), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
for i0, i1, i2 in T.grid(4, 4, 512):
|
||||
with T.sblock("C"):
|
||||
i, j, k = T.axis.remap("SSR", [i0, i1, i2])
|
||||
T.reads(A[i, k], B[k, j])
|
||||
T.writes(C[i, j])
|
||||
with T.init():
|
||||
C[i, j] = T.float32(0)
|
||||
C[i, j] = C[i, j] + A[i, k] * B[k, j]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_matmul_1(
|
||||
A: T.Buffer((4, 512), "float32"),
|
||||
B: T.Buffer((512, 4), "float32"),
|
||||
C: T.Buffer((4, 4), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
C_rf = T.sblock_alloc_buffer([4, 4, 128], dtype="float32")
|
||||
for i0, i1, i2_0, i2_1 in T.grid(4, 4, 4, 128):
|
||||
with T.sblock("C_rf"):
|
||||
vi2_1, i, j, vi2_0 = T.axis.remap("SSSR", [i2_1, i0, i1, i2_0])
|
||||
T.reads(A[i, vi2_0 * 128 + vi2_1], B[vi2_0 * 128 + vi2_1, j])
|
||||
T.writes(C_rf[i, j, vi2_1])
|
||||
with T.init():
|
||||
C_rf[i, j, vi2_1] = T.float32(0)
|
||||
C_rf[i, j, vi2_1] = (
|
||||
C_rf[i, j, vi2_1] + A[i, vi2_0 * 128 + vi2_1] * B[vi2_0 * 128 + vi2_1, j]
|
||||
)
|
||||
for i0, i1, i2_1 in T.grid(4, 4, 128):
|
||||
with T.sblock("C"):
|
||||
vi2_1, i, j = T.axis.remap("RSS", [i2_1, i0, i1])
|
||||
T.reads(C_rf[i, j, vi2_1])
|
||||
T.writes(C[i, j])
|
||||
T.sblock_attr({"meta_schedule.random_compute_producer": 1})
|
||||
with T.init():
|
||||
C[i, j] = T.float32(0)
|
||||
C[i, j] = C[i, j] + C_rf[i, j, vi2_1]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_matmul_2(
|
||||
A: T.Buffer((4, 512), "float32"),
|
||||
B: T.Buffer((512, 4), "float32"),
|
||||
C: T.Buffer((4, 4), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
C_rf = T.sblock_alloc_buffer([4, 4, 4], dtype="float32")
|
||||
for i0, i1, i2_0, i2_1 in T.grid(4, 4, 4, 128):
|
||||
with T.sblock("C_rf"):
|
||||
vi2_0, i, j, vi2_1 = T.axis.remap("SSSR", [i2_0, i0, i1, i2_1])
|
||||
T.reads(A[i, vi2_0 * 128 + vi2_1], B[vi2_0 * 128 + vi2_1, j])
|
||||
T.writes(C_rf[i, j, vi2_0])
|
||||
with T.init():
|
||||
C_rf[i, j, vi2_0] = T.float32(0)
|
||||
C_rf[i, j, vi2_0] = (
|
||||
C_rf[i, j, vi2_0] + A[i, vi2_0 * 128 + vi2_1] * B[vi2_0 * 128 + vi2_1, j]
|
||||
)
|
||||
for i0, i1, i2_0 in T.grid(4, 4, 4):
|
||||
with T.sblock("C"):
|
||||
vi2_0, i, j = T.axis.remap("RSS", [i2_0, i0, i1])
|
||||
T.reads(C_rf[i, j, vi2_0])
|
||||
T.writes(C[i, j])
|
||||
T.sblock_attr({"meta_schedule.random_compute_producer": 1})
|
||||
with T.init():
|
||||
C[i, j] = T.float32(0)
|
||||
C[i, j] = C[i, j] + C_rf[i, j, vi2_0]
|
||||
|
||||
decision_0 = [] # type: ignore
|
||||
decision_1 = [
|
||||
("SamplePerfectTile", [4, 128]),
|
||||
]
|
||||
decision_2 = [
|
||||
("SamplePerfectTile", [4, 128]),
|
||||
]
|
||||
mod = create_prim_func(te_workload.matmul(n=4, m=4, k=512))
|
||||
actual = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=Target({"kind": "llvm", "num-cores": 32}),
|
||||
types=ms.schedule_rule.AddRFactor,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[cpu_matmul_0, cpu_matmul_1, cpu_matmul_2],
|
||||
expected_decisions=[decision_0, decision_1, decision_2],
|
||||
)
|
||||
|
||||
|
||||
def test_cpu_argmax():
|
||||
@T.prim_func(s_tir=True)
|
||||
def argmax(
|
||||
idx: T.Buffer((128, 128), "int32"),
|
||||
val: T.Buffer((128, 128), "float32"),
|
||||
argmax_v0: T.Buffer((128,), "int32"),
|
||||
argmax_v1: T.Buffer((128,), "float32"),
|
||||
) -> None:
|
||||
for i0, i1 in T.grid(128, 128):
|
||||
with T.sblock("argmax"):
|
||||
i = T.axis.spatial(128, i0)
|
||||
k = T.axis.reduce(128, i1)
|
||||
T.reads(idx[i, k], val[i, k])
|
||||
T.writes(argmax_v0[i], argmax_v1[i])
|
||||
with T.init():
|
||||
argmax_v0[i] = -1
|
||||
argmax_v1[i] = T.min_value("float32")
|
||||
v_argmax_v0: T.let[T.int32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v0[i], idx[i, k]
|
||||
)
|
||||
v_argmax_v1: T.let[T.float32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v1[i], val[i, k]
|
||||
)
|
||||
argmax_v0[i] = v_argmax_v0
|
||||
argmax_v1[i] = v_argmax_v1
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def argmax_0(
|
||||
idx: T.Buffer((128, 128), "int32"),
|
||||
val: T.Buffer((128, 128), "float32"),
|
||||
argmax_v0: T.Buffer(128, "int32"),
|
||||
argmax_v1: T.Buffer(128, "float32"),
|
||||
) -> None:
|
||||
for i0, i1 in T.grid(128, 128):
|
||||
with T.sblock("argmax"):
|
||||
i, k = T.axis.remap("SR", [i0, i1])
|
||||
T.reads(idx[i, k], val[i, k])
|
||||
T.writes(argmax_v0[i], argmax_v1[i])
|
||||
with T.init():
|
||||
argmax_v0[i] = -1
|
||||
argmax_v1[i] = T.float32(-3.4028234663852886e38)
|
||||
v_argmax_v0: T.let[T.int32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v0[i], idx[i, k]
|
||||
)
|
||||
v_argmax_v1: T.let[T.float32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v1[i], val[i, k]
|
||||
)
|
||||
argmax_v0[i] = v_argmax_v0
|
||||
argmax_v1[i] = v_argmax_v1
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def argmax_1(
|
||||
idx: T.Buffer((128, 128), "int32"),
|
||||
val: T.Buffer((128, 128), "float32"),
|
||||
argmax_v0: T.Buffer(128, "int32"),
|
||||
argmax_v1: T.Buffer(128, "float32"),
|
||||
) -> None:
|
||||
argmax_v0_rf = T.sblock_alloc_buffer([128, 16], dtype="int32")
|
||||
argmax_v1_rf = T.sblock_alloc_buffer([128, 16], dtype="float32")
|
||||
for i0, i1_0, i1_1 in T.grid(128, 8, 16):
|
||||
with T.sblock("argmax_rf"):
|
||||
vi1_1, i, vi1_0 = T.axis.remap("SSR", [i1_1, i0, i1_0])
|
||||
T.reads(idx[i, vi1_0 * 16 + vi1_1], val[i, vi1_0 * 16 + vi1_1])
|
||||
T.writes(argmax_v0_rf[i, vi1_1], argmax_v1_rf[i, vi1_1])
|
||||
with T.init():
|
||||
argmax_v0_rf[i, vi1_1] = -1
|
||||
argmax_v1_rf[i, vi1_1] = T.float32(-3.4028234663852886e38)
|
||||
v_argmax_v0_rf: T.let[T.int32] = T.Select(
|
||||
argmax_v1_rf[i, vi1_1] >= val[i, vi1_0 * 16 + vi1_1],
|
||||
argmax_v0_rf[i, vi1_1],
|
||||
idx[i, vi1_0 * 16 + vi1_1],
|
||||
)
|
||||
v_argmax_v1_rf: T.let[T.float32] = T.Select(
|
||||
argmax_v1_rf[i, vi1_1] >= val[i, vi1_0 * 16 + vi1_1],
|
||||
argmax_v1_rf[i, vi1_1],
|
||||
val[i, vi1_0 * 16 + vi1_1],
|
||||
)
|
||||
argmax_v0_rf[i, vi1_1] = v_argmax_v0_rf
|
||||
argmax_v1_rf[i, vi1_1] = v_argmax_v1_rf
|
||||
for i0, i1_1 in T.grid(128, 16):
|
||||
with T.sblock("argmax"):
|
||||
vi1_1, i = T.axis.remap("RS", [i1_1, i0])
|
||||
T.reads(argmax_v0_rf[i, vi1_1], argmax_v1_rf[i, vi1_1])
|
||||
T.writes(argmax_v0[i], argmax_v1[i])
|
||||
T.sblock_attr({"meta_schedule.random_compute_producer": 1})
|
||||
with T.init():
|
||||
argmax_v0[i] = -1
|
||||
argmax_v1[i] = T.float32(-3.4028234663852886e38)
|
||||
v_argmax_v0: T.let[T.int32] = T.Select(
|
||||
argmax_v1[i] >= argmax_v1_rf[i, vi1_1], argmax_v0[i], argmax_v0_rf[i, vi1_1]
|
||||
)
|
||||
v_argmax_v1: T.let[T.float32] = T.Select(
|
||||
argmax_v1[i] >= argmax_v1_rf[i, vi1_1], argmax_v1[i], argmax_v1_rf[i, vi1_1]
|
||||
)
|
||||
argmax_v0[i] = v_argmax_v0
|
||||
argmax_v1[i] = v_argmax_v1
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def argmax_2(
|
||||
idx: T.Buffer((128, 128), "int32"),
|
||||
val: T.Buffer((128, 128), "float32"),
|
||||
argmax_v0: T.Buffer(128, "int32"),
|
||||
argmax_v1: T.Buffer(128, "float32"),
|
||||
) -> None:
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
argmax_v0_rf = T.sblock_alloc_buffer([128, 8], dtype="int32")
|
||||
argmax_v1_rf = T.sblock_alloc_buffer([128, 8], dtype="float32")
|
||||
for i0, i1_0, i1_1 in T.grid(128, 8, 16):
|
||||
with T.sblock("argmax_rf"):
|
||||
vi1_0, i, vi1_1 = T.axis.remap("SSR", [i1_0, i0, i1_1])
|
||||
T.reads(idx[i, vi1_0 * 16 + vi1_1], val[i, vi1_0 * 16 + vi1_1])
|
||||
T.writes(argmax_v0_rf[i, vi1_0], argmax_v1_rf[i, vi1_0])
|
||||
with T.init():
|
||||
argmax_v0_rf[i, vi1_0] = -1
|
||||
argmax_v1_rf[i, vi1_0] = T.float32(-3.4028234663852886e38)
|
||||
v_argmax_v0_rf: T.let[T.int32] = T.Select(
|
||||
argmax_v1_rf[i, vi1_0] >= val[i, vi1_0 * 16 + vi1_1],
|
||||
argmax_v0_rf[i, vi1_0],
|
||||
idx[i, vi1_0 * 16 + vi1_1],
|
||||
)
|
||||
v_argmax_v1_rf: T.let[T.float32] = T.Select(
|
||||
argmax_v1_rf[i, vi1_0] >= val[i, vi1_0 * 16 + vi1_1],
|
||||
argmax_v1_rf[i, vi1_0],
|
||||
val[i, vi1_0 * 16 + vi1_1],
|
||||
)
|
||||
argmax_v0_rf[i, vi1_0] = v_argmax_v0_rf
|
||||
argmax_v1_rf[i, vi1_0] = v_argmax_v1_rf
|
||||
for i0, i1_0 in T.grid(128, 8):
|
||||
with T.sblock("argmax"):
|
||||
vi1_0, i = T.axis.remap("RS", [i1_0, i0])
|
||||
T.reads(argmax_v0_rf[i, vi1_0], argmax_v1_rf[i, vi1_0])
|
||||
T.writes(argmax_v0[i], argmax_v1[i])
|
||||
T.sblock_attr({"meta_schedule.random_compute_producer": 1})
|
||||
with T.init():
|
||||
argmax_v0[i] = -1
|
||||
argmax_v1[i] = T.float32(-3.4028234663852886e38)
|
||||
v_argmax_v0: T.let[T.int32] = T.Select(
|
||||
argmax_v1[i] >= argmax_v1_rf[i, vi1_0], argmax_v0[i], argmax_v0_rf[i, vi1_0]
|
||||
)
|
||||
v_argmax_v1: T.let[T.float32] = T.Select(
|
||||
argmax_v1[i] >= argmax_v1_rf[i, vi1_0], argmax_v1[i], argmax_v1_rf[i, vi1_0]
|
||||
)
|
||||
argmax_v0[i] = v_argmax_v0
|
||||
argmax_v1[i] = v_argmax_v1
|
||||
|
||||
decision_0 = [] # type: ignore
|
||||
decision_1 = [
|
||||
("SamplePerfectTile", [8, 16]),
|
||||
]
|
||||
decision_2 = [
|
||||
("SamplePerfectTile", [8, 16]),
|
||||
]
|
||||
mod = argmax
|
||||
actual = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=Target({"kind": "llvm", "num-cores": 32}),
|
||||
types=ms.schedule_rule.AddRFactor,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[argmax_0, argmax_1, argmax_2],
|
||||
expected_decisions=[decision_0, decision_1, decision_2],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_cpu_matmul()
|
||||
test_cpu_argmax()
|
||||
+67
@@ -0,0 +1,67 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.schedule_rule import ApplyCustomRule
|
||||
from tvm.script import tirx as T
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
T.sblock_attr({"schedule_rule": "test_apply_custom_rule"})
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.register_global_func("s_tir.meta_schedule.cpu.test_apply_custom_rule")
|
||||
def sch_fn(sch: tvm.s_tir.Schedule, block: tvm.tirx.SBlock) -> list[tvm.s_tir.Schedule]:
|
||||
raise ValueError("Intended for s_tir.meta_schedule.cpu.test_apply_custom_rule")
|
||||
|
||||
|
||||
@pytest.mark.skip("Tuning test - launches runner")
|
||||
def test_custom_rule():
|
||||
with pytest.raises(ValueError) as e_info:
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
sch_rules = [ApplyCustomRule()]
|
||||
space_gen = ms.space_generator.PostOrderApply(sch_rules=sch_rules)
|
||||
ms.tune_tir(
|
||||
mod=Matmul,
|
||||
target={"kind": "llvm", "num-cores": 1},
|
||||
work_dir=tmpdir,
|
||||
max_trials_global=10,
|
||||
space=space_gen,
|
||||
)
|
||||
assert "Intended for s_tir.meta_schedule.cpu.test_apply_custom_rule" in str(e_info.value)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_custom_rule()
|
||||
@@ -0,0 +1,166 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: F841
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing.space_generation import (
|
||||
check_sketches,
|
||||
generate_design_space,
|
||||
)
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def element_wise(var_A: T.handle, var_B: T.handle) -> None:
|
||||
A = T.match_buffer(var_A, [512, 512], dtype="float32")
|
||||
B = T.match_buffer(var_B, [512, 512], dtype="float32")
|
||||
for i, j in T.grid(512, 512):
|
||||
with T.sblock("C"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = A[vi, vj] + 1.0
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def reduction_loop_only(
|
||||
A: T.Buffer(2, "float32"),
|
||||
B: T.Buffer(2, "float32"),
|
||||
C: T.Buffer((), "float32"),
|
||||
) -> None:
|
||||
for i0 in T.serial(2):
|
||||
with T.sblock("C"):
|
||||
k0 = T.axis.reduce(2, i0)
|
||||
T.reads(A[k0], B[k0])
|
||||
T.writes(C[()])
|
||||
with T.init():
|
||||
C[()] = T.float32(1.0)
|
||||
C[()] = T.min(C[()], A[k0] / B[k0])
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def zero_dim_add(
|
||||
A: T.Buffer((), "float32"),
|
||||
B: T.Buffer((), "float32"),
|
||||
C: T.Buffer((), "float32"),
|
||||
) -> None:
|
||||
with T.sblock("C"):
|
||||
vi = T.axis.spatial(1, 0)
|
||||
C[()] = A[()] + B[()]
|
||||
|
||||
|
||||
def test_cuda_element_wise():
|
||||
@T.prim_func(s_tir=True)
|
||||
def elementwise_0(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
) -> None:
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
for i_j_fused_0 in T.thread_binding(256, thread="blockIdx.x"):
|
||||
for i_j_fused_1 in T.thread_binding(1024, thread="threadIdx.x"):
|
||||
with T.sblock("C"):
|
||||
vi = T.axis.spatial(512, (i_j_fused_0 * 1024 + i_j_fused_1) // 512)
|
||||
vj = T.axis.spatial(512, (i_j_fused_0 * 1024 + i_j_fused_1) % 512)
|
||||
T.reads(A[vi, vj])
|
||||
T.writes(B[vi, vj])
|
||||
B[vi, vj] = A[vi, vj] + T.float32(1)
|
||||
|
||||
decision_0 = [
|
||||
("SampleCategorical", 5),
|
||||
]
|
||||
mod = element_wise
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-3080", host="llvm"),
|
||||
types=ms.schedule_rule.AutoBind,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[elementwise_0],
|
||||
expected_decisions=[decision_0],
|
||||
)
|
||||
|
||||
|
||||
def test_cuda_reduction_loop_only():
|
||||
@T.prim_func(s_tir=True)
|
||||
def reduction_loop_only_0(
|
||||
A: T.Buffer(2, "float32"),
|
||||
B: T.Buffer(2, "float32"),
|
||||
C: T.Buffer((), "float32"),
|
||||
) -> None:
|
||||
for u_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
|
||||
for u_fused_1 in T.thread_binding(1, thread="threadIdx.x"):
|
||||
for i0 in T.serial(2):
|
||||
with T.sblock("C"):
|
||||
k0 = T.axis.reduce(2, i0)
|
||||
T.reads(A[k0], B[k0])
|
||||
T.writes(C[()])
|
||||
with T.init():
|
||||
C[()] = T.float32(1)
|
||||
C[()] = T.min(C[()], A[k0] / B[k0])
|
||||
|
||||
mod = reduction_loop_only
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-3080", host="llvm"),
|
||||
types=ms.schedule_rule.AutoBind,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[reduction_loop_only_0],
|
||||
expected_decisions=[[]],
|
||||
)
|
||||
|
||||
|
||||
def test_cuda_zero_dim_add():
|
||||
@T.prim_func(s_tir=True)
|
||||
def zero_dim_add_0(
|
||||
A: T.Buffer((), "float32"),
|
||||
B: T.Buffer((), "float32"),
|
||||
C: T.Buffer((), "float32"),
|
||||
) -> None:
|
||||
for u_fused_0 in T.thread_binding(1, thread="blockIdx.x"):
|
||||
for u_fused_1 in T.thread_binding(1, thread="threadIdx.x"):
|
||||
with T.sblock("C"):
|
||||
vi = T.axis.spatial(1, 0)
|
||||
T.reads(A[()], B[()])
|
||||
T.writes(C[()])
|
||||
C[()] = A[()] + B[()]
|
||||
|
||||
mod = zero_dim_add
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-3080", host="llvm"),
|
||||
types=ms.schedule_rule.AutoBind,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[zero_dim_add_0],
|
||||
expected_decisions=[[]],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_cuda_element_wise()
|
||||
test_cuda_reduction_loop_only()
|
||||
test_cuda_zero_dim_add()
|
||||
@@ -0,0 +1,567 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: E501, E741, F841
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir.base import assert_structural_equal
|
||||
from tvm.s_tir import Schedule
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing.space_generation import generate_design_space
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# fmt: off
|
||||
# pylint: disable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2DBiasBnReLU:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_X: T.handle, var_W: T.handle, var_B: T.handle, var_bn_scale: T.handle, var_bn_offset: T.handle, var_compute: T.handle) -> None:
|
||||
X = T.match_buffer(var_X, [1, 512, 56, 56], dtype="float32")
|
||||
W = T.match_buffer(var_W, [512, 512, 3, 3], dtype="float32")
|
||||
B = T.match_buffer(var_B, [512, 1, 1], dtype="float32")
|
||||
bn_scale = T.match_buffer(var_bn_scale, [512, 1, 1], dtype="float32")
|
||||
bn_offset = T.match_buffer(var_bn_offset, [512, 1, 1], dtype="float32")
|
||||
compute = T.match_buffer(var_compute, [1, 512, 56, 56], dtype="float32")
|
||||
pad_temp = T.sblock_alloc_buffer([1, 512, 58, 58], dtype="float32")
|
||||
compute_1 = T.sblock_alloc_buffer([1, 512, 56, 56], dtype="float32")
|
||||
bias_add = T.sblock_alloc_buffer([1, 512, 56, 56], dtype="float32")
|
||||
bn_mul = T.sblock_alloc_buffer([1, 512, 56, 56], dtype="float32")
|
||||
bn_add = T.sblock_alloc_buffer([1, 512, 56, 56], dtype="float32")
|
||||
for i0, i1, i2, i3 in T.grid(1, 512, 58, 58):
|
||||
with T.sblock("pad_temp"):
|
||||
i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
pad_temp[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(i2_1 >= 1 and i2_1 < 57 and i3_1 >= 1 and i3_1 < 57, X[i0_1, i1_1, i2_1 - 1, i3_1 - 1], T.float32(0), dtype="float32")
|
||||
for i0, i1, i2, i3, i4, i5, i6 in T.grid(1, 512, 56, 56, 512, 3, 3):
|
||||
with T.sblock("compute"):
|
||||
nn, ff, yy, xx, rc, ry, rx = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
|
||||
with T.init():
|
||||
compute_1[nn, ff, yy, xx] = T.float32(0)
|
||||
compute_1[nn, ff, yy, xx] = compute_1[nn, ff, yy, xx] + pad_temp[nn, rc, yy + ry, xx + rx] * W[ff, rc, ry, rx]
|
||||
for i0, i1, i2, i3 in T.grid(1, 512, 56, 56):
|
||||
with T.sblock("bias_add"):
|
||||
i, j, k, l = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
bias_add[i, j, k, l] = compute_1[i, j, k, l] + B[j, 0, 0]
|
||||
for i0, i1, i2, i3 in T.grid(1, 512, 56, 56):
|
||||
with T.sblock("bn_mul"):
|
||||
i, j, k, l = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
bn_mul[i, j, k, l] = bias_add[i, j, k, l] * bn_scale[j, 0, 0]
|
||||
for i0, i1, i2, i3 in T.grid(1, 512, 56, 56):
|
||||
with T.sblock("bn_add"):
|
||||
i, j, k, l = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
bn_add[i, j, k, l] = bn_mul[i, j, k, l] + bn_offset[j, 0, 0]
|
||||
for i0, i1, i2, i3 in T.grid(1, 512, 56, 56):
|
||||
with T.sblock("compute_1"):
|
||||
i0_2, i1_2, i2_2, i3_2 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
compute[i0_2, i1_2, i2_2, i3_2] = T.max(bn_add[i0_2, i1_2, i2_2, i3_2], T.float32(0))
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2DBiasBnReLUInlined:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_X: T.handle, var_W: T.handle, var_B: T.handle, var_bn_scale: T.handle, var_bn_offset: T.handle, var_compute: T.handle) -> None:
|
||||
X = T.match_buffer(var_X, [1, 512, 56, 56], dtype="float32")
|
||||
W = T.match_buffer(var_W, [512, 512, 3, 3], dtype="float32")
|
||||
B = T.match_buffer(var_B, [512, 1, 1], dtype="float32")
|
||||
bn_scale = T.match_buffer(var_bn_scale, [512, 1, 1], dtype="float32")
|
||||
bn_offset = T.match_buffer(var_bn_offset, [512, 1, 1], dtype="float32")
|
||||
compute = T.match_buffer(var_compute, [1, 512, 56, 56], dtype="float32")
|
||||
pad_temp = T.sblock_alloc_buffer([1, 512, 58, 58], dtype="float32")
|
||||
compute_1 = T.sblock_alloc_buffer([1, 512, 56, 56], dtype="float32")
|
||||
for i0, i1, i2, i3 in T.grid(1, 512, 58, 58):
|
||||
with T.sblock("pad_temp"):
|
||||
i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
pad_temp[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(i2_1 >= 1 and i2_1 < 57 and i3_1 >= 1 and i3_1 < 57, X[i0_1, i1_1, i2_1 - 1, i3_1 - 1], T.float32(0), dtype="float32")
|
||||
for i0, i1, i2, i3, i4, i5, i6 in T.grid(1, 512, 56, 56, 512, 3, 3):
|
||||
with T.sblock("compute"):
|
||||
nn, ff, yy, xx, rc, ry, rx = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
|
||||
with T.init():
|
||||
compute_1[nn, ff, yy, xx] = T.float32(0)
|
||||
compute_1[nn, ff, yy, xx] = compute_1[nn, ff, yy, xx] + pad_temp[nn, rc, yy + ry, xx + rx] * W[ff, rc, ry, rx]
|
||||
for i0, i1, i2, i3 in T.grid(1, 512, 56, 56):
|
||||
with T.sblock("compute_1"):
|
||||
i0_2, i1_2, i2_2, i3_2 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
compute[i0_2, i1_2, i2_2, i3_2] = T.max((compute_1[i0_2, i1_2, i2_2, i3_2] + B[i1_2, 0, 0]) * bn_scale[i1_2, 0, 0] + bn_offset[i1_2, 0, 0], T.float32(0))
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class MultiLevelTiledConv2D:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(var_X: T.handle, var_W: T.handle, var_B: T.handle, var_bn_scale: T.handle, var_bn_offset: T.handle, var_compute: T.handle) -> None:
|
||||
X = T.match_buffer(var_X, [1, 512, 56, 56], dtype="float32")
|
||||
W = T.match_buffer(var_W, [512, 512, 3, 3], dtype="float32")
|
||||
B = T.match_buffer(var_B, [512, 1, 1], dtype="float32")
|
||||
bn_scale = T.match_buffer(var_bn_scale, [512, 1, 1], dtype="float32")
|
||||
bn_offset = T.match_buffer(var_bn_offset, [512, 1, 1], dtype="float32")
|
||||
compute = T.match_buffer(var_compute, [1, 512, 56, 56], dtype="float32")
|
||||
pad_temp = T.sblock_alloc_buffer([1, 512, 58, 58], dtype="float32")
|
||||
compute_1 = T.sblock_alloc_buffer([1, 512, 56, 56], dtype="float32")
|
||||
compute_local = T.sblock_alloc_buffer([1, 512, 56, 56], dtype="float32", scope="local")
|
||||
pad_temp_shared = T.sblock_alloc_buffer([1, 512, 58, 58], dtype="float32", scope="shared")
|
||||
W_shared = T.sblock_alloc_buffer([512, 512, 3, 3], dtype="float32", scope="shared")
|
||||
for i0, i1, i2, i3 in T.grid(1, 512, 58, 58):
|
||||
with T.sblock("pad_temp"):
|
||||
i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
pad_temp[i0_1, i1_1, i2_1, i3_1] = T.if_then_else(i2_1 >= 1 and i2_1 < 57 and i3_1 >= 1 and i3_1 < 57, X[i0_1, i1_1, i2_1 - 1, i3_1 - 1], T.float32(0), dtype="float32")
|
||||
for i0_0_i1_0_i2_0_i3_0_fused in T.thread_binding(0, 224, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_i2_1_i3_1_fused in T.thread_binding(0, 2, thread="vthread.x"):
|
||||
for i0_2_i1_2_i2_2_i3_2_fused in T.thread_binding(0, 8, thread="threadIdx.x"):
|
||||
for i4_0, i5_0, i6_0 in T.grid(1, 3, 1):
|
||||
for ax0_ax1_ax2_ax3_fused_0 in T.serial(0, 40960, annotations={"meta_schedule.cooperative_fetch":1}):
|
||||
for ax0_ax1_ax2_ax3_fused_1 in T.vectorized(0, 3):
|
||||
with T.sblock("pad_temp_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(512, (ax0_ax1_ax2_ax3_fused_0 * 3 + ax0_ax1_ax2_ax3_fused_1) // 30 // 8 % 512)
|
||||
v2 = T.axis.spatial(58, i0_0_i1_0_i2_0_i3_0_fused % 14 // 2 * 8 + i5_0 + (ax0_ax1_ax2_ax3_fused_0 * 3 + ax0_ax1_ax2_ax3_fused_1) // 30 % 8)
|
||||
v3 = T.axis.spatial(58, i0_0_i1_0_i2_0_i3_0_fused % 2 * 28 + (ax0_ax1_ax2_ax3_fused_0 * 3 + ax0_ax1_ax2_ax3_fused_1) % 30)
|
||||
pad_temp_shared[v0, v1, v2, v3] = pad_temp[v0, v1, v2, v3]
|
||||
for ax0_ax1_ax2_ax3_fused_0 in T.serial(0, 12288, annotations={"meta_schedule.cooperative_fetch":1}):
|
||||
for ax0_ax1_ax2_ax3_fused_1 in T.vectorized(0, 4):
|
||||
with T.sblock("W_shared"):
|
||||
v0 = T.axis.spatial(512, i0_0_i1_0_i2_0_i3_0_fused // 14 * 32 + (ax0_ax1_ax2_ax3_fused_0 * 4 + ax0_ax1_ax2_ax3_fused_1) // 1536)
|
||||
v1 = T.axis.spatial(512, (ax0_ax1_ax2_ax3_fused_0 * 4 + ax0_ax1_ax2_ax3_fused_1) // 3 % 512)
|
||||
v2 = T.axis.spatial(3, i5_0)
|
||||
v3 = T.axis.spatial(3, (ax0_ax1_ax2_ax3_fused_0 * 4 + ax0_ax1_ax2_ax3_fused_1) % 3)
|
||||
W_shared[v0, v1, v2, v3] = W[v0, v1, v2, v3]
|
||||
for i4_1, i5_1, i6_1, i0_3, i1_3, i2_3, i3_3, i4_2, i5_2, i6_2, i0_4, i1_4, i2_4, i3_4 in T.grid(32, 1, 1, 1, 1, 1, 1, 16, 1, 3, 1, 8, 2, 28):
|
||||
with T.sblock("compute"):
|
||||
nn = T.axis.spatial(1, 0)
|
||||
ff = T.axis.spatial(512, i0_0_i1_0_i2_0_i3_0_fused // 14 * 32 + i0_2_i1_2_i2_2_i3_2_fused // 2 * 8 + i1_4)
|
||||
yy = T.axis.spatial(56, i0_0_i1_0_i2_0_i3_0_fused // 2 % 7 * 8 + i0_1_i1_1_i2_1_i3_1_fused * 4 + i0_2_i1_2_i2_2_i3_2_fused % 2 * 2 + i2_4)
|
||||
xx = T.axis.spatial(56, i0_0_i1_0_i2_0_i3_0_fused % 2 * 28 + i3_4)
|
||||
rc = T.axis.reduce(512, i4_1 * 16 + i4_2)
|
||||
ry, rx = T.axis.remap("RR", [i5_0, i6_2])
|
||||
with T.init():
|
||||
compute_local[nn, ff, yy, xx] = T.float32(0)
|
||||
compute_local[nn, ff, yy, xx] = compute_local[nn, ff, yy, xx] + pad_temp_shared[nn, rc, yy + ry, xx + rx] * W_shared[ff, rc, ry, rx]
|
||||
for ax0, ax1, ax2, ax3 in T.grid(1, 8, 2, 28):
|
||||
with T.sblock("compute_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(512, i0_0_i1_0_i2_0_i3_0_fused // 14 * 32 + i0_2_i1_2_i2_2_i3_2_fused // 2 * 8 + ax1)
|
||||
v2 = T.axis.spatial(56, i0_0_i1_0_i2_0_i3_0_fused % 14 // 2 * 8 + i0_1_i1_1_i2_1_i3_1_fused * 4 + i0_2_i1_2_i2_2_i3_2_fused % 2 * 2 + ax2)
|
||||
v3 = T.axis.spatial(56, i0_0_i1_0_i2_0_i3_0_fused % 2 * 28 + ax3)
|
||||
compute_1[v0, v1, v2, v3] = compute_local[v0, v1, v2, v3]
|
||||
for i0, i1, i2, i3 in T.grid(1, 512, 56, 56):
|
||||
with T.sblock("compute_1"):
|
||||
i0_2, i1_2, i2_2, i3_2 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
compute[i0_2, i1_2, i2_2, i3_2] = T.max((compute_1[i0_2, i1_2, i2_2, i3_2] + B[i1_2, 0, 0]) * bn_scale[i1_2, 0, 0] + bn_offset[i1_2, 0, 0], T.float32(0))
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class MultiLevelTiledConv2DAfterInline:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(X: T.Buffer((1, 512, 56, 56), "float32"), W: T.Buffer((512, 512, 3, 3), "float32"), B: T.Buffer((512, 1, 1), "float32"), bn_scale: T.Buffer((512, 1, 1), "float32"), bn_offset: T.Buffer((512, 1, 1), "float32"), compute: T.Buffer((1, 512, 56, 56), "float32")) -> None:
|
||||
compute_local = T.sblock_alloc_buffer([1, 512, 56, 56], dtype="float32", scope="local")
|
||||
for i0_0_i1_0_i2_0_i3_0_fused in T.thread_binding(224, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_i2_1_i3_1_fused in T.thread_binding(2, thread="vthread.x"):
|
||||
for i0_2_i1_2_i2_2_i3_2_fused in T.thread_binding(8, thread="threadIdx.x"):
|
||||
for i4_0, i5_0, i6_0, i4_1, i5_1, i6_1, i0_3, i1_3, i2_3, i3_3, i4_2, i5_2, i6_2, i0_4, i1_4, i2_4, i3_4 in T.grid(1, 3, 1, 32, 1, 1, 1, 1, 1, 1, 16, 1, 3, 1, 8, 2, 28):
|
||||
with T.sblock("compute"):
|
||||
nn = T.axis.spatial(1, 0)
|
||||
ff = T.axis.spatial(512, i0_0_i1_0_i2_0_i3_0_fused // 14 * 32 + i0_2_i1_2_i2_2_i3_2_fused // 2 * 8 + i1_4)
|
||||
yy = T.axis.spatial(56, i0_0_i1_0_i2_0_i3_0_fused // 2 % 7 * 8 + i0_1_i1_1_i2_1_i3_1_fused * 4 + i0_2_i1_2_i2_2_i3_2_fused % 2 * 2 + i2_4)
|
||||
xx = T.axis.spatial(56, i0_0_i1_0_i2_0_i3_0_fused % 2 * 28 + i3_4)
|
||||
rc = T.axis.reduce(512, i4_1 * 16 + i4_2)
|
||||
ry, rx = T.axis.remap("RR", [i5_0, i6_2])
|
||||
with T.init():
|
||||
compute_local[nn, ff, yy, xx] = T.float32(0)
|
||||
compute_local[nn, ff, yy, xx] = compute_local[nn, ff, yy, xx] + T.if_then_else(yy + ry >= 1 and yy + ry < 57 and xx + rx >= 1 and xx + rx < 57, X[nn, rc, yy + ry - 1, xx + rx - 1], T.float32(0), dtype="float32") * W[ff, rc, ry, rx]
|
||||
for ax0, ax1, ax2, ax3 in T.grid(1, 8, 2, 28):
|
||||
with T.sblock("compute_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(512, i0_0_i1_0_i2_0_i3_0_fused // 14 * 32 + i0_2_i1_2_i2_2_i3_2_fused // 2 * 8 + ax1)
|
||||
v2 = T.axis.spatial(56, i0_0_i1_0_i2_0_i3_0_fused % 14 // 2 * 8 + i0_1_i1_1_i2_1_i3_1_fused * 4 + i0_2_i1_2_i2_2_i3_2_fused % 2 * 2 + ax2)
|
||||
v3 = T.axis.spatial(56, i0_0_i1_0_i2_0_i3_0_fused % 2 * 28 + ax3)
|
||||
compute[v0, v1, v2, v3] = T.max((compute_local[v0, v1, v2, v3] + B[v1, 0, 0]) * bn_scale[v1, 0, 0] + bn_offset[v1, 0, 0], T.float32(0))
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class SoftmaxBeforeInline:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")) -> None:
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
T_softmax_exp = T.sblock_alloc_buffer([256, 256], dtype="float32")
|
||||
T_softmax_expsum = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
i0_1, k = T.axis.remap("SR", [i0, i1])
|
||||
with T.init():
|
||||
T_softmax_maxelem[i0_1] = T.min_value("float32")
|
||||
T_softmax_maxelem[i0_1] = T.max(T_softmax_maxelem[i0_1], A[i0_1, k])
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_exp"):
|
||||
i0_2, i1_1 = T.axis.remap("SS", [i0, i1])
|
||||
T_softmax_exp[i0_2, i1_1] = T.exp(A[i0_2, i1_1] - T_softmax_maxelem[i0_2], dtype="float32")
|
||||
for i0_3, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
i0_4, k = T.axis.remap("SR", [i0_3, i1])
|
||||
with T.init():
|
||||
T_softmax_expsum[i0_4] = T.float32(0)
|
||||
T_softmax_expsum[i0_4] = T_softmax_expsum[i0_4] + T_softmax_exp[i0_4, k]
|
||||
for i0_5, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
i0_6, i1_2 = T.axis.remap("SS", [i0_5, i1])
|
||||
T_softmax_norm[i0_6, i1_2] = T_softmax_exp[i0_6, i1_2] / T_softmax_expsum[i0_6]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class SoftmaxAfterInline:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")) -> None:
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
T_softmax_expsum = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
i0_1, k = T.axis.remap("SR", [i0, i1])
|
||||
with T.init():
|
||||
T_softmax_maxelem[i0_1] = T.min_value("float32")
|
||||
T_softmax_maxelem[i0_1] = T.max(T_softmax_maxelem[i0_1], A[i0_1, k])
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
i0_2, k = T.axis.remap("SR", [i0, i1])
|
||||
with T.init():
|
||||
T_softmax_expsum[i0_2] = T.float32(0)
|
||||
T_softmax_expsum[i0_2] = T_softmax_expsum[i0_2] + T.exp(A[i0_2, k] - T_softmax_maxelem[i0_2], dtype="float32")
|
||||
for i0_3, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
i0_4, i1_1 = T.axis.remap("SS", [i0_3, i1])
|
||||
T_softmax_norm[i0_4, i1_1] = T.exp(A[i0_4, i1_1] - T_softmax_maxelem[i0_4], dtype="float32") / T_softmax_expsum[i0_4]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class BeforePureSpatial:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
placeholder: T.Buffer((1, 384), "int64"),
|
||||
placeholder_1: T.Buffer((30522, 768), "float32"),
|
||||
placeholder_2: T.Buffer((1, 384, 768), "float32"),
|
||||
T_add: T.Buffer((1, 384, 768), "float32"),
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
compile_engine_const = T.sblock_alloc_buffer([], dtype="int64")
|
||||
T_less = T.sblock_alloc_buffer([1, 384], dtype="bool")
|
||||
compile_engine_const_1 = T.sblock_alloc_buffer([], dtype="int64")
|
||||
T_add_1 = T.sblock_alloc_buffer([1, 384], dtype="int64")
|
||||
T_where = T.sblock_alloc_buffer([1, 384], dtype="int64")
|
||||
T_take = T.sblock_alloc_buffer([1, 384, 768], dtype="float32")
|
||||
with T.sblock("compile_engine_const"):
|
||||
vi = T.axis.spatial(1, 0)
|
||||
T.reads()
|
||||
T.writes(compile_engine_const[()])
|
||||
compile_engine_const[()] = T.int64(0)
|
||||
for i0, i1 in T.grid(1, 384):
|
||||
with T.sblock("T_less"):
|
||||
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(placeholder[ax0, ax1], compile_engine_const[()])
|
||||
T.writes(T_less[ax0, ax1])
|
||||
T_less[ax0, ax1] = placeholder[ax0, ax1] < compile_engine_const[()]
|
||||
with T.sblock("compile_engine_const_1"):
|
||||
vi = T.axis.spatial(1, 0)
|
||||
T.reads()
|
||||
T.writes(compile_engine_const_1[()])
|
||||
compile_engine_const_1[()] = T.int64(30522)
|
||||
for i0, i1 in T.grid(1, 384):
|
||||
with T.sblock("T_add"):
|
||||
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(placeholder[ax0, ax1], compile_engine_const_1[()])
|
||||
T.writes(T_add_1[ax0, ax1])
|
||||
T_add_1[ax0, ax1] = placeholder[ax0, ax1] + compile_engine_const_1[()]
|
||||
for i0, i1 in T.grid(1, 384):
|
||||
with T.sblock("T_where"):
|
||||
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(T_less[ax0, ax1], T_add_1[ax0, ax1], placeholder[ax0, ax1])
|
||||
T.writes(T_where[ax0, ax1])
|
||||
T_where[ax0, ax1] = T.Select(
|
||||
T.cast(T_less[ax0, ax1], "int32") != 0, T_add_1[ax0, ax1], placeholder[ax0, ax1]
|
||||
)
|
||||
for i0, i1, i2 in T.grid(1, 384, 768):
|
||||
with T.sblock("T_take"):
|
||||
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(
|
||||
placeholder_1[T.min(T.max(T.int64(0), T_where[ax0, ax1]), T.int64(30521)), ax2],
|
||||
T_where[ax0, ax1],
|
||||
)
|
||||
T.writes(T_take[ax0, ax1, ax2])
|
||||
T_take[ax0, ax1, ax2] = placeholder_1[
|
||||
T.min(T.max(T.int64(0), T_where[ax0, ax1]), T.int64(30521)), ax2
|
||||
]
|
||||
for i0, i1, i2 in T.grid(1, 384, 768):
|
||||
with T.sblock("T_add_1"):
|
||||
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(T_take[ax0, ax1, ax2], placeholder_2[ax0, ax1, ax2])
|
||||
T.writes(T_add[ax0, ax1, ax2])
|
||||
T_add[ax0, ax1, ax2] = T_take[ax0, ax1, ax2] + placeholder_2[ax0, ax1, ax2]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class AfterPureSpatial:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(placeholder: T.Buffer((1, 384), "int64"), placeholder_1: T.Buffer((30522, 768), "float32"), placeholder_2: T.Buffer((1, 384, 768), "float32"), T_add: T.Buffer((1, 384, 768), "float32")) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
for i0, i1, i2 in T.grid(1, 384, 768):
|
||||
with T.sblock("T_add_1"):
|
||||
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(placeholder[ax0, ax1], placeholder_1[T.min(T.max(T.int64(0), placeholder[ax0, ax1]), T.int64(30521)) : T.min(T.max(T.int64(0), placeholder[ax0, ax1] + T.int64(30522)), T.int64(30521)) + T.int64(1), ax2], placeholder_2[ax0, ax1, ax2])
|
||||
T.writes(T_add[ax0, ax1, ax2])
|
||||
T_add[ax0, ax1, ax2] = placeholder_1[T.min(T.max(T.int64(0), T.Select(T.cast(placeholder[ax0, ax1] < T.int64(0), "int32") != 0, placeholder[ax0, ax1] + T.int64(30522), placeholder[ax0, ax1])), T.int64(30521)), ax2] + placeholder_2[ax0, ax1, ax2]
|
||||
|
||||
@tvm.script.ir_module
|
||||
class ConstConsumer:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(T_full: T.Buffer((1, 12, 4096), "int64")) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
for i0, i1, i2 in T.grid(1, 12, 4096):
|
||||
with T.sblock("T_full"):
|
||||
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads()
|
||||
T.writes(T_full[ax0, ax1, ax2])
|
||||
T_full[ax0, ax1, ax2] = T.int64(0)
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Conv2dInt8:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(p0: T.Buffer((16, 14, 14, 256), "int8"), p1: T.Buffer((1024, 1, 1, 256), "int8"), p2: T.Buffer((1, 1, 1, 1024), "int32"), p3: T.Buffer((1, 1, 1, 1024), "int32"), p4: T.Buffer(1024, "int32"), p5: T.Buffer(1024, "int32"), p6: T.Buffer(1024, "int32"), p7: T.Buffer(1, "int32"), p8: T.Buffer((16, 14, 14, 1024), "int32"), compute: T.Buffer((16, 14, 14, 1024), "int32")) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
compile_engine_const = T.sblock_alloc_buffer([], dtype="int32")
|
||||
pad_temp = T.sblock_alloc_buffer([16, 14, 14, 256], dtype="int8")
|
||||
conv2d_nhwc = T.sblock_alloc_buffer([16, 14, 14, 1024], dtype="int32")
|
||||
T_subtract = T.sblock_alloc_buffer([16, 14, 14, 1024], dtype="int32")
|
||||
T_add = T.sblock_alloc_buffer([16, 14, 14, 1024], dtype="int32")
|
||||
compute_1 = T.sblock_alloc_buffer([16, 14, 14, 1024], dtype="int32")
|
||||
T_add_1 = T.sblock_alloc_buffer([16, 14, 14, 1024], dtype="int32")
|
||||
compute_2 = T.sblock_alloc_buffer([16, 14, 14, 1024], dtype="int32")
|
||||
T_subtract_1 = T.sblock_alloc_buffer([16, 14, 14, 1024], dtype="int32")
|
||||
compute_3 = T.sblock_alloc_buffer([16, 14, 14, 1024], dtype="int32")
|
||||
T_add_2 = T.sblock_alloc_buffer([16, 14, 14, 1024], dtype="int32")
|
||||
with T.sblock("compile_engine_const"):
|
||||
vi = T.axis.spatial(1, 0)
|
||||
T.reads()
|
||||
T.writes(compile_engine_const[()])
|
||||
compile_engine_const[()] = 59
|
||||
for i0, i1, i2, i3 in T.grid(16, 14, 14, 256):
|
||||
with T.sblock("pad_temp"):
|
||||
i0_1, i1_1, i2_1, i3_1 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(p0[i0_1, i1_1, i2_1, i3_1])
|
||||
T.writes(pad_temp[i0_1, i1_1, i2_1, i3_1])
|
||||
pad_temp[i0_1, i1_1, i2_1, i3_1] = p0[i0_1, i1_1, i2_1, i3_1]
|
||||
for i0, i1, i2, i3, i4, i5, i6 in T.grid(16, 14, 14, 1024, 1, 1, 256):
|
||||
with T.sblock("conv2d_nhwc"):
|
||||
nn, yy, xx, ff, ry, rx, rc = T.axis.remap("SSSSRRR", [i0, i1, i2, i3, i4, i5, i6])
|
||||
T.reads(pad_temp[nn, yy + ry, xx + rx, rc], p1[ff, ry, rx, rc])
|
||||
T.writes(conv2d_nhwc[nn, yy, xx, ff])
|
||||
with T.init():
|
||||
conv2d_nhwc[nn, yy, xx, ff] = 0
|
||||
conv2d_nhwc[nn, yy, xx, ff] = conv2d_nhwc[nn, yy, xx, ff] + T.cast(pad_temp[nn, yy + ry, xx + rx, rc], "int32") * T.cast(p1[ff, ry, rx, rc], "int32")
|
||||
for i0, i1, i2, i3 in T.grid(16, 14, 14, 1024):
|
||||
with T.sblock("T_subtract"):
|
||||
ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(conv2d_nhwc[ax0, ax1, ax2, ax3], p2[0, 0, 0, ax3])
|
||||
T.writes(T_subtract[ax0, ax1, ax2, ax3])
|
||||
T_subtract[ax0, ax1, ax2, ax3] = conv2d_nhwc[ax0, ax1, ax2, ax3] - p2[0, 0, 0, ax3]
|
||||
for i0, i1, i2, i3 in T.grid(16, 14, 14, 1024):
|
||||
with T.sblock("T_add"):
|
||||
ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(T_subtract[ax0, ax1, ax2, ax3], p3[0, 0, 0, ax3])
|
||||
T.writes(T_add[ax0, ax1, ax2, ax3])
|
||||
T_add[ax0, ax1, ax2, ax3] = T_subtract[ax0, ax1, ax2, ax3] + p3[0, 0, 0, ax3]
|
||||
for i0, i1, i2, i3 in T.grid(16, 14, 14, 1024):
|
||||
with T.sblock("compute"):
|
||||
i0_2, i1_2, i2_2, i3_2 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(T_add[i0_2, i1_2, i2_2, i3_2], p4[i3_2], p5[i3_2], p6[i3_2])
|
||||
T.writes(compute_1[i0_2, i1_2, i2_2, i3_2])
|
||||
compute_1[i0_2, i1_2, i2_2, i3_2] = T.q_multiply_shift_per_axis(T_add[i0_2, i1_2, i2_2, i3_2], p4[i3_2], p5[i3_2], p6[i3_2], 31, False, True, dtype="int32")
|
||||
for i0_3, i1_3, i2_3, i3_3 in T.grid(16, 14, 14, 1024):
|
||||
with T.sblock("T_add_1"):
|
||||
ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0_3, i1_3, i2_3, i3_3])
|
||||
T.reads(compile_engine_const[()], compute_1[ax0, ax1, ax2, ax3])
|
||||
T.writes(T_add_1[ax0, ax1, ax2, ax3])
|
||||
T_add_1[ax0, ax1, ax2, ax3] = compile_engine_const[()] + compute_1[ax0, ax1, ax2, ax3]
|
||||
for i0_4, i1_4, i2_4, i3_4 in T.grid(16, 14, 14, 1024):
|
||||
with T.sblock("compute_1"):
|
||||
i0_5, i1_5, i2_5, i3_5 = T.axis.remap("SSSS", [i0_4, i1_4, i2_4, i3_4])
|
||||
T.reads(T_add_1[i0_5, i1_5, i2_5, i3_5])
|
||||
T.writes(compute_2[i0_5, i1_5, i2_5, i3_5])
|
||||
compute_2[i0_5, i1_5, i2_5, i3_5] = T.max(T.min(T_add_1[i0_5, i1_5, i2_5, i3_5], 255), 0)
|
||||
for i0_6, i1_6, i2_6, i3_6 in T.grid(16, 14, 14, 1024):
|
||||
with T.sblock("T_subtract_1"):
|
||||
ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0_6, i1_6, i2_6, i3_6])
|
||||
T.reads(compute_2[ax0, ax1, ax2, ax3], p7[0])
|
||||
T.writes(T_subtract_1[ax0, ax1, ax2, ax3])
|
||||
T_subtract_1[ax0, ax1, ax2, ax3] = compute_2[ax0, ax1, ax2, ax3] - p7[0]
|
||||
for i0_7, i1_7, i2_7, i3_7 in T.grid(16, 14, 14, 1024):
|
||||
with T.sblock("compute_2"):
|
||||
i0_8, i1_8, i2_8, i3_8 = T.axis.remap("SSSS", [i0_7, i1_7, i2_7, i3_7])
|
||||
T.reads(T_subtract_1[i0_8, i1_8, i2_8, i3_8])
|
||||
T.writes(compute_3[i0_8, i1_8, i2_8, i3_8])
|
||||
compute_3[i0_8, i1_8, i2_8, i3_8] = T.q_multiply_shift(T_subtract_1[i0_8, i1_8, i2_8, i3_8], 1408572815, 31, 1, dtype="int32")
|
||||
for i0_9, i1_9, i2_9, i3_9 in T.grid(16, 14, 14, 1024):
|
||||
with T.sblock("T_add_2"):
|
||||
ax0, ax1, ax2, ax3 = T.axis.remap("SSSS", [i0_9, i1_9, i2_9, i3_9])
|
||||
T.reads(compute_3[ax0, ax1, ax2, ax3], p8[ax0, ax1, ax2, ax3])
|
||||
T.writes(T_add_2[ax0, ax1, ax2, ax3])
|
||||
T_add_2[ax0, ax1, ax2, ax3] = compute_3[ax0, ax1, ax2, ax3] + p8[ax0, ax1, ax2, ax3]
|
||||
for i0_10, i1_10, i2_10, i3_10 in T.grid(16, 14, 14, 1024):
|
||||
with T.sblock("compute_3"):
|
||||
i0_11, i1_11, i2_11, i3_11 = T.axis.remap("SSSS", [i0_10, i1_10, i2_10, i3_10])
|
||||
T.reads(T_add_2[i0_11, i1_11, i2_11, i3_11])
|
||||
T.writes(compute[i0_11, i1_11, i2_11, i3_11])
|
||||
compute[i0_11, i1_11, i2_11, i3_11] = T.max(T.min(T_add_2[i0_11, i1_11, i2_11, i3_11], 255), 0)
|
||||
|
||||
|
||||
# pylint: enable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
# fmt: on
|
||||
|
||||
|
||||
def test_inline_consumer_chain():
|
||||
mod = Conv2DBiasBnReLU
|
||||
target = Target("llvm")
|
||||
(space,) = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=target,
|
||||
types=ms.schedule_rule.AutoInline,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(lhs=space.mod, rhs=Conv2DBiasBnReLUInlined)
|
||||
|
||||
|
||||
def test_inline_into_cache():
|
||||
mod = MultiLevelTiledConv2D
|
||||
target = Target("cuda", host="llvm")
|
||||
(space,) = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=target,
|
||||
types=ms.schedule_rule.AutoInline,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(lhs=space.mod, rhs=MultiLevelTiledConv2DAfterInline)
|
||||
|
||||
|
||||
def test_inline_into_multiple_consumers():
|
||||
mod = SoftmaxBeforeInline
|
||||
target = Target("cuda", host="llvm")
|
||||
(space,) = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=target,
|
||||
types=ms.schedule_rule.AutoInline,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(lhs=space.mod, rhs=SoftmaxAfterInline)
|
||||
|
||||
|
||||
def test_inline_pure_spatial():
|
||||
mod = BeforePureSpatial
|
||||
target = Target("llvm")
|
||||
(space,) = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=target,
|
||||
types=ms.schedule_rule.AutoInline,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(lhs=space.mod, rhs=AfterPureSpatial)
|
||||
|
||||
|
||||
def test_inline_constant_tensor():
|
||||
mod = ConstConsumer
|
||||
target = Target("cuda", host="llvm")
|
||||
(space,) = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=target,
|
||||
types=ms.schedule_rule.AutoInline,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(lhs=space.mod, rhs=ConstConsumer)
|
||||
|
||||
|
||||
def test_conv2d_int8_inline_constant_scalars():
|
||||
sch = Schedule(Conv2dInt8)
|
||||
|
||||
conv2d = sch.get_sblock("conv2d_nhwc")
|
||||
sch.cache_write(conv2d, 0, "shared")
|
||||
|
||||
with pytest.raises(tvm.s_tir.ScheduleError) as e:
|
||||
sch.reverse_compute_inline(sch.get_sblock("T_add_1"))
|
||||
|
||||
err_msg = "The block is only allowed to read a single buffer region, but it reads 2 region(s)"
|
||||
assert err_msg in str(e)
|
||||
|
||||
ms.schedule_rule.InlineConstantScalars().apply(sch, sch.get_sblock("compile_engine_const"))
|
||||
sch.reverse_compute_inline(sch.get_sblock("T_add_1"))
|
||||
|
||||
|
||||
def test_inline_constant_scalars_skip_output_block():
|
||||
# If the constant scalar block is an output block, it should not be inlined
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Full:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(T_full: T.Buffer((), "float32")):
|
||||
with T.sblock("T_full"):
|
||||
vi = T.axis.spatial(1, 0)
|
||||
T.reads()
|
||||
T.writes(T_full[()])
|
||||
T_full[()] = T.float32(1)
|
||||
|
||||
sch = Schedule(Full)
|
||||
sch = ms.schedule_rule.InlineConstantScalars().apply(sch, sch.get_sblock("T_full"))[0]
|
||||
assert_structural_equal(sch.mod, Full)
|
||||
|
||||
|
||||
def test_no_inline_root_block():
|
||||
@tvm.script.ir_module
|
||||
class MaxReduction:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
data: T.Buffer((8, 8), "float32"),
|
||||
data_red: T.Buffer((), "float32"),
|
||||
):
|
||||
T.func_attr({"tir.noalias": T.bool(True)})
|
||||
with T.sblock("data_red"):
|
||||
T.reads(data[0:8, 0:8])
|
||||
T.writes(data_red[()])
|
||||
with T.init():
|
||||
data_red[()] = T.float32(-3.4e38)
|
||||
for i, j in T.grid(8, 8):
|
||||
with T.sblock("update"):
|
||||
T.reads(data_red[()], data[i, j])
|
||||
T.writes(data_red[()])
|
||||
data_red[()] = T.max(data_red[()], data[i, j])
|
||||
|
||||
target = Target("llvm")
|
||||
(space,) = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=MaxReduction,
|
||||
target=target,
|
||||
types=ms.schedule_rule.AutoInline,
|
||||
)
|
||||
tvm.ir.assert_structural_equal(lhs=space.mod, rhs=MaxReduction)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
+777
@@ -0,0 +1,777 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
|
||||
import tvm
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing import te_workload
|
||||
from tvm.s_tir.meta_schedule.testing.space_generation import (
|
||||
check_sketches,
|
||||
generate_design_space,
|
||||
)
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
from tvm.te import create_prim_func
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Softmax_mn_after_inline:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(
|
||||
A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")
|
||||
) -> None:
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
T_softmax_expsum = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
i0_1, k = T.axis.remap("SR", [i0, i1])
|
||||
with T.init():
|
||||
T_softmax_maxelem[i0_1] = T.min_value("float32")
|
||||
T_softmax_maxelem[i0_1] = T.max(T_softmax_maxelem[i0_1], A[i0_1, k])
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
i0_2, k = T.axis.remap("SR", [i0, i1])
|
||||
with T.init():
|
||||
T_softmax_expsum[i0_2] = T.float32(0)
|
||||
T_softmax_expsum[i0_2] = T_softmax_expsum[i0_2] + T.exp(
|
||||
A[i0_2, k] - T_softmax_maxelem[i0_2], dtype="float32"
|
||||
)
|
||||
for i0_3, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
i0_4, i1_1 = T.axis.remap("SS", [i0_3, i1])
|
||||
T.sblock_attr({"axis": 1})
|
||||
T_softmax_norm[i0_4, i1_1] = (
|
||||
T.exp(A[i0_4, i1_1] - T_softmax_maxelem[i0_4], dtype="float32")
|
||||
/ T_softmax_expsum[i0_4]
|
||||
)
|
||||
|
||||
|
||||
def test_gpu_softmax_mn():
|
||||
@T.prim_func(s_tir=True)
|
||||
def softmax_mn_0(
|
||||
A: T.Buffer((256, 256), "float32"),
|
||||
T_softmax_norm: T.Buffer((256, 256), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
T_softmax_exp = T.sblock_alloc_buffer([256, 256], dtype="float32")
|
||||
T_softmax_expsum = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
i0_1, k = T.axis.remap("SR", [i0, i1])
|
||||
T.reads(A[i0_1, k])
|
||||
T.writes(T_softmax_maxelem[i0_1])
|
||||
with T.init():
|
||||
T_softmax_maxelem[i0_1] = T.float32(-3.4028234663852886e38)
|
||||
T_softmax_maxelem[i0_1] = T.max(T_softmax_maxelem[i0_1], A[i0_1, k])
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_exp"):
|
||||
i0_2, i1_1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(A[i0_2, i1_1], T_softmax_maxelem[i0_2])
|
||||
T.writes(T_softmax_exp[i0_2, i1_1])
|
||||
T_softmax_exp[i0_2, i1_1] = T.exp(
|
||||
A[i0_2, i1_1] - T_softmax_maxelem[i0_2], dtype="float32"
|
||||
)
|
||||
for i0_3, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
i0_4, k = T.axis.remap("SR", [i0_3, i1])
|
||||
T.reads(T_softmax_exp[i0_4, k])
|
||||
T.writes(T_softmax_expsum[i0_4])
|
||||
with T.init():
|
||||
T_softmax_expsum[i0_4] = T.float32(0)
|
||||
T_softmax_expsum[i0_4] = T_softmax_expsum[i0_4] + T_softmax_exp[i0_4, k]
|
||||
for i0_5, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
i0_6, i1_2 = T.axis.remap("SS", [i0_5, i1])
|
||||
T.reads(T_softmax_exp[i0_6, i1_2], T_softmax_expsum[i0_6])
|
||||
T.writes(T_softmax_norm[i0_6, i1_2])
|
||||
T.sblock_attr({"axis": 1})
|
||||
T_softmax_norm[i0_6, i1_2] = T_softmax_exp[i0_6, i1_2] / T_softmax_expsum[i0_6]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def softmax_mn_1(
|
||||
A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
T_softmax_maxelem_shared = T.sblock_alloc_buffer([256], dtype="float32", scope="shared")
|
||||
T_softmax_exp = T.sblock_alloc_buffer([256, 256], dtype="float32")
|
||||
T_softmax_expsum = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
for i0 in T.serial(256):
|
||||
for ax0, ax1_0 in T.grid(1, 1):
|
||||
for ax1_1 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
T.where(ax1_0 * 512 + ax1_1 < 256)
|
||||
i0_1 = T.axis.spatial(256, i0 + ax0)
|
||||
k = T.axis.reduce(256, ax1_0 * 512 + ax1_1)
|
||||
T.reads(A[i0_1, k])
|
||||
T.writes(T_softmax_maxelem_shared[i0_1])
|
||||
with T.init():
|
||||
T_softmax_maxelem_shared[i0_1] = T.float32(-3.4028234663852886e38)
|
||||
T_softmax_maxelem_shared[i0_1] = T.max(
|
||||
T_softmax_maxelem_shared[i0_1], A[i0_1, k]
|
||||
)
|
||||
for i1_0 in T.serial(1):
|
||||
for i1_1 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_exp"):
|
||||
T.where(i1_0 * 512 + i1_1 < 256)
|
||||
i0_2 = T.axis.spatial(256, i0)
|
||||
i1 = T.axis.spatial(256, i1_0 * 512 + i1_1)
|
||||
T.reads(A[i0_2, i1], T_softmax_maxelem_shared[i0_2])
|
||||
T.writes(T_softmax_exp[i0_2, i1])
|
||||
T_softmax_exp[i0_2, i1] = T.exp(
|
||||
A[i0_2, i1] - T_softmax_maxelem_shared[i0_2], dtype="float32"
|
||||
)
|
||||
for i0_3, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
i0_4, k = T.axis.remap("SR", [i0_3, i1])
|
||||
T.reads(T_softmax_exp[i0_4, k])
|
||||
T.writes(T_softmax_expsum[i0_4])
|
||||
with T.init():
|
||||
T_softmax_expsum[i0_4] = T.float32(0)
|
||||
T_softmax_expsum[i0_4] = T_softmax_expsum[i0_4] + T_softmax_exp[i0_4, k]
|
||||
for i0_5, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
i0_6, i1_2 = T.axis.remap("SS", [i0_5, i1])
|
||||
T.reads(T_softmax_exp[i0_6, i1_2], T_softmax_expsum[i0_6])
|
||||
T.writes(T_softmax_norm[i0_6, i1_2])
|
||||
T.sblock_attr({"axis": 1})
|
||||
T_softmax_norm[i0_6, i1_2] = T_softmax_exp[i0_6, i1_2] / T_softmax_expsum[i0_6]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def softmax_mn_2(
|
||||
A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
T_softmax_exp = T.sblock_alloc_buffer([256, 256], dtype="float32")
|
||||
T_softmax_expsum_shared = T.sblock_alloc_buffer([256], dtype="float32", scope="shared")
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
i0_1, k = T.axis.remap("SR", [i0, i1])
|
||||
T.reads(A[i0_1, k])
|
||||
T.writes(T_softmax_maxelem[i0_1])
|
||||
with T.init():
|
||||
T_softmax_maxelem[i0_1] = T.float32(-3.4028234663852886e38)
|
||||
T_softmax_maxelem[i0_1] = T.max(T_softmax_maxelem[i0_1], A[i0_1, k])
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_exp"):
|
||||
i0_2, i1_1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(A[i0_2, i1_1], T_softmax_maxelem[i0_2])
|
||||
T.writes(T_softmax_exp[i0_2, i1_1])
|
||||
T_softmax_exp[i0_2, i1_1] = T.exp(
|
||||
A[i0_2, i1_1] - T_softmax_maxelem[i0_2], dtype="float32"
|
||||
)
|
||||
for i0_3 in T.serial(256):
|
||||
for ax0, ax1_0 in T.grid(1, 32):
|
||||
for ax1_1 in T.thread_binding(8, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
i0_4 = T.axis.spatial(256, i0_3 + ax0)
|
||||
k = T.axis.reduce(256, ax1_0 * 8 + ax1_1)
|
||||
T.reads(T_softmax_exp[i0_4, k])
|
||||
T.writes(T_softmax_expsum_shared[i0_4])
|
||||
with T.init():
|
||||
T_softmax_expsum_shared[i0_4] = T.float32(0)
|
||||
T_softmax_expsum_shared[i0_4] = (
|
||||
T_softmax_expsum_shared[i0_4] + T_softmax_exp[i0_4, k]
|
||||
)
|
||||
for i1_0 in T.serial(32):
|
||||
for i1_1_1 in T.thread_binding(8, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
i0_5 = T.axis.spatial(256, i0_3)
|
||||
i1 = T.axis.spatial(256, i1_0 * 8 + i1_1_1)
|
||||
T.reads(T_softmax_exp[i0_5, i1], T_softmax_expsum_shared[i0_5])
|
||||
T.writes(T_softmax_norm[i0_5, i1])
|
||||
T.sblock_attr({"axis": 1})
|
||||
T_softmax_norm[i0_5, i1] = (
|
||||
T_softmax_exp[i0_5, i1] / T_softmax_expsum_shared[i0_5]
|
||||
)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def softmax_mn_3(
|
||||
A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
T_softmax_maxelem_shared = T.sblock_alloc_buffer([256], dtype="float32", scope="shared")
|
||||
T_softmax_exp = T.sblock_alloc_buffer([256, 256], dtype="float32")
|
||||
T_softmax_expsum_shared = T.sblock_alloc_buffer([256], dtype="float32", scope="shared")
|
||||
for i0 in T.serial(256):
|
||||
for ax0, ax1_0 in T.grid(1, 1):
|
||||
for ax1_1 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
T.where(ax1_0 * 512 + ax1_1 < 256)
|
||||
i0_1 = T.axis.spatial(256, i0 + ax0)
|
||||
k = T.axis.reduce(256, ax1_0 * 512 + ax1_1)
|
||||
T.reads(A[i0_1, k])
|
||||
T.writes(T_softmax_maxelem_shared[i0_1])
|
||||
with T.init():
|
||||
T_softmax_maxelem_shared[i0_1] = T.float32(-3.4028234663852886e38)
|
||||
T_softmax_maxelem_shared[i0_1] = T.max(
|
||||
T_softmax_maxelem_shared[i0_1], A[i0_1, k]
|
||||
)
|
||||
for i1_0 in T.serial(1):
|
||||
for i1_1 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_exp"):
|
||||
T.where(i1_0 * 512 + i1_1 < 256)
|
||||
i0_2 = T.axis.spatial(256, i0)
|
||||
i1 = T.axis.spatial(256, i1_0 * 512 + i1_1)
|
||||
T.reads(A[i0_2, i1], T_softmax_maxelem_shared[i0_2])
|
||||
T.writes(T_softmax_exp[i0_2, i1])
|
||||
T_softmax_exp[i0_2, i1] = T.exp(
|
||||
A[i0_2, i1] - T_softmax_maxelem_shared[i0_2], dtype="float32"
|
||||
)
|
||||
for i0_3 in T.serial(256):
|
||||
for ax0, ax1_0 in T.grid(1, 32):
|
||||
for ax1_1 in T.thread_binding(8, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
i0_4 = T.axis.spatial(256, i0_3 + ax0)
|
||||
k = T.axis.reduce(256, ax1_0 * 8 + ax1_1)
|
||||
T.reads(T_softmax_exp[i0_4, k])
|
||||
T.writes(T_softmax_expsum_shared[i0_4])
|
||||
with T.init():
|
||||
T_softmax_expsum_shared[i0_4] = T.float32(0)
|
||||
T_softmax_expsum_shared[i0_4] = (
|
||||
T_softmax_expsum_shared[i0_4] + T_softmax_exp[i0_4, k]
|
||||
)
|
||||
for i1_0 in T.serial(32):
|
||||
for i1_1 in T.thread_binding(8, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
i0_5 = T.axis.spatial(256, i0_3)
|
||||
i1 = T.axis.spatial(256, i1_0 * 8 + i1_1)
|
||||
T.reads(T_softmax_exp[i0_5, i1], T_softmax_expsum_shared[i0_5])
|
||||
T.writes(T_softmax_norm[i0_5, i1])
|
||||
T.sblock_attr({"axis": 1})
|
||||
T_softmax_norm[i0_5, i1] = (
|
||||
T_softmax_exp[i0_5, i1] / T_softmax_expsum_shared[i0_5]
|
||||
)
|
||||
|
||||
decision_0 = [] # type: ignore
|
||||
decision_1 = [
|
||||
("SampleCategorical", 7),
|
||||
]
|
||||
decision_2 = [
|
||||
("SampleCategorical", 1),
|
||||
]
|
||||
decision_3 = [
|
||||
("SampleCategorical", 1),
|
||||
("SampleCategorical", 7),
|
||||
]
|
||||
mod = create_prim_func(te_workload.softmax_mn(n=256, m=256))
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-3090", host="llvm"),
|
||||
types=ms.schedule_rule.CrossThreadReduction,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[softmax_mn_0, softmax_mn_1, softmax_mn_2, softmax_mn_3],
|
||||
expected_decisions=[decision_0, decision_1, decision_2, decision_3],
|
||||
)
|
||||
|
||||
|
||||
def test_gpu_softmax_mn_after_inline():
|
||||
@T.prim_func(s_tir=True)
|
||||
def softmax_mn_after_inline_0(
|
||||
A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")
|
||||
) -> None:
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
T_softmax_expsum = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
i0_1, k = T.axis.remap("SR", [i0, i1])
|
||||
T.reads(A[i0_1, k])
|
||||
T.writes(T_softmax_maxelem[i0_1])
|
||||
with T.init():
|
||||
T_softmax_maxelem[i0_1] = T.float32(-3.4028234663852886e38)
|
||||
T_softmax_maxelem[i0_1] = T.max(T_softmax_maxelem[i0_1], A[i0_1, k])
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
i0_2, k = T.axis.remap("SR", [i0, i1])
|
||||
T.reads(A[i0_2, k], T_softmax_maxelem[i0_2])
|
||||
T.writes(T_softmax_expsum[i0_2])
|
||||
with T.init():
|
||||
T_softmax_expsum[i0_2] = T.float32(0)
|
||||
T_softmax_expsum[i0_2] = T_softmax_expsum[i0_2] + T.exp(
|
||||
A[i0_2, k] - T_softmax_maxelem[i0_2], dtype="float32"
|
||||
)
|
||||
for i0_3, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
i0_4, i1_1 = T.axis.remap("SS", [i0_3, i1])
|
||||
T.reads(A[i0_4, i1_1], T_softmax_maxelem[i0_4], T_softmax_expsum[i0_4])
|
||||
T.writes(T_softmax_norm[i0_4, i1_1])
|
||||
T.sblock_attr({"axis": 1})
|
||||
T_softmax_norm[i0_4, i1_1] = (
|
||||
T.exp(A[i0_4, i1_1] - T_softmax_maxelem[i0_4], dtype="float32")
|
||||
/ T_softmax_expsum[i0_4]
|
||||
)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def softmax_mn_after_inline_1(
|
||||
A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")
|
||||
) -> None:
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
T_softmax_expsum = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
for i0, i1_0 in T.grid(256, 4):
|
||||
for i1_1 in T.thread_binding(64, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
i0_1 = T.axis.spatial(256, i0)
|
||||
k = T.axis.reduce(256, i1_0 * 64 + i1_1)
|
||||
T.reads(A[i0_1, k])
|
||||
T.writes(T_softmax_maxelem[i0_1])
|
||||
with T.init():
|
||||
T_softmax_maxelem[i0_1] = T.float32(-3.4028234663852886e38)
|
||||
T_softmax_maxelem[i0_1] = T.max(T_softmax_maxelem[i0_1], A[i0_1, k])
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
i0_2, k = T.axis.remap("SR", [i0, i1])
|
||||
T.reads(A[i0_2, k], T_softmax_maxelem[i0_2])
|
||||
T.writes(T_softmax_expsum[i0_2])
|
||||
with T.init():
|
||||
T_softmax_expsum[i0_2] = T.float32(0)
|
||||
T_softmax_expsum[i0_2] = T_softmax_expsum[i0_2] + T.exp(
|
||||
A[i0_2, k] - T_softmax_maxelem[i0_2], dtype="float32"
|
||||
)
|
||||
for i0_3, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
i0_4, i1_1 = T.axis.remap("SS", [i0_3, i1])
|
||||
T.reads(A[i0_4, i1_1], T_softmax_maxelem[i0_4], T_softmax_expsum[i0_4])
|
||||
T.writes(T_softmax_norm[i0_4, i1_1])
|
||||
T.sblock_attr({"axis": 1})
|
||||
T_softmax_norm[i0_4, i1_1] = (
|
||||
T.exp(A[i0_4, i1_1] - T_softmax_maxelem[i0_4], dtype="float32")
|
||||
/ T_softmax_expsum[i0_4]
|
||||
)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def softmax_mn_after_inline_2(
|
||||
A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")
|
||||
) -> None:
|
||||
T_softmax_maxelem = T.sblock_alloc_buffer([256], dtype="float32")
|
||||
T_softmax_expsum_shared = T.sblock_alloc_buffer([256], dtype="float32", scope="shared")
|
||||
for i0, i1 in T.grid(256, 256):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
i0_1, k = T.axis.remap("SR", [i0, i1])
|
||||
T.reads(A[i0_1, k])
|
||||
T.writes(T_softmax_maxelem[i0_1])
|
||||
with T.init():
|
||||
T_softmax_maxelem[i0_1] = T.float32(-3.4028234663852886e38)
|
||||
T_softmax_maxelem[i0_1] = T.max(T_softmax_maxelem[i0_1], A[i0_1, k])
|
||||
for i0_3 in T.serial(256):
|
||||
for ax0, ax1_0 in T.grid(1, 1):
|
||||
for ax1_1 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
T.where(ax1_0 * 512 + ax1_1 < 256)
|
||||
i0_2 = T.axis.spatial(256, i0_3 + ax0)
|
||||
k = T.axis.reduce(256, ax1_0 * 512 + ax1_1)
|
||||
T.reads(A[i0_2, k], T_softmax_maxelem[i0_2])
|
||||
T.writes(T_softmax_expsum_shared[i0_2])
|
||||
with T.init():
|
||||
T_softmax_expsum_shared[i0_2] = T.float32(0)
|
||||
T_softmax_expsum_shared[i0_2] = T_softmax_expsum_shared[i0_2] + T.exp(
|
||||
A[i0_2, k] - T_softmax_maxelem[i0_2], dtype="float32"
|
||||
)
|
||||
for i1_0 in T.serial(1):
|
||||
for i1_1 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
T.where(i1_0 * 512 + i1_1 < 256)
|
||||
i0_4 = T.axis.spatial(256, i0_3)
|
||||
i1_1_1 = T.axis.spatial(256, i1_0 * 512 + i1_1)
|
||||
T.reads(
|
||||
A[i0_4, i1_1_1], T_softmax_maxelem[i0_4], T_softmax_expsum_shared[i0_4]
|
||||
)
|
||||
T.writes(T_softmax_norm[i0_4, i1_1_1])
|
||||
T.sblock_attr({"axis": 1})
|
||||
T_softmax_norm[i0_4, i1_1_1] = (
|
||||
T.exp(A[i0_4, i1_1_1] - T_softmax_maxelem[i0_4], dtype="float32")
|
||||
/ T_softmax_expsum_shared[i0_4]
|
||||
)
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def softmax_mn_after_inline_3(
|
||||
A: T.Buffer((256, 256), "float32"), T_softmax_norm: T.Buffer((256, 256), "float32")
|
||||
) -> None:
|
||||
T_softmax_maxelem_shared = T.sblock_alloc_buffer([256], dtype="float32", scope="shared")
|
||||
T_softmax_expsum_shared = T.sblock_alloc_buffer([256], dtype="float32", scope="shared")
|
||||
for i0_3 in T.serial(256):
|
||||
for ax0, ax1_0 in T.grid(1, 1):
|
||||
for ax1_1 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_maxelem"):
|
||||
T.where(ax1_0 * 512 + ax1_1 < 256)
|
||||
i0_1 = T.axis.spatial(256, i0_3 + ax0)
|
||||
k = T.axis.reduce(256, ax1_0 * 512 + ax1_1)
|
||||
T.reads(A[i0_1, k])
|
||||
T.writes(T_softmax_maxelem_shared[i0_1])
|
||||
with T.init():
|
||||
T_softmax_maxelem_shared[i0_1] = T.float32(-3.4028234663852886e38)
|
||||
T_softmax_maxelem_shared[i0_1] = T.max(
|
||||
T_softmax_maxelem_shared[i0_1], A[i0_1, k]
|
||||
)
|
||||
for ax0, ax1_0 in T.grid(1, 1):
|
||||
for ax1_1 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_expsum"):
|
||||
T.where(ax1_0 * 512 + ax1_1 < 256)
|
||||
i0_2 = T.axis.spatial(256, i0_3 + ax0)
|
||||
k = T.axis.reduce(256, ax1_0 * 512 + ax1_1)
|
||||
T.reads(A[i0_2, k], T_softmax_maxelem_shared[i0_2])
|
||||
T.writes(T_softmax_expsum_shared[i0_2])
|
||||
with T.init():
|
||||
T_softmax_expsum_shared[i0_2] = T.float32(0)
|
||||
T_softmax_expsum_shared[i0_2] = T_softmax_expsum_shared[i0_2] + T.exp(
|
||||
A[i0_2, k] - T_softmax_maxelem_shared[i0_2], dtype="float32"
|
||||
)
|
||||
for i1_0 in T.serial(1):
|
||||
for i1_1 in T.thread_binding(512, thread="threadIdx.x"):
|
||||
with T.sblock("T_softmax_norm"):
|
||||
T.where(i1_0 * 512 + i1_1 < 256)
|
||||
i0_4 = T.axis.spatial(256, i0_3)
|
||||
i1_1_1 = T.axis.spatial(256, i1_0 * 512 + i1_1)
|
||||
T.reads(
|
||||
A[i0_4, i1_1_1],
|
||||
T_softmax_maxelem_shared[i0_4],
|
||||
T_softmax_expsum_shared[i0_4],
|
||||
)
|
||||
T.writes(T_softmax_norm[i0_4, i1_1_1])
|
||||
T.sblock_attr({"axis": 1})
|
||||
T_softmax_norm[i0_4, i1_1_1] = (
|
||||
T.exp(A[i0_4, i1_1_1] - T_softmax_maxelem_shared[i0_4], dtype="float32")
|
||||
/ T_softmax_expsum_shared[i0_4]
|
||||
)
|
||||
|
||||
decision_0 = [] # type: ignore
|
||||
decision_1 = [
|
||||
("SampleCategorical", 4),
|
||||
]
|
||||
decision_2 = [
|
||||
("SampleCategorical", 7),
|
||||
]
|
||||
decision_3 = [
|
||||
("SampleCategorical", 7),
|
||||
("SampleCategorical", 0),
|
||||
]
|
||||
|
||||
mod = Softmax_mn_after_inline
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-3090", host="llvm"),
|
||||
types=ms.schedule_rule.CrossThreadReduction,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[
|
||||
softmax_mn_after_inline_0,
|
||||
softmax_mn_after_inline_1,
|
||||
softmax_mn_after_inline_2,
|
||||
softmax_mn_after_inline_3,
|
||||
],
|
||||
expected_decisions=[decision_0, decision_1, decision_2, decision_3],
|
||||
)
|
||||
|
||||
|
||||
def test_gpu_batch_norm_bmn():
|
||||
@T.prim_func(s_tir=True)
|
||||
def batch_norm_bmn_0(A: T.Buffer((1, 512, 512), "float32"), D: T.Buffer(1, "float32")) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C = T.sblock_alloc_buffer([1], dtype="float32")
|
||||
for i0, i1, i2 in T.grid(1, 512, 512):
|
||||
with T.sblock("C"):
|
||||
b, i, j = T.axis.remap("SRR", [i0, i1, i2])
|
||||
T.reads(A[b, i, j])
|
||||
T.writes(C[b])
|
||||
with T.init():
|
||||
C[b] = T.float32(0)
|
||||
C[b] = C[b] + A[b, i, j] * A[b, i, j]
|
||||
for i0 in T.serial(1):
|
||||
with T.sblock("D"):
|
||||
b = T.axis.spatial(1, i0)
|
||||
T.reads(C[b])
|
||||
T.writes(D[b])
|
||||
D[b] = T.sqrt(C[b], dtype="float32")
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def batch_norm_bmn_1(A: T.Buffer((1, 512, 512), "float32"), D: T.Buffer(1, "float32")) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C_shared = T.sblock_alloc_buffer([1], dtype="float32", scope="shared")
|
||||
for i0_0 in T.serial(1):
|
||||
for ax0, ax1_ax2_fused_0 in T.grid(1, 1024):
|
||||
for ax1_ax2_fused_1 in T.thread_binding(256, thread="threadIdx.x"):
|
||||
with T.sblock("C"):
|
||||
b = T.axis.spatial(1, ax0)
|
||||
i = T.axis.reduce(512, (ax1_ax2_fused_0 * 256 + ax1_ax2_fused_1) // 512)
|
||||
j = T.axis.reduce(512, (ax1_ax2_fused_0 * 256 + ax1_ax2_fused_1) % 512)
|
||||
T.reads(A[b, i, j])
|
||||
T.writes(C_shared[b])
|
||||
with T.init():
|
||||
C_shared[b] = T.float32(0)
|
||||
C_shared[b] = C_shared[b] + A[b, i, j] * A[b, i, j]
|
||||
for i0_1 in T.thread_binding(256, thread="threadIdx.x"):
|
||||
with T.sblock("D"):
|
||||
T.where(i0_0 * 256 + i0_1 < 1)
|
||||
b = T.axis.spatial(1, i0_0 * 256 + i0_1)
|
||||
T.reads(C_shared[b])
|
||||
T.writes(D[b])
|
||||
D[b] = T.sqrt(C_shared[b], dtype="float32")
|
||||
|
||||
decision_0 = [] # type: ignore
|
||||
decision_1 = [
|
||||
("SampleCategorical", 6),
|
||||
]
|
||||
|
||||
mod = create_prim_func(te_workload.norm_bmn(B=1, M=512, N=512))
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-3090", host="llvm"),
|
||||
types=ms.schedule_rule.CrossThreadReduction,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[batch_norm_bmn_0, batch_norm_bmn_1],
|
||||
expected_decisions=[decision_0, decision_1],
|
||||
)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def argmax(
|
||||
idx: T.Buffer((128, 128), "int32"),
|
||||
val: T.Buffer((128, 128), "float32"),
|
||||
argmax_v0: T.Buffer((128,), "int32"),
|
||||
argmax_v1: T.Buffer((128,), "float32"),
|
||||
) -> None:
|
||||
for i0, i1 in T.grid(128, 128):
|
||||
with T.sblock("argmax"):
|
||||
i = T.axis.spatial(128, i0)
|
||||
k = T.axis.reduce(128, i1)
|
||||
T.reads(idx[i, k], val[i, k])
|
||||
T.writes(argmax_v0[i], argmax_v1[i])
|
||||
with T.init():
|
||||
argmax_v0[i] = -1
|
||||
argmax_v1[i] = T.min_value("float32")
|
||||
v_argmax_v0: T.let[T.int32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v0[i], idx[i, k]
|
||||
)
|
||||
v_argmax_v1: T.let[T.float32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v1[i], val[i, k]
|
||||
)
|
||||
argmax_v0[i] = v_argmax_v0
|
||||
argmax_v1[i] = v_argmax_v1
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def argmax_32(
|
||||
idx: T.Buffer((1, 32), "int32"),
|
||||
val: T.Buffer((1, 32), "float32"),
|
||||
argmax_v0: T.Buffer((1,), "int32"),
|
||||
argmax_v1: T.Buffer((1,), "float32"),
|
||||
) -> None:
|
||||
for i0, i1 in T.grid(1, 32):
|
||||
with T.sblock("argmax"):
|
||||
i = T.axis.spatial(1, i0)
|
||||
k = T.axis.reduce(32, i1)
|
||||
T.reads(idx[i, k], val[i, k])
|
||||
T.writes(argmax_v0[i], argmax_v1[i])
|
||||
with T.init():
|
||||
argmax_v0[i] = -1
|
||||
argmax_v1[i] = T.min_value("float32")
|
||||
v_argmax_v0: T.let[T.int32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v0[i], idx[i, k]
|
||||
)
|
||||
v_argmax_v1: T.let[T.float32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v1[i], val[i, k]
|
||||
)
|
||||
argmax_v0[i] = v_argmax_v0
|
||||
argmax_v1[i] = v_argmax_v1
|
||||
|
||||
|
||||
def test_gpu_argmax():
|
||||
@T.prim_func(s_tir=True)
|
||||
def argmax_0(
|
||||
idx: T.Buffer((128, 128), "int32"),
|
||||
val: T.Buffer((128, 128), "float32"),
|
||||
argmax_v0: T.Buffer(128, "int32"),
|
||||
argmax_v1: T.Buffer(128, "float32"),
|
||||
) -> None:
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
for i0, i1 in T.grid(128, 128):
|
||||
with T.sblock("argmax"):
|
||||
i, k = T.axis.remap("SR", [i0, i1])
|
||||
T.reads(idx[i, k], val[i, k])
|
||||
T.writes(argmax_v0[i], argmax_v1[i])
|
||||
with T.init():
|
||||
argmax_v0[i] = -1
|
||||
argmax_v1[i] = T.float32(-3.4028234663852886e38)
|
||||
v_argmax_v0: T.let[T.int32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v0[i], idx[i, k]
|
||||
)
|
||||
v_argmax_v1: T.let[T.float32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v1[i], val[i, k]
|
||||
)
|
||||
argmax_v0[i] = v_argmax_v0
|
||||
argmax_v1[i] = v_argmax_v1
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def argmax_1(
|
||||
idx: T.Buffer((128, 128), "int32"),
|
||||
val: T.Buffer((128, 128), "float32"),
|
||||
argmax_v0: T.Buffer(128, "int32"),
|
||||
argmax_v1: T.Buffer(128, "float32"),
|
||||
) -> None:
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
for i0, i1_0 in T.grid(128, 2):
|
||||
for i1_1 in T.thread_binding(64, thread="threadIdx.x"):
|
||||
with T.sblock("argmax"):
|
||||
i = T.axis.spatial(128, i0)
|
||||
k = T.axis.reduce(128, i1_0 * 64 + i1_1)
|
||||
T.reads(idx[i, k], val[i, k])
|
||||
T.writes(argmax_v0[i], argmax_v1[i])
|
||||
with T.init():
|
||||
argmax_v0[i] = -1
|
||||
argmax_v1[i] = T.float32(-3.4028234663852886e38)
|
||||
v_argmax_v0: T.let[T.int32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v0[i], idx[i, k]
|
||||
)
|
||||
v_argmax_v1: T.let[T.float32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v1[i], val[i, k]
|
||||
)
|
||||
argmax_v0[i] = v_argmax_v0
|
||||
argmax_v1[i] = v_argmax_v1
|
||||
|
||||
decision_0 = [] # type: ignore
|
||||
decision_1 = [
|
||||
("SampleCategorical", 4),
|
||||
]
|
||||
|
||||
mod = argmax
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-3090", host="llvm"),
|
||||
types=ms.schedule_rule.CrossThreadReduction,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[argmax_0, argmax_1],
|
||||
expected_decisions=[decision_0, decision_1],
|
||||
)
|
||||
|
||||
|
||||
def test_gpu_argmax_32():
|
||||
@T.prim_func(s_tir=True)
|
||||
def argmax_0(
|
||||
idx: T.Buffer((1, 32), "int32"),
|
||||
val: T.Buffer((1, 32), "float32"),
|
||||
argmax_v0: T.Buffer((1,), "int32"),
|
||||
argmax_v1: T.Buffer((1,), "float32"),
|
||||
) -> None:
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
for i0, i1 in T.grid(1, 32):
|
||||
with T.sblock("argmax"):
|
||||
i, k = T.axis.remap("SR", [i0, i1])
|
||||
T.reads(idx[i, k], val[i, k])
|
||||
T.writes(argmax_v0[i], argmax_v1[i])
|
||||
with T.init():
|
||||
argmax_v0[i] = -1
|
||||
argmax_v1[i] = T.float32(-3.4028234663852886e38)
|
||||
v_argmax_v0: T.let[T.int32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v0[i], idx[i, k]
|
||||
)
|
||||
v_argmax_v1: T.let[T.float32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v1[i], val[i, k]
|
||||
)
|
||||
argmax_v0[i] = v_argmax_v0
|
||||
argmax_v1[i] = v_argmax_v1
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def argmax_1(
|
||||
idx: T.Buffer((1, 32), "int32"),
|
||||
val: T.Buffer((1, 32), "float32"),
|
||||
argmax_v0: T.Buffer((1,), "int32"),
|
||||
argmax_v1: T.Buffer((1,), "float32"),
|
||||
) -> None:
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
for i0, i1_0 in T.grid(1, 1):
|
||||
for i1_1 in T.thread_binding(64, thread="threadIdx.x"):
|
||||
with T.sblock("argmax"):
|
||||
i = T.axis.spatial(1, i0)
|
||||
k = T.axis.reduce(32, i1_0 * 64 + i1_1)
|
||||
T.where(i1_0 * 64 + i1_1 < 32)
|
||||
T.reads(idx[i, k], val[i, k])
|
||||
T.writes(argmax_v0[i], argmax_v1[i])
|
||||
with T.init():
|
||||
argmax_v0[i] = -1
|
||||
argmax_v1[i] = T.float32(-3.4028234663852886e38)
|
||||
v_argmax_v0: T.let[T.int32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v0[i], idx[i, k]
|
||||
)
|
||||
v_argmax_v1: T.let[T.float32] = T.Select(
|
||||
argmax_v1[i] >= val[i, k], argmax_v1[i], val[i, k]
|
||||
)
|
||||
argmax_v0[i] = v_argmax_v0
|
||||
argmax_v1[i] = v_argmax_v1
|
||||
|
||||
decision_0 = [] # type: ignore
|
||||
decision_1 = [
|
||||
("SampleCategorical", 4),
|
||||
]
|
||||
|
||||
mod = argmax_32
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-3090", host="llvm"),
|
||||
types=ms.schedule_rule.CrossThreadReduction,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[argmax_0, argmax_1],
|
||||
expected_decisions=[decision_0, decision_1],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_gpu_softmax_mn()
|
||||
test_gpu_softmax_mn_after_inline()
|
||||
test_gpu_batch_norm_bmn()
|
||||
test_gpu_argmax()
|
||||
test_gpu_argmax_32()
|
||||
@@ -0,0 +1,889 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: E741, F401
|
||||
import tvm.testing
|
||||
from tvm import te
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing import te_workload
|
||||
from tvm.s_tir.meta_schedule.testing.space_generation import (
|
||||
check_sketches,
|
||||
generate_design_space,
|
||||
print_sketches,
|
||||
)
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def test_cpu_matmul():
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_matmul_0(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
C: T.Buffer((512, 512), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C_global = T.sblock_alloc_buffer([512, 512], dtype="float32")
|
||||
for i0_0, i1_0, i0_1, i1_1 in T.grid(1, 8, 8, 1):
|
||||
for i2_0, i0_2, i1_2, i2_1, i0_3, i1_3 in T.grid(16, 2, 8, 32, 32, 8):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(512, i0_0 * 512 + i0_1 * 64 + i0_2 * 32 + i0_3)
|
||||
j = T.axis.spatial(512, i1_0 * 64 + i1_1 * 64 + i1_2 * 8 + i1_3)
|
||||
k = T.axis.reduce(512, i2_0 * 32 + i2_1)
|
||||
T.reads(A[i, k], B[k, j])
|
||||
T.writes(C_global[i, j])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SSRSRS"})
|
||||
with T.init():
|
||||
C_global[i, j] = T.float32(0)
|
||||
C_global[i, j] = C_global[i, j] + A[i, k] * B[k, j]
|
||||
for ax0, ax1 in T.grid(64, 64):
|
||||
with T.sblock("C_global"):
|
||||
v0 = T.axis.spatial(512, i0_1 * 64 + ax0)
|
||||
v1 = T.axis.spatial(512, i1_0 * 64 + ax1)
|
||||
T.reads(C_global[v0, v1])
|
||||
T.writes(C[v0, v1])
|
||||
C[v0, v1] = C_global[v0, v1]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_matmul_1(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
C: T.Buffer((512, 512), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C_global = T.sblock_alloc_buffer([512, 512], dtype="float32")
|
||||
for i0_0, i1_0 in T.grid(1, 8):
|
||||
for i0_1, i1_1, i2_0, i0_2, i1_2, i2_1, i0_3, i1_3 in T.grid(8, 1, 16, 2, 8, 32, 32, 8):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(512, i0_0 * 512 + i0_1 * 64 + i0_2 * 32 + i0_3)
|
||||
j = T.axis.spatial(512, i1_0 * 64 + i1_1 * 64 + i1_2 * 8 + i1_3)
|
||||
k = T.axis.reduce(512, i2_0 * 32 + i2_1)
|
||||
T.reads(A[i, k], B[k, j])
|
||||
T.writes(C_global[i, j])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SSRSRS"})
|
||||
with T.init():
|
||||
C_global[i, j] = T.float32(0)
|
||||
C_global[i, j] = C_global[i, j] + A[i, k] * B[k, j]
|
||||
for ax0, ax1 in T.grid(512, 64):
|
||||
with T.sblock("C_global"):
|
||||
v0 = T.axis.spatial(512, ax0)
|
||||
v1 = T.axis.spatial(512, i1_0 * 64 + ax1)
|
||||
T.reads(C_global[v0, v1])
|
||||
T.writes(C[v0, v1])
|
||||
C[v0, v1] = C_global[v0, v1]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_matmul_2(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
C: T.Buffer((512, 512), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
for i0_0, i1_0, i0_1, i1_1, i2_0, i0_2, i1_2, i2_1, i0_3, i1_3 in T.grid(
|
||||
1, 8, 8, 1, 16, 2, 8, 32, 32, 8
|
||||
):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(512, i0_0 * 512 + i0_1 * 64 + i0_2 * 32 + i0_3)
|
||||
j = T.axis.spatial(512, i1_0 * 64 + i1_1 * 64 + i1_2 * 8 + i1_3)
|
||||
k = T.axis.reduce(512, i2_0 * 32 + i2_1)
|
||||
T.reads(A[i, k], B[k, j])
|
||||
T.writes(C[i, j])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SSRSRS"})
|
||||
with T.init():
|
||||
C[i, j] = T.float32(0)
|
||||
C[i, j] = C[i, j] + A[i, k] * B[k, j]
|
||||
|
||||
decision_0 = [
|
||||
("SamplePerfectTile", [1, 8, 2, 32]),
|
||||
("SamplePerfectTile", [8, 1, 8, 8]),
|
||||
("SamplePerfectTile", [16, 32]),
|
||||
]
|
||||
decision_1 = [
|
||||
("SamplePerfectTile", [1, 8, 2, 32]),
|
||||
("SamplePerfectTile", [8, 1, 8, 8]),
|
||||
("SamplePerfectTile", [16, 32]),
|
||||
]
|
||||
decision_2 = [
|
||||
("SamplePerfectTile", [1, 8, 2, 32]),
|
||||
("SamplePerfectTile", [8, 1, 8, 8]),
|
||||
("SamplePerfectTile", [16, 32]),
|
||||
]
|
||||
|
||||
mod = te.create_prim_func(te_workload.matmul(512, 512, 512))
|
||||
actual = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=Target("llvm"),
|
||||
types=ms.schedule_rule.MultiLevelTiling,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[cpu_matmul_0, cpu_matmul_1, cpu_matmul_2],
|
||||
expected_decisions=[decision_0, decision_1, decision_2],
|
||||
)
|
||||
|
||||
|
||||
def test_cpu_matmul_relu():
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_matmul_relu_0(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
compute: T.Buffer((512, 512), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C = T.sblock_alloc_buffer([512, 512], dtype="float32")
|
||||
for i0_0, i1_0, i0_1, i1_1, i2_0, i0_2, i1_2, i2_1, i0_3, i1_3 in T.grid(
|
||||
256, 4, 1, 4, 64, 1, 32, 8, 2, 1
|
||||
):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(512, i0_0 * 2 + i0_1 * 2 + i0_2 * 2 + i0_3)
|
||||
j = T.axis.spatial(512, i1_0 * 128 + i1_1 * 32 + i1_2 + i1_3)
|
||||
k = T.axis.reduce(512, i2_0 * 8 + i2_1)
|
||||
T.reads(A[i, k], B[k, j])
|
||||
T.writes(C[i, j])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SSRSRS"})
|
||||
with T.init():
|
||||
C[i, j] = T.float32(0)
|
||||
C[i, j] = C[i, j] + A[i, k] * B[k, j]
|
||||
for i0, i1 in T.grid(512, 512):
|
||||
with T.sblock("compute"):
|
||||
i0_4, i1_4 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(C[i0_4, i1_4])
|
||||
T.writes(compute[i0_4, i1_4])
|
||||
compute[i0_4, i1_4] = T.max(C[i0_4, i1_4], T.float32(0))
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_matmul_relu_1(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
compute: T.Buffer((512, 512), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C = T.sblock_alloc_buffer([512, 512], dtype="float32")
|
||||
for i0_0, i1_0, i0_1, i1_1 in T.grid(256, 4, 1, 4):
|
||||
for i2_0, i0_2, i1_2, i2_1, i0_3, i1_3 in T.grid(64, 1, 32, 8, 2, 1):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(512, i0_0 * 2 + i0_1 * 2 + i0_2 * 2 + i0_3)
|
||||
j = T.axis.spatial(512, i1_0 * 128 + i1_1 * 32 + i1_2 + i1_3)
|
||||
k = T.axis.reduce(512, i2_0 * 8 + i2_1)
|
||||
T.reads(A[i, k], B[k, j])
|
||||
T.writes(C[i, j])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SSRSRS"})
|
||||
with T.init():
|
||||
C[i, j] = T.float32(0)
|
||||
C[i, j] = C[i, j] + A[i, k] * B[k, j]
|
||||
for ax0, ax1 in T.grid(2, 32):
|
||||
with T.sblock("compute"):
|
||||
i0 = T.axis.spatial(512, i0_0 * 2 + ax0)
|
||||
i1 = T.axis.spatial(512, i1_0 * 128 + i1_1 * 32 + ax1)
|
||||
T.reads(C[i0, i1])
|
||||
T.writes(compute[i0, i1])
|
||||
compute[i0, i1] = T.max(C[i0, i1], T.float32(0))
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_matmul_relu_2(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
compute: T.Buffer((512, 512), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C = T.sblock_alloc_buffer([512, 512], dtype="float32")
|
||||
for i0_0, i1_0 in T.grid(256, 4):
|
||||
for i0_1, i1_1, i2_0, i0_2, i1_2, i2_1, i0_3, i1_3 in T.grid(1, 4, 64, 1, 32, 8, 2, 1):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(512, i0_0 * 2 + i0_1 * 2 + i0_2 * 2 + i0_3)
|
||||
j = T.axis.spatial(512, i1_0 * 128 + i1_1 * 32 + i1_2 + i1_3)
|
||||
k = T.axis.reduce(512, i2_0 * 8 + i2_1)
|
||||
T.reads(A[i, k], B[k, j])
|
||||
T.writes(C[i, j])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SSRSRS"})
|
||||
with T.init():
|
||||
C[i, j] = T.float32(0)
|
||||
C[i, j] = C[i, j] + A[i, k] * B[k, j]
|
||||
for ax0, ax1 in T.grid(2, 128):
|
||||
with T.sblock("compute"):
|
||||
i0 = T.axis.spatial(512, i0_0 * 2 + ax0)
|
||||
i1 = T.axis.spatial(512, i1_0 * 128 + ax1)
|
||||
T.reads(C[i0, i1])
|
||||
T.writes(compute[i0, i1])
|
||||
compute[i0, i1] = T.max(C[i0, i1], T.float32(0))
|
||||
|
||||
decision_0 = [
|
||||
("SamplePerfectTile", [256, 1, 1, 2]),
|
||||
("SamplePerfectTile", [4, 4, 32, 1]),
|
||||
("SamplePerfectTile", [64, 8]),
|
||||
]
|
||||
decision_1 = [
|
||||
("SamplePerfectTile", [256, 1, 1, 2]),
|
||||
("SamplePerfectTile", [4, 4, 32, 1]),
|
||||
("SamplePerfectTile", [64, 8]),
|
||||
]
|
||||
decision_2 = [
|
||||
("SamplePerfectTile", [256, 1, 1, 2]),
|
||||
("SamplePerfectTile", [4, 4, 32, 1]),
|
||||
("SamplePerfectTile", [64, 8]),
|
||||
]
|
||||
mod = te.create_prim_func(te_workload.matmul_relu(512, 512, 512))
|
||||
actual = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=Target("llvm"),
|
||||
types=ms.schedule_rule.MultiLevelTiling,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[cpu_matmul_relu_0, cpu_matmul_relu_1, cpu_matmul_relu_2],
|
||||
expected_decisions=[decision_0, decision_1, decision_2],
|
||||
)
|
||||
|
||||
|
||||
def test_cuda_matmul():
|
||||
@T.prim_func(s_tir=True)
|
||||
def cuda_matmul_0(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
C: T.Buffer((512, 512), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C_local = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="local")
|
||||
A_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
B_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
for i0_0_i1_0_fused in T.thread_binding(128, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_fused in T.thread_binding(8, thread="vthread.x"):
|
||||
for i0_2_i1_2_fused in T.thread_binding(4, thread="threadIdx.x"):
|
||||
for i2_0 in T.serial(128):
|
||||
for ax0_ax1_fused in T.serial(256):
|
||||
with T.sblock("A_shared"):
|
||||
v0 = T.axis.spatial(
|
||||
512, i0_0_i1_0_fused // 16 * 64 + ax0_ax1_fused // 4
|
||||
)
|
||||
v1 = T.axis.spatial(512, i2_0 * 4 + ax0_ax1_fused % 4)
|
||||
T.reads(A[v0, v1])
|
||||
T.writes(A_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 2})
|
||||
A_shared[v0, v1] = A[v0, v1]
|
||||
for ax0_ax1_fused in T.serial(128):
|
||||
with T.sblock("B_shared"):
|
||||
v0 = T.axis.spatial(512, i2_0 * 4 + ax0_ax1_fused // 32)
|
||||
v1 = T.axis.spatial(
|
||||
512, i0_0_i1_0_fused % 16 * 32 + ax0_ax1_fused % 32
|
||||
)
|
||||
T.reads(B[v0, v1])
|
||||
T.writes(B_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 1})
|
||||
B_shared[v0, v1] = B[v0, v1]
|
||||
for i2_1, i0_3, i1_3, i2_2, i0_4, i1_4 in T.grid(2, 1, 1, 2, 16, 4):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(
|
||||
512,
|
||||
i0_0_i1_0_fused // 16 * 64
|
||||
+ i0_1_i1_1_fused // 2 * 16
|
||||
+ i0_3 * 16
|
||||
+ i0_4,
|
||||
)
|
||||
j = T.axis.spatial(
|
||||
512,
|
||||
i0_0_i1_0_fused % 16 * 32
|
||||
+ i0_1_i1_1_fused % 2 * 16
|
||||
+ i0_2_i1_2_fused * 4
|
||||
+ i1_3 * 4
|
||||
+ i1_4,
|
||||
)
|
||||
k = T.axis.reduce(512, i2_0 * 4 + i2_1 * 2 + i2_2)
|
||||
T.reads(A_shared[i, k], B_shared[k, j])
|
||||
T.writes(C_local[i, j])
|
||||
T.sblock_attr(
|
||||
{
|
||||
"meta_schedule.thread_extent_high_inclusive": 1024,
|
||||
"meta_schedule.thread_extent_low_inclusive": 32,
|
||||
"meta_schedule.tiling_structure": "SSSRRSRS",
|
||||
}
|
||||
)
|
||||
with T.init():
|
||||
C_local[i, j] = T.float32(0)
|
||||
C_local[i, j] = C_local[i, j] + A_shared[i, k] * B_shared[k, j]
|
||||
for ax0, ax1 in T.grid(16, 4):
|
||||
with T.sblock("C_local"):
|
||||
v0 = T.axis.spatial(
|
||||
512, i0_0_i1_0_fused // 16 * 64 + i0_1_i1_1_fused // 2 * 16 + ax0
|
||||
)
|
||||
v1 = T.axis.spatial(
|
||||
512,
|
||||
i0_0_i1_0_fused % 16 * 32
|
||||
+ i0_1_i1_1_fused % 2 * 16
|
||||
+ i0_2_i1_2_fused * 4
|
||||
+ ax1,
|
||||
)
|
||||
T.reads(C_local[v0, v1])
|
||||
T.writes(C[v0, v1])
|
||||
C[v0, v1] = C_local[v0, v1]
|
||||
|
||||
decision_0 = [
|
||||
("SamplePerfectTile", [8, 4, 1, 1, 16]),
|
||||
("SamplePerfectTile", [16, 2, 4, 1, 4]),
|
||||
("SamplePerfectTile", [128, 2, 2]),
|
||||
("SampleCategorical", 1),
|
||||
("SampleCategorical", 0),
|
||||
]
|
||||
mod = te.create_prim_func(te_workload.matmul(512, 512, 512))
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-2080"), # disable async trace using sm75
|
||||
types=ms.schedule_rule.MultiLevelTiling,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[cuda_matmul_0],
|
||||
expected_decisions=[decision_0],
|
||||
)
|
||||
|
||||
|
||||
def test_cuda_matmul_relu():
|
||||
@T.prim_func(s_tir=True)
|
||||
def cuda_matmul_relu_0(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
compute: T.Buffer((512, 512), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
C = T.sblock_alloc_buffer([512, 512], dtype="float32")
|
||||
C_local = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="local")
|
||||
A_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
B_shared = T.sblock_alloc_buffer([512, 512], dtype="float32", scope="shared")
|
||||
for i0_0_i1_0_fused in T.thread_binding(64, thread="blockIdx.x"):
|
||||
for i0_1_i1_1_fused in T.thread_binding(64, thread="vthread.x"):
|
||||
for i0_2_i1_2_fused in T.thread_binding(8, thread="threadIdx.x"):
|
||||
for i2_0 in T.serial(8):
|
||||
for ax0_ax1_fused in T.serial(4096):
|
||||
with T.sblock("A_shared"):
|
||||
v0 = T.axis.spatial(
|
||||
512, i0_0_i1_0_fused // 8 * 64 + ax0_ax1_fused // 64
|
||||
)
|
||||
v1 = T.axis.spatial(512, i2_0 * 64 + ax0_ax1_fused % 64)
|
||||
T.reads(A[v0, v1])
|
||||
T.writes(A_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 2})
|
||||
A_shared[v0, v1] = A[v0, v1]
|
||||
for ax0_ax1_fused in T.serial(4096):
|
||||
with T.sblock("B_shared"):
|
||||
v0 = T.axis.spatial(512, i2_0 * 64 + ax0_ax1_fused // 64)
|
||||
v1 = T.axis.spatial(
|
||||
512, i0_0_i1_0_fused % 8 * 64 + ax0_ax1_fused % 64
|
||||
)
|
||||
T.reads(B[v0, v1])
|
||||
T.writes(B_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 4})
|
||||
B_shared[v0, v1] = B[v0, v1]
|
||||
for i2_1, i0_3, i1_3, i2_2, i0_4, i1_4 in T.grid(8, 2, 1, 8, 2, 2):
|
||||
with T.sblock("C"):
|
||||
i = T.axis.spatial(
|
||||
512,
|
||||
i0_0_i1_0_fused // 8 * 64
|
||||
+ i0_1_i1_1_fused // 8 * 8
|
||||
+ i0_2_i1_2_fused // 4 * 4
|
||||
+ i0_3 * 2
|
||||
+ i0_4,
|
||||
)
|
||||
j = T.axis.spatial(
|
||||
512,
|
||||
i0_0_i1_0_fused % 8 * 64
|
||||
+ i0_1_i1_1_fused % 8 * 8
|
||||
+ i0_2_i1_2_fused % 4 * 2
|
||||
+ i1_3 * 2
|
||||
+ i1_4,
|
||||
)
|
||||
k = T.axis.reduce(512, i2_0 * 64 + i2_1 * 8 + i2_2)
|
||||
T.reads(A_shared[i, k], B_shared[k, j])
|
||||
T.writes(C_local[i, j])
|
||||
T.sblock_attr(
|
||||
{
|
||||
"meta_schedule.thread_extent_high_inclusive": 1024,
|
||||
"meta_schedule.thread_extent_low_inclusive": 32,
|
||||
"meta_schedule.tiling_structure": "SSSRRSRS",
|
||||
}
|
||||
)
|
||||
with T.init():
|
||||
C_local[i, j] = T.float32(0)
|
||||
C_local[i, j] = C_local[i, j] + A_shared[i, k] * B_shared[k, j]
|
||||
for ax0, ax1 in T.grid(4, 2):
|
||||
with T.sblock("C_local"):
|
||||
v0 = T.axis.spatial(
|
||||
512,
|
||||
i0_0_i1_0_fused // 8 * 64
|
||||
+ i0_1_i1_1_fused // 8 * 8
|
||||
+ i0_2_i1_2_fused // 4 * 4
|
||||
+ ax0,
|
||||
)
|
||||
v1 = T.axis.spatial(
|
||||
512,
|
||||
i0_0_i1_0_fused % 8 * 64
|
||||
+ i0_1_i1_1_fused % 8 * 8
|
||||
+ i0_2_i1_2_fused % 4 * 2
|
||||
+ ax1,
|
||||
)
|
||||
T.reads(C_local[v0, v1])
|
||||
T.writes(C[v0, v1])
|
||||
C[v0, v1] = C_local[v0, v1]
|
||||
for i0, i1 in T.grid(512, 512):
|
||||
with T.sblock("compute"):
|
||||
i0_1, i1_1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(C[i0_1, i1_1])
|
||||
T.writes(compute[i0_1, i1_1])
|
||||
compute[i0_1, i1_1] = T.max(C[i0_1, i1_1], T.float32(0))
|
||||
|
||||
decision_0 = [
|
||||
("SamplePerfectTile", [8, 8, 2, 2, 2]),
|
||||
("SamplePerfectTile", [8, 8, 4, 1, 2]),
|
||||
("SamplePerfectTile", [8, 8, 8]),
|
||||
("SampleCategorical", 1),
|
||||
("SampleCategorical", 3),
|
||||
]
|
||||
mod = te.create_prim_func(te_workload.matmul_relu(512, 512, 512))
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-2080"), # disable async trace using sm75
|
||||
types=ms.schedule_rule.MultiLevelTiling,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[cuda_matmul_relu_0],
|
||||
expected_decisions=[decision_0],
|
||||
)
|
||||
|
||||
|
||||
def test_cuda_sum_with_trivial_block_iter():
|
||||
@T.prim_func(s_tir=True)
|
||||
def sum_with_trivial_block_iter(
|
||||
A: T.Buffer((1, 64, 768), "float32"),
|
||||
B: T.Buffer((1, 64, 1), "float32"),
|
||||
) -> None:
|
||||
for i0, i1, i2, i3 in T.grid(1, 64, 1, 768):
|
||||
with T.sblock("sum"):
|
||||
ax0, ax1, ax2, k2 = T.axis.remap("SSSR", [i0, i1, i2, i3])
|
||||
T.reads(A[ax0, ax1, k2])
|
||||
T.writes(B[ax0, ax1, ax2])
|
||||
with T.init():
|
||||
B[ax0, ax1, ax2] = T.float32(0)
|
||||
B[ax0, ax1, ax2] = B[ax0, ax1, ax2] + A[ax0, ax1, k2]
|
||||
|
||||
# Expect nothing to happen - the rule is not supposed to be applied in this case
|
||||
mod = sum_with_trivial_block_iter
|
||||
(sch,) = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-3080"),
|
||||
types=ms.schedule_rule.MultiLevelTiling,
|
||||
)
|
||||
assert not sch.trace.simplified(remove_postproc=True).insts
|
||||
|
||||
|
||||
def test_multi_level_tiling_hexagon():
|
||||
@T.prim_func(s_tir=True)
|
||||
def cpu_conv2d_nhwc(
|
||||
inputs: T.Buffer((1, 56, 56, 64), "float16"),
|
||||
weight: T.Buffer((3, 3, 64, 64), "float16"),
|
||||
conv2d_nhwc: T.Buffer((1, 56, 56, 64), "float16"),
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
PadInput = T.sblock_alloc_buffer((1, 58, 58, 64), "float16")
|
||||
for i0, i1, i2, i3 in T.grid(1, 58, 58, 64):
|
||||
with T.sblock("PadInput"):
|
||||
v_i0, v_i1, v_i2, v_i3 = T.axis.remap("SSSS", [i0, i1, i2, i3])
|
||||
T.reads(inputs[v_i0, v_i1 - 1, v_i2 - 1, v_i3])
|
||||
T.writes(PadInput[v_i0, v_i1, v_i2, v_i3])
|
||||
PadInput[v_i0, v_i1, v_i2, v_i3] = T.if_then_else(
|
||||
1 <= v_i1 and v_i1 < 57 and 1 <= v_i2 and v_i2 < 57,
|
||||
inputs[v_i0, v_i1 - 1, v_i2 - 1, v_i3],
|
||||
T.float16(0),
|
||||
)
|
||||
for (
|
||||
n_0,
|
||||
h_0,
|
||||
w_0,
|
||||
co_0,
|
||||
rh_0,
|
||||
rw_0,
|
||||
rc_0,
|
||||
n_1,
|
||||
h_1,
|
||||
w_1,
|
||||
co_1,
|
||||
rh_1,
|
||||
rw_1,
|
||||
rc_1,
|
||||
n_2,
|
||||
h_2,
|
||||
w_2,
|
||||
co_2,
|
||||
) in T.grid(1, 1, 2, 1, 3, 3, 16, 1, 14, 2, 1, 1, 1, 4, 1, 4, 14, 64):
|
||||
with T.sblock("conv2d_nhwc"):
|
||||
v_n = T.axis.spatial(1, n_0 + n_1 + n_2)
|
||||
v_h = T.axis.spatial(56, h_0 * 56 + h_1 * 4 + h_2)
|
||||
v_w = T.axis.spatial(56, w_0 * 28 + w_1 * 14 + w_2)
|
||||
v_co = T.axis.spatial(64, co_0 * 64 + co_1 * 64 + co_2)
|
||||
v_rh = T.axis.reduce(3, rh_0 + rh_1)
|
||||
v_rw = T.axis.reduce(3, rw_0 + rw_1)
|
||||
v_rc = T.axis.reduce(64, rc_0 * 4 + rc_1)
|
||||
T.reads(
|
||||
PadInput[v_n, v_h + v_rh, v_w + v_rw, v_co // 64 * 64 + v_rc],
|
||||
weight[v_rh, v_rw, v_rc, v_co],
|
||||
)
|
||||
T.writes(conv2d_nhwc[v_n, v_h, v_w, v_co])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SRSRS"})
|
||||
with T.init():
|
||||
conv2d_nhwc[v_n, v_h, v_w, v_co] = T.float16(0)
|
||||
conv2d_nhwc[v_n, v_h, v_w, v_co] = (
|
||||
conv2d_nhwc[v_n, v_h, v_w, v_co]
|
||||
+ PadInput[v_n, v_h + v_rh, v_w + v_rw, v_co // 64 * 64 + v_rc]
|
||||
* weight[v_rh, v_rw, v_rc, v_co]
|
||||
)
|
||||
|
||||
target_hexagon = Target("qcom/hexagon-v69")
|
||||
|
||||
I = 64
|
||||
O = 64
|
||||
H = 56
|
||||
W = 56
|
||||
|
||||
mod = te.create_prim_func(
|
||||
te_workload.conv2d_nhwc(1, H, W, I, O, 3, 1, 1, 1, in_dtype="float16", out_dtype="float16")
|
||||
)
|
||||
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target(target_hexagon, host=target_hexagon),
|
||||
types=None,
|
||||
sch_rules=[
|
||||
ms.schedule_rule.MultiLevelTilingWideVector(
|
||||
structure="SRSRS",
|
||||
vector_length_in_bits=1024,
|
||||
max_innermost_factor=64,
|
||||
reuse_read=None,
|
||||
reuse_write=None,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
decision_0 = [
|
||||
("SamplePerfectTile", [1, 1, 1]),
|
||||
("SamplePerfectTile", [1, 14, 4]),
|
||||
("SamplePerfectTile", [2, 2, 14]),
|
||||
("SamplePerfectTile", [3, 1]),
|
||||
("SamplePerfectTile", [3, 1]),
|
||||
("SamplePerfectTile", [16, 4]),
|
||||
]
|
||||
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[cpu_conv2d_nhwc],
|
||||
expected_decisions=[decision_0],
|
||||
)
|
||||
|
||||
|
||||
def test_cache_read_specify_consumer():
|
||||
@T.prim_func(s_tir=True)
|
||||
def cache_read_specify_consumer_0(
|
||||
A: T.Buffer((512, 512), "float32"),
|
||||
B: T.Buffer((512, 512), "float32"),
|
||||
T_add: T.Buffer((512, 512), "float32"),
|
||||
):
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
C = T.sblock_alloc_buffer((512, 512))
|
||||
C_local = T.sblock_alloc_buffer((512, 512), scope="local")
|
||||
A_shared = T.sblock_alloc_buffer((512, 512), scope="shared")
|
||||
B_shared = T.sblock_alloc_buffer((512, 512), scope="shared")
|
||||
for i_0_j_0_fused in T.thread_binding(2, thread="blockIdx.x"):
|
||||
for i_1_j_1_fused in T.thread_binding(512, thread="vthread.x"):
|
||||
for i_2_j_2_fused in T.thread_binding(16, thread="threadIdx.x"):
|
||||
for k_0 in range(2):
|
||||
for ax0_ax1_fused in range(131072):
|
||||
with T.sblock("A_shared"):
|
||||
v0 = T.axis.spatial(512, ax0_ax1_fused // 256)
|
||||
v1 = T.axis.spatial(512, k_0 * 256 + ax0_ax1_fused % 256)
|
||||
T.reads(A[v0, v1])
|
||||
T.writes(A_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 2})
|
||||
A_shared[v0, v1] = A[v0, v1]
|
||||
for ax0_ax1_fused in range(65536):
|
||||
with T.sblock("B_shared"):
|
||||
v0 = T.axis.spatial(512, k_0 * 256 + ax0_ax1_fused // 256)
|
||||
v1 = T.axis.spatial(512, i_0_j_0_fused * 256 + ax0_ax1_fused % 256)
|
||||
T.reads(B[v0, v1])
|
||||
T.writes(B_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 3})
|
||||
B_shared[v0, v1] = B[v0, v1]
|
||||
for k_1, i_3, j_3, k_2, i_4, j_4 in T.grid(64, 1, 1, 4, 1, 16):
|
||||
with T.sblock("C"):
|
||||
v_i = T.axis.spatial(
|
||||
512,
|
||||
i_1_j_1_fused // 8 * 8 + i_2_j_2_fused // 2 + i_3 + i_4,
|
||||
)
|
||||
v_j = T.axis.spatial(
|
||||
512,
|
||||
i_0_j_0_fused * 256
|
||||
+ i_1_j_1_fused % 8 * 32
|
||||
+ i_2_j_2_fused % 2 * 16
|
||||
+ j_3 * 16
|
||||
+ j_4,
|
||||
)
|
||||
v_k = T.axis.reduce(512, k_0 * 256 + k_1 * 4 + k_2)
|
||||
T.reads(A_shared[v_i, v_k], B_shared[v_k, v_j])
|
||||
T.writes(C_local[v_i, v_j])
|
||||
T.sblock_attr(
|
||||
{
|
||||
"meta_schedule.thread_extent_high_inclusive": 1024,
|
||||
"meta_schedule.thread_extent_low_inclusive": 32,
|
||||
"meta_schedule.tiling_structure": "SSSRRSRS",
|
||||
}
|
||||
)
|
||||
with T.init():
|
||||
C_local[v_i, v_j] = T.float32(0)
|
||||
C_local[v_i, v_j] = (
|
||||
C_local[v_i, v_j] + A_shared[v_i, v_k] * B_shared[v_k, v_j]
|
||||
)
|
||||
for ax0, ax1 in T.grid(1, 16):
|
||||
with T.sblock("C_local"):
|
||||
v0 = T.axis.spatial(
|
||||
512,
|
||||
i_1_j_1_fused // 8 * 8 + i_2_j_2_fused // 2 + ax0,
|
||||
)
|
||||
v1 = T.axis.spatial(
|
||||
512,
|
||||
i_0_j_0_fused * 256
|
||||
+ i_1_j_1_fused % 8 * 32
|
||||
+ i_2_j_2_fused % 2 * 16
|
||||
+ ax1,
|
||||
)
|
||||
T.reads(C_local[v0, v1])
|
||||
T.writes(C[v0, v1])
|
||||
C[v0, v1] = C_local[v0, v1]
|
||||
for ax0, ax1 in T.grid(512, 512):
|
||||
with T.sblock("T_add"):
|
||||
v_ax0 = T.axis.spatial(512, ax0)
|
||||
v_ax1 = T.axis.spatial(512, ax1)
|
||||
T.reads(C[v_ax0, v_ax1], A[v_ax0, v_ax1])
|
||||
T.writes(T_add[v_ax0, v_ax1])
|
||||
T_add[v_ax0, v_ax1] = C[v_ax0, v_ax1] + A[v_ax0, v_ax1]
|
||||
|
||||
decision_0 = [
|
||||
("SamplePerfectTile", [1, 64, 8, 1, 1]),
|
||||
("SamplePerfectTile", [2, 8, 2, 1, 16]),
|
||||
("SamplePerfectTile", [2, 64, 4]),
|
||||
("SampleCategorical", 1),
|
||||
("SampleCategorical", 2),
|
||||
]
|
||||
A, B, C = te_workload.matmul(512, 512, 512)
|
||||
mod = te.create_prim_func([A, B, C + A])
|
||||
|
||||
space = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target("nvidia/geforce-rtx-2080"), # disable async trace using sm75
|
||||
types=ms.schedule_rule.MultiLevelTiling,
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=space,
|
||||
expected_mods=[cache_read_specify_consumer_0],
|
||||
expected_decisions=[decision_0],
|
||||
)
|
||||
|
||||
|
||||
def test_max_pool_blocked():
|
||||
# fmt off
|
||||
@T.prim_func(s_tir=True)
|
||||
def pool_blocked_cache_read_write(
|
||||
X: T.Buffer((1, 2, 8, 8, 8, 8, 32), "uint8"),
|
||||
pool: T.Buffer((1, 2, 4, 4, 8, 8, 32), "uint8"),
|
||||
):
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
pool_global = T.sblock_alloc_buffer((1, 2, 4, 4, 8, 8, 32), "uint8")
|
||||
X_global = T.sblock_alloc_buffer((1, 2, 8, 8, 8, 8, 32), "uint8")
|
||||
for b_0, c_o_0, h_o_0, w_o_0, h_i_0, w_i_0, c_i_0 in T.grid(1, 2, 4, 1, 8, 1, 4):
|
||||
for ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused in range(896):
|
||||
with T.sblock("X_global"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(2, c_o_0)
|
||||
v2 = T.axis.spatial(8, h_o_0 * 2)
|
||||
v3 = T.axis.spatial(8, ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused // 128)
|
||||
v4 = T.axis.spatial(
|
||||
8, h_i_0 % 4 * 2 + ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused % 128 // 64
|
||||
)
|
||||
v5 = T.axis.spatial(8, ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused % 64 // 8)
|
||||
v6 = T.axis.spatial(32, c_i_0 * 8 + ax0_ax1_ax2_ax3_ax4_ax5_ax6_fused % 8)
|
||||
T.reads(X[v0, v1, v2, v3, v4, v5, v6])
|
||||
T.writes(X_global[v0, v1, v2, v3, v4, v5, v6])
|
||||
X_global[v0, v1, v2, v3, v4, v5, v6] = X[v0, v1, v2, v3, v4, v5, v6]
|
||||
for wh, ww, b_1, c_o_1, h_o_1, w_o_1, h_i_1, w_i_1, c_i_1 in T.grid(
|
||||
2, 2, 1, 1, 1, 4, 1, 8, 8
|
||||
):
|
||||
with T.sblock("pool"):
|
||||
v_b = T.axis.spatial(1, b_0 + b_1)
|
||||
v_c_o = T.axis.spatial(2, c_o_0 + c_o_1)
|
||||
v_h_o = T.axis.spatial(4, h_o_0 + h_o_1)
|
||||
v_w_o = T.axis.spatial(4, w_o_0 * 4 + w_o_1)
|
||||
v_h_i = T.axis.spatial(8, h_i_0 + h_i_1)
|
||||
v_w_i = T.axis.spatial(8, w_i_0 * 8 + w_i_1)
|
||||
v_c_i = T.axis.spatial(32, c_i_0 * 8 + c_i_1)
|
||||
v_wh, v_ww = T.axis.remap("RR", [wh, ww])
|
||||
T.reads(
|
||||
X_global[
|
||||
v_b,
|
||||
v_c_o,
|
||||
v_h_i // 8 * 2 + v_h_o * 2,
|
||||
v_w_i // 8 * 2 + v_w_o * 2,
|
||||
v_h_i % 4 * 2 + v_wh,
|
||||
v_w_i % 4 * 2 + v_ww,
|
||||
v_c_i,
|
||||
]
|
||||
)
|
||||
T.writes(pool_global[v_b, v_c_o, v_h_o, v_w_o, v_h_i, v_w_i, v_c_i])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SRS"})
|
||||
with T.init():
|
||||
pool_global[v_b, v_c_o, v_h_o, v_w_o, v_h_i, v_w_i, v_c_i] = T.uint8(0)
|
||||
pool_global[v_b, v_c_o, v_h_o, v_w_o, v_h_i, v_w_i, v_c_i] = T.max(
|
||||
pool_global[v_b, v_c_o, v_h_o, v_w_o, v_h_i, v_w_i, v_c_i],
|
||||
X_global[
|
||||
v_b,
|
||||
v_c_o,
|
||||
v_h_i // 8 * 2 + v_h_o * 2,
|
||||
v_w_i // 8 * 2 + v_w_o * 2,
|
||||
v_h_i % 4 * 2 + v_wh,
|
||||
v_w_i % 4 * 2 + v_ww,
|
||||
v_c_i,
|
||||
],
|
||||
)
|
||||
for ax0, ax1, ax2, ax3, ax4, ax5, ax6 in T.grid(1, 1, 1, 4, 1, 8, 8):
|
||||
with T.sblock("pool_global"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(2, c_o_0 + ax1)
|
||||
v2 = T.axis.spatial(4, h_o_0 + ax2)
|
||||
v3 = T.axis.spatial(4, ax3)
|
||||
v4 = T.axis.spatial(8, h_i_0 + ax4)
|
||||
v5 = T.axis.spatial(8, ax5)
|
||||
v6 = T.axis.spatial(32, c_i_0 * 8 + ax6)
|
||||
T.reads(pool_global[v0, v1, v2, v3, v4, v5, v6])
|
||||
T.writes(pool[v0, v1, v2, v3, v4, v5, v6])
|
||||
pool[v0, v1, v2, v3, v4, v5, v6] = pool_global[v0, v1, v2, v3, v4, v5, v6]
|
||||
|
||||
# fmt on
|
||||
|
||||
def max_pool_blocked_compute(height, width, channel):
|
||||
ishape = (1, channel // 32, height // 8, width // 8, 8, 8, 32)
|
||||
oshape = (1, channel // 32, height // 8 // 2, width // 8 // 2, 8, 8, 32)
|
||||
X = te.placeholder(ishape, name="X", dtype="uint8")
|
||||
|
||||
window_h = te.reduce_axis((0, 2), name="wh")
|
||||
window_w = te.reduce_axis((0, 2), name="ww")
|
||||
|
||||
out = te.compute(
|
||||
oshape,
|
||||
lambda b, c_o, h_o, w_o, h_i, w_i, c_i: te.max(
|
||||
X[
|
||||
b,
|
||||
c_o,
|
||||
(h_o * 8 + h_i) // 8 * 2,
|
||||
(w_o * 8 + w_i) // 8 * 2,
|
||||
(h_o * 8 + h_i) % 4 * 2 + window_h,
|
||||
(w_o * 8 + w_i) % 4 * 2 + window_w,
|
||||
c_i,
|
||||
],
|
||||
axis=[window_h, window_w],
|
||||
),
|
||||
name="pool",
|
||||
)
|
||||
return [X, out]
|
||||
|
||||
height = width = 64
|
||||
channel = 64
|
||||
|
||||
mod = te.create_prim_func(max_pool_blocked_compute(height, width, channel))
|
||||
|
||||
actual = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=Target("llvm"),
|
||||
types=None,
|
||||
sch_rules=[
|
||||
ms.schedule_rule.MultiLevelTiling(
|
||||
structure="SRS",
|
||||
tile_binds=None,
|
||||
max_innermost_factor=64,
|
||||
vector_load_lens=None,
|
||||
reuse_read=ms.schedule_rule.ReuseType(
|
||||
req="must",
|
||||
levels=[1],
|
||||
scope="global",
|
||||
),
|
||||
reuse_write=ms.schedule_rule.ReuseType(req="must", levels=[1], scope="global"),
|
||||
filter_fn=lambda sch, block_rv: sch.get(block_rv).name_hint == "pool",
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
decision = [
|
||||
("SamplePerfectTile", [1, 1]),
|
||||
("SamplePerfectTile", [2, 1]),
|
||||
("SamplePerfectTile", [4, 1]),
|
||||
("SamplePerfectTile", [1, 4]),
|
||||
("SamplePerfectTile", [8, 1]),
|
||||
("SamplePerfectTile", [1, 8]),
|
||||
("SamplePerfectTile", [4, 8]),
|
||||
]
|
||||
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[pool_blocked_cache_read_write],
|
||||
expected_decisions=[decision],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,427 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: E501, F401, F841
|
||||
from tvm import te
|
||||
from tvm.ir import assert_structural_equal
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing.space_generation import (
|
||||
check_sketches,
|
||||
generate_design_space,
|
||||
print_sketches,
|
||||
)
|
||||
from tvm.s_tir.tensor_intrin.arm_cpu import DP4A_S8S8S32_INTRIN
|
||||
from tvm.s_tir.tensor_intrin.x86 import AVX512_DOT_16x4_INTRIN as AVX512_INTRIN
|
||||
from tvm.s_tir.tensor_intrin.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def test_x86_conv2d_nchwc(
|
||||
intrin=VNNI_INTRIN, target={"kind": "llvm", "mcpu": "cascadelake", "num-cores": 4}
|
||||
):
|
||||
@T.prim_func(s_tir=True)
|
||||
def conv2d_nchwc(
|
||||
placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"),
|
||||
placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"),
|
||||
conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32"),
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
for i0, i1, i2, i3, i4, i5, i6, i7, i8, i9 in T.grid(1, 16, 56, 56, 16, 1, 1, 4, 4, 4):
|
||||
with T.sblock("conv2d_NCHWc_int8"):
|
||||
(
|
||||
n,
|
||||
oc_chunk,
|
||||
oh,
|
||||
ow,
|
||||
oc_block,
|
||||
kh,
|
||||
kw,
|
||||
ic_outer,
|
||||
ic_f_inner,
|
||||
ic_s_inner,
|
||||
) = T.axis.remap("SSSSSRRRRR", [i0, i1, i2, i3, i4, i5, i6, i7, i8, i9])
|
||||
T.reads(
|
||||
placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner],
|
||||
placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
|
||||
)
|
||||
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block])
|
||||
with T.init():
|
||||
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = 0
|
||||
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block] = conv2d_NCHWc_int8[
|
||||
n, oc_chunk, oh, ow, oc_block
|
||||
] + T.cast(
|
||||
placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner], "int32"
|
||||
) * T.cast(
|
||||
placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block, ic_s_inner],
|
||||
"int32",
|
||||
)
|
||||
|
||||
# fmt: off
|
||||
@T.prim_func(s_tir=True)
|
||||
def x86_conv2d_nchwc_0(placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"), placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"), conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32")) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
conv2d_NCHWc_int8_global = T.sblock_alloc_buffer((1, 16, 56, 56, 16), "int32")
|
||||
for i0_0, i1_0, i2_0, i3_0, i4_0_0, i0_1, i1_1, i2_1, i3_1, i4_0_1 in T.grid(1, 8, 28, 56, 1, 1, 2, 1, 1, 1):
|
||||
for i5_0, i6_0, i7_0, i8_0, i9_0_0, i0_2, i1_2, i2_2, i3_2, i4_0_2, i5_1, i6_1, i7_1, i8_1, i9_0_1, i0_3, i1_3, i2_3, i3_3, i4_0_3 in T.grid(1, 1, 1, 4, 1, 1, 1, 2, 1, 1, 1, 1, 4, 1, 1, 1, 1, 1, 1, 1):
|
||||
with T.sblock("conv2d_NCHWc_int8_o"):
|
||||
n = T.axis.spatial(1, i0_0 + i0_1 + i0_2 + i0_3)
|
||||
oc_chunk = T.axis.spatial(16, i1_0 * 2 + i1_1 + i1_2 + i1_3)
|
||||
oh = T.axis.spatial(56, i2_0 * 2 + i2_1 * 2 + i2_2 + i2_3)
|
||||
ow = T.axis.spatial(56, i3_0 + i3_1 + i3_2 + i3_3)
|
||||
oc_block_o = T.axis.spatial(1, i4_0_0 + i4_0_1 + i4_0_2 + i4_0_3)
|
||||
kh = T.axis.reduce(1, i5_0 + i5_1)
|
||||
kw = T.axis.reduce(1, i6_0 + i6_1)
|
||||
ic_outer = T.axis.reduce(4, i7_0 * 4 + i7_1)
|
||||
ic_f_inner = T.axis.reduce(4, i8_0 + i8_1)
|
||||
ic_s_inner_o = T.axis.reduce(1, i9_0_0 + i9_0_1)
|
||||
T.reads(placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4:ic_f_inner * 4 + 4], placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, 0:16, 0:4])
|
||||
T.writes(conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, 0:16])
|
||||
T.sblock_attr({"meta_schedule.auto_tensorize": intrin})
|
||||
with T.init():
|
||||
for i4_1 in range(16):
|
||||
with T.sblock("conv2d_NCHWc_int8_init"):
|
||||
oc_block_i_init = T.axis.spatial(16, i4_1)
|
||||
T.reads()
|
||||
T.writes(conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i_init])
|
||||
conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i_init] = 0
|
||||
for i4_1, i9_1 in T.grid(16, 4):
|
||||
with T.sblock("conv2d_NCHWc_int8"):
|
||||
oc_block_i, ic_s_inner_i = T.axis.remap("SR", [i4_1, i9_1])
|
||||
T.reads(conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i], placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner_i], placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block_i, ic_s_inner_i])
|
||||
T.writes(conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SSRSRS"})
|
||||
conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i] = conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i] + T.Cast("int32", placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner_i]) * T.Cast("int32", placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block_i, ic_s_inner_i])
|
||||
for ax0, ax1, ax2, ax3, ax4 in T.grid(1, 1, 2, 1, 16):
|
||||
with T.sblock("conv2d_NCHWc_int8_global"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(16, i1_0 * 2 + i1_1 + ax1)
|
||||
v2 = T.axis.spatial(56, i2_0 * 2 + ax2)
|
||||
v3 = T.axis.spatial(56, i3_0 + ax3)
|
||||
v4 = T.axis.spatial(16, ax4)
|
||||
T.reads(conv2d_NCHWc_int8_global[v0, v1, v2, v3, v4])
|
||||
T.writes(conv2d_NCHWc_int8[v0, v1, v2, v3, v4])
|
||||
conv2d_NCHWc_int8[v0, v1, v2, v3, v4] = conv2d_NCHWc_int8_global[v0, v1, v2, v3, v4]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def x86_conv2d_nchwc_1(placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"), placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"), conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32")) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
conv2d_NCHWc_int8_global = T.sblock_alloc_buffer((1, 16, 56, 56, 16), "int32")
|
||||
for i0_0, i1_0, i2_0, i3_0, i4_0_0 in T.grid(1, 8, 28, 56, 1):
|
||||
for i0_1, i1_1, i2_1, i3_1, i4_0_1, i5_0, i6_0, i7_0, i8_0, i9_0_0, i0_2, i1_2, i2_2, i3_2, i4_0_2, i5_1, i6_1, i7_1, i8_1, i9_0_1, i0_3, i1_3, i2_3, i3_3, i4_0_3 in T.grid(1, 2, 1, 1, 1, 1, 1, 1, 4, 1, 1, 1, 2, 1, 1, 1, 1, 4, 1, 1, 1, 1, 1, 1, 1):
|
||||
with T.sblock("conv2d_NCHWc_int8_o"):
|
||||
n = T.axis.spatial(1, i0_0 + i0_1 + i0_2 + i0_3)
|
||||
oc_chunk = T.axis.spatial(16, i1_0 * 2 + i1_1 + i1_2 + i1_3)
|
||||
oh = T.axis.spatial(56, i2_0 * 2 + i2_1 * 2 + i2_2 + i2_3)
|
||||
ow = T.axis.spatial(56, i3_0 + i3_1 + i3_2 + i3_3)
|
||||
oc_block_o = T.axis.spatial(1, i4_0_0 + i4_0_1 + i4_0_2 + i4_0_3)
|
||||
kh = T.axis.reduce(1, i5_0 + i5_1)
|
||||
kw = T.axis.reduce(1, i6_0 + i6_1)
|
||||
ic_outer = T.axis.reduce(4, i7_0 * 4 + i7_1)
|
||||
ic_f_inner = T.axis.reduce(4, i8_0 + i8_1)
|
||||
ic_s_inner_o = T.axis.reduce(1, i9_0_0 + i9_0_1)
|
||||
T.reads(placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4:ic_f_inner * 4 + 4], placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, 0:16, 0:4])
|
||||
T.writes(conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, 0:16])
|
||||
T.sblock_attr({"meta_schedule.auto_tensorize": intrin})
|
||||
with T.init():
|
||||
for i4_1 in range(16):
|
||||
with T.sblock("conv2d_NCHWc_int8_init"):
|
||||
oc_block_i_init = T.axis.spatial(16, i4_1)
|
||||
T.reads()
|
||||
T.writes(conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i_init])
|
||||
conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i_init] = 0
|
||||
for i4_1, i9_1 in T.grid(16, 4):
|
||||
with T.sblock("conv2d_NCHWc_int8"):
|
||||
oc_block_i, ic_s_inner_i = T.axis.remap("SR", [i4_1, i9_1])
|
||||
T.reads(conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i], placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner_i], placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block_i, ic_s_inner_i])
|
||||
T.writes(conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SSRSRS"})
|
||||
conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i] = conv2d_NCHWc_int8_global[n, oc_chunk, oh, ow, oc_block_i] + T.Cast("int32", placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner_i]) * T.Cast("int32", placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block_i, ic_s_inner_i])
|
||||
for ax0, ax1, ax2, ax3, ax4 in T.grid(1, 2, 2, 1, 16):
|
||||
with T.sblock("conv2d_NCHWc_int8_global"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(16, i1_0 * 2 + ax1)
|
||||
v2 = T.axis.spatial(56, i2_0 * 2 + ax2)
|
||||
v3 = T.axis.spatial(56, i3_0 + ax3)
|
||||
v4 = T.axis.spatial(16, ax4)
|
||||
T.reads(conv2d_NCHWc_int8_global[v0, v1, v2, v3, v4])
|
||||
T.writes(conv2d_NCHWc_int8[v0, v1, v2, v3, v4])
|
||||
conv2d_NCHWc_int8[v0, v1, v2, v3, v4] = conv2d_NCHWc_int8_global[v0, v1, v2, v3, v4]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def x86_conv2d_nchwc_2(placeholder: T.Buffer((1, 4, 56, 56, 16), "uint8"), placeholder_1: T.Buffer((16, 4, 1, 1, 4, 16, 4), "int8"), conv2d_NCHWc_int8: T.Buffer((1, 16, 56, 56, 16), "int32")) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
for i0_0, i1_0, i2_0, i3_0, i4_0_0, i0_1, i1_1, i2_1, i3_1, i4_0_1, i5_0, i6_0, i7_0, i8_0, i9_0_0, i0_2, i1_2, i2_2, i3_2, i4_0_2, i5_1, i6_1, i7_1, i8_1, i9_0_1, i0_3, i1_3, i2_3, i3_3, i4_0_3 in T.grid(1, 8, 28, 56, 1, 1, 2, 1, 1, 1, 1, 1, 1, 4, 1, 1, 1, 2, 1, 1, 1, 1, 4, 1, 1, 1, 1, 1, 1, 1):
|
||||
with T.sblock("conv2d_NCHWc_int8_o"):
|
||||
n = T.axis.spatial(1, i0_0 + i0_1 + i0_2 + i0_3)
|
||||
oc_chunk = T.axis.spatial(16, i1_0 * 2 + i1_1 + i1_2 + i1_3)
|
||||
oh = T.axis.spatial(56, i2_0 * 2 + i2_1 * 2 + i2_2 + i2_3)
|
||||
ow = T.axis.spatial(56, i3_0 + i3_1 + i3_2 + i3_3)
|
||||
oc_block_o = T.axis.spatial(1, i4_0_0 + i4_0_1 + i4_0_2 + i4_0_3)
|
||||
kh = T.axis.reduce(1, i5_0 + i5_1)
|
||||
kw = T.axis.reduce(1, i6_0 + i6_1)
|
||||
ic_outer = T.axis.reduce(4, i7_0 * 4 + i7_1)
|
||||
ic_f_inner = T.axis.reduce(4, i8_0 + i8_1)
|
||||
ic_s_inner_o = T.axis.reduce(1, i9_0_0 + i9_0_1)
|
||||
T.reads(placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4:ic_f_inner * 4 + 4], placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, 0:16, 0:4])
|
||||
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, 0:16])
|
||||
T.sblock_attr({"meta_schedule.auto_tensorize": intrin})
|
||||
with T.init():
|
||||
for i4_1 in range(16):
|
||||
with T.sblock("conv2d_NCHWc_int8_init"):
|
||||
oc_block_i_init = T.axis.spatial(16, i4_1)
|
||||
T.reads()
|
||||
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block_i_init])
|
||||
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block_i_init] = 0
|
||||
for i4_1, i9_1 in T.grid(16, 4):
|
||||
with T.sblock("conv2d_NCHWc_int8"):
|
||||
oc_block_i, ic_s_inner_i = T.axis.remap("SR", [i4_1, i9_1])
|
||||
T.reads(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block_i], placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner_i], placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block_i, ic_s_inner_i])
|
||||
T.writes(conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block_i])
|
||||
T.sblock_attr({"meta_schedule.tiling_structure": "SSRSRS"})
|
||||
conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block_i] = conv2d_NCHWc_int8[n, oc_chunk, oh, ow, oc_block_i] + T.Cast("int32", placeholder[n, ic_outer, oh + kh, ow + kw, ic_f_inner * 4 + ic_s_inner_i]) * T.Cast("int32", placeholder_1[oc_chunk, ic_outer, kh, kw, ic_f_inner, oc_block_i, ic_s_inner_i])
|
||||
# fmt: on
|
||||
decision_0 = [
|
||||
("SamplePerfectTile", [1, 1, 1, 1]),
|
||||
("SamplePerfectTile", [8, 2, 1, 1]),
|
||||
("SamplePerfectTile", [28, 1, 2, 1]),
|
||||
("SamplePerfectTile", [56, 1, 1, 1]),
|
||||
("SamplePerfectTile", [1, 1, 1, 1]),
|
||||
("SamplePerfectTile", [1, 1]),
|
||||
("SamplePerfectTile", [1, 1]),
|
||||
("SamplePerfectTile", [1, 4]),
|
||||
("SamplePerfectTile", [4, 1]),
|
||||
("SamplePerfectTile", [1, 1]),
|
||||
]
|
||||
decision_1 = [
|
||||
("SamplePerfectTile", [1, 1, 1, 1]),
|
||||
("SamplePerfectTile", [8, 2, 1, 1]),
|
||||
("SamplePerfectTile", [28, 1, 2, 1]),
|
||||
("SamplePerfectTile", [56, 1, 1, 1]),
|
||||
("SamplePerfectTile", [1, 1, 1, 1]),
|
||||
("SamplePerfectTile", [1, 1]),
|
||||
("SamplePerfectTile", [1, 1]),
|
||||
("SamplePerfectTile", [1, 4]),
|
||||
("SamplePerfectTile", [4, 1]),
|
||||
("SamplePerfectTile", [1, 1]),
|
||||
]
|
||||
decision_2 = [
|
||||
("SamplePerfectTile", [1, 1, 1, 1]),
|
||||
("SamplePerfectTile", [8, 2, 1, 1]),
|
||||
("SamplePerfectTile", [28, 1, 2, 1]),
|
||||
("SamplePerfectTile", [56, 1, 1, 1]),
|
||||
("SamplePerfectTile", [1, 1, 1, 1]),
|
||||
("SamplePerfectTile", [1, 1]),
|
||||
("SamplePerfectTile", [1, 1]),
|
||||
("SamplePerfectTile", [1, 4]),
|
||||
("SamplePerfectTile", [4, 1]),
|
||||
("SamplePerfectTile", [1, 1]),
|
||||
]
|
||||
|
||||
mod = conv2d_nchwc
|
||||
actual = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=Target(target),
|
||||
types=None,
|
||||
sch_rules=[
|
||||
ms.schedule_rule.MultiLevelTilingWithIntrin(
|
||||
intrin,
|
||||
structure="SSRSRS",
|
||||
tile_binds=None,
|
||||
max_innermost_factor=64,
|
||||
vector_load_lens=None,
|
||||
reuse_read=None,
|
||||
reuse_write=ms.schedule_rule.ReuseType(req="may", levels=[1, 2], scope="global"),
|
||||
),
|
||||
],
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[x86_conv2d_nchwc_0, x86_conv2d_nchwc_1, x86_conv2d_nchwc_2],
|
||||
expected_decisions=[decision_0, decision_1, decision_2],
|
||||
)
|
||||
|
||||
|
||||
def _check_dp4a_dense(m, n, k, in_dtype, out_dtype, expected_mods, expected_decisions):
|
||||
def _dense(m, n, k, in_dtype, out_dtype):
|
||||
X = te.placeholder((m, k), name="X", dtype=in_dtype)
|
||||
W = te.placeholder((n, k), name="W", dtype=in_dtype)
|
||||
ak = te.reduce_axis((0, k), name="k")
|
||||
matmul = te.compute(
|
||||
(m, n),
|
||||
lambda i, j: te.sum(
|
||||
X[i, ak].astype(out_dtype) * W[j, ak].astype(out_dtype),
|
||||
axis=ak,
|
||||
),
|
||||
name="compute",
|
||||
)
|
||||
return te.create_prim_func([X, W, matmul])
|
||||
|
||||
mod = _dense(m, n, k, in_dtype, out_dtype)
|
||||
actual = generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=Target({"kind": "cuda", "arch": "sm_70"}),
|
||||
types=None,
|
||||
sch_rules=[
|
||||
ms.schedule_rule.MultiLevelTilingWithIntrin(
|
||||
DP4A_S8S8S32_INTRIN,
|
||||
structure="SSSRRSRS",
|
||||
tile_binds=["blockIdx.x", "vthread.x", "threadIdx.x"],
|
||||
max_innermost_factor=64,
|
||||
vector_load_lens=[1, 2, 3, 4],
|
||||
reuse_read=ms.schedule_rule.ReuseType(req="must", levels=[4], scope="shared"),
|
||||
reuse_write=ms.schedule_rule.ReuseType(req="must", levels=[3], scope="local"),
|
||||
)
|
||||
],
|
||||
)
|
||||
if expected_mods is None:
|
||||
assert expected_decisions is None
|
||||
assert len(actual) == 1
|
||||
assert_structural_equal(mod, actual[0].mod["main"])
|
||||
else:
|
||||
check_sketches(mod, actual, expected_mods, expected_decisions)
|
||||
|
||||
|
||||
def test_dp4a_dense():
|
||||
@T.prim_func(s_tir=True)
|
||||
def dp4a_dense_0(
|
||||
X: T.Buffer((128, 128), "int8"),
|
||||
W: T.Buffer((128, 128), "int8"),
|
||||
compute: T.Buffer((128, 128), "int32"),
|
||||
) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
# with T.sblock("root"):
|
||||
compute_local = T.sblock_alloc_buffer((128, 128), "int32", scope="local")
|
||||
X_shared = T.sblock_alloc_buffer((128, 128), "int8", scope="shared")
|
||||
W_shared = T.sblock_alloc_buffer((128, 128), "int8", scope="shared")
|
||||
for i_0_j_0_fused in T.thread_binding(1, thread="blockIdx.x"):
|
||||
for i_1_j_1_fused in T.thread_binding(512, thread="vthread.x"):
|
||||
for i_2_j_2_fused in T.thread_binding(2, thread="threadIdx.x"):
|
||||
for k_0_0 in range(1):
|
||||
for ax0_ax1_fused in range(16384):
|
||||
with T.sblock("X_shared"):
|
||||
v0 = T.axis.spatial(128, ax0_ax1_fused // 128)
|
||||
v1 = T.axis.spatial(128, ax0_ax1_fused % 128)
|
||||
T.reads(X[v0, v1])
|
||||
T.writes(X_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 1})
|
||||
X_shared[v0, v1] = X[v0, v1]
|
||||
for ax0_ax1_fused in range(16384):
|
||||
with T.sblock("W_shared"):
|
||||
v0 = T.axis.spatial(128, ax0_ax1_fused // 128)
|
||||
v1 = T.axis.spatial(128, ax0_ax1_fused % 128)
|
||||
T.reads(W[v0, v1])
|
||||
T.writes(W_shared[v0, v1])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 1})
|
||||
W_shared[v0, v1] = W[v0, v1]
|
||||
for k_0_1, i_3, j_3, k_0_2, i_4, j_4 in T.grid(1, 2, 4, 32, 2, 1):
|
||||
with T.sblock("compute_o"):
|
||||
v_i = T.axis.spatial(
|
||||
128, i_1_j_1_fused // 32 * 8 + i_2_j_2_fused * 4 + i_3 * 2 + i_4
|
||||
)
|
||||
v_j = T.axis.spatial(128, i_1_j_1_fused % 32 * 4 + j_3 + j_4)
|
||||
v_k_o = T.axis.reduce(32, k_0_0 * 32 + k_0_1 * 32 + k_0_2)
|
||||
T.reads(
|
||||
X_shared[v_i, v_k_o * 4 : v_k_o * 4 + 4],
|
||||
W_shared[v_j, v_k_o * 4 : v_k_o * 4 + 4],
|
||||
)
|
||||
T.writes(compute_local[v_i, v_j])
|
||||
T.sblock_attr({"meta_schedule.auto_tensorize": "dp4a_s8s8s32"})
|
||||
with T.init():
|
||||
with T.sblock("compute_init"):
|
||||
T.reads()
|
||||
T.writes(compute_local[v_i, v_j])
|
||||
compute_local[v_i, v_j] = 0
|
||||
for k_1 in range(4):
|
||||
with T.sblock("compute"):
|
||||
v_k_i = T.axis.reduce(4, k_1)
|
||||
T.reads(
|
||||
compute_local[v_i, v_j],
|
||||
X_shared[v_i, v_k_o * 4 + v_k_i],
|
||||
W_shared[v_j, v_k_o * 4 + v_k_i],
|
||||
)
|
||||
T.writes(compute_local[v_i, v_j])
|
||||
T.sblock_attr(
|
||||
{"meta_schedule.tiling_structure": "SSSRRSRS"}
|
||||
)
|
||||
compute_local[v_i, v_j] = compute_local[v_i, v_j] + T.Cast(
|
||||
"int32", X_shared[v_i, v_k_o * 4 + v_k_i]
|
||||
) * T.Cast("int32", W_shared[v_j, v_k_o * 4 + v_k_i])
|
||||
for ax0, ax1 in T.grid(4, 4):
|
||||
with T.sblock("compute_local"):
|
||||
v0 = T.axis.spatial(
|
||||
128, i_1_j_1_fused // 32 * 8 + i_2_j_2_fused * 4 + ax0
|
||||
)
|
||||
v1 = T.axis.spatial(128, i_1_j_1_fused % 32 * 4 + ax1)
|
||||
T.reads(compute_local[v0, v1])
|
||||
T.writes(compute[v0, v1])
|
||||
compute[v0, v1] = compute_local[v0, v1]
|
||||
|
||||
decision_0 = [
|
||||
("SamplePerfectTile", [1, 16, 2, 2, 2]),
|
||||
("SamplePerfectTile", [1, 32, 1, 4, 1]),
|
||||
("SamplePerfectTile", [1, 1, 32]),
|
||||
("SampleCategorical", 0),
|
||||
("SampleCategorical", 0),
|
||||
]
|
||||
_check_dp4a_dense(
|
||||
m=128,
|
||||
n=128,
|
||||
k=128,
|
||||
in_dtype="int8",
|
||||
out_dtype="int32",
|
||||
expected_mods=[dp4a_dense_0],
|
||||
expected_decisions=[decision_0],
|
||||
)
|
||||
|
||||
|
||||
def test_dp4a_dense_no_tensorize_1():
|
||||
_check_dp4a_dense(
|
||||
m=128,
|
||||
n=128,
|
||||
k=128,
|
||||
in_dtype="float32",
|
||||
out_dtype="float32",
|
||||
expected_mods=None,
|
||||
expected_decisions=None,
|
||||
)
|
||||
|
||||
|
||||
def test_dp4a_dense_no_tensorize_2():
|
||||
_check_dp4a_dense(
|
||||
m=127,
|
||||
n=127,
|
||||
k=127,
|
||||
in_dtype="int8",
|
||||
out_dtype="int32",
|
||||
expected_mods=None,
|
||||
expected_decisions=None,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_x86_conv2d_nchwc()
|
||||
test_x86_conv2d_nchwc(AVX512_INTRIN, {"kind": "llvm", "mcpu": "skylake-avx512", "num-cores": 4})
|
||||
test_dp4a_dense()
|
||||
test_dp4a_dense_no_tensorize_1()
|
||||
test_dp4a_dense_no_tensorize_2()
|
||||
File diff suppressed because it is too large
Load Diff
+304
@@ -0,0 +1,304 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
# ruff: noqa: E501
|
||||
import tvm
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing.space_generation import (
|
||||
check_sketches,
|
||||
generate_design_space,
|
||||
)
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# fmt: off
|
||||
# pylint: disable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class ParallelizeVectorizeUnroll:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
with T.sblock("root"):
|
||||
T.reads([])
|
||||
T.writes([])
|
||||
T.sblock_attr({"meta_schedule.parallel": 128, "meta_schedule.vectorize": 16, "meta_schedule.unroll_explicit": 2})
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
# from tvm.script import tirx as T
|
||||
@tvm.script.ir_module
|
||||
class PureSpatial:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(placeholder: T.Buffer((1, 13, 13, 3, 85), "float32"), placeholder_1: T.Buffer((1, 26, 26, 3, 85), "float32"), placeholder_2: T.Buffer((1, 52, 52, 3, 85), "float32"), T_expand_dims: T.Buffer((1, 80, 10647), "float32")) -> None:
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
T_strided_slice_with_axes = T.sblock_alloc_buffer([1, 52, 52, 3, 1], dtype="float32")
|
||||
T_sigmoid = T.sblock_alloc_buffer([1, 52, 52, 3, 1], dtype="float32")
|
||||
T_strided_slice_with_axes_1 = T.sblock_alloc_buffer([1, 52, 52, 3, 80], dtype="float32")
|
||||
T_sigmoid_1 = T.sblock_alloc_buffer([1, 52, 52, 3, 80], dtype="float32")
|
||||
T_multiply = T.sblock_alloc_buffer([1, 52, 52, 3, 80], dtype="float32")
|
||||
T_reshape = T.sblock_alloc_buffer([8112, 80], dtype="float32")
|
||||
T_strided_slice_with_axes_2 = T.sblock_alloc_buffer([1, 26, 26, 3, 1], dtype="float32")
|
||||
T_sigmoid_2 = T.sblock_alloc_buffer([1, 26, 26, 3, 1], dtype="float32")
|
||||
T_strided_slice_with_axes_3 = T.sblock_alloc_buffer([1, 26, 26, 3, 80], dtype="float32")
|
||||
T_sigmoid_3 = T.sblock_alloc_buffer([1, 26, 26, 3, 80], dtype="float32")
|
||||
T_multiply_1 = T.sblock_alloc_buffer([1, 26, 26, 3, 80], dtype="float32")
|
||||
T_reshape_1 = T.sblock_alloc_buffer([2028, 80], dtype="float32")
|
||||
T_strided_slice_with_axes_4 = T.sblock_alloc_buffer([1, 13, 13, 3, 1], dtype="float32")
|
||||
T_sigmoid_4 = T.sblock_alloc_buffer([1, 13, 13, 3, 1], dtype="float32")
|
||||
T_strided_slice_with_axes_5 = T.sblock_alloc_buffer([1, 13, 13, 3, 80], dtype="float32")
|
||||
T_sigmoid_5 = T.sblock_alloc_buffer([1, 13, 13, 3, 80], dtype="float32")
|
||||
T_multiply_2 = T.sblock_alloc_buffer([1, 13, 13, 3, 80], dtype="float32")
|
||||
T_reshape_2 = T.sblock_alloc_buffer([507, 80], dtype="float32")
|
||||
T_concat = T.sblock_alloc_buffer([10647, 80], dtype="float32")
|
||||
T_transpose = T.sblock_alloc_buffer([80, 10647], dtype="float32")
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 52, 52, 3, 1):
|
||||
with T.sblock("T_strided_slice_with_axes"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(placeholder_2[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(4)])
|
||||
T.writes(T_strided_slice_with_axes[ax0, ax1, ax2, ax3, ax4])
|
||||
T_strided_slice_with_axes[ax0, ax1, ax2, ax3, ax4] = placeholder_2[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(4)]
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 52, 52, 3, 1):
|
||||
with T.sblock("T_sigmoid"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(T_strided_slice_with_axes[ax0, ax1, ax2, ax3, ax4])
|
||||
T.writes(T_sigmoid[ax0, ax1, ax2, ax3, ax4])
|
||||
T_sigmoid[ax0, ax1, ax2, ax3, ax4] = T.sigmoid(T_strided_slice_with_axes[ax0, ax1, ax2, ax3, ax4], dtype="float32")
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 52, 52, 3, 80):
|
||||
with T.sblock("T_strided_slice_with_axes_1"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(placeholder_2[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(5)])
|
||||
T.writes(T_strided_slice_with_axes_1[ax0, ax1, ax2, ax3, ax4])
|
||||
T_strided_slice_with_axes_1[ax0, ax1, ax2, ax3, ax4] = placeholder_2[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(5)]
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 52, 52, 3, 80):
|
||||
with T.sblock("T_sigmoid_1"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(T_strided_slice_with_axes_1[ax0, ax1, ax2, ax3, ax4])
|
||||
T.writes(T_sigmoid_1[ax0, ax1, ax2, ax3, ax4])
|
||||
T_sigmoid_1[ax0, ax1, ax2, ax3, ax4] = T.sigmoid(T_strided_slice_with_axes_1[ax0, ax1, ax2, ax3, ax4], dtype="float32")
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 52, 52, 3, 80):
|
||||
with T.sblock("T_multiply"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(T_sigmoid[ax0, ax1, ax2, ax3, 0], T_sigmoid_1[ax0, ax1, ax2, ax3, ax4])
|
||||
T.writes(T_multiply[ax0, ax1, ax2, ax3, ax4])
|
||||
T_multiply[ax0, ax1, ax2, ax3, ax4] = T_sigmoid[ax0, ax1, ax2, ax3, 0] * T_sigmoid_1[ax0, ax1, ax2, ax3, ax4]
|
||||
for i0, i1 in T.grid(8112, 80):
|
||||
with T.sblock("T_reshape"):
|
||||
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(T_multiply[0, (ax1 // 80 + ax0) % 8112 // 156, (ax1 // 80 + ax0) % 156 // 3, (ax1 // 80 + ax0) % 3, ax1 % 80])
|
||||
T.writes(T_reshape[ax0, ax1])
|
||||
T_reshape[ax0, ax1] = T_multiply[0, (ax1 // 80 + ax0) % 8112 // 156, (ax1 // 80 + ax0) % 156 // 3, (ax1 // 80 + ax0) % 3, ax1 % 80]
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 26, 26, 3, 1):
|
||||
with T.sblock("T_strided_slice_with_axes_2"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(placeholder_1[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(4)])
|
||||
T.writes(T_strided_slice_with_axes_2[ax0, ax1, ax2, ax3, ax4])
|
||||
T_strided_slice_with_axes_2[ax0, ax1, ax2, ax3, ax4] = placeholder_1[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(4)]
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 26, 26, 3, 1):
|
||||
with T.sblock("T_sigmoid_2"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(T_strided_slice_with_axes_2[ax0, ax1, ax2, ax3, ax4])
|
||||
T.writes(T_sigmoid_2[ax0, ax1, ax2, ax3, ax4])
|
||||
T_sigmoid_2[ax0, ax1, ax2, ax3, ax4] = T.sigmoid(T_strided_slice_with_axes_2[ax0, ax1, ax2, ax3, ax4], dtype="float32")
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 26, 26, 3, 80):
|
||||
with T.sblock("T_strided_slice_with_axes_3"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(placeholder_1[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(5)])
|
||||
T.writes(T_strided_slice_with_axes_3[ax0, ax1, ax2, ax3, ax4])
|
||||
T_strided_slice_with_axes_3[ax0, ax1, ax2, ax3, ax4] = placeholder_1[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(5)]
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 26, 26, 3, 80):
|
||||
with T.sblock("T_sigmoid_3"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(T_strided_slice_with_axes_3[ax0, ax1, ax2, ax3, ax4])
|
||||
T.writes(T_sigmoid_3[ax0, ax1, ax2, ax3, ax4])
|
||||
T_sigmoid_3[ax0, ax1, ax2, ax3, ax4] = T.sigmoid(T_strided_slice_with_axes_3[ax0, ax1, ax2, ax3, ax4], dtype="float32")
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 26, 26, 3, 80):
|
||||
with T.sblock("T_multiply_1"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(T_sigmoid_2[ax0, ax1, ax2, ax3, 0], T_sigmoid_3[ax0, ax1, ax2, ax3, ax4])
|
||||
T.writes(T_multiply_1[ax0, ax1, ax2, ax3, ax4])
|
||||
T_multiply_1[ax0, ax1, ax2, ax3, ax4] = T_sigmoid_2[ax0, ax1, ax2, ax3, 0] * T_sigmoid_3[ax0, ax1, ax2, ax3, ax4]
|
||||
for i0, i1 in T.grid(2028, 80):
|
||||
with T.sblock("T_reshape_1"):
|
||||
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(T_multiply_1[0, (ax1 // 80 + ax0) % 2028 // 78, (ax1 // 80 + ax0) % 78 // 3, (ax1 // 80 + ax0) % 3, ax1 % 80])
|
||||
T.writes(T_reshape_1[ax0, ax1])
|
||||
T_reshape_1[ax0, ax1] = T_multiply_1[0, (ax1 // 80 + ax0) % 2028 // 78, (ax1 // 80 + ax0) % 78 // 3, (ax1 // 80 + ax0) % 3, ax1 % 80]
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 13, 13, 3, 1):
|
||||
with T.sblock("T_strided_slice_with_axes_4"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(placeholder[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(4)])
|
||||
T.writes(T_strided_slice_with_axes_4[ax0, ax1, ax2, ax3, ax4])
|
||||
T_strided_slice_with_axes_4[ax0, ax1, ax2, ax3, ax4] = placeholder[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(4)]
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 13, 13, 3, 1):
|
||||
with T.sblock("T_sigmoid_4"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(T_strided_slice_with_axes_4[ax0, ax1, ax2, ax3, ax4])
|
||||
T.writes(T_sigmoid_4[ax0, ax1, ax2, ax3, ax4])
|
||||
T_sigmoid_4[ax0, ax1, ax2, ax3, ax4] = T.sigmoid(T_strided_slice_with_axes_4[ax0, ax1, ax2, ax3, ax4], dtype="float32")
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 13, 13, 3, 80):
|
||||
with T.sblock("T_strided_slice_with_axes_5"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(placeholder[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(5)])
|
||||
T.writes(T_strided_slice_with_axes_5[ax0, ax1, ax2, ax3, ax4])
|
||||
T_strided_slice_with_axes_5[ax0, ax1, ax2, ax3, ax4] = placeholder[ax0, ax1, ax2, ax3, T.cast(ax4, "int64") + T.int64(5)]
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 13, 13, 3, 80):
|
||||
with T.sblock("T_sigmoid_5"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(T_strided_slice_with_axes_5[ax0, ax1, ax2, ax3, ax4])
|
||||
T.writes(T_sigmoid_5[ax0, ax1, ax2, ax3, ax4])
|
||||
T_sigmoid_5[ax0, ax1, ax2, ax3, ax4] = T.sigmoid(T_strided_slice_with_axes_5[ax0, ax1, ax2, ax3, ax4], dtype="float32")
|
||||
for i0, i1, i2, i3, i4 in T.grid(1, 13, 13, 3, 80):
|
||||
with T.sblock("T_multiply_2"):
|
||||
ax0, ax1, ax2, ax3, ax4 = T.axis.remap("SSSSS", [i0, i1, i2, i3, i4])
|
||||
T.reads(T_sigmoid_4[ax0, ax1, ax2, ax3, 0], T_sigmoid_5[ax0, ax1, ax2, ax3, ax4])
|
||||
T.writes(T_multiply_2[ax0, ax1, ax2, ax3, ax4])
|
||||
T_multiply_2[ax0, ax1, ax2, ax3, ax4] = T_sigmoid_4[ax0, ax1, ax2, ax3, 0] * T_sigmoid_5[ax0, ax1, ax2, ax3, ax4]
|
||||
for i0, i1 in T.grid(507, 80):
|
||||
with T.sblock("T_reshape_2"):
|
||||
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(T_multiply_2[0, (ax1 // 80 + ax0) % 507 // 39, (ax1 // 80 + ax0) % 39 // 3, (ax1 // 80 + ax0) % 3, ax1 % 80])
|
||||
T.writes(T_reshape_2[ax0, ax1])
|
||||
T_reshape_2[ax0, ax1] = T_multiply_2[0, (ax1 // 80 + ax0) % 507 // 39, (ax1 // 80 + ax0) % 39 // 3, (ax1 // 80 + ax0) % 3, ax1 % 80]
|
||||
for i0, i1 in T.grid(10647, 80):
|
||||
with T.sblock("T_concat"):
|
||||
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(T_reshape[ax0 - 2535, ax1], T_reshape_1[ax0 - 507, ax1], T_reshape_2[ax0, ax1])
|
||||
T.writes(T_concat[ax0, ax1])
|
||||
T_concat[ax0, ax1] = T.if_then_else(2535 <= ax0, T_reshape[ax0 - 2535, ax1], T.if_then_else(507 <= ax0, T_reshape_1[ax0 - 507, ax1], T_reshape_2[ax0, ax1], dtype="float32"), dtype="float32")
|
||||
for i0, i1 in T.grid(80, 10647):
|
||||
with T.sblock("T_transpose"):
|
||||
ax0, ax1 = T.axis.remap("SS", [i0, i1])
|
||||
T.reads(T_concat[ax1, ax0])
|
||||
T.writes(T_transpose[ax0, ax1])
|
||||
T_transpose[ax0, ax1] = T_concat[ax1, ax0]
|
||||
for i0, i1, i2 in T.grid(1, 80, 10647):
|
||||
with T.sblock("T_expand_dims"):
|
||||
ax0, ax1, ax2 = T.axis.remap("SSS", [i0, i1, i2])
|
||||
T.reads(T_transpose[ax1, ax2])
|
||||
T.writes(T_expand_dims[ax0, ax1, ax2])
|
||||
T_expand_dims[ax0, ax1, ax2] = T_transpose[ax1, ax2]
|
||||
|
||||
|
||||
# pylint: enable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
# fmt: on
|
||||
|
||||
|
||||
def test_parallel_vectorize_unroll():
|
||||
@T.prim_func(s_tir=True)
|
||||
def Matmul_0(
|
||||
A: T.Buffer((1024, 1024), "float32"),
|
||||
B: T.Buffer((1024, 1024), "float32"),
|
||||
C: T.Buffer((1024, 1024), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
# body
|
||||
with T.sblock("root"):
|
||||
T.reads()
|
||||
T.writes()
|
||||
T.sblock_attr(
|
||||
{
|
||||
"meta_schedule.parallel": 512,
|
||||
"meta_schedule.unroll_explicit": 16,
|
||||
"meta_schedule.vectorize": 32,
|
||||
}
|
||||
)
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
T.reads(A[vi, vk], B[vk, vj])
|
||||
T.writes(C[vi, vj])
|
||||
with T.init():
|
||||
C[vi, vj] = T.float32(0)
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
decision_0 = [
|
||||
("SampleCategorical", 1),
|
||||
]
|
||||
|
||||
mod = Matmul
|
||||
actual = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=Target({"kind": "llvm", "num-cores": 32}),
|
||||
types=None,
|
||||
sch_rules=[
|
||||
ms.schedule_rule.ParallelizeVectorizeUnroll(
|
||||
max_jobs_per_core=16,
|
||||
max_vectorize_extent=32,
|
||||
unroll_max_steps=[0, 16, 64, 512],
|
||||
unroll_explicit=True,
|
||||
),
|
||||
],
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[Matmul_0],
|
||||
expected_decisions=[decision_0],
|
||||
)
|
||||
|
||||
|
||||
def test_parallel_vectorize_unroll_spatial():
|
||||
mod = PureSpatial
|
||||
actual = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=Target({"kind": "llvm", "num-cores": 32}),
|
||||
types=None,
|
||||
sch_rules=[
|
||||
ms.schedule_rule.ParallelizeVectorizeUnroll(
|
||||
max_jobs_per_core=-1,
|
||||
max_vectorize_extent=-1,
|
||||
unroll_max_steps=[0, 16, 64, 512],
|
||||
unroll_explicit=True,
|
||||
),
|
||||
],
|
||||
)
|
||||
assert len(actual) == 1
|
||||
trace = actual[0].trace.simplified(remove_postproc=True)
|
||||
assert not trace.insts
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_parallel_vectorize_unroll()
|
||||
test_parallel_vectorize_unroll_spatial()
|
||||
+109
@@ -0,0 +1,109 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
|
||||
import tvm
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing.space_generation import (
|
||||
check_sketches,
|
||||
generate_design_space,
|
||||
)
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# fmt: off
|
||||
# pylint: disable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Add:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, [2048, 2048, 2048], dtype="float32")
|
||||
B = T.match_buffer(b, [2048, 2048, 2048], dtype="float32")
|
||||
A_cached = T.sblock_alloc_buffer([2048, 2048, 2048], dtype="float32")
|
||||
# body
|
||||
for i, j, k in T.grid(2048, 2048, 2048):
|
||||
with T.sblock("move"):
|
||||
vi, vj, vk = T.axis.remap("SSS", [i, j, k])
|
||||
T.reads([A[vi, vj, vk]])
|
||||
T.writes([A_cached[vi, vj, vk]])
|
||||
A_cached[vi, vj, vk] = A[vi, vj, vk]
|
||||
for i0, j0, i1, j1, k0, i2, j2, k1 in T.grid(128, 64, 4, 4, 64, 4, 8, 32):
|
||||
with T.sblock("add"):
|
||||
vi = T.axis.spatial(2048, i0 * 16 + i1 * 4 + i2)
|
||||
vj = T.axis.spatial(2048, j0 * 32 + j1 * 8 + j2)
|
||||
vk = T.axis.spatial(2048, k0 * 32 + k1)
|
||||
T.reads([A_cached[vi, vj, vk]])
|
||||
T.writes([B[vi, vj, vk]])
|
||||
B[vi, vj, vk] = A_cached[vi, vj, vk] + T.float32(1)
|
||||
|
||||
# pylint: enable=no-member,invalid-name,unused-variable,no-self-argument,line-too-long,chained-comparison,not-callable,too-many-nested-blocks
|
||||
# fmt: on
|
||||
|
||||
|
||||
def test_random_compute_location():
|
||||
@T.prim_func(s_tir=True)
|
||||
def add_0(
|
||||
A: T.Buffer((2048, 2048, 2048), "float32"),
|
||||
B: T.Buffer((2048, 2048, 2048), "float32"),
|
||||
) -> None:
|
||||
# function attr dict
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
# body
|
||||
# with T.sblock("root")
|
||||
A_cached = T.sblock_alloc_buffer([2048, 2048, 2048], dtype="float32")
|
||||
for i0, j0, i1, j1, k0, i2 in T.grid(128, 64, 4, 4, 64, 4):
|
||||
for ax0, ax1, ax2 in T.grid(1, 8, 32):
|
||||
with T.sblock("move"):
|
||||
vi = T.axis.spatial(2048, i0 * 16 + i1 * 4 + i2 + ax0)
|
||||
vj = T.axis.spatial(2048, j0 * 32 + j1 * 8 + ax1)
|
||||
vk = T.axis.spatial(2048, k0 * 32 + ax2)
|
||||
T.reads(A[vi, vj, vk])
|
||||
T.writes(A_cached[vi, vj, vk])
|
||||
A_cached[vi, vj, vk] = A[vi, vj, vk]
|
||||
for j2, k1 in T.grid(8, 32):
|
||||
with T.sblock("add"):
|
||||
vi = T.axis.spatial(2048, i0 * 16 + i1 * 4 + i2)
|
||||
vj = T.axis.spatial(2048, j0 * 32 + j1 * 8 + j2)
|
||||
vk = T.axis.spatial(2048, k0 * 32 + k1)
|
||||
T.reads(A_cached[vi, vj, vk])
|
||||
T.writes(B[vi, vj, vk])
|
||||
B[vi, vj, vk] = A_cached[vi, vj, vk] + T.float32(1)
|
||||
|
||||
decision_0 = [
|
||||
("SampleComputeLocation", 5),
|
||||
]
|
||||
|
||||
mod = Add
|
||||
actual = generate_design_space(
|
||||
kind="llvm",
|
||||
mod=mod,
|
||||
target=Target("llvm"),
|
||||
types=None,
|
||||
sch_rules=[ms.schedule_rule.RandomComputeLocation()],
|
||||
)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[add_0],
|
||||
expected_decisions=[decision_0],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_random_compute_location()
|
||||
@@ -0,0 +1,420 @@
|
||||
# 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.
|
||||
# ruff: noqa: E501
|
||||
"""Test Meta Schedule SearchStrategy"""
|
||||
|
||||
# pylint: disable=missing-function-docstring
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing.dummy_object import DummyMutator
|
||||
from tvm.s_tir.schedule import Schedule, Trace
|
||||
from tvm.script import tirx as T
|
||||
|
||||
MATMUL_M = 32
|
||||
|
||||
# pylint: disable=missing-class-docstring,invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument, unbalanced-tuple-unpacking
|
||||
# fmt: off
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None: # type: ignore
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (32, 32), "float32")
|
||||
B = T.match_buffer(b, (32, 32), "float32")
|
||||
C = T.match_buffer(c, (32, 32), "float32")
|
||||
for i, j, k in T.grid(32, 32, 32):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0 # type: ignore
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class OtherBlock:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None: # type: ignore
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (32, 32), "float32")
|
||||
B = T.match_buffer(b, (32, 32), "float32")
|
||||
C = T.match_buffer(c, (32, 32), "float32")
|
||||
for i, j, k in T.grid(32, 32, 32):
|
||||
with T.sblock("other"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0 # type: ignore
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
# fmt: on
|
||||
# pylint: enable=missing-class-docstring,invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
|
||||
|
||||
def _is_trace_equal(sch_1: Schedule, sch_2: Schedule, remove_decisions=True) -> bool:
|
||||
if remove_decisions:
|
||||
trace_1 = Trace(sch_1.trace.insts, {})
|
||||
trace_2 = Trace(sch_2.trace.insts, {})
|
||||
else:
|
||||
trace_1 = sch_1.trace
|
||||
trace_2 = sch_2.trace
|
||||
return str(trace_1) == str(trace_2)
|
||||
|
||||
|
||||
def _schedule_matmul(sch: Schedule):
|
||||
block = sch.get_sblock("matmul")
|
||||
i, j, k = sch.get_loops(block=block)
|
||||
i_0, i_1, i_2, i_3 = sch.split(i, sch.sample_perfect_tile(i, n=4))
|
||||
j_0, j_1, j_2, j_3 = sch.split(j, sch.sample_perfect_tile(j, n=4))
|
||||
k_0, k_1 = sch.split(k, sch.sample_perfect_tile(k, n=2))
|
||||
sch.reorder(i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"TestClass",
|
||||
[
|
||||
ms.search_strategy.ReplayFunc,
|
||||
ms.search_strategy.ReplayTrace,
|
||||
],
|
||||
)
|
||||
def test_meta_schedule_replay_func(
|
||||
TestClass: ms.search_strategy.SearchStrategy,
|
||||
): # pylint: disable = invalid-name
|
||||
num_trials_per_iter = 7
|
||||
max_trials_per_task = 20
|
||||
|
||||
context = ms.TuneContext(
|
||||
mod=Matmul,
|
||||
num_threads=1,
|
||||
space_generator=ms.space_generator.ScheduleFn(sch_fn=_schedule_matmul, postprocs=[]),
|
||||
search_strategy=TestClass(),
|
||||
)
|
||||
strategy = context.search_strategy
|
||||
spaces = context.space_generator.generate_design_space(context.mod)
|
||||
strategy.pre_tuning(
|
||||
max_trials=max_trials_per_task,
|
||||
num_trials_per_iter=num_trials_per_iter,
|
||||
design_spaces=spaces,
|
||||
)
|
||||
(correct_sch,) = ms.space_generator.ScheduleFn(sch_fn=_schedule_matmul).generate_design_space(
|
||||
Matmul
|
||||
)
|
||||
num_trials_each_iter: list[int] = []
|
||||
candidates = strategy.generate_measure_candidates()
|
||||
while candidates is not None:
|
||||
num_trials_each_iter.append(len(candidates))
|
||||
runner_results: list[ms.runner.RunnerResult] = []
|
||||
for candidate in candidates:
|
||||
_is_trace_equal(
|
||||
candidate.sch,
|
||||
correct_sch,
|
||||
remove_decisions=(isinstance(strategy, ms.search_strategy.ReplayTrace)),
|
||||
)
|
||||
runner_results.append(
|
||||
ms.runner.RunnerResult(
|
||||
run_secs=[0.11, 0.41, 0.54],
|
||||
error_msg=None,
|
||||
)
|
||||
)
|
||||
strategy.notify_runner_results(candidates, runner_results)
|
||||
candidates = strategy.generate_measure_candidates()
|
||||
strategy.post_tuning()
|
||||
assert num_trials_each_iter == [7, 7, 6]
|
||||
|
||||
|
||||
def test_meta_schedule_evolutionary_search(): # pylint: disable = invalid-name
|
||||
def _schedule_matmul_small(sch: Schedule):
|
||||
block = sch.get_sblock("matmul")
|
||||
_, j, k = sch.get_loops(block=block)
|
||||
_, _ = sch.split(j, sch.sample_perfect_tile(j, n=2))
|
||||
_, _ = sch.split(k, sch.sample_perfect_tile(k, n=2))
|
||||
|
||||
num_trials_per_iter = 10
|
||||
max_trials_per_task = 2000
|
||||
(correct_sch,) = ms.space_generator.ScheduleFn(sch_fn=_schedule_matmul).generate_design_space(
|
||||
Matmul
|
||||
)
|
||||
|
||||
context = ms.TuneContext(
|
||||
mod=Matmul,
|
||||
space_generator=ms.space_generator.ScheduleFn(
|
||||
sch_fn=_schedule_matmul_small,
|
||||
sch_rules=[],
|
||||
postprocs=[],
|
||||
mutator_probs={
|
||||
DummyMutator(): 1.0,
|
||||
},
|
||||
),
|
||||
search_strategy=ms.search_strategy.EvolutionarySearch(
|
||||
population_size=5,
|
||||
init_measured_ratio=0.1,
|
||||
init_min_unmeasured=50,
|
||||
genetic_num_iters=3,
|
||||
genetic_mutate_prob=0.5,
|
||||
genetic_max_fail_count=10,
|
||||
eps_greedy=0.9,
|
||||
),
|
||||
target=tvm.target.Target("llvm"),
|
||||
num_threads=1, # because we are using a mutator from the python side
|
||||
)
|
||||
strategy = context.search_strategy
|
||||
strategy.pre_tuning(
|
||||
max_trials=max_trials_per_task,
|
||||
num_trials_per_iter=num_trials_per_iter,
|
||||
design_spaces=context.space_generator.generate_design_space(context.mod),
|
||||
database=ms.database.MemoryDatabase(),
|
||||
cost_model=ms.cost_model.RandomModel(),
|
||||
)
|
||||
num_trials_each_iter: list[int] = []
|
||||
candidates = strategy.generate_measure_candidates()
|
||||
while candidates is not None:
|
||||
num_trials_each_iter.append(len(candidates))
|
||||
runner_results: list[ms.runner.RunnerResult] = []
|
||||
for candidate in candidates:
|
||||
_is_trace_equal(
|
||||
candidate.sch,
|
||||
correct_sch,
|
||||
remove_decisions=(isinstance(strategy, ms.search_strategy.ReplayTrace)),
|
||||
)
|
||||
runner_results.append(
|
||||
ms.runner.RunnerResult(
|
||||
run_secs=[0.11, 0.41, 0.54],
|
||||
error_msg=None,
|
||||
)
|
||||
)
|
||||
strategy.notify_runner_results(candidates, runner_results)
|
||||
candidates = strategy.generate_measure_candidates()
|
||||
strategy.post_tuning()
|
||||
assert sum(num_trials_each_iter) == 25
|
||||
assert num_trials_each_iter.count(0) < 5
|
||||
|
||||
|
||||
def test_meta_schedule_evolutionary_search_early_stop(): # pylint: disable = invalid-name
|
||||
def _schedule_matmul_empty(sch: Schedule):
|
||||
return sch
|
||||
|
||||
(correct_sch,) = ms.space_generator.ScheduleFn(sch_fn=_schedule_matmul).generate_design_space(
|
||||
Matmul
|
||||
)
|
||||
|
||||
num_trials_per_iter = 10
|
||||
max_trials_per_task = 100
|
||||
|
||||
context = ms.TuneContext(
|
||||
mod=Matmul,
|
||||
search_strategy=ms.search_strategy.EvolutionarySearch(
|
||||
population_size=5,
|
||||
init_measured_ratio=0.1,
|
||||
init_min_unmeasured=50,
|
||||
genetic_num_iters=3,
|
||||
genetic_mutate_prob=0.5,
|
||||
genetic_max_fail_count=10,
|
||||
eps_greedy=0.9,
|
||||
),
|
||||
space_generator=ms.space_generator.ScheduleFn(
|
||||
sch_fn=_schedule_matmul_empty,
|
||||
sch_rules=[],
|
||||
postprocs=[],
|
||||
mutator_probs={
|
||||
DummyMutator(): 1.0,
|
||||
},
|
||||
),
|
||||
target=tvm.target.Target("llvm"),
|
||||
num_threads=1,
|
||||
)
|
||||
strategy = context.search_strategy
|
||||
strategy.pre_tuning(
|
||||
max_trials=max_trials_per_task,
|
||||
num_trials_per_iter=num_trials_per_iter,
|
||||
design_spaces=context.space_generator.generate_design_space(context.mod),
|
||||
database=ms.database.MemoryDatabase(),
|
||||
cost_model=ms.cost_model.RandomModel(),
|
||||
)
|
||||
num_trials_each_iter: list[int] = []
|
||||
candidates = strategy.generate_measure_candidates()
|
||||
while candidates is not None:
|
||||
num_trials_each_iter.append(len(candidates))
|
||||
runner_results: list[ms.runner.RunnerResult] = []
|
||||
for candidate in candidates:
|
||||
_is_trace_equal(
|
||||
candidate.sch,
|
||||
correct_sch,
|
||||
remove_decisions=(isinstance(strategy, ms.search_strategy.ReplayTrace)),
|
||||
)
|
||||
runner_results.append(
|
||||
ms.runner.RunnerResult(
|
||||
run_secs=[0.11, 0.41, 0.54],
|
||||
error_msg=None,
|
||||
),
|
||||
)
|
||||
strategy.notify_runner_results(candidates, runner_results)
|
||||
candidates = strategy.generate_measure_candidates()
|
||||
strategy.post_tuning()
|
||||
assert num_trials_each_iter == [1, 0, 0, 0, 0]
|
||||
|
||||
|
||||
def test_meta_schedule_evolutionary_search_fail_init_population(): # pylint: disable = invalid-name
|
||||
@derived_object
|
||||
class AlwaysFailPostproc(ms.postproc.PyPostproc):
|
||||
"""A postproc that always fails."""
|
||||
|
||||
def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
|
||||
pass
|
||||
|
||||
def apply(self, sch: Schedule) -> bool:
|
||||
return False
|
||||
|
||||
def clone(self) -> "AlwaysFailPostproc":
|
||||
return AlwaysFailPostproc()
|
||||
|
||||
def __str__(self) -> str:
|
||||
return "AlwaysFailPostproc"
|
||||
|
||||
num_trials_per_iter = 10
|
||||
max_trials_per_task = 2000
|
||||
|
||||
context = ms.TuneContext(
|
||||
mod=Matmul,
|
||||
space_generator=ms.space_generator.ScheduleFn(
|
||||
sch_fn=_schedule_matmul,
|
||||
sch_rules=[],
|
||||
postprocs=[AlwaysFailPostproc()],
|
||||
mutator_probs={
|
||||
DummyMutator(): 1.0,
|
||||
},
|
||||
),
|
||||
search_strategy=ms.search_strategy.EvolutionarySearch(
|
||||
population_size=5,
|
||||
init_measured_ratio=0.1,
|
||||
init_min_unmeasured=50,
|
||||
genetic_num_iters=3,
|
||||
genetic_mutate_prob=0.5,
|
||||
genetic_max_fail_count=10,
|
||||
eps_greedy=0.9,
|
||||
),
|
||||
target=tvm.target.Target("llvm"),
|
||||
num_threads=1, # because we are using a mutator from the python side
|
||||
)
|
||||
strategy = context.search_strategy
|
||||
strategy.pre_tuning(
|
||||
max_trials=max_trials_per_task,
|
||||
num_trials_per_iter=num_trials_per_iter,
|
||||
design_spaces=context.space_generator.generate_design_space(context.mod),
|
||||
database=ms.database.MemoryDatabase(),
|
||||
cost_model=ms.cost_model.RandomModel(),
|
||||
)
|
||||
candidates = strategy.generate_measure_candidates()
|
||||
assert candidates is None
|
||||
|
||||
|
||||
def test_meta_schedule_evolutionary_search_skip_invalid_measured_trace(): # pylint: disable = invalid-name
|
||||
# Construct an incompatible measured trace: it references block name "other",
|
||||
# which doesn't exist in Matmul. Replaying this trace should fail and be skipped.
|
||||
wrong_sch = Schedule(OtherBlock)
|
||||
wrong_sch.get_sblock("other")
|
||||
wrong_trace = wrong_sch.trace
|
||||
|
||||
database = ms.database.MemoryDatabase()
|
||||
workload = database.commit_workload(Matmul)
|
||||
database.commit_tuning_record(
|
||||
ms.database.TuningRecord(
|
||||
trace=wrong_trace,
|
||||
workload=workload,
|
||||
run_secs=[0.1],
|
||||
target=tvm.target.Target("llvm"),
|
||||
args_info=ms.arg_info.ArgInfo.from_prim_func(func=Matmul["main"]),
|
||||
)
|
||||
)
|
||||
|
||||
context = ms.TuneContext(
|
||||
mod=Matmul,
|
||||
space_generator=ms.space_generator.ScheduleFn(
|
||||
sch_fn=_schedule_matmul,
|
||||
sch_rules=[],
|
||||
postprocs=[],
|
||||
mutator_probs={
|
||||
DummyMutator(): 1.0,
|
||||
},
|
||||
),
|
||||
search_strategy=ms.search_strategy.EvolutionarySearch(
|
||||
population_size=5,
|
||||
init_measured_ratio=1.0,
|
||||
init_min_unmeasured=1,
|
||||
genetic_num_iters=1,
|
||||
genetic_mutate_prob=0.5,
|
||||
genetic_max_fail_count=4,
|
||||
eps_greedy=0.9,
|
||||
),
|
||||
target=tvm.target.Target("llvm"),
|
||||
num_threads=1,
|
||||
)
|
||||
strategy = context.search_strategy
|
||||
strategy.pre_tuning(
|
||||
max_trials=4,
|
||||
num_trials_per_iter=2,
|
||||
design_spaces=context.space_generator.generate_design_space(context.mod),
|
||||
database=database,
|
||||
cost_model=ms.cost_model.RandomModel(),
|
||||
)
|
||||
|
||||
candidates = strategy.generate_measure_candidates()
|
||||
strategy.post_tuning()
|
||||
|
||||
# Regression assertion: invalid measured trace should be skipped, not crash
|
||||
assert candidates is not None
|
||||
|
||||
|
||||
def test_search_strategy_abstract_class_instantiation():
|
||||
"""Test that directly instantiating abstract SearchStrategy raises TypeError instead of segfault."""
|
||||
from tvm.s_tir.meta_schedule import SearchStrategy, TuneContext
|
||||
from tvm.target import Target
|
||||
|
||||
# Test that direct instantiation raises TypeError
|
||||
# This prevents segfault when SearchStrategy() is called directly
|
||||
with pytest.raises(TypeError, match="Cannot instantiate abstract class SearchStrategy"):
|
||||
SearchStrategy()
|
||||
|
||||
# Test that TuneContext with SearchStrategy() raises TypeError
|
||||
# The error should occur when trying to create SearchStrategy() instance in the function call
|
||||
# Since SearchStrategy() fails in __new__, it will fail before TuneContext.__init__ is called
|
||||
with pytest.raises(TypeError, match="Cannot instantiate abstract class SearchStrategy"):
|
||||
# This will fail when evaluating SearchStrategy() as an argument
|
||||
TuneContext(
|
||||
mod=Matmul, # Use the existing Matmul module from the test file
|
||||
target=Target("llvm"),
|
||||
search_strategy=SearchStrategy(), # This should fail in __new__ before reaching TuneContext
|
||||
)
|
||||
|
||||
# Test that SearchStrategy.create() works correctly
|
||||
strategy = SearchStrategy.create("evolutionary")
|
||||
assert strategy is not None
|
||||
assert isinstance(strategy, SearchStrategy)
|
||||
# Verify it's not the abstract class itself
|
||||
assert type(strategy) is not SearchStrategy
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_meta_schedule_replay_func(ms.search_strategy.ReplayFunc)
|
||||
test_meta_schedule_replay_func(ms.search_strategy.ReplayTrace)
|
||||
test_meta_schedule_evolutionary_search()
|
||||
test_meta_schedule_evolutionary_search_early_stop()
|
||||
test_meta_schedule_evolutionary_search_fail_init_population()
|
||||
test_search_strategy_abstract_class_instantiation()
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,339 @@
|
||||
# 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.
|
||||
# ruff: noqa: E501, F401
|
||||
"""Tests for MetaSchedule search space on CUDA"""
|
||||
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing.space_generation import (
|
||||
check_sketches,
|
||||
generate_design_space,
|
||||
print_sketches,
|
||||
)
|
||||
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def _target():
|
||||
return Target("nvidia/geforce-rtx-3070")
|
||||
|
||||
|
||||
def _design_space(mod):
|
||||
return generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=_target(),
|
||||
types=ms.ScheduleRule,
|
||||
)
|
||||
|
||||
|
||||
def get_c2d_prim_func(stage: int):
|
||||
if stage == 0:
|
||||
# fmt: off
|
||||
@T.prim_func(s_tir=True)
|
||||
def c2d(inputs: T.Buffer((1, 224, 224, 3), "float32"), weight: T.Buffer((7, 7, 3, 64), "float32"), conv2d_nhwc: T.Buffer((1, 112, 112, 64), "float32")):
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
with T.sblock("root"):
|
||||
T.reads()
|
||||
T.writes()
|
||||
T.sblock_attr({"meta_schedule.unroll_explicit": 1024})
|
||||
conv2d_nhwc_local = T.sblock_alloc_buffer((1, 112, 112, 64), scope="local")
|
||||
PadInput_shared = T.sblock_alloc_buffer((1, 230, 230, 3), scope="shared")
|
||||
weight_shared = T.sblock_alloc_buffer((7, 7, 3, 64), scope="shared")
|
||||
for n_0_h_0_w_0_co_0_fused in T.thread_binding(112, thread="blockIdx.x"):
|
||||
for n_1_h_1_w_1_co_1_fused in T.thread_binding(8, thread="vthread.x"):
|
||||
for n_2_h_2_w_2_co_2_fused in T.thread_binding(64, thread="threadIdx.x"):
|
||||
for rh_0, rw_0, rc_0 in T.grid(1, 1, 3):
|
||||
for ax0_ax1_ax2_ax3_fused in range(693):
|
||||
with T.sblock("PadInput_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(230, n_0_h_0_w_0_co_0_fused // 8 * 16 + ax0_ax1_ax2_ax3_fused // 33)
|
||||
v2 = T.axis.spatial(230, n_0_h_0_w_0_co_0_fused % 8 * 28 + ax0_ax1_ax2_ax3_fused % 33)
|
||||
v3 = T.axis.spatial(3, rc_0)
|
||||
T.reads(inputs[v0, v1 - 3, v2 - 3, v3])
|
||||
T.writes(PadInput_shared[v0, v1, v2, v3])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 4})
|
||||
PadInput_shared[v0, v1, v2, v3] = T.if_then_else(3 <= v1 and v1 < 227 and 3 <= v2 and v2 < 227, inputs[v0, v1 - 3, v2 - 3, v3], T.float32(0))
|
||||
for ax0_ax1_ax2_ax3_fused in range(3136):
|
||||
with T.sblock("weight_shared"):
|
||||
v0 = T.axis.spatial(7, ax0_ax1_ax2_ax3_fused // 448)
|
||||
v1 = T.axis.spatial(7, ax0_ax1_ax2_ax3_fused % 448 // 64)
|
||||
v2 = T.axis.spatial(3, rc_0)
|
||||
v3 = T.axis.spatial(64, ax0_ax1_ax2_ax3_fused % 64)
|
||||
T.reads(weight[v0, v1, v2, v3])
|
||||
T.writes(weight_shared[v0, v1, v2, v3])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 3})
|
||||
weight_shared[v0, v1, v2, v3] = weight[v0, v1, v2, v3]
|
||||
for rh_1, rw_1, rc_1, n_3, h_3, w_3, co_3, rh_2, rw_2, rc_2, n_4, h_4, w_4, co_4 in T.grid(7, 1, 1, 1, 1, 14, 1, 1, 7, 1, 1, 1, 1, 1):
|
||||
with T.sblock("conv2d_nhwc"):
|
||||
v_n = T.axis.spatial(1, n_3 + n_4)
|
||||
v_h = T.axis.spatial(112, n_0_h_0_w_0_co_0_fused // 8 * 8 + n_1_h_1_w_1_co_1_fused // 4 * 4 + n_2_h_2_w_2_co_2_fused // 16 + h_3 + h_4)
|
||||
v_w = T.axis.spatial(112, n_0_h_0_w_0_co_0_fused % 8 * 14 + w_3 + w_4)
|
||||
v_co = T.axis.spatial(64, n_1_h_1_w_1_co_1_fused % 4 * 16 + n_2_h_2_w_2_co_2_fused % 16 + co_3 + co_4)
|
||||
v_rh = T.axis.reduce(7, rh_0 * 7 + rh_1 + rh_2)
|
||||
v_rw = T.axis.reduce(7, rw_0 * 7 + rw_1 * 7 + rw_2)
|
||||
v_rc = T.axis.reduce(3, rc_0 + rc_1 + rc_2)
|
||||
T.reads(PadInput_shared[v_n, v_h * 2 + v_rh, v_w * 2 + v_rw, v_co // 64 * 3 + v_rc], weight_shared[v_rh, v_rw, v_rc, v_co])
|
||||
T.writes(conv2d_nhwc_local[v_n, v_h, v_w, v_co])
|
||||
T.sblock_attr({"meta_schedule.thread_extent_high_inclusive": 1024, "meta_schedule.thread_extent_low_inclusive": 32, "meta_schedule.tiling_structure": "SSSRRSRS"})
|
||||
with T.init():
|
||||
conv2d_nhwc_local[v_n, v_h, v_w, v_co] = T.float32(0)
|
||||
conv2d_nhwc_local[v_n, v_h, v_w, v_co] = conv2d_nhwc_local[v_n, v_h, v_w, v_co] + PadInput_shared[v_n, v_h * 2 + v_rh, v_w * 2 + v_rw, v_co // 64 * 3 + v_rc] * weight_shared[v_rh, v_rw, v_rc, v_co]
|
||||
for ax0, ax1, ax2, ax3 in T.grid(1, 1, 14, 1):
|
||||
with T.sblock("conv2d_nhwc_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(112, n_0_h_0_w_0_co_0_fused // 8 * 8 + n_1_h_1_w_1_co_1_fused // 4 * 4 + n_2_h_2_w_2_co_2_fused // 16 + ax1)
|
||||
v2 = T.axis.spatial(112, n_0_h_0_w_0_co_0_fused % 8 * 14 + ax2)
|
||||
v3 = T.axis.spatial(64, n_1_h_1_w_1_co_1_fused % 4 * 16 + n_2_h_2_w_2_co_2_fused % 16 + ax3)
|
||||
T.reads(conv2d_nhwc_local[v0, v1, v2, v3])
|
||||
T.writes(conv2d_nhwc[v0, v1, v2, v3])
|
||||
conv2d_nhwc[v0, v1, v2, v3] = conv2d_nhwc_local[v0, v1, v2, v3]
|
||||
# fmt: on
|
||||
else:
|
||||
# fmt: off
|
||||
@T.prim_func(s_tir=True)
|
||||
def c2d(inputs: T.Buffer((1, 224, 224, 3), "float32"), weight: T.Buffer((7, 7, 3, 64), "float32"), conv2d_nhwc: T.Buffer((1, 112, 112, 64), "float32")):
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
with T.sblock("root"):
|
||||
T.reads()
|
||||
T.writes()
|
||||
T.sblock_attr({"meta_schedule.unroll_explicit": 1024})
|
||||
conv2d_nhwc_local = T.sblock_alloc_buffer((1, 112, 112, 64), scope="local")
|
||||
PadInput_shared = T.sblock_alloc_buffer((1, 230, 230, 3), scope="shared")
|
||||
weight_shared = T.sblock_alloc_buffer((7, 7, 3, 64), scope="shared")
|
||||
for n_0_h_0_w_0_co_0_fused in T.thread_binding(112, thread="blockIdx.x"):
|
||||
for n_1_h_1_w_1_co_1_fused in T.thread_binding(8, thread="vthread.x"):
|
||||
for n_2_h_2_w_2_co_2_fused in T.thread_binding(64, thread="threadIdx.x"):
|
||||
for rh_0_rw_0_rc_0_fused in T.serial(3, annotations={"software_pipeline_async_stages": [0], "software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, stage - 2]}):
|
||||
for ax0_ax1_ax2_ax3_fused in range(693):
|
||||
with T.sblock("PadInput_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(230, n_0_h_0_w_0_co_0_fused // 8 * 16 + ax0_ax1_ax2_ax3_fused // 33)
|
||||
v2 = T.axis.spatial(230, n_0_h_0_w_0_co_0_fused % 8 * 28 + ax0_ax1_ax2_ax3_fused % 33)
|
||||
v3 = T.axis.spatial(3, rh_0_rw_0_rc_0_fused)
|
||||
T.reads(inputs[v0, v1 - 3, v2 - 3, v3])
|
||||
T.writes(PadInput_shared[v0, v1, v2, v3])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 4})
|
||||
PadInput_shared[v0, v1, v2, v3] = T.if_then_else(3 <= v1 and v1 < 227 and 3 <= v2 and v2 < 227, inputs[v0, v1 - 3, v2 - 3, v3], T.float32(0))
|
||||
for ax0_ax1_ax2_ax3_fused in range(3136):
|
||||
with T.sblock("weight_shared"):
|
||||
v0 = T.axis.spatial(7, ax0_ax1_ax2_ax3_fused // 448)
|
||||
v1 = T.axis.spatial(7, ax0_ax1_ax2_ax3_fused % 448 // 64)
|
||||
v2 = T.axis.spatial(3, rh_0_rw_0_rc_0_fused)
|
||||
v3 = T.axis.spatial(64, ax0_ax1_ax2_ax3_fused % 64)
|
||||
T.reads(weight[v0, v1, v2, v3])
|
||||
T.writes(weight_shared[v0, v1, v2, v3])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 3})
|
||||
weight_shared[v0, v1, v2, v3] = weight[v0, v1, v2, v3]
|
||||
for rh_1, rw_1, rc_1, n_3, h_3, w_3, co_3, rh_2, rw_2, rc_2, n_4, h_4, w_4, co_4 in T.grid(7, 1, 1, 1, 1, 14, 1, 1, 7, 1, 1, 1, 1, 1):
|
||||
with T.sblock("conv2d_nhwc"):
|
||||
v_n = T.axis.spatial(1, n_3 + n_4)
|
||||
v_h = T.axis.spatial(112, n_0_h_0_w_0_co_0_fused // 8 * 8 + n_1_h_1_w_1_co_1_fused // 4 * 4 + n_2_h_2_w_2_co_2_fused // 16 + h_3 + h_4)
|
||||
v_w = T.axis.spatial(112, n_0_h_0_w_0_co_0_fused % 8 * 14 + w_3 + w_4)
|
||||
v_co = T.axis.spatial(64, n_1_h_1_w_1_co_1_fused % 4 * 16 + n_2_h_2_w_2_co_2_fused % 16 + co_3 + co_4)
|
||||
v_rh = T.axis.reduce(7, rh_1 + rh_2)
|
||||
v_rw = T.axis.reduce(7, rw_1 * 7 + rw_2)
|
||||
v_rc = T.axis.reduce(3, rh_0_rw_0_rc_0_fused + rc_1 + rc_2)
|
||||
T.reads(PadInput_shared[v_n, v_h * 2 + v_rh, v_w * 2 + v_rw, v_co // 64 * 3 + v_rc], weight_shared[v_rh, v_rw, v_rc, v_co])
|
||||
T.writes(conv2d_nhwc_local[v_n, v_h, v_w, v_co])
|
||||
T.sblock_attr({"meta_schedule.thread_extent_high_inclusive": 1024, "meta_schedule.thread_extent_low_inclusive": 32, "meta_schedule.tiling_structure": "SSSRRSRS"})
|
||||
with T.init():
|
||||
conv2d_nhwc_local[v_n, v_h, v_w, v_co] = T.float32(0)
|
||||
conv2d_nhwc_local[v_n, v_h, v_w, v_co] = conv2d_nhwc_local[v_n, v_h, v_w, v_co] + PadInput_shared[v_n, v_h * 2 + v_rh, v_w * 2 + v_rw, v_co // 64 * 3 + v_rc] * weight_shared[v_rh, v_rw, v_rc, v_co]
|
||||
for ax0, ax1, ax2, ax3 in T.grid(1, 1, 14, 1):
|
||||
with T.sblock("conv2d_nhwc_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(112, n_0_h_0_w_0_co_0_fused // 8 * 8 + n_1_h_1_w_1_co_1_fused // 4 * 4 + n_2_h_2_w_2_co_2_fused // 16 + ax1)
|
||||
v2 = T.axis.spatial(112, n_0_h_0_w_0_co_0_fused % 8 * 14 + ax2)
|
||||
v3 = T.axis.spatial(64, n_1_h_1_w_1_co_1_fused % 4 * 16 + n_2_h_2_w_2_co_2_fused % 16 + ax3)
|
||||
T.reads(conv2d_nhwc_local[v0, v1, v2, v3])
|
||||
T.writes(conv2d_nhwc[v0, v1, v2, v3])
|
||||
conv2d_nhwc[v0, v1, v2, v3] = conv2d_nhwc_local[v0, v1, v2, v3]
|
||||
# fmt: on
|
||||
return c2d
|
||||
|
||||
|
||||
def test_cuda_c2d():
|
||||
c2d_decision = [
|
||||
("SamplePerfectTile", [1, 1, 1, 1, 1]),
|
||||
("SamplePerfectTile", [14, 2, 4, 1, 1]),
|
||||
("SamplePerfectTile", [8, 1, 1, 14, 1]),
|
||||
("SamplePerfectTile", [1, 4, 16, 1, 1]),
|
||||
("SamplePerfectTile", [1, 7, 1]),
|
||||
("SamplePerfectTile", [1, 1, 7]),
|
||||
("SamplePerfectTile", [3, 1, 1]),
|
||||
("SampleCategorical", 3),
|
||||
("SampleCategorical", 2),
|
||||
("SampleCategorical", 7),
|
||||
]
|
||||
|
||||
mod = create_te_workload("C2D", 0)
|
||||
actual = _design_space(mod)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[
|
||||
get_c2d_prim_func(stage=0),
|
||||
get_c2d_prim_func(stage=4),
|
||||
get_c2d_prim_func(stage=5),
|
||||
],
|
||||
expected_decisions=[c2d_decision, c2d_decision, c2d_decision],
|
||||
)
|
||||
|
||||
|
||||
def get_gmm_prim_func(stage: int):
|
||||
if stage == 0:
|
||||
# fmt: off
|
||||
@T.prim_func(s_tir=True)
|
||||
def gmm(X: T.Buffer((1, 1024, 1024), "float32"), Y: T.Buffer((1, 1024, 1024), "float32"), Z: T.Buffer((1, 1024, 1024), "float32")):
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
with T.sblock("root"):
|
||||
T.reads()
|
||||
T.writes()
|
||||
T.sblock_attr({"meta_schedule.unroll_explicit": 16})
|
||||
Z_local = T.sblock_alloc_buffer((1, 1024, 1024), scope="local")
|
||||
X_shared = T.sblock_alloc_buffer((1, 1024, 1024), scope="shared")
|
||||
Y_shared = T.sblock_alloc_buffer((1, 1024, 1024), scope="shared")
|
||||
for b_0_i_0_j_0_fused in T.thread_binding(256, thread="blockIdx.x"):
|
||||
for b_1_i_1_j_1_fused in T.thread_binding(32, thread="vthread.x"):
|
||||
for b_2_i_2_j_2_fused in T.thread_binding(64, thread="threadIdx.x"):
|
||||
for k_0 in range(64):
|
||||
for ax0_ax1_ax2_fused in range(1024):
|
||||
with T.sblock("X_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(1024, b_0_i_0_j_0_fused // 16 * 64 + ax0_ax1_ax2_fused // 16)
|
||||
v2 = T.axis.spatial(1024, k_0 * 16 + ax0_ax1_ax2_fused % 16)
|
||||
T.reads(X[v0, v1, v2])
|
||||
T.writes(X_shared[v0, v1, v2])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 4})
|
||||
X_shared[v0, v1, v2] = X[v0, v1, v2]
|
||||
for ax0_ax1_ax2_fused in range(1024):
|
||||
with T.sblock("Y_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(1024, k_0 * 16 + ax0_ax1_ax2_fused // 64)
|
||||
v2 = T.axis.spatial(1024, b_0_i_0_j_0_fused % 16 * 64 + ax0_ax1_ax2_fused % 64)
|
||||
T.reads(Y[v0, v1, v2])
|
||||
T.writes(Y_shared[v0, v1, v2])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 4})
|
||||
Y_shared[v0, v1, v2] = Y[v0, v1, v2]
|
||||
for k_1, b_3, i_3, j_3, k_2, b_4, i_4, j_4 in T.grid(2, 1, 1, 1, 8, 1, 1, 2):
|
||||
with T.sblock("Z"):
|
||||
v_b = T.axis.spatial(1, b_3 + b_4)
|
||||
v_i = T.axis.spatial(1024, b_0_i_0_j_0_fused // 16 * 64 + b_1_i_1_j_1_fused // 4 * 8 + b_2_i_2_j_2_fused // 8 + i_3 + i_4)
|
||||
v_j = T.axis.spatial(1024, b_0_i_0_j_0_fused % 16 * 64 + b_1_i_1_j_1_fused % 4 * 16 + b_2_i_2_j_2_fused % 8 * 2 + j_3 * 2 + j_4)
|
||||
v_k = T.axis.reduce(1024, k_0 * 16 + k_1 * 8 + k_2)
|
||||
T.reads(X_shared[v_b, v_i, v_k], Y_shared[v_b, v_k, v_j])
|
||||
T.writes(Z_local[v_b, v_i, v_j])
|
||||
T.sblock_attr({"meta_schedule.thread_extent_high_inclusive": 1024, "meta_schedule.thread_extent_low_inclusive": 32, "meta_schedule.tiling_structure": "SSSRRSRS"})
|
||||
with T.init():
|
||||
Z_local[v_b, v_i, v_j] = T.float32(0)
|
||||
Z_local[v_b, v_i, v_j] = Z_local[v_b, v_i, v_j] + X_shared[v_b, v_i, v_k] * Y_shared[v_b, v_k, v_j]
|
||||
for ax0, ax1, ax2 in T.grid(1, 1, 2):
|
||||
with T.sblock("Z_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(1024, b_0_i_0_j_0_fused // 16 * 64 + b_1_i_1_j_1_fused // 4 * 8 + b_2_i_2_j_2_fused // 8 + ax1)
|
||||
v2 = T.axis.spatial(1024, b_0_i_0_j_0_fused % 16 * 64 + b_1_i_1_j_1_fused % 4 * 16 + b_2_i_2_j_2_fused % 8 * 2 + ax2)
|
||||
T.reads(Z_local[v0, v1, v2])
|
||||
T.writes(Z[v0, v1, v2])
|
||||
Z[v0, v1, v2] = Z_local[v0, v1, v2]
|
||||
# fmt: on
|
||||
else:
|
||||
# fmt: off
|
||||
@T.prim_func(s_tir=True)
|
||||
def gmm(X: T.Buffer((1, 1024, 1024), "float32"), Y: T.Buffer((1, 1024, 1024), "float32"), Z: T.Buffer((1, 1024, 1024), "float32")):
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
with T.sblock("root"):
|
||||
T.reads()
|
||||
T.writes()
|
||||
T.sblock_attr({"meta_schedule.unroll_explicit": 16})
|
||||
Z_local = T.sblock_alloc_buffer((1, 1024, 1024), scope="local")
|
||||
X_shared = T.sblock_alloc_buffer((1, 1024, 1024), scope="shared")
|
||||
Y_shared = T.sblock_alloc_buffer((1, 1024, 1024), scope="shared")
|
||||
for b_0_i_0_j_0_fused in T.thread_binding(256, thread="blockIdx.x"):
|
||||
for b_1_i_1_j_1_fused in T.thread_binding(32, thread="vthread.x"):
|
||||
for b_2_i_2_j_2_fused in T.thread_binding(64, thread="threadIdx.x"):
|
||||
for k_0_fused in T.serial(64, annotations={"software_pipeline_async_stages": [0], "software_pipeline_order": [0, 1, 2], "software_pipeline_stage": [0, 0, stage - 2]}):
|
||||
for ax0_ax1_ax2_fused in range(1024):
|
||||
with T.sblock("X_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(1024, b_0_i_0_j_0_fused // 16 * 64 + ax0_ax1_ax2_fused // 16)
|
||||
v2 = T.axis.spatial(1024, k_0_fused * 16 + ax0_ax1_ax2_fused % 16)
|
||||
T.reads(X[v0, v1, v2])
|
||||
T.writes(X_shared[v0, v1, v2])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 4})
|
||||
X_shared[v0, v1, v2] = X[v0, v1, v2]
|
||||
for ax0_ax1_ax2_fused in range(1024):
|
||||
with T.sblock("Y_shared"):
|
||||
v0 = T.axis.spatial(1, 0)
|
||||
v1 = T.axis.spatial(1024, k_0_fused * 16 + ax0_ax1_ax2_fused // 64)
|
||||
v2 = T.axis.spatial(1024, b_0_i_0_j_0_fused % 16 * 64 + ax0_ax1_ax2_fused % 64)
|
||||
T.reads(Y[v0, v1, v2])
|
||||
T.writes(Y_shared[v0, v1, v2])
|
||||
T.sblock_attr({"meta_schedule.cooperative_fetch": 4})
|
||||
Y_shared[v0, v1, v2] = Y[v0, v1, v2]
|
||||
for k_1, b_3, i_3, j_3, k_2, b_4, i_4, j_4 in T.grid(2, 1, 1, 1, 8, 1, 1, 2):
|
||||
with T.sblock("Z"):
|
||||
v_b = T.axis.spatial(1, b_3 + b_4)
|
||||
v_i = T.axis.spatial(1024, b_0_i_0_j_0_fused // 16 * 64 + b_1_i_1_j_1_fused // 4 * 8 + b_2_i_2_j_2_fused // 8 + i_3 + i_4)
|
||||
v_j = T.axis.spatial(1024, b_0_i_0_j_0_fused % 16 * 64 + b_1_i_1_j_1_fused % 4 * 16 + b_2_i_2_j_2_fused % 8 * 2 + j_3 * 2 + j_4)
|
||||
v_k = T.axis.reduce(1024, k_0_fused * 16 + k_1 * 8 + k_2)
|
||||
T.reads(X_shared[v_b, v_i, v_k], Y_shared[v_b, v_k, v_j])
|
||||
T.writes(Z_local[v_b, v_i, v_j])
|
||||
T.sblock_attr({"meta_schedule.thread_extent_high_inclusive": 1024, "meta_schedule.thread_extent_low_inclusive": 32, "meta_schedule.tiling_structure": "SSSRRSRS"})
|
||||
with T.init():
|
||||
Z_local[v_b, v_i, v_j] = T.float32(0)
|
||||
Z_local[v_b, v_i, v_j] = Z_local[v_b, v_i, v_j] + X_shared[v_b, v_i, v_k] * Y_shared[v_b, v_k, v_j]
|
||||
for ax0, ax1, ax2 in T.grid(1, 1, 2):
|
||||
with T.sblock("Z_local"):
|
||||
v0 = T.axis.spatial(1, ax0)
|
||||
v1 = T.axis.spatial(1024, b_0_i_0_j_0_fused // 16 * 64 + b_1_i_1_j_1_fused // 4 * 8 + b_2_i_2_j_2_fused // 8 + ax1)
|
||||
v2 = T.axis.spatial(1024, b_0_i_0_j_0_fused % 16 * 64 + b_1_i_1_j_1_fused % 4 * 16 + b_2_i_2_j_2_fused % 8 * 2 + ax2)
|
||||
T.reads(Z_local[v0, v1, v2])
|
||||
T.writes(Z[v0, v1, v2])
|
||||
Z[v0, v1, v2] = Z_local[v0, v1, v2]
|
||||
# fmt: on
|
||||
return gmm
|
||||
|
||||
|
||||
def test_cuda_gmm():
|
||||
gmm_decision = [
|
||||
("SamplePerfectTile", [1, 1, 1, 1, 1]),
|
||||
("SamplePerfectTile", [16, 8, 8, 1, 1]),
|
||||
("SamplePerfectTile", [16, 4, 8, 1, 2]),
|
||||
("SamplePerfectTile", [64, 2, 8]),
|
||||
("SampleCategorical", 3),
|
||||
("SampleCategorical", 3),
|
||||
("SampleCategorical", 1),
|
||||
]
|
||||
|
||||
mod = create_te_workload("GMM", 3)
|
||||
actual = _design_space(mod)
|
||||
check_sketches(
|
||||
mod,
|
||||
sketches=actual,
|
||||
expected_mods=[
|
||||
get_gmm_prim_func(stage=0),
|
||||
get_gmm_prim_func(stage=4),
|
||||
get_gmm_prim_func(stage=5),
|
||||
],
|
||||
expected_decisions=[gmm_decision, gmm_decision, gmm_decision],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_cuda_c2d()
|
||||
test_cuda_gmm()
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
# 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.
|
||||
"""Tests for MetaSchedule search space on CUDA"""
|
||||
|
||||
# isort: off
|
||||
from typing import Literal
|
||||
|
||||
# isort: on
|
||||
from tvm.ir import IRModule
|
||||
from tvm.s_tir import Schedule
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing.space_generation import get_rules
|
||||
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
|
||||
from tvm.target import Target
|
||||
|
||||
|
||||
def generate_design_space(
|
||||
kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"],
|
||||
mod: IRModule,
|
||||
target: Target,
|
||||
types: type | tuple[type, ...],
|
||||
sch_rules: list[ms.ScheduleRule] | None = None,
|
||||
initialize_time: int = 1,
|
||||
) -> list[Schedule]:
|
||||
if sch_rules is None:
|
||||
sch_rules = get_rules(kind, types)
|
||||
else:
|
||||
assert types is None
|
||||
ctx = ms.TuneContext(
|
||||
mod=mod,
|
||||
target=target,
|
||||
space_generator=ms.space_generator.PostOrderApply(
|
||||
sch_rules=sch_rules,
|
||||
postprocs=[],
|
||||
mutator_probs={},
|
||||
),
|
||||
task_name="test",
|
||||
)
|
||||
# each time cloning will trigger one more initialization
|
||||
for _ in range(initialize_time - 1):
|
||||
ctx = ctx.clone()
|
||||
return ctx.generate_design_space()
|
||||
|
||||
|
||||
def _target():
|
||||
return Target("nvidia/geforce-rtx-3070")
|
||||
|
||||
|
||||
def _design_space(mod):
|
||||
return generate_design_space(
|
||||
kind="cuda",
|
||||
mod=mod,
|
||||
target=_target(),
|
||||
types=ms.ScheduleRule,
|
||||
initialize_time=100,
|
||||
)
|
||||
|
||||
|
||||
def test_c2d():
|
||||
mod = create_te_workload("C2D", 0)
|
||||
actual = _design_space(mod)
|
||||
assert len(actual) == 3
|
||||
|
||||
|
||||
def test_gmm():
|
||||
mod = create_te_workload("GMM", 0)
|
||||
actual = _design_space(mod)
|
||||
assert len(actual) == 3
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_c2d()
|
||||
test_gmm()
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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.
|
||||
"""Test Meta Schedule SpaceGenerator"""
|
||||
# pylint: disable=missing-function-docstring
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.s_tir.meta_schedule.space_generator import (
|
||||
PySpaceGenerator,
|
||||
ScheduleFn,
|
||||
SpaceGeneratorUnion,
|
||||
)
|
||||
from tvm.s_tir.meta_schedule.tune_context import TuneContext
|
||||
from tvm.s_tir.schedule import Schedule
|
||||
from tvm.script import tirx as T
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
# fmt: off
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
T.func_attr({"global_symbol": "main"})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
# fmt: on
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,no-self-argument
|
||||
|
||||
|
||||
def schedule_matmul(sch: Schedule):
|
||||
block = sch.get_sblock("matmul")
|
||||
i, j, k = sch.get_loops(block=block)
|
||||
# TODO(@zxybazh): Change to `sample_perfect_tile` after upstreaming
|
||||
i_0, i_1, i_2, i_3 = sch.split(loop=i, factors=[2, 4, 64, 2])
|
||||
j_0, j_1, j_2, j_3 = sch.split(loop=j, factors=[4, 64, 2, 2])
|
||||
k_0, k_1 = sch.split(loop=k, factors=[32, 32])
|
||||
sch.reorder(i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3)
|
||||
|
||||
|
||||
def _check_correct(schedule: Schedule):
|
||||
trace = schedule.trace
|
||||
for inst in trace.decisions:
|
||||
assert math.prod(trace.decisions[inst]) == 1024
|
||||
|
||||
|
||||
def test_meta_schedule_space_generator_schedule_fn():
|
||||
mod = Matmul
|
||||
space_generator = ScheduleFn(sch_fn=schedule_matmul)
|
||||
design_spaces = space_generator.generate_design_space(mod)
|
||||
assert len(design_spaces) == 1
|
||||
(schedule,) = design_spaces
|
||||
_check_correct(schedule)
|
||||
|
||||
|
||||
def test_meta_schedule_design_space_generator_union():
|
||||
mod = Matmul
|
||||
space_generator = ScheduleFn(sch_fn=schedule_matmul)
|
||||
space_generator_union = SpaceGeneratorUnion([space_generator, space_generator])
|
||||
design_spaces = space_generator_union.generate_design_space(mod)
|
||||
assert len(design_spaces) == 2
|
||||
for design_space in design_spaces:
|
||||
_check_correct(design_space)
|
||||
|
||||
|
||||
def test_meta_schedule_design_space_generator_NIE():
|
||||
@derived_object
|
||||
class TestPySpaceGenerator(PySpaceGenerator):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.sch_rules = []
|
||||
self.postprocs = []
|
||||
self.mutator_probs = {}
|
||||
|
||||
with pytest.raises(
|
||||
RuntimeError, match="PySpaceGenerator's InitializeWithTuneContext method not implemented!"
|
||||
):
|
||||
generator = TestPySpaceGenerator()
|
||||
generator._initialize_with_tune_context(TuneContext())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,121 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-docstring,no-member,invalid-name,unused-variable
|
||||
# ruff: noqa: F401
|
||||
import logging
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.runner.config import EvaluatorConfig
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
from tvm.testing import env
|
||||
|
||||
logging.basicConfig()
|
||||
logging.getLogger("tvm.s_tir.meta_schedule").setLevel(logging.DEBUG)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
|
||||
@pytest.mark.skip("Integration test")
|
||||
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
||||
def test_tune_matmul_cpu():
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
target = Target({"kind": "llvm", "num-cores": 16})
|
||||
database = ms.tir_integration.tune_tir(
|
||||
mod=matmul,
|
||||
target=target,
|
||||
work_dir=work_dir,
|
||||
max_trials_global=32,
|
||||
num_trials_per_iter=16,
|
||||
post_optimization=True,
|
||||
)
|
||||
trials = 32
|
||||
database = ms.tune_tir(
|
||||
mod=matmul,
|
||||
target=target,
|
||||
max_trials_global=trials,
|
||||
num_trials_per_iter=64,
|
||||
work_dir=work_dir,
|
||||
runner=ms.runner.LocalRunner(
|
||||
evaluator_config=EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=100,
|
||||
)
|
||||
),
|
||||
cost_model=ms.cost_model.XGBModel(
|
||||
extractor=ms.feature_extractor.PerStoreFeature(),
|
||||
adaptive_training=False,
|
||||
),
|
||||
strategy=ms.search_strategy.EvolutionarySearch(),
|
||||
post_optimization=True, # testing post optmization
|
||||
)
|
||||
# +1 because of post optmization
|
||||
assert len(database) == trials + 1
|
||||
|
||||
|
||||
@pytest.mark.skip("Integration test")
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_tune_matmul_cuda():
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
target = Target("nvidia/geforce-rtx-3070")
|
||||
trials = 32
|
||||
database = ms.tune_tir(
|
||||
mod=matmul,
|
||||
target=target,
|
||||
max_trials_global=trials,
|
||||
num_trials_per_iter=64,
|
||||
work_dir=work_dir,
|
||||
runner=ms.runner.LocalRunner(
|
||||
evaluator_config=EvaluatorConfig(
|
||||
number=1,
|
||||
repeat=1,
|
||||
min_repeat_ms=100,
|
||||
)
|
||||
),
|
||||
cost_model=ms.cost_model.XGBModel(
|
||||
extractor=ms.feature_extractor.PerStoreFeature(),
|
||||
adaptive_training=False,
|
||||
),
|
||||
strategy=ms.search_strategy.EvolutionarySearch(),
|
||||
post_optimization=True, # testing post optmization
|
||||
)
|
||||
# +1 because of post optmization
|
||||
assert len(database) == trials + 1
|
||||
|
||||
|
||||
if __name__ == """__main__""":
|
||||
test_tune_matmul_cpu()
|
||||
test_tune_matmul_cuda()
|
||||
@@ -0,0 +1,461 @@
|
||||
# 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.
|
||||
# ruff: noqa: F821, RUF012
|
||||
"""Test Meta Schedule Task Scheduler"""
|
||||
|
||||
import random
|
||||
import weakref
|
||||
|
||||
import pytest
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.s_tir import Schedule
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing.dummy_object import DummyBuilder, DummyRunner
|
||||
from tvm.script import tirx as T
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class MatmulModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main( # type: ignore
|
||||
a: T.handle,
|
||||
b: T.handle,
|
||||
c: T.handle,
|
||||
) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0 # type: ignore
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class MatmulReluModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main( # type: ignore
|
||||
a: T.handle,
|
||||
b: T.handle,
|
||||
d: T.handle,
|
||||
) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
D = T.match_buffer(d, (1024, 1024), "float32")
|
||||
C = T.sblock_alloc_buffer((1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0 # type: ignore
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
for i, j in T.grid(1024, 1024):
|
||||
with T.sblock("relu"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
D[vi, vj] = T.max(C[vi, vj], 0.0) # type: ignore
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class BatchMatmulModule:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main( # type: ignore
|
||||
a: T.handle,
|
||||
b: T.handle,
|
||||
c: T.handle,
|
||||
) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, [16, 128, 128])
|
||||
B = T.match_buffer(b, [16, 128, 128])
|
||||
C = T.match_buffer(c, [16, 128, 128])
|
||||
for n, i, j, k in T.grid(16, 128, 128, 128):
|
||||
with T.sblock("matmul"):
|
||||
vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k])
|
||||
with T.init():
|
||||
C[vn, vi, vj] = 0.0 # type: ignore
|
||||
C[vn, vi, vj] = C[vn, vi, vj] + A[vn, vi, vk] * B[vn, vj, vk]
|
||||
|
||||
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks
|
||||
|
||||
|
||||
def _schedule_matmul(sch: Schedule):
|
||||
block = sch.get_sblock("matmul")
|
||||
i, j, k = sch.get_loops(block=block)
|
||||
i_0, i_1, i_2, i_3 = sch.split(loop=i, factors=[2, 4, 64, 2])
|
||||
j_0, j_1, j_2, j_3 = sch.split(loop=j, factors=[4, 64, 2, 2])
|
||||
k_0, k_1 = sch.split(loop=k, factors=[32, 32])
|
||||
sch.reorder(i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3)
|
||||
|
||||
|
||||
def _schedule_batch_matmul(sch: Schedule):
|
||||
block = sch.get_sblock("matmul")
|
||||
i, j, k, t = sch.get_loops(block=block)
|
||||
i_0, i_1, i_2, i_3 = sch.split(loop=i, factors=[2, 2, 2, 2])
|
||||
j_0, j_1, j_2, j_3 = sch.split(loop=j, factors=[2, 4, 64, 2])
|
||||
k_0, k_1 = sch.split(loop=k, factors=[32, 32])
|
||||
t_0, t_1 = sch.split(loop=t, factors=[2, 512])
|
||||
sch.reorder(i_0, j_0, i_1, j_1, k_0, i_2, j_2, k_1, i_3, j_3, t_0, t_1)
|
||||
|
||||
|
||||
@derived_object
|
||||
class MyTaskScheduler(ms.task_scheduler.PyTaskScheduler):
|
||||
done: set = set()
|
||||
|
||||
def next_task_id(self) -> int:
|
||||
tasks = self._outer().tasks_
|
||||
while len(self.done) != len(tasks):
|
||||
x = random.randint(0, len(tasks) - 1)
|
||||
task = tasks[x]
|
||||
if not task.is_terminated:
|
||||
"""Calling base func via following route:
|
||||
Python side:
|
||||
PyTaskScheduler does not have `_touch_task`
|
||||
Call TaskScheduler's `touch_task`, which calls ffi
|
||||
C++ side:
|
||||
The ffi calls TaskScheduler's `touch_task`
|
||||
But it is overridden in PyTaskScheduler
|
||||
PyTaskScheduler checks if the function is overridden in python
|
||||
If not, it returns the TaskScheduler's vtable, calling
|
||||
TaskScheduler::TouchTask
|
||||
"""
|
||||
if task.runner_futures is not None:
|
||||
self.join_running_task(x)
|
||||
return x
|
||||
self.done.add(x)
|
||||
return -1
|
||||
|
||||
|
||||
def test_meta_schedule_task_scheduler_single():
|
||||
num_trials_per_iter = 3
|
||||
max_trials_per_task = 10
|
||||
database = ms.database.MemoryDatabase()
|
||||
round_robin = ms.task_scheduler.RoundRobin()
|
||||
round_robin.tune(
|
||||
[
|
||||
ms.TuneContext(
|
||||
MatmulModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="Test",
|
||||
rand_state=42,
|
||||
)
|
||||
],
|
||||
[1.0],
|
||||
max_trials_global=num_trials_per_iter,
|
||||
max_trials_per_task=max_trials_per_task,
|
||||
num_trials_per_iter=64,
|
||||
builder=DummyBuilder(),
|
||||
runner=DummyRunner(),
|
||||
database=database,
|
||||
measure_callbacks=[ms.measure_callback.AddToDatabase()],
|
||||
cost_model=None,
|
||||
)
|
||||
assert len(database) == max_trials_per_task
|
||||
|
||||
|
||||
def test_meta_schedule_task_scheduler_multiple():
|
||||
num_trials_per_iter = 6
|
||||
max_trials_per_task = 101
|
||||
tasks = [
|
||||
ms.TuneContext(
|
||||
MatmulModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="Matmul",
|
||||
rand_state=42,
|
||||
),
|
||||
ms.TuneContext(
|
||||
MatmulReluModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="MatmulRelu",
|
||||
rand_state=0xDEADBEEF,
|
||||
),
|
||||
ms.TuneContext(
|
||||
BatchMatmulModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_batch_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="BatchMatmul",
|
||||
rand_state=0x114514,
|
||||
),
|
||||
]
|
||||
database = ms.database.MemoryDatabase()
|
||||
round_robin = ms.task_scheduler.RoundRobin()
|
||||
round_robin.tune(
|
||||
tasks,
|
||||
[1.0, 1.0, 1.0],
|
||||
builder=DummyBuilder(),
|
||||
runner=DummyRunner(),
|
||||
database=database,
|
||||
measure_callbacks=[ms.measure_callback.AddToDatabase()],
|
||||
max_trials_global=max_trials_per_task * len(tasks),
|
||||
max_trials_per_task=max_trials_per_task,
|
||||
num_trials_per_iter=num_trials_per_iter,
|
||||
cost_model=None,
|
||||
)
|
||||
assert len(database) == max_trials_per_task * len(tasks)
|
||||
for task in tasks:
|
||||
assert (
|
||||
len(
|
||||
database.get_top_k(
|
||||
database.commit_workload(task.mod),
|
||||
100000,
|
||||
)
|
||||
)
|
||||
== max_trials_per_task
|
||||
)
|
||||
|
||||
|
||||
def test_meta_schedule_task_scheduler_NIE(): # pylint: disable=invalid-name
|
||||
@derived_object
|
||||
class NIETaskScheduler(ms.task_scheduler.PyTaskScheduler):
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError, match="next_task_id is not defined"):
|
||||
scheduler = NIETaskScheduler()
|
||||
scheduler.next_task_id()
|
||||
|
||||
|
||||
def test_meta_schedule_task_scheduler_avoid_cyclic(): # pylint: disable=invalid-name
|
||||
scheduler = MyTaskScheduler()
|
||||
test = weakref.ref(scheduler) # test if it can be destructed successfully
|
||||
del scheduler
|
||||
assert test() is None
|
||||
|
||||
|
||||
def test_meta_schedule_task_scheduler_override_next_task_id_only(): # pylint: disable=invalid-name
|
||||
max_trials_per_task = 101
|
||||
tasks = [
|
||||
ms.TuneContext(
|
||||
MatmulModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="Matmul",
|
||||
rand_state=42,
|
||||
),
|
||||
ms.TuneContext(
|
||||
MatmulReluModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="MatmulRelu",
|
||||
rand_state=0xDEADBEEF,
|
||||
),
|
||||
ms.TuneContext(
|
||||
BatchMatmulModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_batch_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="BatchMatmul",
|
||||
rand_state=0x114514,
|
||||
),
|
||||
]
|
||||
database = ms.database.MemoryDatabase()
|
||||
scheduler = MyTaskScheduler()
|
||||
scheduler.tune(
|
||||
tasks,
|
||||
task_weights=[1.0] * len(tasks),
|
||||
builder=DummyBuilder(),
|
||||
runner=DummyRunner(),
|
||||
database=database,
|
||||
measure_callbacks=[ms.measure_callback.AddToDatabase()],
|
||||
max_trials_global=max_trials_per_task * len(tasks),
|
||||
max_trials_per_task=max_trials_per_task,
|
||||
num_trials_per_iter=6,
|
||||
cost_model=None,
|
||||
)
|
||||
assert len(database) == max_trials_per_task * len(tasks)
|
||||
for task in tasks:
|
||||
assert (
|
||||
len(
|
||||
database.get_top_k(
|
||||
database.commit_workload(task.mod),
|
||||
100000,
|
||||
)
|
||||
)
|
||||
== max_trials_per_task
|
||||
)
|
||||
|
||||
|
||||
def test_meta_schedule_task_scheduler_multiple_gradient_based():
|
||||
max_trials_per_task = 101
|
||||
tasks = [
|
||||
ms.TuneContext(
|
||||
MatmulModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="Matmul",
|
||||
rand_state=42,
|
||||
),
|
||||
ms.TuneContext(
|
||||
MatmulReluModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="MatmulRelu",
|
||||
rand_state=0xDEADBEEF,
|
||||
),
|
||||
ms.TuneContext(
|
||||
BatchMatmulModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_batch_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="BatchMatmul",
|
||||
rand_state=0x114514,
|
||||
),
|
||||
]
|
||||
database = ms.database.MemoryDatabase()
|
||||
gradient_based = ms.task_scheduler.GradientBased()
|
||||
gradient_based.tune(
|
||||
tasks,
|
||||
task_weights=[1.0, 1.0, 1.0],
|
||||
builder=DummyBuilder(),
|
||||
runner=DummyRunner(),
|
||||
database=database,
|
||||
measure_callbacks=[ms.measure_callback.AddToDatabase()],
|
||||
max_trials_global=max_trials_per_task * len(tasks),
|
||||
max_trials_per_task=max_trials_per_task,
|
||||
num_trials_per_iter=6,
|
||||
cost_model=None,
|
||||
)
|
||||
assert len(database) == max_trials_per_task * len(tasks)
|
||||
for task in tasks:
|
||||
assert (
|
||||
len(database.get_top_k(database.commit_workload(task.mod), 10000))
|
||||
== max_trials_per_task
|
||||
)
|
||||
|
||||
|
||||
def test_meta_schedule_task_scheduler_gradient_based_with_null_search_strategy():
|
||||
"""
|
||||
When search strategy of one task returns empty list of candidates or None,
|
||||
the scheduler should continue working as normal for other tasks
|
||||
"""
|
||||
|
||||
@derived_object
|
||||
class NullSearchStrategy(ms.search_strategy.PySearchStrategy):
|
||||
def __init__(self, rounds_with_empty_candidates):
|
||||
self.rounds_with_empty_candidates = rounds_with_empty_candidates
|
||||
|
||||
def _initialize_with_tune_context(self, context: "TuneContext") -> None:
|
||||
pass
|
||||
|
||||
def pre_tuning(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def post_tuning(self):
|
||||
pass
|
||||
|
||||
def generate_measure_candidates(self):
|
||||
"""
|
||||
Returns empty list to indicate there is no result from search, while
|
||||
the search isn't ended.
|
||||
"""
|
||||
if self.rounds_with_empty_candidates:
|
||||
self.rounds_with_empty_candidates -= 1
|
||||
return []
|
||||
return None
|
||||
|
||||
def notify_runner_results(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def clone(self):
|
||||
return NullSearchStrategy(n=self.n)
|
||||
|
||||
tasks = [
|
||||
ms.TuneContext(
|
||||
MatmulModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_matmul,
|
||||
search_strategy=NullSearchStrategy(rounds_with_empty_candidates=5),
|
||||
task_name="Matmul",
|
||||
rand_state=42,
|
||||
),
|
||||
ms.TuneContext(
|
||||
BatchMatmulModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_batch_matmul,
|
||||
search_strategy=NullSearchStrategy(rounds_with_empty_candidates=0),
|
||||
task_name="BatchMatmul",
|
||||
rand_state=0x114514,
|
||||
),
|
||||
ms.TuneContext(
|
||||
MatmulReluModule,
|
||||
num_threads=1,
|
||||
target=tvm.target.Target("llvm"),
|
||||
space_generator=_schedule_matmul,
|
||||
search_strategy=ms.search_strategy.ReplayTrace(),
|
||||
task_name="MatmulRelu",
|
||||
rand_state=0xDEADBEEF,
|
||||
),
|
||||
]
|
||||
database = ms.database.MemoryDatabase()
|
||||
gradient_based = ms.task_scheduler.GradientBased()
|
||||
gradient_based.tune(
|
||||
tasks,
|
||||
task_weights=[1.0, 1.0, 1.0],
|
||||
builder=DummyBuilder(),
|
||||
runner=DummyRunner(),
|
||||
database=database,
|
||||
measure_callbacks=[ms.measure_callback.AddToDatabase()],
|
||||
max_trials_global=30,
|
||||
max_trials_per_task=10,
|
||||
num_trials_per_iter=6,
|
||||
cost_model=None,
|
||||
)
|
||||
|
||||
assert len(database) == 10
|
||||
assert len(database.get_top_k(database.commit_workload(MatmulModule), 100)) == 0
|
||||
assert len(database.get_top_k(database.commit_workload(BatchMatmulModule), 100)) == 0
|
||||
assert len(database.get_top_k(database.commit_workload(MatmulReluModule), 100)) == 10
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_meta_schedule_task_scheduler_single()
|
||||
test_meta_schedule_task_scheduler_multiple()
|
||||
test_meta_schedule_task_scheduler_NIE()
|
||||
test_meta_schedule_task_scheduler_avoid_cyclic()
|
||||
test_meta_schedule_task_scheduler_override_next_task_id_only()
|
||||
test_meta_schedule_task_scheduler_multiple_gradient_based()
|
||||
test_meta_schedule_task_scheduler_gradient_based_with_null_search_strategy()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,63 @@
|
||||
# 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.
|
||||
# ruff: noqa: F401
|
||||
"""Test the tune context of meta schedule."""
|
||||
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.s_tir.meta_schedule import TuneContext
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
|
||||
# pylint: disable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring
|
||||
|
||||
|
||||
@tvm.script.ir_module
|
||||
class Matmul:
|
||||
@T.prim_func(s_tir=True)
|
||||
def main(a: T.handle, b: T.handle, c: T.handle) -> None: # pylint: disable=no-self-argument
|
||||
T.func_attr({"global_symbol": "main", "tirx.noalias": True})
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.match_buffer(b, (1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j, k in T.grid(1024, 1024, 1024):
|
||||
with T.sblock("matmul"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]
|
||||
|
||||
|
||||
# pylint: enable=invalid-name,no-member,line-too-long,too-many-nested-blocks,missing-docstring
|
||||
|
||||
|
||||
def test_tune_context_create():
|
||||
mod = Matmul
|
||||
context = TuneContext(mod=mod, target=Target("llvm"), task_name="Test Task")
|
||||
assert context.num_threads > 0
|
||||
assert context.rand_state != -1
|
||||
assert context.task_name == "Test Task"
|
||||
assert context.mod == mod or tvm_ffi.structural_equal(context.mod, mod)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
tvm.testing.main()
|
||||
@@ -0,0 +1,195 @@
|
||||
# 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.
|
||||
# pylint: disable=missing-docstring,no-member,invalid-name,unused-variable
|
||||
import logging
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("tornado") # tvm.rpc.tracker (LocalRPC) requires tornado
|
||||
|
||||
import tvm
|
||||
import tvm.testing
|
||||
from tvm.ir.utils import derived_object
|
||||
from tvm.s_tir import meta_schedule as ms
|
||||
from tvm.s_tir.meta_schedule.testing.custom_builder_runner import run_module_via_rpc
|
||||
from tvm.s_tir.meta_schedule.testing.local_rpc import LocalRPC
|
||||
from tvm.s_tir.schedule import SBlockRV, Schedule
|
||||
from tvm.script import tirx as T
|
||||
from tvm.target import Target
|
||||
from tvm.testing import env
|
||||
|
||||
logging.basicConfig()
|
||||
logging.getLogger("tvm.s_tir.meta_schedule").setLevel(logging.DEBUG)
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, [128, 128])
|
||||
B = T.match_buffer(b, [128, 128])
|
||||
C = T.match_buffer(c, [128, 128])
|
||||
for i, j, k in T.grid(128, 128, 128):
|
||||
with T.sblock("update"):
|
||||
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
|
||||
with T.init():
|
||||
C[vi, vj] = 0.0
|
||||
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
|
||||
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def two_step(a: T.handle, c: T.handle) -> None:
|
||||
A = T.match_buffer(a, (1024, 1024), "float32")
|
||||
B = T.sblock_alloc_buffer((1024, 1024), "float32")
|
||||
C = T.match_buffer(c, (1024, 1024), "float32")
|
||||
for i, j in T.grid(1024, 1024):
|
||||
with T.sblock("A"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
B[vi, vj] = A[vi, vj] * 2.0
|
||||
for i, j in T.grid(1024, 1024):
|
||||
with T.sblock("B"):
|
||||
vi, vj = T.axis.remap("SS", [i, j])
|
||||
C[vi, vj] = B[vi, vj] + 3.0
|
||||
|
||||
|
||||
@pytest.mark.skip("Integration test")
|
||||
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
|
||||
def test_tune_matmul_cpu():
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
target = Target({"kind": "llvm", "num-cores": 16})
|
||||
database = ms.tir_integration.tune_tir(
|
||||
mod=matmul,
|
||||
target=target,
|
||||
work_dir=work_dir,
|
||||
max_trials_global=32,
|
||||
num_trials_per_iter=16,
|
||||
)
|
||||
sch = ms.tir_integration.compile_tir(database, matmul, target)
|
||||
if sch is None:
|
||||
print("No valid schedule found!")
|
||||
else:
|
||||
sch.mod.show()
|
||||
sch.trace.show()
|
||||
|
||||
|
||||
@pytest.mark.skip("Integration test")
|
||||
@pytest.mark.gpu
|
||||
@pytest.mark.skipif(not env.has_cuda(), reason="need cuda")
|
||||
def test_tune_matmul_cuda():
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
target = Target("nvidia/geforce-rtx-3070")
|
||||
database = ms.tir_integration.tune_tir(
|
||||
mod=matmul,
|
||||
target=target,
|
||||
work_dir=work_dir,
|
||||
max_trials_global=32,
|
||||
num_trials_per_iter=16,
|
||||
)
|
||||
sch = ms.tir_integration.compile_tir(database, matmul, target)
|
||||
if sch is None:
|
||||
print("No valid schedule found!")
|
||||
else:
|
||||
sch.mod.show()
|
||||
sch.trace.show()
|
||||
|
||||
|
||||
@pytest.mark.skip("Integration test")
|
||||
def test_tune_run_module_via_rpc():
|
||||
target = tvm.target.Target("llvm")
|
||||
rt_mod = tvm.compile(matmul, target)
|
||||
|
||||
# construct the input
|
||||
input_data = {}
|
||||
input_shape = (128, 128)
|
||||
input_dtype = "float32"
|
||||
a_np = np.random.uniform(size=input_shape).astype(input_dtype)
|
||||
b_np = np.random.uniform(size=input_shape).astype(input_dtype)
|
||||
c_np = np.zeros(input_shape).astype(input_dtype)
|
||||
for i in range(128):
|
||||
for j in range(128):
|
||||
for k in range(128):
|
||||
c_np[i, j] = c_np[i, j] + a_np[i, k] * b_np[j, k]
|
||||
input_data["a"] = a_np
|
||||
input_data["b"] = b_np
|
||||
input_data["c"] = np.zeros(input_shape).astype(input_dtype)
|
||||
|
||||
with LocalRPC() as rpc:
|
||||
rpc_config = ms.runner.RPCConfig(
|
||||
tracker_host=rpc.tracker_host,
|
||||
tracker_port=rpc.tracker_port,
|
||||
tracker_key=rpc.tracker_key,
|
||||
session_priority=1,
|
||||
session_timeout_sec=100,
|
||||
)
|
||||
|
||||
def f_timer(rt_mod, dev, input_data):
|
||||
rt_mod(input_data["a"], input_data["b"], input_data["c"])
|
||||
return input_data["c"]
|
||||
|
||||
result = run_module_via_rpc(
|
||||
rpc_config=rpc_config,
|
||||
lib=rt_mod,
|
||||
dev_type=target.kind.name,
|
||||
args=input_data,
|
||||
continuation=f_timer,
|
||||
)
|
||||
tvm.testing.assert_allclose(result.numpy(), c_np, rtol=1e-3)
|
||||
|
||||
|
||||
@pytest.mark.skip("Integration test")
|
||||
def test_tune_block_cpu():
|
||||
@derived_object
|
||||
class RemoveBlock(ms.schedule_rule.PyScheduleRule):
|
||||
def _initialize_with_tune_context(self, context: ms.TuneContext) -> None:
|
||||
pass
|
||||
|
||||
def apply(self, sch: Schedule, block: SBlockRV):
|
||||
if sch.get(block).name_hint == "root":
|
||||
return [sch]
|
||||
sch = sch.copy()
|
||||
sch.compute_inline(block)
|
||||
return [sch]
|
||||
|
||||
def clone(self) -> "RemoveBlock":
|
||||
return RemoveBlock()
|
||||
|
||||
with tempfile.TemporaryDirectory() as work_dir:
|
||||
target = Target({"kind": "llvm", "num-cores": 16})
|
||||
database = ms.tir_integration.tune_tir(
|
||||
mod=two_step,
|
||||
target=target,
|
||||
work_dir=work_dir,
|
||||
max_trials_global=32,
|
||||
num_trials_per_iter=16,
|
||||
space=ms.space_generator.PostOrderApply(
|
||||
f_block_filter=lambda block: block.name_hint == "A",
|
||||
sch_rules=[RemoveBlock()],
|
||||
postprocs=[],
|
||||
mutator_probs={},
|
||||
),
|
||||
)
|
||||
sch = ms.tir_integration.compile_tir(database, two_step, target)
|
||||
assert sch is not None
|
||||
sch.mod.show()
|
||||
sch.trace.show()
|
||||
|
||||
|
||||
if __name__ == """__main__""":
|
||||
test_tune_matmul_cpu()
|
||||
test_tune_matmul_cuda()
|
||||
test_tune_run_module_via_rpc()
|
||||
test_tune_block_cpu()
|
||||
Reference in New Issue
Block a user