Support multi-dimensional tensor prints in CPU runtime.

[ienkovich]

This patch adds recursion to existing printer to handle N-dim tensors. Example kernel:

@triton.jit
def kernel_print(a):
    a_block_ptr = tl.make_block_ptr(base=a, shape=(2, 2, 2, 16), strides=(64, 32, 16, 1), offsets=(0, 0, 0, 0),
                                    block_shape=(2, 2, 2, 16), order=(3, 2, 1, 0))
    tl.device_print("1d", tl.load(a + tl.arange(0, 16)))
    tl.device_print("2d", tl.load(a_block_ptr).reshape(8, 16))
    tl.device_print("4d", tl.load(a_block_ptr))

a = torch.randn((2, 2, 2, 16), dtype=torch.float32, device=device)
kernel_print[(1,)](a)

Output:

(0, 0, 0) 1d: [-0.5524, -0.4098, -1.7870,  0.2918, -1.1225, -0.0241,  0.0628,  0.2458,
               -0.5359,  0.3338, -0.3889, -0.9465,  3.1933,  0.9299,  0.3630,  0.4745]
(0, 0, 0) 2d: [[-0.5524, -0.4098, -1.7870,  0.2918, -1.1225, -0.0241,  0.0628,  0.2458,
                -0.5359,  0.3338, -0.3889, -0.9465,  3.1933,  0.9299,  0.3630,  0.4745],
               [-0.5122, -0.0762,  0.0591,  0.1670, -1.8423, -1.9155,  1.2343, -0.2337,
                -1.5709, -1.3229,  0.0674, -0.4463,  1.6680,  0.9151, -0.4779,  1.2519],
               [-0.7417, -0.2565, -0.3983,  1.0564,  0.4838,  0.2603, -0.1971,  1.4788,
                -0.6320, -1.8124,  0.0076,  0.0049, -0.8248, -0.0163, -0.0900,  0.0212],
               [-0.2815, -1.2157,  1.9806,  1.1678,  0.2476,  0.3623, -0.3772, -0.4139,
                 0.4568, -1.5914,  0.1012, -0.8285,  1.0688,  0.2999, -0.0060,  0.2093],
               [ 2.9888,  0.4598,  0.1646, -1.1593,  0.8276, -0.7299,  0.4085, -0.6787,
                 0.5134, -0.0486,  0.5971, -0.6865,  0.6573, -0.7823, -0.0347, -1.4166],
               [-0.8509, -0.3272, -0.2053,  0.7272, -1.5200, -0.1023,  1.7000,  0.9346,
                -0.2285, -0.3346,  1.2642, -1.1812,  0.4497, -1.1152,  0.2743,  0.4295],
               [-1.3411,  0.1982, -2.3510, -0.0179,  0.7062,  0.5876, -1.2185, -1.4348,
                -0.9621,  2.1502,  0.8021, -0.7234, -1.1524, -0.1315, -1.0360, -0.0446],
               [ 0.6504, -0.4704, -0.4960,  0.1469, -0.6344, -0.3554, -0.6259,  0.8477,
                -2.0600,  0.6924,  1.2824, -0.9905,  0.1201,  0.4415,  0.1252,  0.6513]]
(0, 0, 0) 4d: [[[[-0.5524, -0.4098, -1.7870,  0.2918, -1.1225, -0.0241,  0.0628,  0.2458,
                  -0.5359,  0.3338, -0.3889, -0.9465,  3.1933,  0.9299,  0.3630,  0.4745],
                 [-0.5122, -0.0762,  0.0591,  0.1670, -1.8423, -1.9155,  1.2343, -0.2337,
                  -1.5709, -1.3229,  0.0674, -0.4463,  1.6680,  0.9151, -0.4779,  1.2519]],
                [[-0.7417, -0.2565, -0.3983,  1.0564,  0.4838,  0.2603, -0.1971,  1.4788,
                  -0.6320, -1.8124,  0.0076,  0.0049, -0.8248, -0.0163, -0.0900,  0.0212],
                 [-0.2815, -1.2157,  1.9806,  1.1678,  0.2476,  0.3623, -0.3772, -0.4139,
                   0.4568, -1.5914,  0.1012, -0.8285,  1.0688,  0.2999, -0.0060,  0.2093]]],
               [[[ 2.9888,  0.4598,  0.1646, -1.1593,  0.8276, -0.7299,  0.4085, -0.6787,
                   0.5134, -0.0486,  0.5971, -0.6865,  0.6573, -0.7823, -0.0347, -1.4166],
                 [-0.8509, -0.3272, -0.2053,  0.7272, -1.5200, -0.1023,  1.7000,  0.9346,
                  -0.2285, -0.3346,  1.2642, -1.1812,  0.4497, -1.1152,  0.2743,  0.4295]],
                [[-1.3411,  0.1982, -2.3510, -0.0179,  0.7062,  0.5876, -1.2185, -1.4348,
                  -0.9621,  2.1502,  0.8021, -0.7234, -1.1524, -0.1315, -1.0360, -0.0446],
                 [ 0.6504, -0.4704, -0.4960,  0.1469, -0.6344, -0.3554, -0.6259,  0.8477,
                  -2.0600,  0.6924,  1.2824, -0.9905,  0.1201,  0.4415,  0.1252,  0.6513]]]]

[minjang]

Looks pretty legit!