wehub-resource-sync
126 lines
4.5 KiB
Python
126 lines
4.5 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.
|
|
# pylint: disable=unused-argument, invalid-name, no-else-return, abstract-method, arguments-differ
|
|
"""Relax transformation passes for testing"""
|
|
|
|
import logging
|
|
import os
|
|
|
|
import tvm_ffi
|
|
|
|
import tvm
|
|
from tvm.ir import Call
|
|
from tvm.ir.module import IRModule
|
|
from tvm.relax.expr import DataflowBlock, Var
|
|
from tvm.runtime import Object
|
|
|
|
|
|
def ApplyEmptyCppMutator() -> tvm.ir.transform.Pass:
|
|
"""Create empty cpp mutator"""
|
|
packed_func = tvm.get_global_func("relax.testing.transform.ApplyEmptyCppMutator")
|
|
return packed_func()
|
|
|
|
|
|
def dataflow_liveness_analysis(block: DataflowBlock) -> dict[Var, tuple[int, int]]:
|
|
"""
|
|
Inner function for the dataflow inplace transformation exposed for testing.
|
|
"""
|
|
if "PYTEST_CURRENT_TEST" not in os.environ:
|
|
logging.warning("The function dataflow_liveness_analysis is exposed for testing only.")
|
|
|
|
live_ranges = tvm.get_global_func("relax.testing.transform.DataflowLivenessAnalysis")(block) # type: ignore
|
|
ret = {}
|
|
for var, live_range in live_ranges.items():
|
|
ret[var] = tuple(live_range)
|
|
return ret # type: ignore
|
|
|
|
|
|
def dataflow_alias_analysis(
|
|
block: DataflowBlock, inputs: list[Var]
|
|
) -> tuple[dict[Var, set[int]], dict[int, list[set[int]]]]:
|
|
"""
|
|
Inner function for the dataflow inplace transformation exposed for testing.
|
|
"""
|
|
if "PYTEST_CURRENT_TEST" not in os.environ:
|
|
logging.warning("The function dataflow_alias_analysis is exposed for testing only.")
|
|
|
|
alias_sets, tuple_map = tvm.get_global_func("relax.testing.transform.DataflowAliasAnalysis")(
|
|
block,
|
|
inputs,
|
|
) # type: ignore
|
|
res_alias_sets = {}
|
|
res_tuple_map = {}
|
|
for var, alias_set in alias_sets.items():
|
|
res_alias_sets[var] = set(alias_set)
|
|
for idx, elem_alias_sets in tuple_map.items():
|
|
res_tuple_map[idx] = [set(alias_set) for alias_set in elem_alias_sets]
|
|
return res_alias_sets, res_tuple_map # type: ignore
|
|
|
|
|
|
@tvm_ffi.register_object("relax.transform.InplaceOpportunity")
|
|
class InplaceOpportunity(Object):
|
|
"""
|
|
Represents an opportunity to make a binding in-place. Exposed only for testing;
|
|
the constructor is not exposed.
|
|
|
|
Parameters:
|
|
-----------
|
|
binding_idx: int
|
|
Index of the binding within its block
|
|
|
|
arg_idxs: List[int]
|
|
Indices of arguments that are eligible to be used as in-place targets.
|
|
"""
|
|
|
|
def __init__(self, _binding_idx, _arg_idxs):
|
|
raise NotImplementedError("Constructor for InplaceOpportunity not exposed!")
|
|
|
|
|
|
def dataflow_inplace_analysis(
|
|
block: DataflowBlock, inputs: list[Var], mod: IRModule
|
|
) -> tuple[list[tuple[int, set[int]]], list[tuple[int, set[int]]]]:
|
|
"""
|
|
Inner function for the dataflow inplace transformation exposed for testing.
|
|
"""
|
|
if "PYTEST_CURRENT_TEST" not in os.environ:
|
|
logging.warning("The function dataflow_inplace_analysis is exposed for testing only.")
|
|
index_lists = tvm.get_global_func("relax.testing.transform.DataflowInplaceAnalysis")(
|
|
block, inputs, mod
|
|
) # type: ignore
|
|
|
|
def convert(opp_list):
|
|
return list(map(lambda opp: (int(opp.binding_idx), set(map(int, opp.arg_idxs))), opp_list))
|
|
|
|
return (convert(index_lists[0]), convert(index_lists[1])) # type: ignore
|
|
|
|
|
|
def dataflow_single_inplace_call(
|
|
mod: IRModule, call: Call, inplace_indices: list[int]
|
|
) -> tuple[Call, IRModule]:
|
|
"""
|
|
Inner function for the dataflow inplace transformation exposed for testing.
|
|
"""
|
|
if "PYTEST_CURRENT_TEST" not in os.environ:
|
|
logging.warning("The function dataflow_single_inplace_call is exposed for testing only.")
|
|
|
|
ret = tvm.get_global_func("relax.testing.transform.SingleInplaceCall")(
|
|
mod,
|
|
call,
|
|
inplace_indices,
|
|
) # type: ignore
|
|
return (ret[0], ret[1]) # type: ignore
|