kernel seg fault when processing large tensors

[yiansu]

Hello, I got seg faults using the kernel generated by taco when processing large tensors. I used taco compiler to generate the following kernel computation

taco "A(i, j) = B(i, j, k) * c(k)" -f=A:dd -f=B:dss -f=c:d -write-assembly=assemble.hpp -write-compute=compute.hpp

and use the following main function

#include <random>
#include "taco.h"
#include "assemble.hpp"
#include "compute.hpp"

using namespace taco;
int main(int argc, char* argv[]) {
  std::default_random_engine gen(0);
  std::uniform_real_distribution<double> unif(0.0, 1.0);

  Format dd({Dense, Dense});
  Format dss({Dense, Sparse, Sparse});
  Format d({Dense});

  // Tensor<double> B = read("nell-1.tns", dss);
  Tensor<double> B = read("nell-2.tns", dss);

  Tensor<double> c({B.getDimension(2)}, d);
  for (int i = 0; i < c.getDimension(0); ++i) {
    c.insert({i}, unif(gen));
  }
  c.pack();

  Tensor<double> A({B.getDimension(0), B.getDimension(1)}, dd);
  A.pack();

  assemble(A.getTacoTensorT(), B.getTacoTensorT(), c.getTacoTensorT());
  compute(A.getTacoTensorT(), B.getTacoTensorT(), c.getTacoTensorT());

  return 0;
}

I replaced the restrict keyword with __restrict__ in the generated two *.hpp files and compile it using clang from LLVM15 with the following command

clang++ -std=c++17 -O3 -DNDEBUG -DTACO -I ../../include -L../../build/lib main.cpp -o main -ltaco

The sparse tensors I used are nell-1 and nell-2 from FROSTT, which are also used for the original taco paper. However, the kernel successfully run with the tensor nell-2 but not for the larger one nell-1 (seg fault). Is this a known issue or it's a bug such as data overflow within taco? The issue happens to other large tensors as well.

[rohany]

TACO currently (and it's not a trivial fix) generates code using 32 bit integers. The dimension sizes of nell-1 will cause integer overflow for your dense output tensor, leading to the segfault that you're seeing.