tensor-compiler / taco

The Tensor Algebra Compiler (taco) computes sparse tensor expressions on CPUs and GPUs
http://tensor-compiler.org
Other
1.25k stars 188 forks source link

Question: Does python API support examine the generated c++ kernel #375

Open ByzanTine opened 3 years ago

ByzanTine commented 3 years ago

Hi, I am wondering if there is a way to view the c++ kernel for a python generated index expression/einsum expression.

When I mean the c++ kernel, I mean the ones generated like https://tensor-compiler.org/codegen.html. We want to make sure the python wrapper generated kernel matches our expectation.

Infinoid commented 3 years ago

Taco generates code in a temp directory and compiles it from there. Debug builds of taco leave the temp directory around after the application ends, so you can take a look at the code that was generated.

Try running a script that uses pytaco to calculate something, and then go look for the tmpdir. Like this:

$ ls -ltr /tmp | tail -n1
drwx------ 2 infinoid infinoid    4096 Jan 21 07:55 taco_tmp_MBllBB

$ ls /tmp/taco_tmp_MBllBB/
qyix7mt8kaz9.c  qyix7mt8kaz9.h  qyix7mt8kaz9.so

$ less /tmp/taco_tmp_MBllBB/qyix7mt8kaz9.c
[snip]
int compute(taco_tensor_t *A6, taco_tensor_t *A0) {
  float* restrict A6_vals = (float*)(A6->vals);
  int A01_dimension = (int)(A0->dimensions[0]);
  int A02_dimension = (int)(A0->dimensions[1]);
  int64_t* restrict A0_vals = (int64_t*)(A0->vals);

  #pragma omp parallel for schedule(runtime)
  for (int32_t x = 0; x < A01_dimension; x++) {
    for (int32_t y = 0; y < A02_dimension; y++) {
      int32_t yA6 = x * 4 + y;
      int32_t yA0 = x * 4 + y;
      A6_vals[yA6] = -A0_vals[yA0];
    }
  }
  return 0;
}
[snip]