aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h
diff options
context:
space:
mode:
Diffstat (limited to 'tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h')
-rw-r--r--tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h74
1 files changed, 43 insertions, 31 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 6f7a9d94e3..b00f903d56 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h
+++ b/tensorflow/compiler/xla/service/llvm_ir/kernel_support_library.h
@@ -30,14 +30,14 @@ namespace xla {
// flow more readable.
class KernelSupportLibrary {
public:
- // `ir_builder` is the llvm::IRBuilder instance used to generate LLVM IR.
+ // `b` is the llvm::IRBuilder instance used to generate LLVM IR.
// `unroll_mode` specifies the desired LLVM unrolling behavior for every loop
// generated by this instance of KernelSupportLibrary.
explicit KernelSupportLibrary(
- llvm::IRBuilder<>* ir_builder,
+ llvm::IRBuilder<>* b,
llvm_ir::UnrollMode unroll_mode = llvm_ir::UnrollMode::kNoUnroll,
bool prevent_vectorization = true)
- : ir_builder_(ir_builder),
+ : b_(b),
unroll_mode_(unroll_mode),
prevent_vectorization_(prevent_vectorization) {}
@@ -71,18 +71,18 @@ class KernelSupportLibrary {
const std::function<Status(llvm::Value* ind_var,
bool is_first_iteration)>&
for_body_generator) {
- return For(name, /*start=*/ir_builder_->getInt64(start),
- /*end=*/ir_builder_->getInt64(end),
- /*step=*/ir_builder_->getInt64(step), for_body_generator);
+ return For(name, /*start=*/b_->getInt64(start),
+ /*end=*/b_->getInt64(end),
+ /*step=*/b_->getInt64(step), for_body_generator);
}
void ForReturnVoid(
tensorflow::StringPiece name, int64 start, int64 end, int64 step,
const std::function<void(llvm::Value* ind_var, bool is_first_iteration)>&
for_body_generator) {
- ForReturnVoid(name, /*start=*/ir_builder_->getInt64(start),
- /*end=*/ir_builder_->getInt64(end),
- /*step=*/ir_builder_->getInt64(step), for_body_generator);
+ ForReturnVoid(name, /*start=*/b_->getInt64(start),
+ /*end=*/b_->getInt64(end),
+ /*step=*/b_->getInt64(step), for_body_generator);
}
// Generates the following control flow structure if `peel_first_iteration` is
@@ -184,17 +184,17 @@ class KernelSupportLibrary {
Status For(
tensorflow::StringPiece name, int64 start, int64 end, int64 step,
const std::function<Status(llvm::Value* ind_var)>& for_body_generator) {
- return For(name, /*start=*/ir_builder_->getInt64(start),
- /*end=*/ir_builder_->getInt64(end),
- /*step=*/ir_builder_->getInt64(step), for_body_generator);
+ return For(name, /*start=*/b_->getInt64(start),
+ /*end=*/b_->getInt64(end),
+ /*step=*/b_->getInt64(step), for_body_generator);
}
void ForReturnVoid(
tensorflow::StringPiece name, int64 start, int64 end, int64 step,
const std::function<void(llvm::Value* ind_var)>& for_body_generator) {
- ForReturnVoid(name, /*start=*/ir_builder_->getInt64(start),
- /*end=*/ir_builder_->getInt64(end),
- /*step=*/ir_builder_->getInt64(step), for_body_generator);
+ ForReturnVoid(name, /*start=*/b_->getInt64(start),
+ /*end=*/b_->getInt64(end),
+ /*step=*/b_->getInt64(step), for_body_generator);
}
// Generates the following control flow structure:
@@ -203,16 +203,30 @@ class KernelSupportLibrary {
// `true_block_generator()`;
// else
// `false_block_generator()`;
- Status If(llvm::Value* condition,
+ Status If(tensorflow::StringPiece name, llvm::Value* condition,
const std::function<Status()>& true_block_generator,
const std::function<Status()>& false_block_generator =
[]() -> Status { return Status::OK(); });
+ Status If(llvm::Value* condition,
+ const std::function<Status()>& true_block_generator,
+ const std::function<Status()>& false_block_generator =
+ []() -> Status { return Status::OK(); }) {
+ return If("", condition, true_block_generator, false_block_generator);
+ }
+
void IfReturnVoid(llvm::Value* condition,
const std::function<void()>& true_block_generator,
const std::function<void()>& false_block_generator = []() {
}) {
- TF_CHECK_OK(If(condition,
+ IfReturnVoid("", condition, true_block_generator, false_block_generator);
+ }
+
+ void IfReturnVoid(tensorflow::StringPiece name, llvm::Value* condition,
+ const std::function<void()>& true_block_generator,
+ const std::function<void()>& false_block_generator = []() {
+ }) {
+ TF_CHECK_OK(If(name, condition,
[&]() {
true_block_generator();
return Status::OK();
@@ -244,41 +258,39 @@ class KernelSupportLibrary {
// in a nullptr llvm::Value* in its position to `kernel_body_generator`.
// Currently we only support at most one nullptr value in `arguments`.
static void EmitAndCallOutlinedKernel(
- bool enable_fast_math, bool optimize_for_size,
- llvm::IRBuilder<>* ir_builder, tensorflow::StringPiece kernel_name,
- ArgumentVector arguments,
+ bool enable_fast_math, bool optimize_for_size, llvm::IRBuilder<>* b,
+ tensorflow::StringPiece kernel_name, ArgumentVector arguments,
const std::function<void(ArgumentVector)>& kernel_body_generator);
// Thin wrappers around the more general EmitAndCallOutlinedKernel above.
static void EmitAndCallOutlinedKernel(
- bool enable_fast_math, bool optimize_for_size,
- llvm::IRBuilder<>* ir_builder, tensorflow::StringPiece kernel_name,
- llvm::Value* arg0, llvm::Value* arg1, llvm::Value* arg2,
+ bool enable_fast_math, bool optimize_for_size, llvm::IRBuilder<>* b,
+ tensorflow::StringPiece kernel_name, llvm::Value* arg0, llvm::Value* arg1,
+ llvm::Value* arg2,
const std::function<void(llvm::Value*, llvm::Value*, llvm::Value*)>&
kernel_body_generator) {
EmitAndCallOutlinedKernel(
- enable_fast_math, optimize_for_size, ir_builder, kernel_name,
- {arg0, arg1, arg2}, [&](ArgumentVector args) {
+ enable_fast_math, optimize_for_size, b, kernel_name, {arg0, arg1, arg2},
+ [&](ArgumentVector args) {
kernel_body_generator(args[0], args[1], args[2]);
});
}
static void EmitAndCallOutlinedKernel(
- bool enable_fast_math, bool optimize_for_size,
- llvm::IRBuilder<>* ir_builder, tensorflow::StringPiece kernel_name,
- llvm::Value* arg0, llvm::Value* arg1, llvm::Value* arg2,
- llvm::Value* arg3,
+ bool enable_fast_math, bool optimize_for_size, llvm::IRBuilder<>* b,
+ tensorflow::StringPiece kernel_name, llvm::Value* arg0, llvm::Value* arg1,
+ llvm::Value* arg2, llvm::Value* arg3,
const std::function<void(llvm::Value*, llvm::Value*, llvm::Value*,
llvm::Value*)>& kernel_body_generator) {
EmitAndCallOutlinedKernel(
- enable_fast_math, optimize_for_size, ir_builder, kernel_name,
+ enable_fast_math, optimize_for_size, b, kernel_name,
{arg0, arg1, arg2, arg3}, [&](ArgumentVector args) {
kernel_body_generator(args[0], args[1], args[2], args[3]);
});
}
private:
- llvm::IRBuilder<>* ir_builder_;
+ llvm::IRBuilder<>* b_;
llvm_ir::UnrollMode unroll_mode_;
bool prevent_vectorization_;
};
generated by cgit v1.2.3 (git 2.39.1) at 2024-05-29 09:15:51 +0000