Packages

exla

0.7.2
0.13.0 0.12.0 0.11.0 0.10.0 0.9.2 0.9.1 0.9.0 0.8.0 0.7.3 0.7.2 0.7.1 0.7.0 0.6.4 0.6.3 0.6.1 0.6.0 0.5.3 0.5.2 0.5.1 0.5.0 0.4.2 0.4.1 0.4.0 0.3.0 0.2.3 0.2.2 0.2.1 0.2.0

Google's XLA (Accelerated Linear Algebra) compiler/backend for Nx

HexDocs

Current section

Files

Jump to
Readme
28 Versions
5 Dependencies
Dependants
Files
Activity
exla c_src exla mlir builder.h
Raw

c_src/exla/mlir/builder.h 

#ifndef EXLA_MLIR_BUILDER_H_

#define EXLA_MLIR_BUILDER_H_



#include <stack>



#include "../exla_nif_util.h"

#include "mlir/Dialect/Func/IR/FuncOps.h"

#include "mlir/IR/Builders.h"

#include "mlir/IR/BuiltinOps.h"

#include "mlir/IR/BuiltinTypes.h"

#include "mlir/IR/MLIRContext.h"

#include "mlir/IR/OwningOpRef.h"

#include "stablehlo/dialect/StablehloOps.h"

#include "stablehlo/reference/Types.h"

#include "xla/shape.h"

#include "xla/types.h"



namespace exla {



class MLIRModule;



class MLIRFunction {

 public:

  MLIRFunction(MLIRModule *module, std::unique_ptr<mlir::func::FuncOp> func);



  mlir::Value AddOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value SubtractOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value TupleOp(std::vector<mlir::Value> vals);

  mlir::Value GetTupleElementOp(mlir::Value tuple, tsl::int64 index);

  mlir::Value MulOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value MinOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value MaxOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value RemOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value PowOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value DivOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value Atan2Op(mlir::Value lhs, mlir::Value rhs);

  mlir::Value EqualOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value NotEqualOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value LessOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value LessEqualOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value GreaterOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value GreaterEqualOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value BitwiseAndOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value BitwiseOrOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value BitwiseXorOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value BitwiseNotOp(mlir::Value operand);

  mlir::Value ShiftLeftOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value ShiftRightLogicalOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value ShiftRightArithmeticOp(mlir::Value lhs, mlir::Value rhs);

  mlir::Value ConvertOp(mlir::Value operand, mlir::Type type);

  mlir::Value BitcastConvertOp(mlir::Value operand, xla::Shape shape);

  mlir::Value PadOp(mlir::Value op, mlir::Value pad, std::vector<int64_t> padding_low, std::vector<int64_t> padding_high, std::vector<int64_t> padding_mid);

  mlir::Value AbsOp(mlir::Value operand);

  mlir::Value RealOp(mlir::Value operand);

  mlir::Value ImagOp(mlir::Value operand);

  mlir::Value ConjOp(mlir::Value operand);

  mlir::Value ExpOp(mlir::Value operand);

  mlir::Value Expm1Op(mlir::Value operand);

  mlir::Value FloorOp(mlir::Value operand);

  mlir::Value CeilOp(mlir::Value operand);

  mlir::Value RoundOp(mlir::Value operand);

  mlir::Value LogOp(mlir::Value operand);

  mlir::Value LogisticOp(mlir::Value operand);

  mlir::Value Log1pOp(mlir::Value operand);

  mlir::Value SignOp(mlir::Value operand);

  mlir::Value CosOp(mlir::Value operand);

  mlir::Value SinOp(mlir::Value operand);

  mlir::Value TanOp(mlir::Value operand);

  mlir::Value AcosOp(mlir::Value operand);

  mlir::Value AsinOp(mlir::Value operand);

  mlir::Value AtanOp(mlir::Value operand);

  mlir::Value CoshOp(mlir::Value operand);

  mlir::Value SinhOp(mlir::Value operand);

  mlir::Value TanhOp(mlir::Value operand);

  mlir::Value AcoshOp(mlir::Value operand);

  mlir::Value AsinhOp(mlir::Value operand);

  mlir::Value AtanhOp(mlir::Value operand);

  mlir::Value SqrtOp(mlir::Value operand);

  mlir::Value CbrtOp(mlir::Value operand);

  mlir::Value NegateOp(mlir::Value operand);

  mlir::Value ErfOp(mlir::Value operand);

  mlir::Value ErfInvOp(mlir::Value operand);

  mlir::Value ErfcOp(mlir::Value operand);

  mlir::Value IsFiniteOp(mlir::Value operand);

  mlir::Value IsInfOp(mlir::Value operand);

  mlir::Value IsNanOp(mlir::Value operand);

  mlir::Value RsqrtOp(mlir::Value operand);

  mlir::Value ClzOp(mlir::Value operand);

  mlir::Value PopulationCountOp(mlir::Value operand);

  mlir::Value IotaOp(xla::Shape shape, int64_t dimension);

  mlir::Value TransposeOp(mlir::Value operand, std::vector<int64_t> axes);

  mlir::Value ReshapeOp(mlir::Value operand, std::vector<int64_t> target_shape);

  mlir::Value ReverseOp(mlir::Value operand, std::vector<int64_t> dims);

  mlir::Value SliceOp(mlir::Value operand, std::vector<int64_t> starts, std::vector<int64_t> limites, std::vector<int64_t> strides);

  std::vector<mlir::Value> TopKOp(mlir::Value operand, int64_t k);

  std::vector<mlir::Value> SortOp(MLIRFunction *comparator, std::vector<mlir::Value> operand, int64_t dim, bool status);

  mlir::Value DynamicSliceOp(mlir::Value operand, std::vector<mlir::Value> starts, std::vector<int64_t> lengths);

  mlir::Value BroadcastInDimOp(mlir::Value operand, xla::Shape result_shape, std::vector<int64_t> axes);

  mlir::Value DotGeneralOp(xla::Shape output_shape, mlir::Value lhs, mlir::Value rhs, xla::DotDimensionNumbers dnums, xla::PrecisionConfig config);

  mlir::Value ConcatenateOp(std::vector<mlir::Value> operands, int64_t dimension);

  mlir::Value OptimizationBarrierOp(mlir::Value operand);

  mlir::Value ClampOp(mlir::Value min, mlir::Value operand, mlir::Value max);

  mlir::Value SelectOp(mlir::Value pred, mlir::Value on_true, mlir::Value on_false);

  mlir::Value ScatterOp(mlir::Value target, mlir::Value indices, mlir::Value updates, bool add_or_put, int64_t indices_rank, std::vector<int64_t> update_window_dims, std::vector<int64_t> inserted_window_dims, std::vector<int64_t> index_dims_to_window_dims);

  mlir::Value SelectAndScatterOp(mlir::Value target, mlir::Value source, mlir::Value init_value, bool gt_or_lt, std::vector<int64_t> window_dimensions, std::vector<int64_t> window_strides, std::vector<int64_t> padding);

  mlir::Value GatherOp(mlir::Value source, mlir::Value indices, std::vector<int64_t> offset_dims, std::vector<int64_t> collapsed_slice_dims, std::vector<int64_t> start_index_map, std::vector<int64_t> slice_sizes, int64_t index_vector_dim);

  mlir::Value FFTOp(mlir::Value tensor, bool forward_fft, std::vector<int64_t> fft_length);

  mlir::Value ConvOp(mlir::Value tensor, mlir::Value kernel, std::vector<int64_t> window_strides, std::vector<int64_t> padding, std::vector<int64_t> tensor_dilation, std::vector<int64_t> kernel_dilation, xla::ConvolutionDimensionNumbers dimension_numbers, uint64_t feature_group_count, uint64_t batch_group_count, uint64_t precision_config, std::vector<int64_t> output_dims);

  mlir::Value CreateTokenOp();

  mlir::Value TriangularSolveOp(mlir::Value a, mlir::Value b, bool left_side, bool lower, bool transpose_a);

  mlir::Value DynamicUpdateSliceOp(mlir::Value operand, mlir::Value update, std::vector<mlir::Value> start_indices);

  std::vector<mlir::Value> ReduceOp(MLIRFunction *function, std::vector<mlir::Value> init_values, std::vector<mlir::Value> inputs, std::vector<int64_t> dimensions);

  std::vector<mlir::Value> WindowReduceOp(MLIRFunction *function, std::vector<mlir::Value> init_values, std::vector<mlir::Value> inputs, std::vector<int64_t> window_dimensions, std::vector<int64_t> window_strides, std::vector<int64_t> input_dilations, std::vector<int64_t> window_dilations, std::vector<std::pair<int64_t, int64_t>> padding);

  mlir::Value MapOp(MLIRFunction *function, std::vector<mlir::Value> inputs, std::vector<int64_t> dimensions);

  std::pair<std::vector<mlir::Value>, std::pair<mlir::Region *, mlir::Region *>> IfOp(mlir::Value pred, std::vector<xla::Shape> output_shape);

  ERL_NIF_TERM ConstantOp(mlir::Type type, ErlNifEnv *env, ERL_NIF_TERM value_ptr, std::optional<std::vector<int64_t>> dims = std::nullopt);

  std::pair<mlir::Value, std::vector<mlir::Value>> InfeedOp(mlir::Value token, std::vector<xla::Shape> shapes);

  mlir::Value OutfeedOp(std::vector<mlir::Value> inputs, mlir::Value token);

  std::vector<mlir::Value> CallOp(std::vector<mlir::Value> inputs, MLIRFunction *computation);

  std::pair<std::vector<mlir::Value>, std::pair<mlir::Region *, mlir::Region *>> WhileOp(std::vector<mlir::Value> initial);

  std::vector<mlir::Value> ReturnOp(std::vector<mlir::Value> values);

  int get_mlir_type(ErlNifEnv *env, ERL_NIF_TERM term, mlir::Type *type);

  std::vector<mlir::Value> PushRegion(mlir::Region *region);

  std::pair<mlir::Value, mlir::Value> QRCpuCustomCall(mlir::Value operand, std::vector<int64_t> q_shape, std::vector<int64_t> r_shape);

  void PopRegion();

  void Build(mlir::Value root);



  llvm::MutableArrayRef<mlir::BlockArgument> get_arguments() { return func_->getBody().front().getArguments(); }



  mlir::func::FuncOp *function() { return func_.get(); }



 private:

  std::shared_ptr<MLIRModule> module_;

  std::unique_ptr<mlir::func::FuncOp> func_;



  std::stack<mlir::Region *> regions;



  void dump_mlir_module();

  void setInsertionPoint();

};



class MLIRModule {

 public:

  MLIRModule(mlir::MLIRContext *context);



  MLIRFunction *CreateFunction(

      std::string name,

      std::vector<xla::Shape *> arg_shapes,

      std::vector<xla::Shape *> ret_shape,

      bool is_public);



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

  mlir::OpBuilder *builder() { return builder_.get(); }

  mlir::MLIRContext *context() { return context_; }



  void LowerPatterns();



 private:

  mlir::MLIRContext *context_;

  mlir::OwningOpRef<mlir::ModuleOp> module_;

  std::unique_ptr<mlir::OpBuilder> builder_;

};



mlir::Type

TypeIntToMLIRType(mlir::OpBuilder *builder, xla::PrimitiveType type_int);



xla::PrimitiveType MLIRTypeToPrimitiveType(mlir::Type);

}  // namespace exla



#endif