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 :)

[jeremylt]

just occurred to me - changing the JiT compile flags to -O0 -g would probably help see what the error is

[zatkins-dev]

I found it

[zatkins-dev]

rather, I found the cause of the seg fault -- there's still at least something wrong

[jeremylt]

Awesome. I figure there are multiple bugs - I got excited and started coding way too fast.

[zatkins-dev]

aha! it was the tensor contractions

[jeremylt]

Ok, t500 rens now, so it compiles and executes the operator function without issues. But there's a memory leak somewhere, and t501 fails, so there's a correctness issue. And AtPoints still needs to be piped in

[jeremylt]

Awesome, hopefully our fixes complement each other

it should be relatively easy to setup the even-odd stuff and to replace this cantraction with XSMM, I think

[zatkins-dev]

Hmmm this is working locally...I wonder what the issue is

[jeremylt]

If dlopen is too fiddly, we may want to use LLVM instead

[zatkins-dev]

We're going to need to be careful -- having multiple functions with the same name seems to result in unpredictable behavior (see t506)

[jrwrigh]

I'm guessing y'all have probably already thought about this, but having the compiler flags be editable/overwritable would be nice. There are times when I've noticed that -02 is actually faster than -O3 in some cases (most of which I never fully understand).