code: def convert_blockmask(blockmask, causal):
      """Convert from the 0-1 format to the format used by the CUDA code.
      0 means the block is skipped.
      nonzero means the block is not skipped.
      Argument:
          blockmask: (row, col): a 0-1 tensor
      Return:
          blockmask_converted: (col, row), dtype torch.int32: for each column, it contains the row
              indices of the nonzero blocks, padded with -1 to reach length @row.
              The indices are multiplied by 4, with the smallest bit used to encode whether
              it is the first nonzero in its row, and the 2nd smallest bit to encode whether it is
              the last nonzero in its row..
      """
      assert not causal
      # TD [2022-05-13]: The indexing and sorting is very tricky
      nrow, ncol = blockmask.shape
      # Sort does not support bool on CUDA
      blockmask = blockmask.to(dtype=torch.uint8)
      nonzero_val, nonzero_sorted_rowidx = blockmask.sort(
          dim=0, stable=True, descending=True
      )
      nonzero_unsorted_rowidx = nonzero_sorted_rowidx.argsort(dim=0)
      last_nonzero_col_per_row = blockmask.sort(
          dim=-1, stable=True
      ).indices[:, -1]
      last_nonzero_col_per_row_after_sort = nonzero_unsorted_rowidx[
          torch.arange(nrow, device=blockmask.device),
          last_nonzero_col_per_row,
      ]
      first_nonzero_col_per_row = blockmask.sort(
          dim=-1, stable=True, descending=True
      ).indices[:, 0]
      first_nonzero_col_per_row_after_sort = nonzero_unsorted_rowidx[
          torch.arange(nrow, device=blockmask.device),
          first_nonzero_col_per_row,
      ]
      nonzero_idx = nonzero_sorted_rowidx * 4
      nonzero_idx[
          last_nonzero_col_per_row_after_sort, last_nonzero_col_per_row
      ] += 2
      nonzero_idx[
          first_nonzero_col_per_row_after_sort,
          first_nonzero_col_per_row,
      ] += 1
      nonzero_idx[nonzero_val == 0] = -1
      return nonzero_idx.T.contiguous().to(dtype=torch.int32)
  
file_path: flash-attention/flash_attn/flash_blocksparse_attn_interface.py
imports: import flash_attn_cuda
  import torch
  import torch.nn
label_desc: Use image node_id 1 for calling a global function with example usage: convert_blockmask(blockmask, causal) and returns: nonzero_idx