CPU JiT Backend

[jeremylt]

Purpose:

This PR is for a CPU JiT backend. At the moment it doesn't run, but the bones and basic vision should (I hope) be clear. I am wondering if we can make some performance gains by compiling to the specific basis and restriction parameters.

In principle, once this works, we can drop in XSMM support easily.

Closes: #1239

LLM/GenAI Disclosure:

Describe any LLM and GenAI usage here.

None

[jeremylt]

Sample generated code:

#define BLOCK_SIZE 8

#include <ceed/jit-source/cpu-gen/cpu-jit.h>

#include <ceed/jit-source/cpu-gen/cpu-gen-utils.h>

// Ceed object templates
#include <ceed/jit-source/cpu-gen/cpu-gen-restriction-templates.h>
#include <ceed/jit-source/cpu-gen/cpu-gen-basis-templates.h>

// User QFunction source
#include "/home/jeremy/Dev/libCEED/tests/t500-operator.h"

static int CeedCpuGenOperator_setup(void *ctx, const FieldData_Cpu_Gen *inputs, FieldData_Cpu_Gen *outputs) {
  // Operator constants
  const CeedInt Q = 8;
  const CeedInt Q_1d = 8;
  const CeedInt num_elem = 15;
  const CeedInt num_blocks = (num_elem / BLOCK_SIZE) + !!(num_elem % BLOCK_SIZE);

  // Field constants
  // -- Input Fields
  // ---- Input Field 0: weight
  const CeedInt num_q_comp_in_0 = 1;
  const CeedInt dim_in_0 = 1;
  const CeedInt P_1d_in_0 = 2;
  // ---- Input Field 1: dx
  const CeedInt elem_size_in_1 = 2;
  const CeedInt num_comp_in_1 = 1;
  const CeedInt num_q_comp_in_1 = 1;
  const CeedInt dim_in_1 = 1;
  const CeedInt P_1d_in_1 = 2;
  
  // -- Output Fields
  // ---- Output Field 0: rho
  const CeedInt elem_size_out_0 = 8;
  const CeedInt num_comp_out_0 = 1;
  const CeedInt num_q_comp_out_0 = num_comp_out_0;
  const CeedInt dim_out_0 = 1;
  const CeedInt P_out_0 = elem_size_out_0;

  // Scratch restriction buffer space
  const CeedInt max_e_vec_buffer_size = 2;
  CeedScalar e_vec_scratch[max_e_vec_buffer_size * BLOCK_SIZE];

  // Loop over blocks
  for (CeedInt block = 0; block < num_blocks; block++) {
    // -- Input ElemRestrictions and Bases
    // ---- Input Field 0: weight
    // ------ Restriction Type: none
    // ------ Evaluation Mode: quadrature weights
    CeedScalar q_vec_in_0[Q * BLOCK_SIZE];
    CeedCall(CeedBasis_Apply_Weight_Tensor_1D<BLOCK_SIZE, Q_1d>(inputs[0].weights, q_vec_in_0));
    // ---- Input Field 1: dx
    // ------ Restriction Type: offset
    CeedScalar *e_vec_in_1 = e_vec_scratch;
    {
      const CeedInt comp_stride = 1;
      
      CeedCall(CeedElemRestriction_Apply_NoTranspose_Offset<BLOCK_SIZE, num_comp_in_1, elem_size_in_1, num_elem, comp_stride>(block, inputs[1].offsets, inputs[1].array, e_vec_in_1));
    }
    // ------ Evaluation Mode: gradient
    CeedScalar q_vec_in_1[num_q_comp_in_1 * Q * BLOCK_SIZE];
    CeedCall(CeedBasis_Apply_NoTranspose_Grad_Tensor_1D<BLOCK_SIZE, num_comp_in_1, P_1d_in_1, Q_1d>(inputs[1].interp, inputs[1].grad, e_vec_in_1, q_vec_in_1));
    
    // -- QFunction
    // ---- QFunction inputs
    const CeedScalar* q_vecs_in[2] = {
      q_vec_in_0,
      q_vec_in_1,
    };

    // ---- QFunction outputs
    CeedScalar q_vec_out_0[num_q_comp_out_0 * Q * BLOCK_SIZE];
    CeedScalar* q_vecs_out[1] = {
      q_vec_out_0,
    };

    // ---- Call User QFunction
    CeedCall(setup(ctx, Q * BLOCK_SIZE, q_vecs_in, q_vecs_out));

    // -- Output Bases and ElemRestrictions
    // ---- Output Field 0: rho
    // ------ Evaluation Mode: none
    CeedScalar *e_vec_out_0 = q_vec_out_0;
    // ------ Restriction Type: strided
    {
      const CeedInt strides_0 = 1, strides_1 = 8, strides_2 = 8;
      
      CeedCall(CeedElemRestriction_ApplyAdd_Transpose_Strided<BLOCK_SIZE, num_comp_out_0, elem_size_out_0, num_elem, strides_0, strides_1, strides_2>(block, e_vec_out_0, outputs[0].array));
    }
  }
  return CEED_ERROR_SUCCESS;
}

[jeremylt]

@zatkins-dev there's bugs somewhere in my templated code, but this compiles and runs now, so I think the idea is workable

[zatkins-dev]

i'll take a look, i like templates

[jeremylt]

As with [cuda,hip]/gen, CEED_DEBUG=1 turn on printing out the generated code, but it also gets dumped into the folder temp (same folder for CUDA + Rust)

[zatkins-dev]

  const CeedInt Q = 8;
  const CeedInt Q_1d = 8;
  const CeedInt num_elem = 15;
  const CeedInt num_blocks = (num_elem / BLOCK_SIZE) + !!(num_elem % BLOCK_SIZE);

  // Field constants
  // -- Input Fields
  // ---- Input Field 0: weight
  const CeedInt num_q_comp_in_0 = 1;
  const CeedInt dim_in_0 = 1;
  const CeedInt P_1d_in_0 = 2;
  // ---- Input Field 1: dx
  const CeedInt elem_size_in_1 = 2;
  const CeedInt num_comp_in_1 = 1;
  const CeedInt num_q_comp_in_1 = 1;
  const CeedInt dim_in_1 = 1;
  const CeedInt P_1d_in_1 = 2;
  
  // -- Output Fields
  // ---- Output Field 0: rho
  const CeedInt elem_size_out_0 = 8;
  const CeedInt num_comp_out_0 = 1;
  const CeedInt num_q_comp_out_0 = num_comp_out_0;
  const CeedInt dim_out_0 = 1;
  const CeedInt P_out_0 = elem_size_out_0;

If we're using C++ to compile this, these should really be constexpr to give the compiler extra hints

[jeremylt]

Feel free to push direct to the branch - my puppy needs exercise now, according to him :)