gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[PYTHON][TENSORFLOW] More bugfixes for forward/backward signatures

Browse Source
This commit is contained in:
Philippe Tillet
2019-10-31 01:49:30 -04:00
parent 93a86d4fc6
commit 91a2fd463b
6 changed files with 74 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
python/triton/function.py
View File

@@ -1,10 +1,14 @@
import triton.frameworks as fw
import triton.utils
import triton.utils as utils
class OpContext(object):
def __init__(self):
self.to_save = []
def save_for_backward(self, *tensors):
self.to_save = tensors
self.to_save = [x.to_tensor() if isinstance(x, utils.tf_empty_proxy) else x
for x in tensors]
@property
def saved_tensors(self):
@@ -16,7 +20,7 @@ class function_meta(type):
cls.registered = False
return super(function_meta, cls).__init__(name, bases, attrs)
ctx_registry = triton.utils.id_dict()
ctx_registry = utils.id_dict()
class function(metaclass = function_meta):
@@ -43,13 +47,39 @@ class function(metaclass = function_meta):
@classmethod
def extract_tf_tensors(cls, lst, err):
ret = []
for x in lst:
if x and not isinstance(x, triton.utils.tf_empty_proxy):
raise ValueError('Results of ' + err + ' must be created using triton.empty()')
if x and x.tensor is None:
raise ValueError('Empty tensor never filled during ' + err)
return [x.tensor if x else None for x in lst]
if x is None:
ret += [None]
elif isinstance(x, fw.tensorflow.Tensor):
ret += [x]
elif isinstance(x, utils.tf_empty_proxy):
if x.tensor is None:
raise ValueError('Empty tensor never filled during ' + err)
else:
ret += [x.tensor]
else:
raise ValueError('Unsupported return type', type(x))
return ret
@classmethod
def map_in_to_args(cls, op, args):
ret = dict()
for i, ix in enumerate(op.inputs):
for j, jx in enumerate(args):
if ix is jx:
ret[j] = i
return ret
@classmethod
def map_res_to_out(cls, op, result):
ret = []
for i, ix in enumerate(result):
for j, jx in enumerate(op.outputs):
if ix is jx:
ret.append(j)
return ret
@classmethod
def apply_tensorflow(cls, *args, **kwargs):
ctx = OpContext()
@@ -60,30 +90,21 @@ class function(metaclass = function_meta):
result = result if isinstance(result, tuple) else (result, )
result = function.extract_tf_tensors(result, 'forward')
# Find a mapping between ::forward arguments and tensorflow op arguments
op = result[0].op
remap_in = dict()
for i, ix in enumerate(op.inputs):
for j, jx in enumerate(args):
if ix is jx:
remap_in[j] = i
remap_out = []
for i, ix in enumerate(result):
for j, jx in enumerate(op.outputs):
if ix is jx:
remap_out.append(j)
# Register backward pass
ctx_registry[op] = ctx
key = result[0]
op = result[0].op
ctx_registry[key] = ctx
remap_in = cls.map_in_to_args(op, args)
remap_out = cls.map_res_to_out(op, result)
name = op.op_def.name
if not cls.registered:
@fw.tensorflow.RegisterGradient(name)
def gradient(op, *dy):
dy = dy if len(dy) > 1 else dy[0]
# Remap gradient inputs in the right order
grads = [dy[i] for i in remap_out]
dy = [dy[i] for i in remap_out]
dy = dy if len(dy) > 1 else dy[0]
# Execute gradient function
grad = cls.backward(ctx_registry[op], grads)
grad = cls.backward(ctx_registry[key], dy)
grad = function.extract_tf_tensors(grad, 'backward')
# Remap gradient in the right order
ret = [None] * len(op.inputs)
@@ -95,7 +116,7 @@ class function(metaclass = function_meta):
cls.registered = True
# Return tensor
return result
return result[0] if len(result)==1 else result
@classmethod
def apply(cls, *args, **kwargs):

16
python/triton/kernel.py
View File

@@ -237,14 +237,20 @@ class kernel:
kwargs['T' + str(i)] = x.dtype
# launch
ret = self.fw_op(*operands, **kwargs)
# fill empty tensors with corresponding values
for j, y in enumerate(ret[0].op.op_def.output_arg):
for i, x in enumerate(ret[0].op.op_def.input_arg):
ret = [ret] if isinstance(ret, fw.tensorflow.Tensor) else ret
op_def = ret[0].op.op_def
# fill empty tensors with corresponding values
for j, y in enumerate(op_def.output_arg):
found = False
for i, x in enumerate(op_def.input_arg):
if y.name + '_shape' == x.name:
empty[i].tensor = ret[j]
args[i].tensor = ret[j]
found = True
assert found
# store timing information
if bench > 0:
bench_registry[ret] = triton.utils.id_dict.lazy_entry(bench_id)
for y in ret:
bench_registry[y] = triton.utils.id_dict.lazy_entry(bench_id)
elif fw.has_torch():
args = [x.contiguous() if isinstance(x, fw.torch.Tensor) else x for x in args[:-1]]
self.fw_op(op_id, bench, bench_id, *args)

2
python/triton/ops/batchnorm.py
View File

@@ -104,7 +104,7 @@ void bwdbatchnorm(float *DX, float *DG, float *DB,
lambda opt: [1, C],
TM = 128)
# save
ctx.save_for_backward(x, gamma, beta, mean.tensor, var.tensor)
ctx.save_for_backward(x, gamma, beta, mean, var)
ctx.eps = eps
return y

9
python/triton/ops/dot.py
View File

@@ -76,10 +76,11 @@ void dot(TYPE * A, TYPE * B, TYPE * C,
'BROADCAST_BK': 'newaxis, :' if transpose_b else ':, newaxis',
'BROADCAST_BN': ':, newaxis' if transpose_b else 'newaxis, :',
'SHAPE_B' : 'TN, TK' if transpose_b else 'TK, TN'}
return _dot.kernel(a, b, c, M, N, Ka, lda, ldb, ldc,
grid, bench=bench,
AT = transpose_a, BT = transpose_b, TYPE = dtype,
TM = [64, 128], TN = [64, 128], TK = [8], **macros)
_dot.kernel(a, b, c, M, N, Ka, lda, ldb, ldc,
grid, bench=bench,
AT = transpose_a, BT = transpose_b, TYPE = dtype,
TM = [64, 128], TN = [64, 128], TK = [8], **macros)
return c
@staticmethod
def forward(ctx, a, b, transpose_a = False, transpose_b = False, bench = 0):

17
python/triton/ops/einsum.py
View File

@@ -169,14 +169,15 @@ void einsumk(TYPE * A, TYPE * B, TYPE * C,
TN = [2**i for i in range(5, max(6, min(8, int(math.log2(bmnk[2]) + 1 ))))]
TB = [2**i for i in range(0, max(1, min(3, int(math.log2(bmnk[0]) + 1 ))))]
TK = [bmnk[2]] if bmnk[2] < 16 else [8, 16]
return _einsum.kernel(a, b, c,
bmnk[1], bmnk[2], bmnk[3],
std0[0], std0[1], std0[2],
std1[0], std1[1], std1[2],
grid, bench=bench,
**macros,
TYPE=dtype, TM=TM, TN=TN, TK=TK, TB=TB)
_einsum.kernel(a, b, c,
bmnk[1], bmnk[2], bmnk[3],
std0[0], std0[1], std0[2],
std1[0], std1[1], std1[2],
grid, bench=bench,
**macros,
TYPE=dtype, TM=TM, TN=TN, TK=TK, TB=TB)
return c
@staticmethod
def forward(ctx, subscripts, a, b, bench = 0):

2
python/triton/utils.py
View File

@@ -13,7 +13,7 @@ class tf_empty_proxy:
self.tensor = None
def to_tensor(self):
assert self.tensor
assert self.tensor is not None
return self.tensor
def empty(shape, dtype):