[FEA] [CuTeDSL] `print_latex` in `CuTeDSL`

[simveit]

Is your feature request related to a problem? Please describe. It would be nice to have utility function in CuTeDSL like print_latex in C++ API

Describe the solution you'd like A similar functionality as in C++ API

Describe alternatives you've considered Using C++ API

[thakkarV]

this is not a priority atm, contributions are welcome

[Chillee]

I wrote one here that only uses CuteDSL (as well as numpy/matplotlib): https://gist.github.com/Chillee/e2b07157caeade8c6b0bdf463d10f833

I could clean it up and land it if folks have opinions about what API it should support.

[thakkarV]

@Chillee this looks awesome :D

@ccecka @fengxie we can critique this with the intention of Horace upstreaming it into the DSL as a first class utility?

[ccecka]

I've found matplotlib to not scale well and suffer greatly for larger layouts. Instead, I wrote mine with SVG vector graphics output.

import svgwrite
import numpy as np

def layout_svg(layout):
  # 8 RBG-255 Greyscale colors
  rgb_255_colors = [(255, 255, 255), (230, 230, 230), (205, 205, 205), (180, 180, 180),
                    (155, 155, 155), (130, 130, 130), (105, 105, 105), ( 80,  80,  80)]

  # Cell size in pixels
  cell_size = 20

  # Grid size
  M, N = size(layout[0]), size(layout[1])

  # Create SVG canvas
  dwg = svgwrite.Drawing("layout.svg", size=(N*cell_size, M*cell_size))

  # Draw grid cells
  for i in range(M):
    for j in range(N):
      idx = layout(i, j)
      x = j * cell_size
      y = i * cell_size

      # Draw rectangle
      dwg.add(dwg.rect(
        insert=(x, y), size=(cell_size, cell_size),
        fill=svgwrite.rgb(*rgb_255_colors[idx % len(rgb_255_colors)], mode='RGB'), stroke='black'))

      # Add label text
      dwg.add(dwg.text(
        str(idx), insert=(x + cell_size//2, y + cell_size//2),
        text_anchor="middle", alignment_baseline="central", font_size="8px"))

  dwg.save()

def tvlayout_svg(layout, tile_mn=None):
  tile_mn = coshape(layout) if tile_mn is None else tile_mn
  assert rank(layout)  == 2, "Expected a rank-2 TV Layout"
  assert rank(tile_mn) == 2, "Expected a rank-2 MN Tile"
  assert congruent(coprofile(layout), (0,0)), "Expected a 2D codomain (tid,vid) -> (m,n)"

  # 8 RGB-255 colors, TODO Generalize
  rgb_255_colors = [(175, 175, 255), (175, 255, 175), (255, 255, 175), (255, 175, 175),
                    (210, 210, 255), (210, 255, 210), (255, 255, 210), (255, 210, 210)]

  # Cell size in pixels
  cell_size = 20

  # Grid size
  M, N = size(tile_mn[0]), size(tile_mn[1])
  filled = np.zeros((M,N), dtype=bool)

  # Create SVG canvas
  dwg = svgwrite.Drawing("tvlayout.svg", size=(N*cell_size, M*cell_size))

  # Fill in grid
  for i in range(M):
    for j in range(N):
      dwg.add(dwg.rect(
        insert=(j*cell_size, i*cell_size), size=(cell_size, cell_size),
        fill='white', stroke='black'))

  # Draw TV cells
  for tid in range(size(layout[0])):
    for vid in range(size(layout[1])):
      i,j = layout(tid, vid)
      x = j * cell_size
      y = i * cell_size

      if filled[i,j]:
        continue
      filled[i,j] = True

      # Draw rectangle
      dwg.add(dwg.rect(
        insert=(x, y), size=(cell_size, cell_size),
        fill=svgwrite.rgb(*rgb_255_colors[tid % len(rgb_255_colors)], mode='RGB'), stroke='black'))

      # Add label text
      dwg.add(dwg.text(
        f"T{tid}", insert=(x + cell_size//2, y + 1*cell_size//4),
        text_anchor="middle", alignment_baseline="central", font_size="8px"))
      dwg.add(dwg.text(
        f"V{vid}", insert=(x + cell_size//2, y + 3*cell_size//4),
        text_anchor="middle", alignment_baseline="central", font_size="8px"))

  dwg.save()

Not fully featured, but the vector graphics are consistently good. These are essentially a transliteration of the C++ print_latex/print_svg to Python.

[Chillee]

@ccecka where are the dependencies coming from? (e.g. congruent, coprofile, etc.)

and yeah, in my testing I found that matplotlib ends up being a bottleneck at larger sizes. It's about a minute for a 256x256 layout, but that's not that bad :P

[fengxie]

@ccecka where are the dependencies coming from? (e.g. congruent, coprofile, etc.)

and yeah, in my testing I found that matplotlib ends up being a bottleneck at larger sizes. It's about a minute for a 256x256 layout, but that's not that bad :P

Do we really want to visualize huge tile like this ;-) I was thinking about bring my old version to work "yesterday". But glad you two are so quick :D

[Chillee]

@fengxie haha probably not. Just making the point that I think matplotlib actually scales "ok-ish" even to fairly large sizes.

[github-actions[bot]]

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.

[NTT123]

@ccecka's code is almost working for me. So, here is a slightly modified version in case it is helpful for anyone: https://gist.github.com/NTT123/05f603c748e5f26bbc07bb636883e9ee

Expected output:

[NTT123]

I spent some time vibe coding this package cute-viz with Claude Code. It supports swizzle, thread-value, copy, and MMA layouts. All of them are simple translations from the C++ version. Here is an example MMA layout:

[thakkarV]

@NTT123 you are welcome to open a PR :)

[github-actions[bot]]

This issue has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with a comment indicating any updates or changes to the original issue and/or confirm this issue still needs to be addressed. This issue will be labeled inactive-90d if there is no activity in the next 60 days.