Inside Mirage (1) - Superoptimizing Tensor Programs via MuGraph and Search

class Graph {
public:
  // ...
  // Create a new input tensor (resides in Global Memory)
  DTensor new_input(std::vector<int> const &dims, ...);

  // Create a standard operator (e.g., Matmul)
  DTensor matmul(DTensor const &A, DTensor const &B);

  // Create a "Customized" operator defined by a Thread Block Graph
  // This is where kernel fusion happens!
  std::vector<DTensor> customized(std::vector<DTensor> const &inputs,
                                  mirage::threadblock::Graph const &_graph);
  
  std::vector<mirage::kernel::KNOperator *> operators;
  // ...
};

class Graph {
public:
  // ...
  // Create an input operator (Loads data from Global -> Shared Memory)
  STensor new_input(mirage::kernel::DTensor const &dtensor, ...);

  // Perform computation in Shared Memory
  STensor matmul(STensor const &A, STensor const &B);
  STensor reduction(STensor const &A, int dim);

  // Mark output (Stores data from Shared -> Global Memory)
  mirage::kernel::DTensor mark_output(STensor const &stensor, ...);

  std::vector<mirage::threadblock::TBOperator *> operators;
  // ...
};

// TODO

int cython_search(mirage::kernel::Graph const *input_graph,
                  char const *backend,
                  int max_num_graphs,
                  mirage::kernel::Graph **new_graphs,
                  std::vector<MInt3> imap_to_explore,
                  std::vector<MInt3> omap_to_explore,
                  std::vector<MDim3> grid_dim_to_explore,
                  std::vector<MDim3> block_dim_to_explore,
                  std::vector<int> fmap_to_explore,
                  std::vector<int> frange_to_explore,
                  char const *filename,
                  bool verbose,
                  char const *default_config,
                  bool is_formal_verified);

int cython_search(...) {
  // 1. If a checkpoint file exists, load graphs from it instead of searching
  if (filename) {
    //..
    new_graphs[num] = new kernel::Graph();
    from_json(graph, *new_graphs[num]);
  }

  {
    // 2. Initialize Search Configuration
    search::GeneratorConfig config =
        search::GeneratorConfig::get_default_config();
    //..

    // 3. Set Exploration Parameters (Search Space Pruning)
    // These vectors limit the search to specific configurations if provided.
    // .. imap, omap, grid_dim, block_dim, fmap, frange

    // 4. Run the Search
    // Instantiate the Generator
    search::KernelGraphGenerator gen(
        *input_graph, config, result_filename, verbose);
    
    // Start the search process
    gen.generate_kernel_graphs();

    // 5. Collect and Return Results
    //..
  }
}

void KernelGraphGenerator::preprocess(kernel::Graph const &computation_graph) {
  // 1. Get input attributes: enumerate computation_graph.operators, check op_type == KN_INPUT_OP, and store in computation_graph_input_attrs
  //..

  // 2. Get possible abstract expressions, showing the kernel graph's computation tasks.
  //    They will be used for pruning illegal graphs during search.
  //    
  abstract_expr_eval(computation_graph, computation_graph_exprs);

  // 3. Some ”range“ things. I'm not sure what they are for.
  // ..

  // 4. Get output expressions (abstract)
  for (kernel::KNOperator *op : computation_graph.operators) {
    if (op->op_type == type::KNOperatorType::KN_OUTPUT_OP) {
      computation_graph_output_exprs.push_back(
          computation_graph_exprs.at(op->input_tensors[0].guid));
    }
  }
  // 4.5  Initalize final expressions
  //      THIS CONSTRUCTS AN E-GRAPH
  for (auto const &final_expr : computation_graph_output_exprs) {
    initialize_final_expr(final_expr);
    // ⤷ get_egraph(expr->to_egg().c_str()) // This build an e-graph!
  }
  
  // 5. Initialize verifier, formal / probabilistic
  // ..
}

void KernelGraphGenerator::generate_kernel_graphs() {
  
  // Create a SearchContext
  SearchContext c;
  c.level = SearchLevel::LV_KERNEL;
  c.kn_graph = std::make_shared<kernel::Graph>();

  std::vector<SerializedSearchContext> verified_graphs;

  // This is the main backtracking search function
  generate_next_operator(
      c,
      [this](SearchContext const &c) {
        return c.level == SearchLevel::LV_KERNEL &&
                this->verify(*c.kn_graph);
      },
      verified_graphs,
      /*search_depth=*/0,
      /*is_a_new_thread_start=*/true);


  save_results();
}

void KernelGraphGenerator::generate_next_operator(SearchContext &c, ...) {
    // 1. Check if we found a valid graph
    if (verify(c)) {
        verified_graphs.push_back(SerializedSearchContext(c));
        return;
    }

    // 2. TODO: Something about abstract expr
    std::unordered_map<type::GuidType, std::shared_ptr<AbstractExpr const>> algebraic_expr;
    abstract_expr_eval(*c.kn_graph, algebraic_expr);
    if (c.tb_graph) {
      abstract_expr_eval(*c.tb_graph, algebraic_expr);
    }

    // 2.5 Lambda: check if adding an operator yields a valid abstract expression
    //     What is valid? check_abstract_expr checks if the expression is an subexpr to any final exprs
  
    auto infer_and_check_abstract_expr = [&](auto const &input_tensors,
                                           auto op_type) {
      std::vector<std::shared_ptr<AbstractExpr const>> input_exprs =
          vector_map(input_tensors,
                    [&](auto const &t) { return algebraic_expr.at(t.guid); });
      std::shared_ptr<AbstractExpr const> expr =
          get_abstract_expr(op_type, input_tensors, input_exprs);
      return check_abstract_expr(expr);
    };


    // 2. Try adding a Kernel Operator (KNOperator)
    for (type::KNOperatorType op_type : dim_strategy.get_knop_cand()) {
        // Case K1: finish and verify the current graph
        if (op_type != type::KNOperatorType::KN_CUSTOMIZED_OP) {
          // Case K2: generate a pre-defined kernel operator

          // .. Some checking
          // Add to the graph
          KNOperator *new_op = create_op(*c.kn_graph, op_type, input_tensors);
          if (new_op) {
            c.kn_graph->operators.push_back(new_op);
            // Search deeper
            generate_next_operator(
                c, verify, verified_graphs, search_depth + 1);
          }
          // Backtrack
          while (c.kn_graph->operators.back() != old_last_op) {
            delete c.kn_graph->operators.back();
            c.kn_graph->operators.pop_back();
          }
        } else {
          // Case K3: generate a graph-def kernel operator
          // .. Some checking
          
          // Enumerate grid_dim, block_dim, input_map, forloop_dim, forloop_range
          for (dim3 grid_dim : dim_strategy.get_grid_dim_cand(..)) {
            for (dim3 block_dim :
                 dim_strategy.get_block_dim_cand(..)) {
              for (std::vector<int3> const &input_map :
                   dim_strategy.get_input_map_cand(..)) {
                for (std::vector<int> const &forloop_dim :
                     dim_strategy.get_forloop_dim_cand(..)) {
                  for (int forloop_range :
                       dim_strategy.get_forloop_range_cand(..)) {
                    c.tb_graph = std::make_shared<threadblock::Graph>(
                        grid_dim,
                        block_dim,
                        forloop_range,
                        config.reduction_dimx);
                    
                    // Create input tensors
                    // ..

                    c.level = SearchLevel::LV_THREADBLOCK;
                    // Search deeper
                    generate_next_operator(
                        c, verify, verified_graphs, search_depth + 1);
                    
                    // Backtrack
                    c.level = SearchLevel::LV_KERNEL;
                    c.tb_graph = nullptr;
                  }
                }
              }
            }
          }
        }

    }
    
    // 3. Try adding a Customized Operator (Thread Block Graph Search)
    //    I think this looks very similar to the Kernel Search, except no graph-def kernels
    // ...
}

std::vector<type::KNOperatorType> DimStrategy::get_knop_cand() {
  // Returns a list of operators to explore (e.g., MATMUL, EXP, ADD), random_shuffled
  return config.knop_to_explore;
}

std::vector<dim3> DimStrategy::get_grid_dim_cand(std::vector<DTensor> const &tensors) {
    // Proposes grid dimensions (e.g., {128, 1, 1}, {256, 1, 1}) based on input shapes
    // ...
}

bool KernelGraphGenerator::check_abstract_expr(std::shared_ptr<AbstractExpr const> expr, ..) {
  for (auto const &final_expr : computation_graph_output_exprs) {
    if (subexpr_to_final_expr(expr)) {
      return true;
    }
  }
  return false;
}

std::vector<bool> subexpr_to_final_expr(
    std::vector<std::shared_ptr<AbstractExpr const>> const &exprs) {
  // The `egg_equiv` functions use e-graph's pattern matching to get the expressions that appear in the e-graph, indicating equivalance to any sub-expr of final expressions.
  bool *results_in_raw_array =
      egg_equiv(exprs_c_str, static_cast<int>(exprs.size()));
  // ..
}