change package name akg to _akg

pull/223/head
lizhenyu 5 years ago
parent 268d358a1d
commit 8cd3308a4c

@ -18,7 +18,7 @@ import sys
import os import os
def AKGAddPath(): def AKGAddPath():
"""akg add path.""" """_akg add path."""
pwd = os.path.dirname(os.path.realpath(__file__)) pwd = os.path.dirname(os.path.realpath(__file__))
tvm_path = os.path.realpath(pwd) tvm_path = os.path.realpath(pwd)
if tvm_path not in sys.path: if tvm_path not in sys.path:
@ -32,12 +32,12 @@ class AKGMetaPathFinder:
"""class AKGMetaPath finder.""" """class AKGMetaPath finder."""
def find_module(self, fullname, path=None): def find_module(self, fullname, path=None):
"""method akg find module.""" """method _akg find module."""
if fullname.startswith("akg.tvm"): if fullname.startswith("_akg.tvm"):
rname = fullname[4:] rname = fullname[5:]
return AKGMetaPathLoader(rname) return AKGMetaPathLoader(rname)
if fullname.startswith("akg.topi"): if fullname.startswith("_akg.topi"):
rname = fullname[4:] rname = fullname[5:]
return AKGMetaPathLoader(rname) return AKGMetaPathLoader(rname)
return None return None

@ -14,9 +14,9 @@
"""cast""" """cast"""
import logging import logging
import akg.tvm import _akg.tvm
from akg.ops.math import cast from _akg.ops.math import cast
from akg.topi.generic import schedule_elemwise from _akg.topi.generic import schedule_elemwise
def Cast(x, dst_type): def Cast(x, dst_type):
"""cast.""" """cast."""
@ -34,10 +34,10 @@ def gpu_schedule_Cast(outs):
sch (schedule.Schedule): The created schedule. sch (schedule.Schedule): The created schedule.
""" """
device = 'cuda' device = 'cuda'
ctx = akg.tvm.context(device, 0) ctx = _akg.tvm.context(device, 0)
if not ctx.exist: if not ctx.exist:
logging.info("Skip because %s is not enabled", device) logging.info("Skip because %s is not enabled", device)
return None return None
with akg.tvm.target.create(device): with _akg.tvm.target.create(device):
sch = schedule_elemwise(outs) sch = schedule_elemwise(outs)
return sch return sch

@ -15,7 +15,7 @@
"""default schedule function for GPU""" """default schedule function for GPU"""
from queue import Queue from queue import Queue
import akg.tvm as tvm import _akg.tvm as tvm
DEFAULT_GPU_THREAD = 1024 DEFAULT_GPU_THREAD = 1024
@ -31,7 +31,7 @@ def default_schedule(outs):
sch (schedule.Schedule): The created schedule. sch (schedule.Schedule): The created schedule.
""" """
if not isinstance(outs, tvm.tensor.Tensor) and not isinstance(outs, list): if not isinstance(outs, tvm.tensor.Tensor) and not isinstance(outs, list):
raise ValueError("outs should be list of akg.tvm.tensor.Tensor or akg.tvm.tensor.Tensor") raise ValueError("outs should be list of _akg.tvm.tensor.Tensor or _akg.tvm.tensor.Tensor")
device = 'cuda' device = 'cuda'
ctx = tvm.context(device, 0) ctx = tvm.context(device, 0)
if not ctx.exist: if not ctx.exist:

@ -12,9 +12,9 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""equal""" """equal"""
import akg.tvm import _akg.tvm
from akg.ops.math import equal from _akg.ops.math import equal
from akg.topi.generic import schedule_elemwise from _akg.topi.generic import schedule_elemwise
def Equal(x, y): def Equal(x, y):
"""equal.""" """equal."""
@ -32,9 +32,9 @@ def gpu_schedule_Equal(outs):
sch (schedule.Schedule): The created schedule. sch (schedule.Schedule): The created schedule.
""" """
device = 'cuda' device = 'cuda'
ctx = akg.tvm.context(device, 0) ctx = _akg.tvm.context(device, 0)
if not ctx.exist: if not ctx.exist:
raise SystemError("Skip because %s is not enabled" % device) raise SystemError("Skip because %s is not enabled" % device)
with akg.tvm.target.create(device): with _akg.tvm.target.create(device):
sch = schedule_elemwise(outs) sch = schedule_elemwise(outs)
return sch return sch

@ -13,8 +13,8 @@
# limitations under the License. # limitations under the License.
"""mean op compute and schedule""" """mean op compute and schedule"""
import akg.tvm as tvm import _akg.tvm as tvm
from akg.ops.math.mean import mean from _akg.ops.math.mean import mean
from .default_schedule import DEFAULT_GPU_THREAD from .default_schedule import DEFAULT_GPU_THREAD
def Mean(x, axis=None, keepdims=True): def Mean(x, axis=None, keepdims=True):

@ -13,9 +13,9 @@
# limitations under the License. # limitations under the License.
"""mean_grad""" """mean_grad"""
import akg.tvm as tvm import _akg.tvm as tvm
import akg import _akg
from akg.ops.math import mean from _akg.ops.math import mean
from .default_schedule import DEFAULT_GPU_THREAD from .default_schedule import DEFAULT_GPU_THREAD
@ -30,7 +30,7 @@ def mean_ad(head, input_shape, axis, keepdims):
if tensor_b.op.name == "mean_output": if tensor_b.op.name == "mean_output":
tensor_b = tensor_b.op.input_tensors[0] tensor_b = tensor_b.op.input_tensors[0]
jacs = list(akg.differentiate(tensor_b, [tensor_a], head)) jacs = list(_akg.differentiate(tensor_b, [tensor_a], head))
return jacs[0] return jacs[0]

@ -13,9 +13,9 @@
# limitations under the License. # limitations under the License.
"""mul""" """mul"""
import akg.topi as topi import _akg.topi as topi
import akg.tvm as tvm import _akg.tvm as tvm
from akg.ops.math import mul from _akg.ops.math import mul
def Mul(x, y): def Mul(x, y):
"""mul.""" """mul."""

@ -13,9 +13,9 @@
# limitations under the License. # limitations under the License.
"""relu6""" """relu6"""
import akg.topi as topi import _akg.topi as topi
import akg.tvm as tvm import _akg.tvm as tvm
from akg.topi import tag from _akg.topi import tag
@tvm.tag_scope(tag=tag.ELEMWISE) @tvm.tag_scope(tag=tag.ELEMWISE)
def topi_nn_relu6(x): def topi_nn_relu6(x):

@ -13,8 +13,8 @@
# limitations under the License. # limitations under the License.
"""relu6 grad""" """relu6 grad"""
import akg.topi as topi import _akg.topi as topi
import akg.tvm as tvm import _akg.tvm as tvm
def ReLU6Grad(y_grad, x): def ReLU6Grad(y_grad, x):
""" """

@ -13,8 +13,8 @@
# limitations under the License. # limitations under the License.
"""squeeze""" """squeeze"""
import akg.topi as topi import _akg.topi as topi
import akg.tvm as tvm import _akg.tvm as tvm
def Squeeze(x, axis=None): def Squeeze(x, axis=None):
""" """

@ -13,7 +13,7 @@
# limitations under the License. # limitations under the License.
"""squeeze grad""" """squeeze grad"""
import akg.topi as topi import _akg.topi as topi
def SqueezeGrad(y_grad, x_shape, axis=None): def SqueezeGrad(y_grad, x_shape, axis=None):
""" """

@ -12,9 +12,9 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
"""tile""" """tile"""
import akg.tvm import _akg.tvm
from akg.ops.array import tile from _akg.ops.array import tile
from akg.topi.generic import schedule_elemwise from _akg.topi.generic import schedule_elemwise
def Tile(x, multiples): def Tile(x, multiples):
"""tile.""" """tile."""
@ -31,9 +31,9 @@ def gpu_schedule_Tile(outs):
sch (schedule.Schedule): The created schedule. sch (schedule.Schedule): The created schedule.
""" """
device = 'cuda' device = 'cuda'
ctx = akg.tvm.context(device, 0) ctx = _akg.tvm.context(device, 0)
if not ctx.exist: if not ctx.exist:
raise SystemError("Skip because %s is not enabled" % device) raise SystemError("Skip because %s is not enabled" % device)
with akg.tvm.target.create(device): with _akg.tvm.target.create(device):
s = schedule_elemwise(outs) s = schedule_elemwise(outs)
return s return s

@ -20,9 +20,9 @@ import logging
import traceback import traceback
import os.path import os.path
from pathlib import Path from pathlib import Path
import akg.tvm import _akg.tvm
from akg.utils import validation_check as vc_util from _akg.utils import validation_check as vc_util
from akg.utils.dsl_create import TensorUtils from _akg.utils.dsl_create import TensorUtils
from . import gpu from . import gpu
from . import op_build from . import op_build
@ -67,7 +67,7 @@ def compilewithjson(json_str):
tensor_shape = input_desc[0]['shape'] tensor_shape = input_desc[0]['shape']
tensor_shape = (1,) if not tensor_shape else tensor_shape tensor_shape = (1,) if not tensor_shape else tensor_shape
vc_util.shape_dtype_max_size_check(tensor_shape) vc_util.shape_dtype_max_size_check(tensor_shape)
args[input_desc[0]['name']] = akg.tvm.placeholder( args[input_desc[0]['name']] = _akg.tvm.placeholder(
shape=tensor_shape, name=input_desc[0]['tensor_name'], dtype=input_desc[0]['data_type']) shape=tensor_shape, name=input_desc[0]['tensor_name'], dtype=input_desc[0]['data_type'])
tsr.append(args[input_desc[0]['name']]) tsr.append(args[input_desc[0]['name']])
else: else:
@ -76,7 +76,7 @@ def compilewithjson(json_str):
tensor_shape = tmp_desc['shape'] tensor_shape = tmp_desc['shape']
tensor_shape = (1,) if not tensor_shape else tensor_shape tensor_shape = (1,) if not tensor_shape else tensor_shape
vc_util.shape_dtype_max_size_check(tensor_shape) vc_util.shape_dtype_max_size_check(tensor_shape)
tmp_input.append(akg.tvm.placeholder( tmp_input.append(_akg.tvm.placeholder(
shape=tensor_shape, name=tmp_desc['tensor_name'], dtype=tmp_desc['data_type'])) shape=tensor_shape, name=tmp_desc['tensor_name'], dtype=tmp_desc['data_type']))
args[input_desc[0]['name']] = tmp_input args[input_desc[0]['name']] = tmp_input
tsr = tsr + tmp_input tsr = tsr + tmp_input

@ -19,10 +19,10 @@ import types
import typing import typing
import logging import logging
import traceback import traceback
import akg.tvm import _akg.tvm
import akg import _akg
from akg import save_gpu_param as gpu_utils from _akg import save_gpu_param as gpu_utils
from akg.utils import validation_check as vc_util from _akg.utils import validation_check as vc_util
MS_CUDA_KERNEL_PATH = "/tmp/cuda_meta/" MS_CUDA_KERNEL_PATH = "/tmp/cuda_meta/"
@ -38,21 +38,21 @@ def op_build(opnames, computes, args, custom_schedule, device, kernel_name, attr
return None return None
schedule_name = 'gpu_schedule_' + opnames[0] schedule_name = 'gpu_schedule_' + opnames[0]
schedule_func = getattr(akg.gpu, schedule_name) schedule_func = getattr(_akg.gpu, schedule_name)
if not isinstance(schedule_func, (types.FunctionType, typing.Callable)): if not isinstance(schedule_func, (types.FunctionType, typing.Callable)):
logging.error("no schedule func found %s", str(schedule_name)) logging.error("no schedule func found %s", str(schedule_name))
return None return None
ptx_file = os.path.realpath(MS_CUDA_KERNEL_PATH + kernel_name + ".ptx") ptx_file = os.path.realpath(MS_CUDA_KERNEL_PATH + kernel_name + ".ptx")
if os.path.exists(ptx_file): if os.path.exists(ptx_file):
os.remove(ptx_file) os.chmod(ptx_file, 0o600)
try: try:
with open(ptx_file, 'at') as file: with open(ptx_file, 'at') as file:
fcntl.flock(file.fileno(), fcntl.LOCK_EX) fcntl.flock(file.fileno(), fcntl.LOCK_EX)
file.seek(0, 2) file.seek(0, 2)
if file.tell() == 0: if file.tell() == 0:
s = schedule_func(computes) s = schedule_func(computes)
foo = akg.tvm.build(s, args, device, name=kernel_name) foo = _akg.tvm.build(s, args, device, name=kernel_name)
ptx_code = foo.imported_modules[0].get_source("ptx") ptx_code = foo.imported_modules[0].get_source("ptx")
file.write(ptx_code) file.write(ptx_code)
json_file = os.path.realpath(MS_CUDA_KERNEL_PATH + kernel_name + ".json") json_file = os.path.realpath(MS_CUDA_KERNEL_PATH + kernel_name + ".json")

@ -13,12 +13,12 @@
# limitations under the License. # limitations under the License.
"""operator dsl function: tile""" """operator dsl function: tile"""
import akg.tvm import _akg.tvm
import akg.topi import _akg.topi
from akg.utils import validation_check as vc_util from _akg.utils import validation_check as vc_util
@vc_util.check_input_type(akg.tvm.tensor.Tensor, (list, tuple)) @vc_util.check_input_type(_akg.tvm.tensor.Tensor, (list, tuple))
def tile(data, multiples): def tile(data, multiples):
""" """
Repeats the data in the specified dimensions according to the multiples. Repeats the data in the specified dimensions according to the multiples.
@ -32,5 +32,5 @@ def tile(data, multiples):
""" """
vc_util.check_shape(data.shape) vc_util.check_shape(data.shape)
vc_util.check_int_list(multiples, "multiples") vc_util.check_int_list(multiples, "multiples")
output = akg.topi.tile(data, multiples) output = _akg.topi.tile(data, multiples)
return output return output

@ -13,12 +13,12 @@
# limitations under the License. # limitations under the License.
"""operator dsl function: cast""" """operator dsl function: cast"""
import akg.tvm import _akg.tvm
import akg.topi import _akg.topi
from akg.utils import validation_check as vc_util from _akg.utils import validation_check as vc_util
@vc_util.check_input_type(akg.tvm.tensor.Tensor, str) @vc_util.check_input_type(_akg.tvm.tensor.Tensor, str)
def cast(data, dst_type): def cast(data, dst_type):
""" """
cast data to target type. cast data to target type.
@ -31,6 +31,6 @@ def cast(data, dst_type):
tvm.tensor.Tensor, type is dst_type. tvm.tensor.Tensor, type is dst_type.
""" """
vc_util.check_shape(data.shape) vc_util.check_shape(data.shape)
out = akg.topi.cast(data, dst_type) out = _akg.topi.cast(data, dst_type)
return out return out

@ -13,13 +13,13 @@
# limitations under the License. # limitations under the License.
"""operator dsl function: equal""" """operator dsl function: equal"""
import akg.tvm import _akg.tvm
import akg.topi import _akg.topi
from akg.utils.dsl_create import produce_shapes from _akg.utils.dsl_create import produce_shapes
from akg.utils import validation_check as vc_util from _akg.utils import validation_check as vc_util
@vc_util.check_input_type(akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor) @vc_util.check_input_type(_akg.tvm.tensor.Tensor, _akg.tvm.tensor.Tensor)
def equal(input1, input2): def equal(input1, input2):
""" """
check whether input1 equals to input2. check whether input1 equals to input2.
@ -42,13 +42,13 @@ def equal(input1, input2):
dtype = input1.dtype dtype = input1.dtype
# get equal compute # get equal compute
t_value = akg.tvm.compute(shape, lambda *indice: akg.tvm.const(1, dtype), "T") t_value = _akg.tvm.compute(shape, lambda *indice: _akg.tvm.const(1, dtype), "T")
f_value = akg.tvm.compute(shape, lambda *indice: akg.tvm.const(0, dtype), "F") f_value = _akg.tvm.compute(shape, lambda *indice: _akg.tvm.const(0, dtype), "F")
input1_bro = akg.topi.broadcast_to(input1, shape) input1_bro = _akg.topi.broadcast_to(input1, shape)
input2_bro = akg.topi.broadcast_to(input2, shape) input2_bro = _akg.topi.broadcast_to(input2, shape)
c_out = akg.tvm.compute(shape, lambda *indice: akg.tvm.expr.Select(input1_bro[indice] == input2_bro[indice], c_out = _akg.tvm.compute(shape, lambda *indice: _akg.tvm.expr.Select(input1_bro[indice] == input2_bro[indice],
t_value[indice], f_value[indice]), name="C") t_value[indice], f_value[indice]), name="C")
res = akg.tvm.compute(shape, lambda *indice: c_out(*indice).astype("bool"), name="res") res = _akg.tvm.compute(shape, lambda *indice: c_out(*indice).astype("bool"), name="res")
return res return res

@ -13,14 +13,14 @@
# limitations under the License. # limitations under the License.
"""operator dsl function: mean""" """operator dsl function: mean"""
import akg.topi import _akg.topi
import akg.tvm import _akg.tvm
from akg.utils import format_transform as ft_util from _akg.utils import format_transform as ft_util
from akg.utils import validation_check as vc_util from _akg.utils import validation_check as vc_util
from akg.ops.math import sum from _akg.ops.math import sum
@vc_util.check_input_type(akg.tvm.tensor.Tensor, (list, tuple, int, type(None)), (bool, type(None))) @vc_util.check_input_type(_akg.tvm.tensor.Tensor, (list, tuple, int, type(None)), (bool, type(None)))
def mean(data, axis=None, keepdims=False): def mean(data, axis=None, keepdims=False):
""" """
Computes the mean of the values of a Tensor over the whole dataset. Computes the mean of the values of a Tensor over the whole dataset.
@ -42,6 +42,6 @@ def mean(data, axis=None, keepdims=False):
for i in axis: for i in axis:
count *= shape[i] count *= shape[i]
output, _ = sum.sum_value(data, axis, keepdims) output, _ = sum.sum_value(data, axis, keepdims)
res = akg.topi.divide(output, count) res = _akg.topi.divide(output, count)
return res return res

@ -13,11 +13,11 @@
# limitations under the License. # limitations under the License.
"""operator dsl function: mul""" """operator dsl function: mul"""
import akg.topi import _akg.topi
from akg.utils import validation_check as vc_util from _akg.utils import validation_check as vc_util
@vc_util.check_input_type(akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor) @vc_util.check_input_type(_akg.tvm.tensor.Tensor, _akg.tvm.tensor.Tensor)
def mul(l_input, r_input): def mul(l_input, r_input):
""" """
Calculate x * y element-wise. Calculate x * y element-wise.
@ -38,6 +38,6 @@ def mul(l_input, r_input):
vc_util.check_shape(shape2) vc_util.check_shape(shape2)
vc_util.auto_broadcast_check(shape1, shape2) vc_util.auto_broadcast_check(shape1, shape2)
vc_util.elemwise_dtype_check(l_input.dtype, r_input.dtype) vc_util.elemwise_dtype_check(l_input.dtype, r_input.dtype)
output = akg.topi.multiply(l_input, r_input) output = _akg.topi.multiply(l_input, r_input)
return output return output

@ -13,12 +13,12 @@
# limitations under the License. # limitations under the License.
"""operator dsl function: sub""" """operator dsl function: sub"""
import akg.topi import _akg.topi
import akg.tvm import _akg.tvm
from akg.utils import validation_check as vc_util from _akg.utils import validation_check as vc_util
@vc_util.check_input_type(akg.tvm.tensor.Tensor, akg.tvm.tensor.Tensor) @vc_util.check_input_type(_akg.tvm.tensor.Tensor, _akg.tvm.tensor.Tensor)
def sub(data1, data2): def sub(data1, data2):
""" """
Computes data1 - data2 elementwise, broadcast is supported. Computes data1 - data2 elementwise, broadcast is supported.
@ -35,6 +35,6 @@ def sub(data1, data2):
vc_util.check_shape(data2.shape) vc_util.check_shape(data2.shape)
vc_util.auto_broadcast_check(data1.shape, data2.shape) vc_util.auto_broadcast_check(data1.shape, data2.shape)
res = akg.topi.subtract(data1, data2) res = _akg.topi.subtract(data1, data2)
return res return res

@ -14,13 +14,13 @@
"""operator dsl function: sum""" """operator dsl function: sum"""
import akg.topi import _akg.topi
import akg.tvm import _akg.tvm
from akg.utils import format_transform as ft_util from _akg.utils import format_transform as ft_util
from akg.utils import validation_check as vc_util from _akg.utils import validation_check as vc_util
@vc_util.check_input_type(akg.tvm.tensor.Tensor, (list, tuple, int, type(None)), (bool, type(None))) @vc_util.check_input_type(_akg.tvm.tensor.Tensor, (list, tuple, int, type(None)), (bool, type(None)))
def sum_value(inputs, axis=None, keepdims=False): def sum_value(inputs, axis=None, keepdims=False):
""" """
Compute the sum of elements across dimensions of a tensor. Compute the sum of elements across dimensions of a tensor.
@ -38,8 +38,8 @@ def sum_value(inputs, axis=None, keepdims=False):
vc_util.check_shape(inputs.shape) vc_util.check_shape(inputs.shape)
if not axis: if not axis:
output = akg.topi.identity(inputs) output = _akg.topi.identity(inputs)
else: else:
output = akg.topi.sum(inputs, axis=axis, keepdims=keepdims) output = _akg.topi.sum(inputs, axis=axis, keepdims=keepdims)
return output return output

Some files were not shown because too many files have changed in this diff Show More

Loading…
Cancel
Save