def test_tc_autotune_reinforce(self): with tempfile.NamedTemporaryFile() as cache_file: group_normalization = """ def moments(float(N, K) I) -> (mean, var) { # var = E(x^2) - mean^2. mean(n) +=! I(n, r_k) var(n) +=! I(n, r_k) * I(n, r_k) mean(n) = mean(n) / (K) var(n) = var(n) / (K) - mean(n) * mean(n) } def group_normalization( float(N, G, D, H, W) I, float(G, D) gamma, float(G, D) beta, float(N, G) mean, float(N, G) var) -> (O) { O(n, g, d, h, w) = gamma(g, d) * ( I(n, g, d, h, w) - mean(n, g) ) * rsqrt( var(n, g) + 1e-5 ) + beta(g, d) } """ N, G, D, H, W = 32, 32, 4, 56, 56 I, gamma, beta = (torch.randn(N, G, D, H, W, device='cuda'), torch.randn(G, D, device='cuda').fill_(1.0), torch.randn(G, D, device='cuda').zero_()) T = tc.define( group_normalization, tc.make_autotuned_options_factory( starting_options='naive', tuner_config=tuner_config, cache_filename=cache_file.name, store_to_cache=True)) # First occurrence triggers tuning mean, var = T.moments(I.view((N * G, -1))) out = T.group_normalization(I, gamma, beta, mean.view((N, G)), var.view((N, G))) # Create a new TC object to retrigger tuning T = tc.define( group_normalization, tc.make_autotuned_options_factory( tuner_config=tuner_config, cache_filename=cache_file.name, load_from_cache=True, store_to_cache=True)) mean, var = T.moments(I.view((N * G, -1))) out = T.group_normalization(I, gamma, beta, mean.view((N, G)), var.view((N, G))) from torch.nn.modules.normalization import GroupNorm GN = GroupNorm(G, G * D).cuda() ref = GN.forward(I.view((N, G * D, H, W))) tc.assert_almost_equal(ref, out.view((N, G * D, H, W)), I, operations=D * H * W)
def test_conv_with_backward_2kernels(self): conv = """ def convolution(float(N,C,H,W) I, float(M,C,KH,KW) W1, float(M) Bias) -> (O) { O(n, m, h, w) +=! I(n, r_c, h + r_kh, w + r_kw) * W1(m, r_c, r_kh, r_kw) O(n, m, h, w) = O(n, m, h, w) + Bias(m) } def convolution_igrad(float(M,C,KH,KW) W1, float(N,M,H,W) d_O) -> (d_I) { d_I(n, c, h, w) +=! d_O( n, r_m, h - r_kh, w - r_kw) * W1(r_m, c, r_kh, r_kw) } def convolution_wgrad(float(N,C,H,W) I, float(N,M,H,W) d_O) -> (d_W1) { d_W1(m, c, kh, kw) +=! d_O(r_n, m, r_h - kh, r_w - kw) * I(r_n, c, r_h, r_w) } def convolution_biasgrad(float(M) Bias) -> (d_Bias) { # TODO: Bias incorrect + check d_Bias(m) = Bias(m) } """ N, C, H, W, O, kH, kW = 32, 4, 56, 56, 16, 1, 1 T = tc.define( conv, tc.make_autotuned_options_factory(starting_options='naive', tuner_config=tuner_config)) I = torch.randn(N, C, H, W, device='cuda', requires_grad=True) # Reference from torch.nn.modules.conv import Conv2d Conv = Conv2d(C, O, 1, stride=1).cuda() ref = Conv.forward(I) W = Conv.weight.clone() Bias = Conv.bias.clone() def convolution_backward(I, W, Bias, d_O): d_I = T.convolution_igrad(W, d_O) d_O = T.convolution_wgrad(I, d_O) d_Bias = T.convolution_biasgrad(Bias) return (d_I, d_O, d_Bias) convolution_function = tc.make_autograd(T.convolution, convolution_backward) # First occurrence triggers tuning out = convolution_function(I, W, Bias) out.sum().backward() # Subsequent occurrences do not out = convolution_function(I, W, Bias) out.sum().backward() tc.assert_almost_equal(ref, out, I, operations=C * kH * kW)
def test_group_norm_fused(self): group_normalization = """ def group_normalization( float(N, G, D, H, W) I, float(G, D) gamma, float(G, D) beta) -> (Sum, SumSq, O) { Sum(n, g) +=! I(n, g, r_d, r_h, r_w) SumSq(n, g) +=! I(n, g, r_d, r_h, r_w) * I(n, g, r_d, r_h, r_w) O(n, g, d, h, w) = gamma(g, d) * ( I(n, g, d, h, w) - Sum(n, g) / (D * H * W)) * rsqrt( (SumSq(n, g) - Sum(n, g) * Sum(n, g) / (D * H * W)) / (D * H * W) + 1e-5) + beta(g, d) } """ N, G, D, H, W = 32, 32, 4, 56, 56 T = tc.define( group_normalization, tc.make_autotuned_options_factory(starting_options='naive', tuner_config=tuner_config)) I, gamma, beta = (torch.randn(N, G, D, H, W, device='cuda'), torch.randn(G, D, device='cuda').fill_(1.0), torch.randn(G, D, device='cuda').zero_()) Sum, SumSq, O = T.group_normalization(I, gamma, beta) from torch.nn.modules.normalization import GroupNorm GN = GroupNorm(G, G * D).cuda() ref = GN.forward(I.view((N, G * D, H, W))) tc.assert_almost_equal(ref, O.view((N, G * D, H, W)), I, operations=D * H * W)
def generate_options(tc_str: str, entry_point: str, *inputs: torch.Tensor) -> tc.MappingOptions: global reinforce # TODO: comment the line below which serves the purpose of not blowing up # CI time return tc.make_naive_options_factory()(tc_str, entry_point, *inputs) if entry_point == 'make_idx': return tc.make_naive_options_factory()(tc_str, entry_point, *inputs) loaded = tc.make_load_from_cache_options_factory(args.tuner_cache_file)( tc_str, entry_point, *inputs) if loaded is None or entry_point in reinforce_list or '*' in reinforce_list: start = loaded if loaded is not None else 'naive' return tc.make_autotuned_options_factory( starting_options=start, tuner_config=tuner_config, cache_filename=args.tuner_cache_file, store_to_cache=True, )(tc_str, entry_point, *inputs) assert loaded is not None, 'None found' return loaded
def __init__(self, I, C, K, groups=1, padding=0, bias=False, from_cache=False, cache_file='tc_group3d.pt', tuner_config=None): ''' Module providing grouped 3d convolution using tensor comprehensions :param I: Number of input channels :type I: int :param C: Number of output channels :type C: int :param K: Kernel size :type K: tuple or int :param groups: Number of groups :type groups: int :param from_cache: If True load from specified cache file, If False, perform autotuning :type from_cache: bool :param cache_file: Path and name of cache file :type cache_file: string :param padding: Amount of input padding :type padding: tuple or int :param bias: Not implemented :type bias: bool :param tuner_config: Tuner config object to use for auto-tuning :type tuner_config: tensor_comprehensions.TunerConfig ''' import torch.nn.functional as F super().__init__() K = self.int_to_tuple(K) padding = self.int_to_tuple(padding) group_convolution = self.tc_string() if not from_cache: if tuner_config is None: tuner_config = tc.TunerConfig().generations(25).pop_size(100).number_elites(15) conv_option = tc.tclib.MappingOptions('naive').tile([1,1]).mapToThreads([4,16,4]).mapToBlocks([256,256]).unroll(1) TC = tc.define(group_convolution, tc.make_autotuned_options_factory( starting_options=conv_option, tuner_config=tuner_config, cache_filename=cache_file, store_to_cache=True, load_from_cache=False )) else: TC = tc.define(group_convolution, tc.make_load_from_cache_options_factory(cache_file)) self.convolution_grouped = tc.make_autograd(TC.group_convolution, TC.convolution_grad) self.W = torch.nn.Parameter(torch.rand(groups, C/groups, I/groups, K[0], K[1], K[2])) self.pad = F.pad self.groups = groups self.padding = padding self.K = K
cache = MappingOptionsCache(cache_file.name) top10 = cache.load(mm, "matmul", (A, B), 10) assert top1.__str__() == top10[0].__str__() # Compile and run with the new options compilation_cache.compile("matmul", (A, B), top1) time_tc(100, "raw unchecked_run tuned options\t", lambda name, ins: compilation_cache.unchecked_run(name, ins), "matmul", (A, B)) ################################################################################ # 4. Simple torch.autograd.Function ################################################################################ T = tc.define( mm, tc.make_autotuned_options_factory(starting_options='naive', tuner_config=tuner_config)) def backward(A, B, d_C): d_A = T.matmul_agrad(B, d_C, unchecked=True) d_B = T.matmul_bgrad(A, d_C, unchecked=True) return d_A, d_B matmul_function = tc.make_autograd(lambda A, B: T.matmul(A, B, unchecked=True), backward) # Run once to trigger automatic tuning and compilation then time # For example purposes, use retain_graph # retain_graph = True prevents freeing the buffers when performing backward # see e.g. https://stackoverflow.com/questions/46774641/what-does-the-parameter-retain-graph-mean-in-the-variables-backward-method
def main(): parser = argparse.ArgumentParser( "compile + tune + test tensor comp kernels...") parser.add_argument("--kernel_name", default=r"kernel_*") parser.add_argument("--list", const=True, action="store_const", default=False) parser.add_argument("--tune", const=True, action="store_const", default=False) parser.add_argument("--exact", const=True, action="store_const", default=False) parser.add_argument("--float32", const=True, action="store_const", default=False) parser.add_argument("--load_cache", const=True, action="store_const", default=False) parser.add_argument("--generations", default=10, type=int) parser.add_argument("--cache_filename", default="tc_cache", type=str) parser.add_argument("--init", default="naive", type=str) parser.add_argument("--threads", default=16, type=int) parser.add_argument("--pop_size", default=100, type=int) parser.add_argument("--crossover_rate", default=80, type=int) parser.add_argument("--mutation_rate", default=7, type=int) parser.add_argument("--number_elites", default=10, type=int) parser.add_argument("--height", default=32, type=int) parser.add_argument("--width", default=32, type=int) parser.add_argument("--N", default=8, type=int) parser.add_argument("--channels", default=3, type=int) parser.add_argument("--num_gpus", default=1, type=int) args = parser.parse_args() matched_kernels = [] gpus = ",".join([str(x) for x in range(args.num_gpus)]) print("devices: ", gpus) tuner_config = tc.TunerConfig().threads(args.threads).generations( args.generations).pop_size(args.pop_size).crossover_rate( args.crossover_rate).mutation_rate( args.mutation_rate).number_elites(args.number_elites) for k in all_kernel_strs: if not args.exact: if re.match(re.compile(args.kernel_name), k): matched_kernels.append(k) else: if k == args.kernel_name: matched_kernels.append(k) if args.list: print("Kernels available:") for k in matched_kernels: print("\t" + k) if args.init not in ["naive", "pointwise", "mlp"]: assert False start_options = tc.MappingOptions(args.init) if args.tune: if not args.load_cache: opts = tc.make_autotuned_options_factory( starting_options=start_options, cache_filename=args.cache_filename, store_to_cache=True, tuner_config=tuner_config) else: print("loading from cache...") opts = tc.make_autotuned_options_factory( load_from_cache=True, cache_filename=args.cache_filename, store_to_cache=True, tuner_config=tuner_config) else: if not args.load_cache: opts = tc.make_naive_options_factory() else: opts = tc.make_load_from_cache_options_factory( cache_filename=args.cache_filename) kernel_fn_map = {} N = args.N H = args.height W = args.width torch.manual_seed(0) x = torch.randn(N, H, W, args.channels).double().cuda() k_input = torch.randn(N, N, H, W, H, W).double().cuda() z = torch.tensor(1.0).double().cuda() y = x if args.float32: x = x.float() y = y.float() k_input = k_input.float() z = z.float() for k in matched_kernels: print(f"Tuning {k}") kernel_fn = tc.define(all_kernel_strs[k], opts) kernel_fn_map[k] = kernel_fn if "float" in k: k_call = getattr(kernel_fn, k.replace("kernel_float_", "")) else: k_call = getattr(kernel_fn, k.replace("kernel_", "")) if "input" in k: kxy = k_call(x, y) print("output: ", kxy) else: if "exponential_shifted" in k: print("calling exponential shifted") kxy = k_call(k_input, k_input, k_input, z) else: kxy = k_call(k_input, k_input, k_input)