def get_options_from_kwargs_and_tuner_cache(name, cache_file, options_cache,
*inputs, **kwargs):
options = None
if "options" in kwargs and kwargs["options"] is not None:
options = kwargs["options"]
assert "type" in kwargs, "tuning layer type not specified: forward/backward"
# if we pass separate options for forward/backward, we use them otherwise
# use the same options
if isinstance(options, list) and len(options) == 2:
options = options[0] if kwargs["type"] == "forward" else options[1]
elif cache_file and isinstance(kwargs["cache"], str):
options = get_options_from_cache_file(name, *inputs, **kwargs)
elif options_cache and kwargs["type"] in options_cache and options_cache[
kwargs["type"]] is not None:
options = options_cache[kwargs["type"]]
logger.info(
"Kernel was previously tuned, seeding the current tuning with those mapping options"
)
if options is None:
options = Options("naive")
logger.warning(
"Using 'naive' type mapping options for autotuning. See help(your_layer.autotune) for how to set mapping options."
)
if not isinstance(options, Options):
options = Options(options)
return options
def get_options_from_kwargs(name, *inputs, **kwargs):
# now the options can be a tuple (if training) or it will be just options
# (only forward)
options = None
if "options" in kwargs and kwargs["options"] is not None:
options = kwargs["options"]
assert "type" in kwargs, "layer type not specified: forward/backward"
if isinstance(options, list) and len(options) == 2:
options = options[0] if kwargs["type"] == "forward" else options[1]
elif "training" in kwargs and kwargs["training"] and kwargs[
"type"] == "backward":
logger.warning(
'Same mapping options will be used to run backward layer, please pass backward mapping options for better performance.'
)
elif "cache" in kwargs and kwargs["cache"]:
options = get_options_from_cache_file(name, *inputs, **kwargs)
elif "options_cache" in kwargs and kwargs["options_cache"]:
options_cache = kwargs["options_cache"]
assert kwargs[
"type"] is not None, "layer type not specified: forward/backward"
options = options_cache[kwargs["type"]]
logger.info("Tuned kernel options found, using those options")
if options is None:
options = Options("naive")
logger.warning(
"No mapping options passed, 'naive' type mapping options will be used and will likely have bad performance. See help(your_layer.__call__) for setting mapping options."
)
if not isinstance(options, Options):
options = Options(options)
return options
def test_kru_backward_pass_options(self):
lang = KRU3_1_TRAINING
forward_options = Options("naive")
backward_options = Options("naive")
KRU3_1 = tc.define(lang,
training=True,
name="KRU3_1",
backward="KRU3_1_GRAD")
X = Variable(torch.randn(256, 16, 16, 16).cuda(), requires_grad=True)
W2 = Parameter(torch.randn(32, 16).cuda())
out = KRU3_1(W2, X, options=[forward_options, backward_options])
out[0].sum().backward()
def test_kru_train_autotune_manual_options_both(self):
lang = KRU3_1_TRAINING
forward_options, backward_options = Options("mlp"), Options("mlp")
KRU3_1 = tc.define(lang,
training=True,
name="KRU3_1",
backward="KRU3_1_GRAD")
W2, X = torch.randn(32, 16).cuda(), torch.randn(256, 16, 16, 16).cuda()
options = KRU3_1.autotune(W2,
X,
options=[forward_options, backward_options],
cache=True,
**tc.autotuner_default_options)
def tune_and_store(self, tc_name, inputs, mapping_options, cache_file=""):
options = mapping_options
if not isinstance(options, Options):
options = Options(options)
try:
best_options = self.autotuner.tune(cache_file, tc_name, inputs,
options, [options])
return best_options
except RuntimeError:
return options
def test_autotuner_start_options_and_run_kernel(self):
lang = MATMUL_LANG
matmul = tc.define(lang, name="matmul")
mat1, mat2 = torch.randn(100, 400).cuda(), torch.randn(400, 500).cuda()
options = Options("mlp")
best_options = matmul.autotune(mat1,
mat2,
cache=True,
options=options,
**tc.autotuner_default_options)
out = matmul(mat1, mat2, options=best_options)
def tune_and_store(self, tc_name, inputs, mapping_options, cache_file=""):
options = mapping_options
if not isinstance(options, Options):
options = Options(options)
try:
best_options = self.autotuner.tune(cache_file, tc_name, inputs,
options, [options])
return best_options
except Exception as e:
logger.error('Raised exception: {}'.format(e))
return options
def test_manual_options(self):
lang = MATMUL_LANG
matmul = tc.define(lang, name="matmul")
mat1, mat2 = torch.randn(3, 4).cuda(), torch.randn(4, 5).cuda()
options = Options("naive")
out = matmul(mat1, mat2, options=options)
def test_manual_options(self):
lang = MATMUL_KRU1_LANG
KRU3_1 = tc.define(lang, name="KRU3_1")
W2, X = torch.randn(32, 16).cuda(), torch.randn(256, 16, 16, 16).cuda()
options = Options("naive")
out = KRU3_1(W2, X, options=options)
def test_KRU3(self):
# define TC
lang = """
def KRU3_1(float(D2, N2) W2, float(M, N0, N1, N2) X) -> (XW2) {
XW2(m, n0, n1, d2) +=! X(m, n0, n1, n2_red) * W2(d2, n2_red)
}
def KRU3_2(float(D1, N1) W1, float(M, N0, N1, D2) XW2) -> (XW2W1) {
XW2W1(m, n0, d1, d2) +=! XW2(m, n0, n1_red, d2) * W1(d1, n1_red)
}
def KRU3_3(float(D0, N0) W0, float(M, N0, D1, D2) XW2W1) -> (Y) {
Y(m, d0, d1, d2) +=! XW2W1(m, n0_red, d1, d2) * W0(d0, n0_red)
}
"""
# create input tensors
M, D0, D1, D2, N0, N1, N2, max_factors = 256, 32, 32, 32, 16, 16, 16, 3
W0 = torch.randn(D0, N0).cuda()
W1 = torch.randn(D1, N1).cuda()
W2 = torch.randn(D2, N2).cuda()
X = torch.randn(M, N0, N1, N2).cuda()
# define the mapping_options
options = Options("naive")
options.useSharedMemory(True)
options.usePrivateMemory(True)
options.tile([4, 1, 1, 8, 16])
options.mapToBlocks([64, 16, 16])
options.mapToThreads([8, 4, 8])
options.unroll(128)
# create TC compilation unit object and define the TC language
cu = TcCompilationUnit()
cu.define(lang)
print("Running KRU3_1")
inputs1 = [W2, X]
outputs1 = cu.compile_and_run("KRU3_1", inputs1, options=options)
print("Running KRU3_2")
XW2 = outputs1[0]
inputs2 = [W1, XW2]
outputs2 = cu.compile_and_run("KRU3_2", inputs2, options=options)
print("Running KRU3_3")
XW2W1 = outputs2[0]
inputs3 = [W0, XW2W1]
outputs3 = cu.compile_and_run("KRU3_3", inputs3, options=options)
def test_batchmatmul(self):
# define TC
lang = """
def batch_matmul(float(B, N, M) X, float(B, M, K) Y) -> (Z) {
Z(b, n, k) +=! X(b, n, mm) * Y(b, mm, k)
}
"""
# create input tensors
B, K, M, N = 500, 26, 72, 26
X = torch.randn(B, N, M).cuda()
Y = torch.randn(B, M, K).cuda()
inputs = [X, Y]
# define the mapping_options
options = Options("naive")
options.useSharedMemory(True)
options.usePrivateMemory(True)
options.unrollCopyShared(True)
options.outerScheduleFusionStrategy("Preserve3Coincident")
options.fixParametersBeforeScheduling(True)
options.tile([1])
options.tileImperfectlyNested(False)
options.mapToBlocks([72, 16, 1])
options.mapToThreads([7, 26])
options.unroll(128)
# run with TC, get the outputs and check against reference implementation
outputs = self.check(lang, "batch_matmul", options, inputs)
def test_tmm(self):
# define TC
lang = """
def tmm(float(M,K) A, float(N,K) B) -> (C) {
C(m, n) +=! A(m, kk) * B(n, kk)
}
"""
# create input tensors
M, N, K = 128, 256, 32
A = torch.randn(M, K).cuda()
B = torch.randn(N, K).cuda()
inputs = [A, B]
# define the mapping_options
options = Options("naive")
options.useSharedMemory(True)
options.usePrivateMemory(True)
options.unrollCopyShared(False)
options.outerScheduleFusionStrategy("Preserve3Coincident")
options.fixParametersBeforeScheduling(False)
options.tile([4, 32])
options.tileImperfectlyNested(False)
options.mapToBlocks([64, 128])
options.mapToThreads([1, 32])
options.unroll(4)
# run with TC, get the outputs and check against reference implementation
outputs = self.check(lang, "tmm", options, inputs)
expected = torch.mm(A, torch.transpose(B, 0, 1))
diff = outputs[0] - expected
self.assert_almost_equal(diff, inputs, M * N, 3e-7)
def test_options(self):
print('\nCreating mapping_options')
options = Options("naive")
options.useSharedMemory(True)
options.usePrivateMemory(False)
options.unrollCopyShared(False)
options.mapToBlocks([256, 8])
options.mapToThreads([4, 16, 4])
options.tile([2, 8, 64, 128])
options.unroll(128)
options.tileImperfectlyNested(False)
options.fixParametersBeforeScheduling(False)
options.scheduleFusionStrategy("Max")
options.outerScheduleFusionStrategy("Preserve3Coincident")
print('Mapping options created successfully')
def test_mlp(self):
# define TC
lang = """
def mlp3(float(B,N) I, float(O,N) W2, float(O) B2, float(P,O) W3, float(P) B3, float(Q,P) W4, float(Q) B4) -> (O2, O3, O4) {
O2(b, o) +=! I(b, n) * W2(o, n)
O2(b, o) = O2(b, o) + B2(o)
O2(b, o) = fmax(O2(b, o), 0)
O3(b, p) +=! O2(b, o) * W3(p, o)
O3(b, p) = O3(b, p) + B3(p)
O3(b, p) = fmax(O3(b, p), 0)
O4(b, q) +=! O3(b, p) * W4(q, p)
O4(b, q) = O4(b, q) + B4(q)
O4(b, q) = fmax(O4(b, q), 0)
}
"""
# create input tensors
B, N, O, P, Q = 128, 128, 64, 32, 2
I = torch.randn(B, N).cuda()
W2 = torch.randn(O, N).cuda()
B2 = torch.randn(O).cuda()
W3 = torch.randn(P, O).cuda()
B3 = torch.randn(P).cuda()
W4 = torch.randn(Q, P).cuda()
B4 = torch.randn(Q).cuda()
inputs = [I, W2, B2, W3, B3, W4, B4]
# define the mapping_options
options = Options("naive")
options.useSharedMemory(False)
options.usePrivateMemory(False)
options.unrollCopyShared(True)
options.outerScheduleFusionStrategy("Max")
options.fixParametersBeforeScheduling(False)
options.tile([4])
options.tileImperfectlyNested(False)
options.mapToBlocks([128])
options.mapToThreads([64])
options.unroll(128)
# run with TC, get the outputs and check against reference implementation
outputs = self.check(lang, "mlp3", options, inputs)
def test_C3(self):
# define TC
lang = """
def _C3(float(B, WX) I, float(WY, WX) W) -> (C3) {
C3(b, wy) +=! I(b, wxx) * W(wy, wxx)
}
"""
# create input tensors
B, WX, WY = 128, 1000, 1024
I = torch.randn(B, WX).cuda()
W = torch.randn(WY, WX).cuda()
inputs = [I, W]
# define the mapping_options
options = Options("naive")
options.useSharedMemory(True)
options.usePrivateMemory(True)
options.unrollCopyShared(True)
options.outerScheduleFusionStrategy("Preserve3Coincident")
options.fixParametersBeforeScheduling(True)
options.tile([8, 32, 32])
options.tileImperfectlyNested(False)
options.mapToBlocks([128, 128])
options.mapToThreads([1, 32])
options.unroll(256)
# run with TC, get the outputs and check against reference implementation
outputs = self.check(lang, "_C3", options, inputs)
def test_group_convolution(self):
# define TC
lang = """
def group_convolution(float(N,G,C,H,W) I, float(G,F,C,KH,KW) W1, float(G,F) B)
-> (O)
{
O(n, g, f, h, w) +=! I(n, g, c, h + kh, w + kw) * W1(g, f, c, kh, kw)
O(n, g, f, h, w) = O(n, g, f, h, w) + B(g, f)
}
"""
# create input tensors
N, G, C, F, H, W, KH, KW = 32, 32, 32, 32, 7, 7, 3, 3
tI = torch.randn(N, G, C, H, W).cuda()
tW = torch.randn(G, F, C, KH, KW).cuda()
tB = torch.randn(G, F).cuda()
inputs = [tI, tW, tB]
# define the mapping_options
options = Options("naive")
options.useSharedMemory(True)
options.usePrivateMemory(False)
options.unrollCopyShared(True)
options.outerScheduleFusionStrategy("Preserve3Coincident")
options.fixParametersBeforeScheduling(False)
options.tile([1, 1])
options.tileImperfectlyNested(False)
options.mapToBlocks([32, 32, 3])
options.mapToThreads([8, 7, 7])
options.unroll(256)
# run with TC, get the outputs and check against reference implementation
outputs = self.check(lang, "group_convolution", options, inputs)