android-14.0.0_r21/s

/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#ifndef XLA_MLIR_RUNTIME_CUSTOM_CALL_ENCODING_H_
#define XLA_MLIR_RUNTIME_CUSTOM_CALL_ENCODING_H_

#include <functional>
#include <memory>
#include <string>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>

#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/ImplicitLocOpBuilder.h"  // from @llvm-project
#include "mlir/IR/SymbolTable.h"  // from @llvm-project
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/xla/runtime/custom_call.h"
#include "tensorflow/compiler/xla/runtime/type_id.h"

namespace xla {
namespace runtime {

//===----------------------------------------------------------------------===//
// Helper classes to build XLA custom calls' lowering to the LLVM dialect.
//===----------------------------------------------------------------------===//
//
// Arguments to the custom call API intrinsic are encoded as an array of opaque
// pointers and at the runtime side available as `void**`. Runtime decodes
// opaque pointers to the C++ data structures (see runtime/custom_call.h), and
// passes them to the registered callback. Argument encoding/decoding must be
// compatible, otherwise it's very easy to get a segfault because of an illegal
// memory access.
//
// Attributes are encoded into a separate opaque storage together with names,
// so the runtime side can decode the attributes it needs and check that all
// required attributes were passed to the custom call handler.
//
// Custom call attributes are encoded as module global constants, and at run
// time we only need to pass a pointer to the constant section.
//
// Custom call arguments are encoded as an array of pointers allocated on the
// stack. Each individual argument is also encoded on the stack, because
// arguments are run time values and we can't encode them in the constant
// section.

// Forward declare class declared below.
class Globals;

//===----------------------------------------------------------------------===//
// Custom call arguments encoding.
//===----------------------------------------------------------------------===//

// Encodes argument into stack allocated storage according to the ABI. If
// argument is a constant, then it can be packed as a global constant.
class CustomCallArgEncoding {
 public:
  struct Encoded {
    mlir::Value type_id;  // !llvm.ptr<i64>
    mlir::Value value;    // !llvm.ptr<ArgType>
  };

  virtual ~CustomCallArgEncoding() = default;

  virtual mlir::LogicalResult Match(mlir::Value value,
                                    mlir::Value conterted) const = 0;

  virtual mlir::FailureOr<Encoded> Encode(Globals &g,
                                          mlir::ImplicitLocOpBuilder &b,
                                          mlir::Value value,
                                          mlir::Value converted) const = 0;
};

// A set of registered custom call arguments encodings.
class CustomCallArgEncodingSet {
 public:
  using Encoded = CustomCallArgEncoding::Encoded;

  // Finds matching argument encoding and tries to encode the values. Returns
  // failure if didn't match values to any of the argument encodings.
  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::Value value,
                                  mlir::Value converted) const;

  template <typename... Ts, typename = std::enable_if_t<sizeof...(Ts) != 0>>
  CustomCallArgEncodingSet &Add() {
    (encodings_.emplace_back(std::make_unique<Ts>()), ...);
    return *this;
  }

 private:
  std::vector<std::unique_ptr<CustomCallArgEncoding>> encodings_;
};

//===----------------------------------------------------------------------===//
// Custom call attributes encoding.
//===----------------------------------------------------------------------===//

// Attributes encoding packs attribute name, data type and a value into the
// module global constant, and returns values pointing to the encoded data.
struct CustomCallAttrEncoding {
  static constexpr char kAttrName[] = "__rt_attr_name";
  static constexpr char kAttrValue[] = "__rt_attr_value";

  struct Encoded {
    mlir::Value name;     // !llvm.ptr<i8>
    mlir::Value type_id;  // !llvm.ptr<i64>
    mlir::Value value;    // !llvm.ptr<EncodedAttrType>
  };

  virtual ~CustomCallAttrEncoding() = default;

  virtual mlir::LogicalResult Match(llvm::StringRef name,
                                    mlir::Attribute attr) const = 0;

  virtual mlir::FailureOr<Encoded> Encode(Globals &g,
                                          mlir::ImplicitLocOpBuilder &b,
                                          llvm::StringRef name,
                                          mlir::Attribute attr) const = 0;
};

// A set of registered custom call attributes encodings.
class CustomCallAttrEncodingSet {
 public:
  using Encoded = CustomCallAttrEncoding::Encoded;

  // Finds matching attribute encoding and tries to encode the attribute.
  // Returns failure if didn't match attribute to any of the encodings.
  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  llvm::StringRef name,
                                  mlir::Attribute attr) const;

  template <typename... Ts, typename = std::enable_if_t<sizeof...(Ts) != 0>>
  CustomCallAttrEncodingSet &Add() {
    (encodings_.emplace_back(std::make_unique<Ts>()), ...);
    return *this;
  }

  template <typename... Ts, typename ConstructorArg,
            typename... ConstructorArgs,
            typename = std::enable_if_t<sizeof...(Ts) != 0>>
  CustomCallAttrEncodingSet &Add(ConstructorArg &&arg,
                                 ConstructorArgs &&...args) {
    (encodings_.emplace_back(std::make_unique<Ts>(arg, args...)), ...);
    return *this;
  }

 private:
  std::vector<std::unique_ptr<CustomCallAttrEncoding>> encodings_;
};

//===----------------------------------------------------------------------===//
// A set of helper functions for packing primitive attributes.
//===----------------------------------------------------------------------===//

// Packs TypeID as `i64` constant value and casts it to the `!llvm.ptr<i8>`,
// because type id internally is implemented as an opaque pointer.
mlir::Value PackTypeId(Globals &g, mlir::ImplicitLocOpBuilder &b,
                       mlir::TypeID type_id);

// Packs string as a module global null-terminated string constant. We reuse
// the encoding scheme for arrays to store sting with its size, to avoid
// computing the length of the null-terminated string at run tine.
//
// Returns `!llvm.ptr<EncodedArray<char>>`.
mlir::Value PackString(Globals &g, mlir::ImplicitLocOpBuilder &b,
                       llvm::StringRef strref, llvm::StringRef symbol_base);

// Packs scalar attribute as a global constant. Returns `!llvm.ptr<AttrType>`.
mlir::Value PackScalarAttribute(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                mlir::Attribute value,
                                mlir::StringRef symbol_base);

//===----------------------------------------------------------------------===//
// A helper class to create global constants in the module.
//===----------------------------------------------------------------------===//

class Globals {
 public:
  // Global value initializer that build the initialization region.
  using GlobalInitializer =
      std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Attribute)>;

  // Global value initializer that can return failure if it can't initialize the
  // global value from the given attribute.
  using FailureOrGlobalInitializer = std::function<mlir::LogicalResult(
      mlir::ImplicitLocOpBuilder &, mlir::Attribute)>;

  Globals(mlir::ModuleOp module, TypeIDNameRegistry type_id_names)
      : module_(module),
        sym_table_(module_),
        type_id_names_(std::move(type_id_names)) {}

  // Creates a global external variable for the type id.
  mlir::LLVM::GlobalOp GetOrCreate(mlir::ImplicitLocOpBuilder &b,
                                   mlir::TypeID type_id);

  // Creates a global null-terminated string constant.
  mlir::LLVM::GlobalOp GetOrCreate(mlir::ImplicitLocOpBuilder &b,
                                   llvm::StringRef strref,
                                   llvm::StringRef symbol_base);

  // Creates a global constant value from the attribute. Attribute type must be
  // a valid type compatible with LLVM globals.
  mlir::LLVM::GlobalOp GetOrCreate(mlir::ImplicitLocOpBuilder &b,
                                   mlir::TypedAttr attr,
                                   llvm::StringRef symbol_base);

  // Creates a global constant value of the given type from the attribute, using
  // optional user-provided global constant initialization.
  mlir::LLVM::GlobalOp GetOrCreate(
      mlir::ImplicitLocOpBuilder &b, mlir::Attribute attr, mlir::Type type,
      llvm::StringRef symbol_base, GlobalInitializer initialize = {},
      mlir::LLVM::Linkage linkage = mlir::LLVM::Linkage::Internal);

  // Creates a global constant value of the given type from the attribute, using
  // optional user-provided global constant initialization. Returns failure if
  // user-provided initialization failed to initialize the global value.
  mlir::FailureOr<mlir::LLVM::GlobalOp> TryGetOrCreate(
      mlir::ImplicitLocOpBuilder &b, mlir::Attribute attr, mlir::Type type,
      llvm::StringRef symbol_base, FailureOrGlobalInitializer initialize = {},
      mlir::LLVM::Linkage linkage = mlir::LLVM::Linkage::Internal);

  // Returns the address of the global value.
  static mlir::Value AddrOf(mlir::ImplicitLocOpBuilder &b,
                            mlir::LLVM::GlobalOp global);

  // Return the address of the global value casted to `!llvm.ptr<i8>`.
  static mlir::Value OpaqueAddrOf(mlir::ImplicitLocOpBuilder &b,
                                  mlir::LLVM::GlobalOp global);

  mlir::ModuleOp module() { return module_; }

 private:
  // Globals key: {attribute, encoded-type, sym-name}. We can only have global
  // constants of one of the LLVM types, and there could be multiple ways to
  // encode an attribute as an LLVM type, e.g. strings can be stored as null
  // terminated array of bytes, or a pair of string size and and array of bytes.
  using Key = std::tuple<mlir::Attribute, mlir::Type, mlir::StringAttr>;

  mlir::LLVM::GlobalOp Find(Key key);

  mlir::ModuleOp module_;
  mlir::SymbolTable sym_table_;  // symbol table for the `module_`
  llvm::DenseMap<Key, mlir::LLVM::GlobalOp> globals_;

  // A mapping from the TypeID to the unique type name for encoding external
  // globals corresponding to types ids.
  TypeIDNameRegistry type_id_names_;
};

//===----------------------------------------------------------------------===//
// Custom call attributes encoding.
//===----------------------------------------------------------------------===//

// Encodes attribute using a scheme compatible with run time attributes decoding
// (see `internal::DecodedAttrs` in the custom call header file).
//
// Returns a value of `!llvm.ptr<ptr<i8>>` (void**) type pointing to the encoded
// attributes array (array of pointers).
//
// This function is used to encode:
//
//   1. Struct attributes as aggregates of nested attributes, where the order of
//      attributes matches the order defined with the `AggregateAttrDef` schema
//      defined below.
//
//   2. Custom call attributes, where the attributes sorted lexicographically by
//      name, to be able to efficiently decode named attributes.
//
mlir::FailureOr<mlir::Value> EncodeAttributes(
    Globals &g, mlir::ImplicitLocOpBuilder &b,
    const CustomCallAttrEncodingSet &encoding, llvm::StringRef symbol_base,
    llvm::ArrayRef<mlir::NamedAttribute> attrs);

struct StringAttrEncoding : public CustomCallAttrEncoding {
  mlir::LogicalResult Match(llvm::StringRef, mlir::Attribute) const final;
  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::StringRef, mlir::Attribute) const final;
};

struct ScalarAttrEncoding : public CustomCallAttrEncoding {
  mlir::LogicalResult Match(llvm::StringRef, mlir::Attribute) const final;
  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::StringRef, mlir::Attribute) const final;
};

struct DenseElementsAttrEncoding : public CustomCallAttrEncoding {
  mlir::LogicalResult Match(llvm::StringRef, mlir::Attribute) const final;
  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::StringRef, mlir::Attribute) const final;
};

struct ArrayAttrEncoding : public CustomCallAttrEncoding {
  mlir::LogicalResult Match(llvm::StringRef, mlir::Attribute) const final;
  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::StringRef, mlir::Attribute) const final;
};

struct DenseArrayAttrEncoding : public CustomCallAttrEncoding {
  mlir::LogicalResult Match(llvm::StringRef, mlir::Attribute) const final;
  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::StringRef, mlir::Attribute) const final;
};

struct EmptyArrayAttrEncoding : public CustomCallAttrEncoding {
  mlir::LogicalResult Match(llvm::StringRef, mlir::Attribute) const final;
  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::StringRef, mlir::Attribute) const final;
};

// Custom call attribute encoding that encodes enums using their underlying
// scalar type. Type id is based on the enum type passed to the runtime.
//
// This encoding can convert enum types defined in the compiler (e.g. dialect
// enums defined in MLIR) to the enum types used at run time.
template <typename AttrType, typename EnumType,
          typename RuntimeEnumType = EnumType>
struct EnumAttrEncoding : public CustomCallAttrEncoding {
  static_assert(std::is_enum<RuntimeEnumType>::value, "must be an enum class");

  // Convert from the compile time enum to the run time enum.
  using Converter = std::function<RuntimeEnumType(EnumType)>;

  EnumAttrEncoding() {
    static_assert(std::is_same<EnumType, RuntimeEnumType>::value,
                  "requires enum converter");
    convert = [](EnumType value) { return value; };
  }

  explicit EnumAttrEncoding(Converter convert) : convert(std::move(convert)) {}

  mlir::LogicalResult Match(llvm::StringRef, mlir::Attribute attr) const final {
    return mlir::success(attr.isa<AttrType>());
  }

  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::StringRef name,
                                  mlir::Attribute attr) const final {
    // Convert enum underlying integral value to an attribute.
    EnumType compile_time_enum = attr.cast<AttrType>().getValue();
    RuntimeEnumType run_time_enum = convert(compile_time_enum);

    using T = std::underlying_type_t<RuntimeEnumType>;
    T underlying_value = static_cast<T>(run_time_enum);

    mlir::TypeID type_id = mlir::TypeID::get<Tagged<RuntimeEnumType>>();
    mlir::Attribute underlying_attr = AsAttr(b, underlying_value);

    Encoded encoded;
    encoded.name = PackString(g, b, name, kAttrName);
    encoded.type_id = PackTypeId(g, b, type_id);
    encoded.value = PackScalarAttribute(g, b, underlying_attr, kAttrValue);

    return encoded;
  }

  static mlir::Attribute AsAttr(mlir::ImplicitLocOpBuilder &b, uint32_t value) {
    return b.getI32IntegerAttr(value);
  }

  Converter convert;
};

// A helper type to define `AttrType` encoding scheme.
template <typename AttrType>
struct AggregateAttrDef {
  template <typename T>
  using Extract = T (AttrType::*)() const;

  template <typename T, typename Attr = mlir::Attribute>
  using Encode = Attr (mlir::Builder::*)(T);

  template <typename T, typename Attr = mlir::Attribute>
  AggregateAttrDef &Add(std::string name, Extract<T> extract,
                        Encode<T, Attr> encode) {
    bindings.emplace_back([=](AttrType attr, mlir::Builder &b) {
      auto encoded = std::invoke(encode, b, std::invoke(extract, attr));
      return mlir::NamedAttribute(b.getStringAttr(name), encoded);
    });
    return *this;
  }

  AggregateAttrDef &Add(std::string name, Extract<bool> extract) {
    return Add(name, extract, &mlir::Builder::getBoolAttr);
  }

  AggregateAttrDef &Add(std::string name, Extract<int64_t> extract) {
    return Add(name, extract, &mlir::Builder::getI64IntegerAttr);
  }

  AggregateAttrDef &Add(std::string name,
                        Extract<llvm::ArrayRef<int64_t>> extract) {
    return Add(name, extract, &mlir::Builder::getI64TensorAttr);
  }

  // A list of functions to destruct `AttrType` attribute into the aggregate
  // attributes that will be used for encoding.
  using Bind = std::function<mlir::NamedAttribute(AttrType, mlir::Builder &)>;
  llvm::SmallVector<Bind> bindings;
};

// Custom call attribute encoding for the user-defined attributes which encodes
// them as an aggregate of primitive attributes. It uses the encoding scheme
// compatible with the custom call attributes decoding.
template <typename AttrType, typename RuntimeType = AttrType>
struct AggregateAttrEncoding : public CustomCallAttrEncoding {
  using AttrDef = AggregateAttrDef<AttrType>;

  AggregateAttrEncoding(const CustomCallAttrEncodingSet &encoding,
                        AttrDef attrdef)
      : encoding(encoding), attrdef(std::move(attrdef)) {}

  mlir::LogicalResult Match(llvm::StringRef, mlir::Attribute attr) const final {
    return mlir::success(attr.isa<AttrType>());
  }

  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::StringRef name,
                                  mlir::Attribute attr) const final {
    // Extract aggregate attributes from the user-defined attributes.
    llvm::SmallVector<mlir::NamedAttribute> attrs;
    for (auto &bind : attrdef.bindings)
      attrs.emplace_back(bind(attr.cast<AttrType>(), b));

    // Encode extracted attributes as an aggregate.
    auto type_id = mlir::TypeID::get<Tagged<RuntimeType>>();
    auto sym = "__rt_aggregate_" + AttrType::getMnemonic();
    auto aggregate = EncodeAttributes(g, b, encoding, sym.str(), attrs);
    if (mlir::failed(aggregate)) return mlir::failure();

    Encoded encoded;
    encoded.name = PackString(g, b, name, kAttrName);
    encoded.type_id = PackTypeId(g, b, type_id);
    encoded.value = *aggregate;
    return encoded;
  }

  const CustomCallAttrEncodingSet &encoding;
  AttrDef attrdef;
};

//===----------------------------------------------------------------------===//
// Custom call arguments encoding.
//===----------------------------------------------------------------------===//

// Encodes scalar operands.
class ScalarArgEncoding : public CustomCallArgEncoding {
 public:
  mlir::LogicalResult Match(mlir::Value, mlir::Value) const final;
  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::Value, mlir::Value) const final;
};

// Encodes MemRef operands according to the (Strided)MemrefView ABI.
class MemrefArgEncoding : public CustomCallArgEncoding {
 public:
  mlir::LogicalResult Match(mlir::Value, mlir::Value) const final;
  mlir::FailureOr<Encoded> Encode(Globals &g, mlir::ImplicitLocOpBuilder &b,
                                  mlir::Value, mlir::Value) const final;

 private:
  // Encodes memref as LLVM struct value:
  //
  //   { i8: dtype, i8: rank, ptr<i8>: data,
  //     array<2*rank x i64>: sizes_and_strides }
  //
  // This is a type erased version of the MLIR memref descriptor without base
  // pointer. We pack sizes and strides as a single array member, so that on
  // the runtime side we can read it back using C flexible array member.
  mlir::Value EncodeMemRef(mlir::ImplicitLocOpBuilder &b,
                           mlir::MemRefType memref_ty,
                           mlir::Value descriptor) const;
};

//===----------------------------------------------------------------------===//
// Default encodings for arguments and attributes.
//===----------------------------------------------------------------------===//

CustomCallArgEncodingSet DefaultArgEncodings();
CustomCallAttrEncodingSet DefaultAttrEncodings();

}  // namespace runtime
}  // namespace xla

#endif  // XLA_MLIR_RUNTIME_CUSTOM_CALL_ENCODING_H_