code: def specialize_to_ptx(lang, function_name):
      param_cnt = 0
      ins = []
      alu = {
          BinaryOps.ADD: "add",
          BinaryOps.SUB: "sub",
          BinaryOps.MUL: "mul",
          BinaryOps.DIV: "div",
          BinaryOps.MAX: "max",
          BinaryOps.MOD: "rem",
          BinaryOps.CMPLT: "setp.lt",
          UnaryOps.SQRT: "sqrt.approx",
          UnaryOps.NOOP: "mov",
          UnaryOps.NEG: "neg",
          UnaryOps.SIN: "sin.approx",
          UnaryOps.LOG2: "lg2.approx",
          UnaryOps.EXP2: "ex2.approx.ftz",
          TernaryOps.MULACC: "fma.rn",
          TernaryOps.WHERE: "selp",
      }
      for uop, out, vin, arg in lang.ins:
          if uop == UOps.ENDLOOP:
              ins.append("bar.sync 0;")
          elif uop == UOps.DEFINE_LOCAL:
              ins.append(f".shared .align 4 .b8 {arg[0]}[{arg[1]*4}];")
          elif uop == UOps.SPECIAL:
              if arg.startswith("data"):
                  param_cnt += 1
                  ins.append(f"ld.param.u64 {out}, [{arg}];")
                  # TODO: we sometimes want this to be local, nvcc converts to global most of the time, not sure when we would need to?
                  # ins.append(f"cvta.to.global.u64 {out}, {out};")
              elif arg.startswith("gid"):
                  ins.append(
                      f"mov.u32 {out}, %ctaid.{'xyz'[int(arg[3:])]};"
                  )
              elif arg.startswith("lid"):
                  ins.append(
                      f"mov.u32 {out}, %tid.{'xyz'[int(arg[3:])]};"
                  )
          elif uop == UOps.ALU:
              if arg == BinaryOps.MUL and out.dtype == dtypes.bool:
                  ins.append(
                      f"and.pred {out}, {', '.join(str(x) for x in vin)};"
                  )
              else:
                  otype = (
                      vin[0].dtype
                      if arg in [BinaryOps.CMPLT]
                      else out.dtype
                  )
                  if arg == TernaryOps.WHERE:
                      if vin[0].dtype == dtypes.bool:
                          reg = vin[0]
                      else:
                          reg = lang.newreg(
                              (vin[0], "bool"), dtypes.bool
                          )
                          ins.append(
                              f"setp.ne.{dtype_to_nvtype[vin[0].dtype]} {reg}, {'0f00000000' if dtypes.is_float(vin[0].dtype) else '0'}, {vin[0]};"
                          )
                      vin = vin[1:] + [reg]
                  ins.append(
                      f"{alu[arg]}{'.lo' if arg == BinaryOps.MUL and out.dtype != dtypes.float32 else ''}{'.rn' if arg == BinaryOps.DIV and out.dtype == dtypes.float32 else ''}.{dtype_to_nvtype[otype]} {out}, {', '.join(str(x) for x in vin)};"
                  )
          elif uop == UOps.LOAD:
              if arg.__class__ in (int, float):
                  ins.append(
                      f"mov.{dtype_to_nvtype[out.dtype]} {out}, {'0f'+float_to_hex(arg) if dtypes.is_float(out.dtype) else int(arg)};"
                  )
              elif arg[2] is not None and (
                  arg[2] == dtypes.bool or arg[2] != out.dtype
              ):
                  dt = (
                      ("u16", dtypes.uint16)
                      if arg[2] == dtypes.bool == out.dtype
                      else ("u8", dtypes.uint8)
                      if arg[2] == dtypes.bool
                      else ("b16", dtypes.float16)
                      if arg[2] == dtypes.half
                      else (dtype_to_nvtype[arg[2]], arg[2])
                  )
                  reg = lang.newreg((out, dt[0]), dtype=dt[1])
                  ins.append(
                      f"ld.{arg[1]}.{dt[0]} {reg}, [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}];"
                  )
                  render_cast(ins, reg, out)
              else:
                  ins.append(
                      f"ld.{arg[1]}.{dtype_to_nvtype[dtypes.float if arg[2] is None else arg[2]]} {out}, [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}];"
                  )
          elif uop == UOps.STORE:
              if (
                  ptx_needs_cast(
                      dtypes.float if arg[2] is None else arg[2],
                      vin[1].dtype,
                  )
                  or arg[2] == dtypes.bool
              ):
                  if arg[2] == dtypes.bool != vin[1].dtype:
                      prereg = lang.newreg(
                          (vin[1], "bool"), dtype=dtypes.bool
                      )
                      render_cast(ins, vin[1], prereg)
                  else:
                      prereg = vin[1]
                  reg = lang.newreg(
                      (
                          prereg,
                          dtypes.uint16
                          if arg[2] == dtypes.bool
                          else arg[2],
                      ),
                      dtype=dtypes.uint16
                      if arg[2] == dtypes.bool
                      else dtypes.float
                      if arg[2] is None
                      else arg[2],
                  )
                  render_cast(ins, prereg, reg)
                  ins.append(
                      f"st.{arg[1]}.{dtype_to_nvtype['bits16' if arg[2] == dtypes.float16 else dtypes.uint8 if arg[2] == dtypes.bool else dtypes.float if arg[2] is None else arg[2]]} [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}], {reg};"
                  )
              else:
                  ins.append(
                      f"st.{arg[1]}.{dtype_to_nvtype[dtypes.float if arg[2] is None else arg[2]]} [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}], {vin[1]};"
                  )
          elif uop == UOps.CAST:
              render_cast(ins, vin[0], out)
          elif uop == UOps.LABEL:
              ins.append(f"{arg}:")
          elif uop == UOps.COND_BRANCH:
              ins.append(
                  f"@{'!' if not arg[1] else ''}{vin[0]} bra {arg[0]};"
              )
  
      ins_prefix = [
          ".version 7.8",
          ".target " + arch(),
          ".address_size 64",
          f".visible .entry {function_name}({', '.join(f'.param .u64 data{i}' for i in range(param_cnt))}) {{",
      ]
      for arg in [
          (dtype, lang.type_to_letter(dtype), c)
          for dtype, c in lang.cnts.items()
      ]:
          ins_prefix.append(
              f".reg .{dtype_to_nvtype[arg[0][0]]} %{arg[1]}<{arg[2]}>;",
          )
      ins = ins_prefix + ins
      ins += ["ret;", "}"]
      return "\n".join(ins)
  
file_path: tinygrad/extra/assembly/assembly_ptx.py
imports: from typing import List
  import struct
  from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
  from tinygrad.codegen.linearizer import UOps, UOp
  from tinygrad import dtypes
  from tinygrad.ops import BinaryOps, UnaryOps, TernaryOps
  from tinygrad.runtime.ops_cuda import arch
label_desc: Use image node_id 4 for calling a global function with example usage: specialize_to_ptx(lang, function_name) and returns: str