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

Integrating generated kernel for another apps #204

Open awnguyen1559 opened 5 years ago

awnguyen1559 commented 5 years ago

I am still a bit unsure on how to integrate my generated kernel using the taco command line tool (with a new assembly/compute loop) for another computation that I have (that is similar to tensor_times_vector.cpp). Could someone point me to the document that I can find out more about this?

RSenApps commented 5 years ago

What exactly do you mean by integrate? You can use the command line tool if you want just the kernel, but it is typically easier to use the library (like in the tensor_times_vector example). The library will compile the same kernel as the command line tool and wraps it in the Tensor class for the rest of your application to use.

Right now the Tensor API is fairly limited, but new features that will make taco integration easier are being worked on in #189

For now you can iterate over values of computed tensors and use the values in the rest of your application as desired.

awnguyen1559 commented 5 years ago

So I used the command line tool in build/bin/taco to generate my own loop of A(i,k) = B(i,j) * C(j,k) to my own taco-generated-kernel.h. I have another file as the tensor_times_vector.cpp in the apps folder to test the computation. I just didn’t know how to use both of those files (that I generated) within the framework?

RSenApps commented 5 years ago

Instead of using the command line tool, you want to add a new IndexStmt to do your other computation.

For example you can edit the tensor_times_vector example like the following

A(i,j) = B(i,j,k) * c(k);
A.compile();
...
A2(i,k) = B2(i,j) * C2(j,k);
A2.compile();
...

Or even combine the expressions to generate a more efficient kernel for your final result.

awnguyen1559 commented 5 years ago

I see, that’s what I have been doing instead. On that note, what should I do if I want to put some of my code inside the compile and compute functions then? From what I’ve known so far, those functions are in the taco_tensor.h/cpp file. However, it doesn’t show my changes (e.g some print statements) if I put them in those functions.

RSenApps commented 5 years ago

What do you mean by put some code inside these functions? Specifically, what type of state does this code need to access and how does the code change the behavior of the program. There are various options depending on what you want to do. You could wrap the compile and compute methods with your own methods or you could integrate your code into the compiled kernels by modifying codegen_c.cpp, if your code doesn't care about the actual tensor data. Otherwise, there are various options for either modifying the generate code during lowering or performing a rewriting pass on one of our intermediate IRs. If you let me know your exact use case, I can better point you in the right direction.

oicirtap commented 5 years ago

Correct me if I'm wrong @anhwnguyen, but it seems like you want to experiment with changing the source code of the compiled assemble and compute methods for a tensor.

How the taco library works is that calling compile on a tensor will generate the appropriate assemble and compute methods for the specified IndexExpression (as exemplified by @RSenApps above). If you wish to experiment with changing the source code for these expressions, you can use the compileSource method available to the tensor class. compileSource will allow you to write arbitrary code for assemble and compute and run it within the taco framework for the specified tensor. You can get an idea of how to format such source code by calling the tensor method getSource, which will print out the exact source code used to generate assemble and compute given the specified expression.

To partially answer your original question, the code generated by the command line tool assumes that your tensors are stored as taco_tensor_t objects. getSource will also be useful to provide examples of how you can use taco_tensor_t together with the assemble and compute functions.

If you decide to use the c++ library, I would recommend copying the code generated by getSource on an expression similar to what you want, and then modify the output source code to suit your needs. Then, you can call compileSource on your new source code and run it by calling assemble and compute. I don't know the exact details of taco's code compilation, but since the compilation environment is not the same as regular code compilation, calls such as print statements unfortunately might not show up properly in your command line when executing code in this manner.

I believe we currently do not have proper documentation on how to integrate code generated by our command line tool with other source code. While it'd be good to have such documentation, I am not sure we'll be writing it in the near future. For the time being, it might be much easier to use the c++ library as @RSenApps suggested.

I hope this helps answer some of your questions.

awnguyen1559 commented 5 years ago

Yes! That was great, thank you so much!

awnguyen1559 commented 5 years ago

I have one more question if you happen to know the answer or could direct me to a good document/location in the code to start: After generating the new source code with my changes, how can I use that code with some of the existing sample matrices in the Data folder?

oicirtap commented 5 years ago

We have some file IO functions which allow us to read in mtx, ttx, and tns file formats and get Tensor and TensorBase objects. You can find read() and write() functions for this in tensor.h. You'd probably have to read a tensor and then compile your code into the returned Tensor object with compileSource.

read() and write() are not very fast (we have a pending task to speed up these functions), but they are the simplest option to get you started.

nt1tov commented 5 years ago

+1: I really do not understand how to use generated kernels