gds_render.py

# Inspo from https://github.com/TinyTapeout/tt-support-tools
import argparse
import gdstk
import subprocess
import os

# Not secure against maliciously constructed data, assuming lyp file is trusted
import xml.etree.ElementTree as ET 

MET4_LABEL = (71, 5)
AREA_ID_STDCELL = (81, 4)
AREA_ID_DIODE = (81, 23)
PR_BOUNDARY = (235, 4)

SHAPE_STYLE = {
  AREA_ID_STDCELL: {"fill": "none", "stroke": "red"},
}

LABEL_STYLE = {
  MET4_LABEL: {
    "fill": "#e25822", 
    "paint-order": "stroke",
    # "stroke": "#db8e70",
    "stroke": "#4a1f0e",
    "stroke-width": "1px",
    "stroke-linecap": "butt",
    "stroke-linejoin": "miter",
    }
}

IO_VGA_LABELS = {
  "uo_out[0]": "R[1]",
  "uo_out[1]": "G[1]",
  "uo_out[2]": "B[1]",
  "uo_out[3]": "vsync",
  "uo_out[4]": "R[0]",
  "uo_out[5]": "G[0]",
  "uo_out[6]": "B[0]",
  "uo_out[7]": "hsync",
}

# From https://github.com/TinyTapeout/GDS2glTF
# TODO: Ideally would follow stackup-up measurements
# https://skywater-pdk.readthedocs.io/en/main/rules/assumptions.html
BW_LAYER_NAMES = [
  "licon", "li1", "mcon", "met1", "via", "met2",
  "via2", "met3", "via3", "met4", "via4", "met5",
]

class Layer:
  """Parses the layer tag from the lyp file.
  """
  def __init__(self, layer_tag: str):
    # In the form: "prBoundary.boundary - 235/4"
    self.name, layer_data_no = layer_tag.split(" - ")
    self.layer_no, self.data_no = map(int, layer_data_no.split("/"))
  
  def to_tuple(self) -> tuple[int, int]:
    return (self.layer_no, self.data_no)


class LayerStack:
  """A layer stack derived by reading the `lyp` KLayout layout properties XML
  file used to assign names to layer numbers and determine drawing order.
  """
  def __init__(self, lyp_path):
    self.layers = []
    self.layer_drawing_order = {}
    self.missing_layers = set()
    
    lyp_tree = ET.parse(lyp_path)
    # https://www.klayout.de/lyp_format.html
    # TODO: Multi-Tab Layer Files
    props = lyp_tree.getroot()
    assert props.tag == "layer-properties"
    
    i = 0
    for prop in props:
      # TODO: if too many, skip valid = false
      if prop.tag != "properties":
        continue
      
      name_tag = prop.find("name")
      assert name_tag is not None
      layer = Layer(name_tag.text)
      
      self.layer_drawing_order[layer.to_tuple()] = i
      self.layers.append(layer)
      i += 1

  def get_layers_matching(self, patterns: list[str]) -> list[tuple[int, int]]:
    """Returns a list matching gdstk's (layer, data) for every layer containing
    at least one of the given substrings.

    Args:
        patterns (list[str]): A list of substrings to check

    Returns:
        list[tuple[int, int]]: A list of (layer, data) layers that match
    """
    def layer_to_tuple(l: Layer):
      return l.to_tuple()
    def layer_name_contains(l: Layer):
      return any([p in l.name for p in patterns])
    
    return list(map(layer_to_tuple, filter(layer_name_contains, self.layers)))

  def poly_sort(self, a: gdstk.Polygon, b: gdstk.Polygon) -> bool:
    """Returns true if polygon a is below polygon b in the drawing order
    according to the lyp layer stack, also noting any missing layers.
    This places unknown keys towards the bottom, comparing values if both are
    missing.

    Args:
        a (gdstk.Polygon): Polygon A
        b (gdstk.Polygon): Polygon B

    Returns:
        bool: True if polygon a is below polygon b
    """
    key_a, key_b = (a.layer, a.datatype), (b.layer, b.datatype)
    key_a_missing = key_a not in self.layer_drawing_order
    key_b_missing = key_b not in self.layer_drawing_order
    # Place unknown keys on the bottom
    if key_a_missing and key_b_missing:
      self.missing_layers.add(key_a)
      self.missing_layers.add(key_b)
      return key_a < key_b
    elif key_a_missing:
      self.missing_layers.add(key_a)
      return True
    elif key_b_missing:
      self.missing_layers.add(key_b)
      return False
    
    order_a = self.layer_drawing_order[key_a]
    order_b = self.layer_drawing_order[key_b]
    return order_a < order_b

def read_top_cell(gds_path):
  library = gdstk.read_gds(gds_path)
  return library.top_level()[0]

def filter_layers_recursively(cell: gdstk.Cell, spec, remove=True):
  cell.filter(spec, remove=remove)
  for subcell in cell.dependencies(recursive=True):
    subcell.filter(spec, remove=remove)

def is_fill_decap_tap(cell: str) -> bool:
  return "__fill" in cell or "__decap" in cell or "__tap" in cell

def filter_fill_decap_tap(cell: gdstk.Cell):
  for ref in cell.references:
    if is_fill_decap_tap(ref.cell.name):
      cell.remove(ref)

def format_io_crop(cell: gdstk.Cell, chop_min_x: int, chop_min_y: int, chop_max_x: int) -> gdstk.Cell:
  """Formats the polygons and labels into crop displaying IO, removing all
  references outside of the given area, slicing and selecting polygons in the
  given area, formatting labels, and adding labels corresponding to the PMOD
  VGA output.

  Args:
      cell (gdstk.Cell): Original cell
      chop_min_x (int): Left x-value
      chop_min_y (int): Bottom y-value
      chop_max_x (int): Right x-value

  Returns:
      gdstk.Cell: New cell containing polygons inside the given area and newly
        formatted labels
  """
  
  out_cell = gdstk.Cell("out")
  
  def get_rough_polys(refs: list[gdstk.Reference]):
    """Helper function to take all cell references and return polygons of all
    cells overlapping specified area.
    """
    temp_cell = gdstk.Cell("temp")
    # Rough cut
    for ref in refs:
      (min_x, _), (max_x, max_y) = ref.bounding_box()
      if max_y > chop_min_y and max_x > chop_min_x and min_x < chop_max_x:
        temp_cell.add(ref)
    
    temp_cell.flatten()
    return temp_cell.polygons
      
  # slice returns a list of N+1 lists, each containing polys part of a slice position (very cool and efficient)
  result = gdstk.slice(cell.polygons + get_rough_polys(cell.references), chop_min_y, "y")
  polys = result[1] # All polys greater than slice position
  result = gdstk.slice(polys, [chop_min_x, chop_max_x], "x")
  polys = result[1] # All polys greater than min and less than max
  out_cell.add(*polys)

  for label in cell.labels:
    label_x, label_y = label.origin
    if label_x < chop_min_x or label_x > chop_max_x:
      continue
    label.magnification *= 10
    if label.text == "VPWR" or label.text == "VGND":
      # Shift up and center
      label.origin = (label_x + 0.5, 220)
    else:
      if label.text in IO_VGA_LABELS:
        vga_label = label.copy()
        vga_label.text = IO_VGA_LABELS[label.text]
        vga_label.origin = (label_x, label_y - 0.3)
        vga_label.magnification *= 1.8
        vga_label.rotation = 0
        out_cell.add(vga_label)
        
      label.text = f"({label.text})"
      label.origin = (label_x, label_y - 1.1)
      label.magnification *= 1.4
      label.rotation = 0

    out_cell.add(label)
  
  return out_cell

def get_bw_layer_style(layer_stack: LayerStack) -> dict[tuple[int, int], dict[str, str]]:
  """Generates the B&W layer style dictionary, mapping a brighter value to parts
  of BW_LAYER_NAMES that are closer to the top.

  Args:
      layer_stack (LayerStack): Description of layers to get all layer numbers
        given a name.

  Returns:
      dict[tuple[int, int], dict[str, str]]: Style dictionary for use in gdstk's
        write_svg
  """
  bw_layers = set()
  styles = {}
  for i, layer_name in enumerate(BW_LAYER_NAMES):
    value = tuple(min(int((0.12 + 0.06 * i) * 255), 255) for _ in range(3))
    layers = layer_stack.get_layers_matching([layer_name])
    bw_layers.update(layers)
    for layer in layers:
      styles[layer] = {"fill": f"rgb{value}"}

  styles[PR_BOUNDARY] = {"fill": f"black"}

  return styles
    
def write_images(img_path: str, cell: gdstk.Cell, layer_stack: LayerStack, scaling = 10, shape_style=SHAPE_STYLE):
  """Outputs an SVG to `img_path`.svg using gdstk's write_svg, then converts
  to a PNG using rsvg-convert, printing layers compared that are missing
  from the layer_stack.

  Args:
      img_path (str): Output path for the SVG/PNG images
      cell (gdstk.Cell): Cell (preferrebly flattened) to render
      layer_stack (LayerStack): Used to determine drawing order
      scaling (int, optional): Scales the output SVG. Defaults to 10.
      shape_style (dict[tuple[int, int], dict[str, str]], optional): Style by 
        layer passed to write_svg. Defaults to SHAPE_STYLE.
  """
  print(f"Creating {img_path}...")
  # Sort function alone does not work across cells, must flatten first
  cell.flatten()
  cell.write_svg(
    f"{img_path}.svg", pad=0, scaling=scaling,
    shape_style=shape_style,
    label_style=LABEL_STYLE,
    sort_function=layer_stack.poly_sort
  )
  print("SVG created!")
  if len(layer_stack.missing_layers) > 0:
    print("LayerStack missing:", layer_stack.missing_layers)
    layer_stack.missing_layers.clear()
  
  cmd = f"rsvg-convert --unlimited {img_path}.svg -o {img_path}.png --no-keep-image-data"
  p = subprocess.run(cmd, shell=True, capture_output=True)
  if p.returncode == 127:
    print('ERROR: rsvg-convert not found; is package "librsvg2-bin" installed?')
  elif p.returncode != 0 and b"cannot load more than" in p.stderr:
    print( f'ERROR: Too many SVG elements ("{p.stderr.decode().strip()}")')
  else:
    print("PNG created!")

def main(gds_path, lyp_path, out_path):
  # TODO: Bound by the calls to rsvg-convert, so multithread or async if an issue
  os.makedirs(out_path, exist_ok=True)
  layer_stack = LayerStack(lyp_path)
  met_layers = layer_stack.get_layers_matching(["mcon", "met", "via"])
  label_layers = layer_stack.get_layers_matching(["label"])
  label_and_std_layers = label_layers + [AREA_ID_STDCELL]
  
  top_cell = read_top_cell(gds_path)
  # TODO: edits still affecting original even with a deep_copy? 
  # Reading is quick so do that instead
  # top_cell_met = top_cell.copy("top_cell_met", deep_copy=True)

  ### Everything ###
  filter_layers_recursively(top_cell, label_and_std_layers)
  write_images(f"{out_path}/full", top_cell.flatten(), layer_stack)
  
  ### No decap/fill or metal ###
  top_cell = read_top_cell(gds_path)
  filter_fill_decap_tap(top_cell)
  filter_layers_recursively(top_cell, met_layers)
  write_images(f"{out_path}/min", top_cell.flatten(), layer_stack)

  ### Only metal ###
  top_cell_met = read_top_cell(gds_path)
  filter_layers_recursively(top_cell_met, label_layers)
  filter_layers_recursively(top_cell_met, met_layers + [PR_BOUNDARY], remove=False)
  write_images(f"{out_path}/met", top_cell_met.flatten(), layer_stack, shape_style={PR_BOUNDARY: {"fill": "black"}})
  
  ### B&W Depth Everything ###
  top_cell_bw = read_top_cell(gds_path)
  filter_layers_recursively(top_cell_bw, label_and_std_layers + [AREA_ID_DIODE])
  bw_style = get_bw_layer_style(layer_stack)
  filter_layers_recursively(top_cell_bw, list(bw_style.keys()) + [PR_BOUNDARY], remove=False)
  write_images(f"{out_path}/bw", top_cell_bw.flatten(), layer_stack, shape_style=bw_style)

  ### I/O ###
  top_cell_io = read_top_cell(gds_path)
  filter_layers_recursively(top_cell_io, set(label_and_std_layers) - set([MET4_LABEL]))
  filter_layers_recursively(top_cell_io, met_layers, remove=False)
  top_cell_io_chopped = format_io_crop(top_cell_io, 73, 215, 102)
  write_images(f"{out_path}/io_met", top_cell_io_chopped.flatten(), layer_stack, scaling=100)
  
  """Merged with: https://www.imagemagick.org/Usage/montage/
  montage {full,min}.png -tile 2x1 -geometry +0+0 full_min.png
  montage {met,bw}.png -tile 2x1 -geometry +0+0 met_bw.png
  """


if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument("gds_path")
  parser.add_argument("--lyp-file", default="sky130.lyp")
  parser.add_argument("--out", default="gds_renders")
  args = parser.parse_args()
  
  main(args.gds_path, args.lyp_file, args.out)