wehub-resource-sync
130 lines
6.4 KiB
Python
130 lines
6.4 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: E501, F401, F841
|
|
"""Integration test for MetaSchedule"""
|
|
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.runtime import cpu as tvm_cpu
|
|
from tvm.runtime import tensor as tvm_tensor
|
|
from tvm.runtime.vm import VirtualMachine
|
|
from tvm.s_tir import meta_schedule as ms
|
|
from tvm.script import ir as I
|
|
from tvm.script import relax as R
|
|
from tvm.script import tirx as T
|
|
|
|
|
|
# fmt: off
|
|
@I.ir_module
|
|
class Module0:
|
|
@R.function
|
|
def main(data: R.Tensor((1, 8, 8, 4), dtype="int32")) -> R.Tensor((1, 8, 8, 4), dtype="int32"):
|
|
cls = Module0
|
|
with R.dataflow():
|
|
c = R.const([[[[-171701247],[-1719837685],[1801664104],[-634316588]],[[920159370],[-132073802],[2142531563],[1465185701]],[[-1505608067],[1737948828],[1581089391],[-1986167320]]],[[[-1449581822],[35714587],[496324563],[-1430879015]],[[-1615680873],[1198514997],[1494683955],[1567376558]],[[1319924884],[-380548171],[296785437],[-1546305981]]],[[[-398644701],[-2004794585],[-1850413687],[2072643657]],[[847950121],[-544212073],[-199532669],[-343273682]],[[953721562],[-1930209358],[1573600108],[-577689853]]]], "int32")
|
|
lv: R.Tensor((1, 8, 8, 4), dtype="int32") = R.nn.conv2d(data, c, strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=4, data_layout="NHWC", kernel_layout="HWOI", out_layout="NHWC", out_dtype="int32")
|
|
b = R.const([[[[1, 1, 1, 1]]]], "int32")
|
|
lv1: R.Tensor((1, 8, 8, 4), dtype="int32") = R.add(lv, b)
|
|
c1 = R.const([[[[2042349344],[-2076067063],[1528163722],[-1156452837]],[[-2097172051],[1137787079],[-601389657],[1907495997]],[[987801941],[1073738593],[-1410339796],[-689755358]]],[[[90351522],[-44886952],[-1914103775],[-691553659]],[[-1288505112],[-1376578817],[-2067933148],[-1413101824]],[[1261422027],[-156976862],[-1185734459],[1608778622]]],[[[-664209483],[1907479806],[1838595152],[464942526]],[[877953160],[415131837],[-2010736511],[1218242769]],[[-1440127632],[112931],[521745784],[-1931145893]]]], "int32")
|
|
lv2: R.Tensor((1, 8, 8, 4), dtype="int32") = R.nn.conv2d(lv1, c1, strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=4, data_layout="NHWC", kernel_layout="HWOI", out_layout="NHWC", out_dtype="int32")
|
|
c2 = R.const([[[[687940110],[-910571705],[-901609800],[-500525928]],[[506872399],[1070176297],[-305936110],[1625439784]],[[-1565626954],[-1705688881],[-866370805],[-1750740826]]],[[[300497007],[-626864803],[390295545],[222549121]],[[319224543],[-2003064970],[657992492],[2014175448]],[[653278589],[-768810984],[-294555581],[-1197167662]]],[[[1703154671],[-1540759805],[-568817430],[-1729755444]],[[-275458074],[2078945571],[1683298006],[-1029327874]],[[1315093181],[159010501],[875694807],[-223655381]]]], "int32")
|
|
lv3: R.Tensor((1, 8, 8, 4), dtype="int32") = R.nn.conv2d(lv2, c2, strides=[1, 1], padding=[1, 1, 1, 1], dilation=[1, 1], groups=4, data_layout="NHWC", kernel_layout="HWOI", out_layout="NHWC", out_dtype="int32")
|
|
gv: R.Tensor((1, 8, 8, 4), dtype="int32") = lv3
|
|
R.output(gv)
|
|
return gv
|
|
|
|
# fmt: on
|
|
|
|
|
|
def test_extracting_tasks():
|
|
target = {"kind": "llvm", "mcpu": "core-avx2", "num-cores": 1}
|
|
|
|
relax_mod = Module0
|
|
relax_mod = relax.transform.LegalizeOps()(relax_mod)
|
|
relax_mod = relax.transform.AnnotateTIROpPattern()(relax_mod)
|
|
relax_mod = relax.transform.FuseOps()(relax_mod)
|
|
relax_mod = relax.transform.FoldConstant()(relax_mod)
|
|
relax_mod = relax.transform.FuseTIR()(relax_mod)
|
|
|
|
relax_expectation = {
|
|
"structural": 2, # The relax constants do not reach the tirx at the lowering.
|
|
"ignore-tensor": 2,
|
|
"anchor-block": 1,
|
|
}
|
|
for module_equality, count in relax_expectation.items():
|
|
extracted_tasks = ms.relax_integration.extract_tasks(
|
|
relax_mod,
|
|
target,
|
|
{},
|
|
module_equality=module_equality,
|
|
)
|
|
assert len(extracted_tasks) == count
|
|
|
|
|
|
def test_compile_relax_with_database():
|
|
"""End-to-end test: tune with MetaSchedule then compile_relax with the database.
|
|
|
|
Verifies that the pipeline ordering in compile_relax is correct: tasks are
|
|
extracted and tuned against fused-TIR keys, and compile_relax produces those
|
|
same keys (by running LegalizeOps + FuseOps + FuseTIR before applying the
|
|
database), so the scheduled kernels are actually picked up.
|
|
"""
|
|
pytest.importorskip("cloudpickle") # needed by meta_schedule popen workers
|
|
|
|
target = tvm.target.Target({"kind": "llvm", "num-cores": 1})
|
|
|
|
# Prepare the fused module whose TIR keys will populate the database.
|
|
fused_mod = Module0
|
|
fused_mod = relax.transform.LegalizeOps()(fused_mod)
|
|
fused_mod = relax.transform.AnnotateTIROpPattern()(fused_mod)
|
|
fused_mod = relax.transform.FuseOps()(fused_mod)
|
|
fused_mod = relax.transform.FoldConstant()(fused_mod)
|
|
fused_mod = relax.transform.FuseTIR()(fused_mod)
|
|
|
|
with tempfile.TemporaryDirectory() as work_dir:
|
|
database = ms.relax_integration.tune_relax(
|
|
fused_mod,
|
|
params={},
|
|
target=target,
|
|
work_dir=work_dir,
|
|
max_trials_global=4,
|
|
)
|
|
# compile_relax takes the raw module and builds the fused-TIR pipeline
|
|
# internally; the database keys must therefore match the ones above.
|
|
exe = ms.relax_integration.compile_relax(
|
|
database=database,
|
|
mod=Module0,
|
|
target=target,
|
|
params=None,
|
|
)
|
|
|
|
dev = tvm_cpu()
|
|
vm = VirtualMachine(exe.jit(), dev)
|
|
data = tvm_tensor(np.zeros((1, 8, 8, 4), dtype="int32"), device=dev)
|
|
result = vm["main"](data)
|
|
assert result.numpy().shape == (1, 8, 8, 4)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|