Skip to content
Permalink
Browse files

[PYTHON] Removed dead code for alloc_empty and register_scalar

  • Loading branch information
ptillet committed Oct 30, 2019
1 parent f4fcaf8 commit bf3dc63858faca9922a48bb65461d2e36d00889d
@@ -34,18 +34,6 @@ if(BUILD_PYTHON_MODULE)
# PyBind11 wrapper source file
file(GLOB_RECURSE PYTHON_SRC python/src/bindings.cc)
include_directories(python/src/ ${PYTHON_INCLUDE_DIRS})
if(TF_LIBS)
# extra tensorflow ops (e.g., alloc_empty)
# update directories
link_directories(${TF_LIB_DIRS})
include_directories(${TF_INCLUDE_DIRS})
# get sources
file(GLOB_RECURSE EXTRA_TF_OPS_SRC python/src/tensorflow/*.cc)
add_library(extra_tf_ops SHARED ${EXTRA_TF_OPS_SRC})
# create target
target_link_libraries(extra_tf_ops triton ${TF_LIBS})
target_compile_definitions(extra_tf_ops PRIVATE "-D_GLIBCXX_USE_CXX11_ABI=${TF_ABI}")
endif()
endif()


@@ -24,7 +24,7 @@ class MODE(Enum):
# Matmul
cases += [[[4, 1024, 1024], [1024, 1024], [4, 1024, 1024], "btc,ck->btk"]]
# Attention
cases += [[[4, 256, 8, 2, 64], [8, 2, 512, 64], [4, 256, 8, 2, 512], "bchak,hank->bchan"]]
# cases += [[[4, 256, 8, 2, 64], [8, 2, 512, 64], [4, 256, 8, 2, 512], "bchak,hank->bchan"]]

if mode == MODE.TF:
sess = tf.InteractiveSession()

This file was deleted.

This file was deleted.

@@ -4,8 +4,6 @@

torch = None
tensorflow = None
tf_extra_ops = None
gen_resource_variable_ops = None

def _import_torch():
global torch
@@ -14,24 +12,13 @@ def _import_torch():

def _import_tensorflow():
global tensorflow
global gen_resource_variable_ops
if tensorflow is None:
import tensorflow
from tensorflow.python.ops import gen_resource_variable_ops

def _import_tf_extra_ops():
global tf_extra_ops
if tf_extra_ops is None:
path = os.path.dirname(libtriton.__file__)
path = os.path.join(path, 'libextra_tf_ops.so')
_import_tensorflow()
tf_extra_ops = tensorflow.load_op_library(path)

def has_tensorflow():
result = 'tensorflow' in sys.modules
if result:
_import_tensorflow()
_import_tf_extra_ops()
return result

def has_torch():
@@ -219,7 +219,7 @@ def __call__(self, *args, **kwargs):
# retrieve framework op
op_id = self.fw_id[key]
# register grid
libtriton.register_grid(op_id, _make_grid(args))
libtriton.register_grid(op_id, args[-1])
# id for the benchmark result
bench_id = libtriton.make_scalar_id() if bench > 0 else -1
# call framework function
@@ -13,90 +13,22 @@ def __init__(self, shape, dtype):

def empty(shape, dtype):
if fw.has_tensorflow():
shape = [x.handle if isinstance(x, scalar) else fw.tensorflow.constant(x) for x in shape]
shape = [fw.tensorflow.constant(x) for x in shape]
shape = fw.tensorflow.stack(shape)
return tf_empty_proxy(shape, dtype)
#return fw.tf_extra_ops.alloc_empty(args, T = dtype)
elif fw.has_torch():
return fw.torch.empty(*shape).cuda()

class lazy_shape:

def __init__(self, shape):
self.shape = shape

def __getitem__(self, key):
return scalar(self.shape[key])

def shape(A) :
if fw.has_tensorflow():
return A.shape.as_list()
#return lazy_shape(fw.tensorflow.shape(A))
elif fw.has_torch():
return A.shape
else:
assert False


class scalar:

def __init__(self, x):
self.id = libtriton.make_scalar_id()
self.handle = fw.tf_extra_ops.register_scalar(x, id=self.id)
self.assume_initialized = False

def set_assume_initialized(self):
self.assume_initialized = True

def unset_assume_initialized(self):
self.assume_initialized = False

def get_value(self):
if self.assume_initialized:
return libtriton.retrieve_scalar(self.id)
else:
return self.handle

def __add__(self, other):
return self.get_value() + other

def __radd__(self, other):
return other + self.get_value()

def __sub__(self, other):
return self.get_value() - other

def __rsub(self, other):
return other - self.get_value()

def __mul__(self, other):
return self.get_value() * other

def __rmul(self, other):
return other * self.get_value()

def __floordiv__(self, other):
return self.get_value() // other

def __rfloordiv__(self, other):
return other // self.get_value()

def __div__(self, other):
return self.get_value() / other

def __rdiv__(self, other):
return other / self.get_value()

def __truediv__(self, other):
self.get_value().__truediv__(other)

def __rtruediv__(self, other):
other.__truediv__(self.get_value())

def __neg__(self):
return -self.get_value()


class id_dict:

# Lazy entry for e.g., tensorflow, when value of benchmark is

0 comments on commit bf3dc63

Please sign in to comment.
You can’t perform that action at this time.