Packages

xla

0.10.0
0.10.0 0.9.1 0.9.0 0.8.0 0.7.1 0.7.0 0.6.0 0.5.1 0.5.0 0.4.4 0.4.3 0.4.2 retired 0.4.1 0.4.0 0.3.0 0.2.0 0.1.1 0.1.0

Precompiled XLA binaries

HexDocs

Current section

Files

Jump to
Readme
18 Versions
1 Dependency
Dependants
Files
Activity
xla extension custom_calls lu_bf16.cc
Raw

extension/custom_calls/lu_bf16.cc 

#include "lu.h"



#include "xla/tsl/platform/types.h"



ffi::Error lu_cpu_custom_call_bf16_impl(ffi::Buffer<ffi::BF16> operand,

                                        ffi::ResultBuffer<ffi::U8> p,

                                        ffi::ResultBuffer<ffi::BF16> l,

                                        ffi::ResultBuffer<ffi::BF16> u) {

  return lu_cpu_custom_call_impl<tsl::bfloat16, ffi::Buffer<ffi::BF16>>(

      operand, p, l, u);

}



XLA_FFI_DEFINE_HANDLER_SYMBOL(lu_cpu_custom_call_bf16,

                              lu_cpu_custom_call_bf16_impl,

                              ffi::Ffi::Bind()

                                  .Arg<ffi::Buffer<ffi::BF16>>()

                                  .Ret<ffi::Buffer<ffi::U8>>()

                                  .Ret<ffi::Buffer<ffi::BF16>>()

                                  .Ret<ffi::Buffer<ffi::BF16>>());



XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), "lu_cpu_custom_call_bf16", "Host",

                         lu_cpu_custom_call_bf16);