Release 1.0.2

Author: Pavel Sukhachev

Backend/Benchmarks.py

@@ -9,6 +9,10 @@
 mulMatrixOnMatrix = None
 sumOnMatrix = None
 
+mulTensorOnVecGroup = None
+sumOnTensorGroup = None
+mulTensorBatch = None
+
 
 def autoinit():
 	if not Config.shouldInit():
@@ -37,7 +41,7 @@ def initHip():
 
 
 def initGPU(Backend):
-	backend = Backend.getBackend(Config.deviceIdx, initmode=2)
+	backend = Backend.getBackend(Config.deviceIdx, initmode=2, logger=Config.getLogger())
 	GPUArray, memoryPool, blas, matmod = backend.GPUArray, backend.memoryPool, backend.blas, backend.matmod
 
 	def wrapToVectorAddVector(y, x, alpha=1.0):
@@ -70,6 +74,33 @@ def wrapSumOnMatrix(A, out=None, cols=True, alpha=1.0, beta=0.0):
 	mulMatrixOnMatrix = wrapGemm
 	sumOnMatrix = wrapSumOnMatrix
 
+	formats = {
+		"gbp": backend.GroupFormat.gbp.value,
+		"bgp": backend.GroupFormat.bgp.value
+	}
+
+	def wrapMulTensorOnVecGroup(tensor, vecs, out=None, formatT="bgp", transpT=False, alpha=1.0, beta=0.0):
+		assert tensor.ndim == 3 and formatT == "gbp"
+		axis = 0 if transpT else 1
+
+		return matmod.matvec(tensor, vecs, axis, out, alpha, beta, memoryPool)
+
+	def wrapSumOnTensorGroup(tensor, out=None, formatT="bgp", cols=True, alpha=1.0, beta=0.0):
+		assert tensor.ndim == 3
+		axis = (1 if formatT == "gbp" else 0) if cols else 2
+
+		return matmod.matsum(tensor, axis, out, alpha, beta, memoryPool)
+
+	def wrapMulTensorBatch(A, B, formatA="bgp", formatB="bgp", out=None, formatOut="bgp", transpA=False, transpB=False,
+						   alpha=1.0, beta=0.0):
+		formatA, formatB, formatOut = formats[formatA], formats[formatB], formats[formatOut]
+		return blas.gemmBatched(A, B, formatA, formatB, formatOut, transpA, transpB, alpha, beta, out, memoryPool)
+
+	global mulTensorOnVecGroup, sumOnTensorGroup, mulTensorBatch
+	mulTensorOnVecGroup = wrapMulTensorOnVecGroup
+	sumOnTensorGroup = wrapSumOnTensorGroup
+	mulTensorBatch = wrapMulTensorBatch
+
 
 def initCPU():
 	from PuzzleLib.CPU.Wrappers import NumpyBlas

Backend/Blas.py

@@ -34,6 +34,29 @@
 crossMapLRNBackward = None
 
 
+instanceNorm2d = None
+instanceNorm2dBackward = None
+
+
+spatialTf = None
+spatialTfBackward = None
+
+
+RNNMode = None
+DirectionMode = None
+
+createRnn = None
+
+acquireRnnParams = None
+updateRnnParams = None
+
+forwardRnn = None
+backwardDataRnn = None
+backwardParamsRnn = None
+
+deviceSupportsBatchHint = None
+
+
 def autoinit():
 	if not Config.shouldInit():
 		return
@@ -87,10 +110,19 @@ def wrapCrossMapLRNBackward(data, outdata, grad, _, N, alpha, beta, K):
 	crossMapLRN = wrapCrossMapLRN
 	crossMapLRNBackward = wrapCrossMapLRNBackward
 
+	def wrapSpatialTf(data, transform, outshape, getGrid):
+		return dnn.spatialTf(data, transform, outshape, getGrid, allocator=memoryPool)
+
+	def wrapSpatialTfBackward(grad, data, grid):
+		return dnn.spatialTfBackward(grad, data, grid, allocator=memoryPool)
+
+	global spatialTf, spatialTfBackward
+	spatialTf = wrapSpatialTf
+	spatialTfBackward = wrapSpatialTfBackward
+
 
 def initHip():
 	from PuzzleLib.Hip import Backend
-	initGPU(Backend)
 
 	backend = initGPU(Backend)
 	memoryPool, dnn = backend.memoryPool, backend.dnn
@@ -125,9 +157,17 @@ def wrapMapLRNBackward(data, outdata, grad, means, workspace, N, alpha, beta, K)
 
 
 def initGPU(Backend):
-	import numpy as np
+	backend = Backend.getBackend(Config.deviceIdx, initmode=1, logger=Config.getLogger())
+
+	initBaseGPU(backend)
+	initInstanceNormGPU(backend)
+	initRnnGPU(backend)
+
+	return backend
+
 
-	backend = Backend.getBackend(Config.deviceIdx, initmode=1)
+def initBaseGPU(backend):
+	import numpy as np
 	memoryPool, dnn = backend.memoryPool, backend.dnn
 
 	global ConvFwdAlgo, ConvBwdDataAlgo, ConvBwdFilterAlgo
@@ -143,7 +183,7 @@ def wrapConvNd(data, W, bias, stride, pad, dilation, groups, algo):
 
 	def wrapConvNdBackwardData(grad, W, data, stride, pad, dilation, groups, algo):
 		return dnn.convNdBackwardData(
-			grad, W, None, data, stride, pad, dilation, groups, algo.value, None, memoryPool
+			grad, W, None, data, stride, pad, dilation, None, groups, algo.value, None, memoryPool
 		)
 
 	def wrapConvNdBackwardParams(data, grad, W, bias, stride, pad, dilation, groups,
@@ -168,9 +208,9 @@ def wrapConvNdbenchmark(datashape, Wshape, stride, pad, dilation, groups, transp
 	global convNdbenchmark
 	convNdbenchmark = wrapConvNdbenchmark
 
-	def wrapDeconvNd(data, W, bias, stride, pad, dilation, groups, algo):
+	def wrapDeconvNd(data, W, bias, stride, pad, dilation, postpad, groups, algo):
 		return dnn.convNdBackwardData(
-			data, W, bias.ravel() if bias is not None else None, None, stride, pad, dilation, groups,
+			data, W, bias.ravel() if bias is not None else None, None, stride, pad, dilation, postpad, groups,
 			algo.value, None, memoryPool
 		)
 
@@ -231,15 +271,79 @@ def wrapSoftmaxNdBackward(outdata, grad):
 	softmaxNd = wrapSoftmaxNd
 	softmaxNdBackward = wrapSoftmaxNdBackward
 
+
+def initInstanceNormGPU(backend):
+	memoryPool = backend.memoryPool
+
+	def wrapInstanceNorm2d(data, scale, bias, epsilon=1e-5):
+		return backend.instanceNorm2d(data, scale.ravel(), bias.ravel(), epsilon, allocator=memoryPool)
+
+	def wrapInstanceNorm2dBackward(grad, data, extscale, savemean, saveinvvar, epsilon, affine):
+		return backend.instanceNorm2dBackward(
+			grad, data, extscale, savemean, saveinvvar, epsilon, affine, allocator=memoryPool
+		)
+
+	global instanceNorm2d, instanceNorm2dBackward
+	instanceNorm2d = wrapInstanceNorm2d
+	instanceNorm2dBackward = wrapInstanceNorm2dBackward
+
+
+def initRnnGPU(backend):
+	import numpy as np
+	memoryPool = backend.memoryPool
+
+	global RNNMode, DirectionMode
+	RNNMode = backend.RNNMode
+	DirectionMode = backend.DirectionMode
+
+	def wrapCreateRnn(insize, hsize, layers, mode, direction, dropout, seed, batchsize):
+		rnn, W, params = backend.createRnn(
+			insize, hsize, np.float32, layers, mode=mode, direction=direction, dropout=dropout,
+			seed=seed, batchsize=0 if batchsize is None else batchsize
+		)
+
+		return rnn, W, {i: layer for i, layer in enumerate(params)}
+
+	def wrapAcquireRnnParams(descRnn, w):
+		params = backend.acquireRnnParams(descRnn, w)
+		return w, params
+
+	def wrapUpdateRnnParams(descRnn, w, params):
+		params = [params[layer] for layer in sorted(params.keys())]
+		backend.updateRnnParams(descRnn, w, params)
+
+	global createRnn, acquireRnnParams, updateRnnParams
+	createRnn = wrapCreateRnn
+	acquireRnnParams = wrapAcquireRnnParams
+	updateRnnParams = wrapUpdateRnnParams
+
+	def wrapForwardRnn(data, W, descRnn, test=False):
+		return descRnn.forward(data, W, test=test, allocator=memoryPool)
+
+	def wrapBackwardDataRnn(grad, outdata, W, reserve, descRnn):
+		ingrad, _, _ = descRnn.backwardData(grad, outdata, W, reserve, allocator=memoryPool)
+		return ingrad, reserve
+
+	def wrapBackwardParamsRnn(data, outdata, _, reserve, descRnn):
+		return descRnn.backwardParams(data, outdata, reserve, allocator=memoryPool)
+
+	global forwardRnn, backwardDataRnn, backwardParamsRnn
+	forwardRnn = wrapForwardRnn
+	backwardDataRnn = wrapBackwardDataRnn
+	backwardParamsRnn = wrapBackwardParamsRnn
+
+	global deviceSupportsBatchHint
+	deviceSupportsBatchHint = backend.deviceSupportsBatchHint
+
 	return backend
 
 
 def initCPU():
 	from PuzzleLib.CPU.Wrappers import NumpyDnn
 
 	def wrapConvNd(data, W, bias, stride, pad, dilation, groups, algo):
-		assert dilation == (1, 1) and groups == 1
-		return NumpyDnn.conv2d(data, W, bias, stride, pad)
+		assert groups == 1
+		return NumpyDnn.conv2d(data, W, bias, stride, pad, dilation)
 
 	global convNd, convNdBackwardData, convNdBackwardParams
 	convNd = wrapConvNd
@@ -263,9 +367,12 @@ def wrapBatchNormNd(data, scale, bias, mean, var, epsilon, factor, test, mode=No
 	BatchNormMode = ProxyBatchNormMode
 	batchNormNd = wrapBatchNormNd
 
+	global deviceSupportsBatchHint
+	deviceSupportsBatchHint = lambda: False
+
 
 def initIntel():
-	from PuzzleLib.Intel.Wrappers import DNNL
+	from PuzzleLib.Intel.Wrappers import DNNL, DNNLInstanceNorm
 
 	global ConvFwdAlgo, ConvBwdDataAlgo, ConvBwdFilterAlgo
 	ConvFwdAlgo = DNNL.ConvAlgo
@@ -299,9 +406,9 @@ def wrapConvNdbenchmark(datashape, Wshape, stride, pad, dilation, groups, transp
 	global convNdbenchmark
 	convNdbenchmark = wrapConvNdbenchmark
 
-	def wrapDeconvNd(data, W, bias, stride, pad, dilation, groups, algo):
+	def wrapDeconvNd(data, W, bias, stride, pad, dilation, groups, postpad, algo):
 		assert groups == 1
-		return DNNL.convNd(data, W, bias, stride, pad, dilation, algo=algo, transpose=True)
+		return DNNL.convNd(data, W, bias, stride, pad, dilation, postpad, algo=algo, transpose=True)
 
 	def wrapDeconvNdBackwardData(grad, W, data, stride, pad, dilation, groups, algo):
 		assert groups == 1
@@ -360,5 +467,24 @@ def wrapCrossMapLRNBackward(data, outdata, grad, workspace, N, alpha, beta, K):
 	crossMapLRN = wrapCrossMapLRN
 	crossMapLRNBackward = wrapCrossMapLRNBackward
 
+	def wrapInstanceNorm2d(data, scale, bias, epsilon):
+		result = DNNLInstanceNorm.instanceNorm2d(data, scale, bias, epsilon)
+
+		outdata, savemean, savevar, extscale, extbias, desc = result
+		return outdata, savemean, savevar, (extscale, extbias, desc)
+
+	def wrapInstanceNorm2dBackward(grad, data, exts, savemean, savevar, epsilon, affine):
+		extscale, extbias, desc = exts
+		return DNNLInstanceNorm.instanceNorm2dBackward(
+			grad, data, extscale, extbias, savemean, savevar, epsilon, desc, affine
+		)
+
+	global instanceNorm2d, instanceNorm2dBackward
+	instanceNorm2d = wrapInstanceNorm2d
+	instanceNorm2dBackward = wrapInstanceNorm2dBackward
+
+	global deviceSupportsBatchHint
+	deviceSupportsBatchHint = lambda: False
+
 
 autoinit()