Skip to content

Instantly share code, notes, and snippets.

View psrivas2's full-sized avatar

Prakalp Srivastava psrivas2

  • OctoML
  • United States
View GitHub Profile

Here is how the TIR Primfunc would like before hoisting

    @T.prim_func
    def fused_primitives(x: T.Buffer[25, "float32"], y_crop_1: T.Buffer[25, "float32"]) -> None:
        x_pad_0 = T.alloc_buffer([28], dtype="float32")
        x_transform = T.alloc_buffer([4, 7], dtype="float32")
        x_pad_1 = T.alloc_buffer([4, 8], dtype="float32")
        y_crop_0 = T.alloc_buffer([4, 7], dtype="float32")
        y_transform = T.alloc_buffer([28], dtype="float32")