Inside Mirage (2) - Transpiler from MuGraph to CUDA

TranspileResult generate_code() {
  this->resolve_distributed_config();  // Multi-GPU configuration
  this->resolve_dtensor_meta();        // Gather tensor metadata
  this->resolve_tb_fusion();           // Plan operator fusion
  this->resolve_tensor_layout();       // Decide tensor layouts
  this->plan_dtensor_memory();         // Allocate global memory
  return this->transpile_ugraph();     // Generate CUDA code
}

// Currently we only fuse elementwise unary operators and
// forloop_accum_no_red with the previous operator
for (tb::TBOperator *const op : tb_graph.operators) {
  if ((type::is_threadblock_element_unary(op->op_type)) ||
      (op->op_type == type::TB_FORLOOP_ACCUM_NO_RED_OP)) {
    tb::STensor const &input0 = op->input_tensors.at(0);
    tb::TBOperator *prev_op = input0.owner_op;
    // Check various conditions...
    if (num_consumers.at(input0.guid) == 1 &&
        prev_op->op_type != type::TB_FORLOOP_ACCUM_NO_RED_OP &&
        prev_op->op_type != type::TB_INPUT_OP &&
        /* ... other conditions ... */) {
      is_fused_with_prev[op] = true;
      is_fused_with_next[prev_op] = true;
    }
  }
}

// Construct `fusion_chain`
for (tb::TBOperator *const last_op : tb_graph.operators) {
  if (is_fused_with_next[last_op]) continue;
  // Now last_op is the tail of a fusion chain
  std::vector<tb::TBOperator const *> fused_ops;
  tb::TBOperator *cur_op = last_op;
  while (true) {
    fused_ops.push_back(cur_op);
    if (is_fused_with_prev[cur_op]) {
      cur_op = cur_op->input_tensors.at(0).owner_op;
    } else {
      break;
    }
  }
  std::reverse(fused_ops.begin(), fused_ops.end());
  fusion_chain[leading_op] = fused_ops;
}

for (kn::DTensor const &dtensor : all_dtensors) {
  // Every DTensor can only have 1 innermost dim
  z3::expr_vector innermost_exprs(ctx);
  for (int i = 0; i < num_dims; ++i) {
    innermost_exprs.push_back(d_is_innermost[dtensor.guid][i]);
  }
  opt.add(z3::atmost(innermost_exprs, 1));
  opt.add(z3::atleast(innermost_exprs, 1));
}

// For matmul, innermost must be in last two dims
for (kn::DTensor const &tensor : {lhs, rhs, output}) {
  int num_dims = tensor.num_dims;
  opt.add(d_is_innermost[tensor.guid][num_dims - 1] ||
          d_is_innermost[tensor.guid][num_dims - 2]);
}

namespace cost {
cost_t KN_REDUCTION_INNERMOST_EQ_REDUC_DIM = 4000;
cost_t TB_INPUT_NO_WIDE_COPY = 4000;
cost_t TB_OUTPUT_NO_WIDE_COPY = 4000;
cost_t TB_MATMUL_NO_LDMATRIX = 10000;
cost_t TB_INPUT_NO_CP_ASYNC = 20000;
cost_t SWIZZLE_DIM = 1000;
cost_t SMEM_FACTOR = 1;
}

void calc_tensor_strides(size_t *strides, size_t &num_phy_elems,
                         int num_dims, int const dims[],
                         int innermost_dim, int datatype_size) {
  // Order: [innermost, N-1, N-2, ..., innermost+1, innermost-1, ..., 0]
  vector<int> dim_order = {innermost_dim};
  for (int i = num_dims - 1; i >= 0; --i) {
    if (i != innermost_dim) dim_order.push_back(i);
  }

  size_t alignment = std::max(16 / datatype_size, 1);
  size_t cur_stride = 1;
  for (int dim_idx : dim_order) {
    strides[dim_idx] = cur_stride;
    if (dims[dim_idx] != 1) {
      // Pad first non-1 dimension to 16-byte alignment
      cur_stride *= round_to_multiple((size_t)dims[dim_idx], alignment);
    }
  }
  num_phy_elems = cur_stride;
}

// Calculate depth using dynamic programming
for (tb::TBOperator *const op : tb_graph.operators) {
  if (op->op_type == type::TB_INPUT_OP) {
    op2chaining_meta[op] = {0, next_chain_idx++, 0};
  }
}
for (tb::TBOperator *const op : tb_graph.operators) {
  if (op->op_type == type::TB_INPUT_OP) continue;
  if (is_fused_with_prev[op]) {
    // Same depth as predecessor
    OpChainingMeta prev_meta = op2chaining_meta.at(prev_op);
    op2chaining_meta[op] = {prev_meta.level, prev_meta.fuse_chain_idx,
                            prev_meta.pos_in_fuse_chain + 1};
  } else {
    // max(input depths) + 1
    int res = 0;
    for (tb::STensor const &input : op->input_tensors) {
      res = std::max(res, op2chaining_meta.at(input.owner_op).level);
    }
    op2chaining_meta[op] = {res + 1, next_chain_idx++, 0};
  }
}

struct TBSched {
  std::vector<TBSchedNode> pre_loop_nodes;   // Before the forloop
  std::vector<TBSchedNode> loop_nodes;       // Inside the forloop
  std::vector<TBSchedNode> post_loop_nodes;  // After the forloop
};

struct TBSchedOpMeta {
  bool is_chunked_input;    // Use 128-bit wide copies?
  bool is_pipelined_input;  // Use async cp.async?
  bool is_chunked_output;   // Use 128-bit wide writes?
  bool is_accum_in_reg;     // Keep accumulator in registers?
  // ...
};

// Allow accumulators to take up to 192 registers per thread
// (Each thread has 255 32-bit registers max)
if (per_thread_accum_numel_tot + per_thr_accum_numel <= 192) {
  op_meta.is_accum_in_reg = true;
  per_thread_accum_numel_tot += per_thr_accum_numel;
} else {
  op_meta.is_accum_in_reg = false;
}

if (is_power_of_2(num_chunks_in_inner_dim)) {
  // XOR-based swizzling
  meta.is_xor_swizzled = true;
  meta.xor_swizzle_b = log2(num_chunks_in_128B);
  meta.xor_swizzle_m = log2(chunk_size_num_elems);
  meta.xor_swizzle_s = log2(num_chunks_in_inner_dim);
} else {
  // Shift-based swizzling
  int new_num_chunks = num_chunks_in_inner_dim;
  while (gcd(new_num_chunks, num_chunks_in_128B) != 1) {
    new_num_chunks++;
  }
  // Update strides to reflect padding
  meta.strides[meta.innermost_dim] = 1;
  size_t cur_stride = new_num_chunks * chunk_size_num_elems;
  // ... update other strides
}

class MemoryPlanner {
  size_t cur_addr = 0;

  size_t allocate(size_t size) {
    size_t addr = cur_addr;
    cur_addr += size;
    cur_addr = round_to_multiple(cur_addr, 128);  // 128B alignment
    return addr;
  }
};

// Pre-loop inputs: allocated at pre_loop time, freed at 2T or post_loop usage
// Loop intermediates: allocated/freed within each iteration
// Accumulators: allocated before loop, freed after post_loop usage
// Post-loop intermediates: allocated/freed within post_loop phase

vector<std::shared_ptr<AbstractMemoryPlanner>> planners = {
  std::make_shared<FirstFitMemoryPlanner>(),
  std::make_shared<BestFitMemoryPlanner>(),
  std::make_shared<WorseFitMemoryPlanner>()
};

class FirstFitMemoryPlanner : public OnlineAllocMemoryPlannerBase {
  Range select_range(size_t size) override {
    auto it = std::find_if(free_ranges.begin(), free_ranges.end(),
      [&](Range const &range) {
        return range.second - range.first >= size;
      });
    return {it->first, it->first + size};
  }
};

exec.e("kn::gemm<$>($,$,$, $,$,$, ...);",
       compute_type, out0_ptr, in0_ptr, in1_ptr, m, n, k, ...);

exec.e("using kernel = kn::ElementUnaryKernel<$, "
       "kn::ElementUnaryOpType::$, $, $>;",
       datatype, op_type, in_layout, out_layout);
exec.e("kernel::run($, $);", out_ptr, in_ptr);

__global__ void custom_kernel_0(half* output, half const* input) {
  int thread_idx = threadIdx.x;
  extern __shared__ char buf[];

  // Define STensor pointers
  half *stensor0_ptr = (half*)(buf + 0);
  half *stensor1_ptr = (half*)(buf + 1024);

  // Pre-loop: Load static inputs
  STensor0InputAtom::run(stensor0_ptr, dtensor_tile_ptr, thread_idx);
  __syncthreads();

  // Main loop
  for (int for_idx = 0; for_idx < FORLOOP_RANGE; for_idx++) {
    // Load inputs for this iteration
    STensor1InputAtom::run(stensor1_ptr, dtensor_tile_ptr + offset*for_idx, thread_idx);
    __syncthreads();

    // Compute
    Matmul0Kernel::run(matmul_accum, stensor0_ptr, stensor1_ptr, thread_idx);
    // ... more operators
  }
  __syncthreads();

  // Post-loop: reductions, outputs
  OutputAtom::run(dtensor_output, stensor_result, thread_idx);
}

string transpile_fusion_epilogue(chain, dtype) {
  string res = "tb::EpilogueStore<half>";
  for (op in chain[1:]) {
    if (op->op_type == TB_EXP_OP)
      res = fmt("tb::EpilogueExp<half, $>", res);
    else if (op->op_type == TB_FORLOOP_ACCUM_NO_RED_OP)
      res = "tb::EpilogueStoreAccum<half>";
    // ...
  }
  return res;
}

using Epilogue = tb::EpilogueExp<half, tb::EpilogueStore<half>>;

__global__ void custom_kernel_rmsnorm_linear(...) {
  extern __shared__ char buf[];

  // Load weight (pre-loop)
  WNormInputAtom::run(wnorm_ptr, wnorm_tile, thread_idx);
  __syncthreads();

  for (int for_idx = 0; for_idx < K_TILES; for_idx++) {
    // Load X tile
    XInputAtom::run(x_ptr, x_tile + for_idx * TILE_K, thread_idx);
    __syncthreads();

    // RMSNorm computation
    SquareKernel::run(sq_ptr, x_ptr, thread_idx);
    __syncthreads();
    ReductionKernel::run(sum_ptr, sq_ptr, thread_idx);
    __syncthreads();
    SqrtKernel::run(rms_ptr, sum_ptr, thread_idx);
    DivKernel::run(norm_ptr, x_ptr, rms_ptr, thread_idx);
    __syncthreads();

    // Matmul with accumulation
    MatmulKernel::run(accum, norm_ptr, wlinear_ptr, thread_idx);
  }
  __syncthreads();

  // Output
  OutputAtom::run(output_tile, accum, thread_idx);
}