diff options
| author |  Sanjoy Das <sanjoy@google.com> | 2018-05-21 11:11:48 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-05-21 11:14:06 -0700 |
| commit | a0e4081cf5e11556c8e3d3e022a17afca991b3fe (patch) | |
| tree | 31b3e6f4d087b58ae7e4b3cc4d96963fbdd00c07 /tensorflow/compiler/xla/service/llvm_ir | |
| parent | 0f192f9b0ab0f5c30b7284d5b8eff86993aebb3e (diff) | |
Add a kernel usable as a GEBP inner loop for an LLVM IR GEMM
This is not used in any real code path, but I've added an escape hatch that runs
regular matrix multiplies through this kernel for testing purposes.
As far as I can tell this is functionally correct, but I don't yet have a proper
apples-to-apples performance comparison -- that'll have to wait till the
implementation is complete.
PiperOrigin-RevId: 197422075
Diffstat (limited to 'tensorflow/compiler/xla/service/llvm_ir')
| -rw-r--r-- | tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h | 9 |
1 files changed, 9 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h index 1c00b2aabd..64b935bbf1 100644 --- a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h +++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h @@ -101,6 +101,15 @@ class KernelSupportLibrary { } void For( + tensorflow::StringPiece name, llvm::Value* start, llvm::Value* end, + int64 step, + const std::function<void(llvm::Value* ind_var)>& for_body_generator) { + For(name, start, end, ir_builder_->getInt64(step), + /*peel_first_iteration=*/false, + [&](llvm::Value* indvar, llvm::Value*) { for_body_generator(indvar); }); + } + + void For( tensorflow::StringPiece name, int64 start, int64 end, int64 step, const std::function<void(llvm::Value* ind_var)>& for_body_generator) { For(name, /*start=*/ir_builder_->getInt64(start), |