2D tiling of a 3d physical mesh.

  The "outer" mesh dimension corresponds to which tile.
  The "inner" mesh dimension corresponds to the position within a tile
  of processors.

  Optionally, if cores_name is specified, then a 3 dimensional logical mesh
  is returned, with the third dimension representing the two different
  cores within a chip.  If cores_name is not specified, then the
  cores-in-a-chip dimension is folded into the inner dimension.

  TODO(noam): explain this better.

  Example:

  tile_2d(physical_shape=[8, 16, 2], tile_shape=[4, 4])

  The "inner" dimension has size 4x4x2=32 and corresponds to the position
  within a 4x4 tile of processors.

  The "outer" dimension has size 8/4 * 16/4 = 8, and corresponds to the 8
  tiles in the mesh.

  Args:
    physical_shape: a triple of integers [X, Y, cores]
    tile_shape: a pair
    outer_name: a string
    inner_name: a string
    cores_name: an optional string

  Returns:
    mesh_shape: a mtf.Shape
    logical_to_physical: a list

def tile_2d(physical_shape, tile_shape,
            outer_name="outer",
            inner_name="inner",
            cores_name=None):
  """2D tiling of a 3d physical mesh.

  The "outer" mesh dimension corresponds to which tile.
  The "inner" mesh dimension corresponds to the position within a tile
  of processors.

  Optionally, if cores_name is specified, then a 3 dimensional logical mesh
  is returned, with the third dimension representing the two different
  cores within a chip.  If cores_name is not specified, then the
  cores-in-a-chip dimension is folded into the inner dimension.

  TODO(noam): explain this better.

  Example:

  tile_2d(physical_shape=[8, 16, 2], tile_shape=[4, 4])

  The "inner" dimension has size 4x4x2=32 and corresponds to the position
  within a 4x4 tile of processors.

  The "outer" dimension has size 8/4 * 16/4 = 8, and corresponds to the 8
  tiles in the mesh.

  Args:
    physical_shape: a triple of integers [X, Y, cores]
    tile_shape: a pair
    outer_name: a string
    inner_name: a string
    cores_name: an optional string

  Returns:
    mesh_shape: a mtf.Shape
    logical_to_physical: a list
  """
  logical_to_physical = []
  p0, p1, p2 = physical_shape
  t0, t1 = tile_shape
  tile_ring = _ring_2d(t0, t1)
  tiles_ring = _ring_2d(p0 // t0, p1 // t1)
  for logical_pnum in range(p0 * p1 * p2):
    core_on_chip = logical_pnum % p2
    logical_chip_num = logical_pnum // p2
    logical_pos_in_tile = logical_chip_num % (t0 * t1)
    logical_tile_num = logical_chip_num // (t0 * t1)
    tile_i, tile_j = tile_ring[logical_pos_in_tile]
    tiles_i, tiles_j = tiles_ring[logical_tile_num]
    physical_pnum = core_on_chip + p2 * (
        tile_i * p1 + tile_j +
        tiles_i * p1 * t0 + tiles_j * t1)
    logical_to_physical.append(physical_pnum)
  assert sorted(logical_to_physical) == list(range(p0 * p1  * p2))
  tile_size = t0 * t1 * p2
  num_tiles = p0 * p1 // (t0 * t1)
  if cores_name:
    mesh_shape = mtf.Shape(
        [mtf.Dimension(outer_name, int(num_tiles)),
         mtf.Dimension(inner_name, int(t0 * t1)),
         mtf.Dimension(cores_name, int(p2))])
  else:
    mesh_shape = mtf.Shape(
        [mtf.Dimension(outer_name, int(num_tiles)),
         mtf.Dimension(inner_name, int(tile_size))])
  return mesh_shape, logical_to_physical