IrArray represents an XLA array at the LLVM IR level. This class encapsulates a base pointer to the buffer holding the array (as an LLVM Value) and the shape of the array. The class includes methods for emitting LLVM IR sequences which access elements of the array at a multidimensional index (eg, [x, y, z] in a 3-dimensional array). Arbitrary shape and layouts are supported.
| 42 | // index (eg, [x, y, z] in a 3-dimensional array). Arbitrary shape and layouts |
| 43 | // are supported. |
| 44 | class IrArray { |
| 45 | public: |
| 46 | // A multidimensional index into an IrArray. All the runtime indices |
| 47 | // (multidim) and dimensions (Shape::dimensions(), absl::Span<const int64>) |
| 48 | // are major-first. |
| 49 | // |
| 50 | // This may also keep a linear index and the layout and dimensions it was |
| 51 | // emitted for; if the shape where this `Index` is used matches, the linear |
| 52 | // index may be used, potentially sparing the cost of computing the |
| 53 | // multidimensional index, which LLVM DCE can delete. |
| 54 | class Index { |
| 55 | public: |
| 56 | // Constructs an index for a scalar shape. |
| 57 | explicit Index(llvm::Type* index_ty) : index_type_(index_ty) { |
| 58 | CHECK(index_ty->isIntegerTy()); |
| 59 | } |
| 60 | |
| 61 | // Constructs an index from linear index "linear" and computes the |
| 62 | // multi-dimensional index from "linear" and "shape". "b" is the IR |
| 63 | // builder to emit the index of each dimension in the multi-dimensional |
| 64 | // index. |
| 65 | // |
| 66 | // Precondition: "shape" has a layout. |
| 67 | Index(llvm::Value* linear, const Shape& shape, llvm::IRBuilder<>* b); |
| 68 | |
| 69 | // Constructs an index from a multi-dimensional index. 'shape' is the shape |
| 70 | // for which the multi-dimensional index is used. 'index_type' is the type |
| 71 | // of the index. |
| 72 | // |
| 73 | // Precondition: "shape" has a layout. |
| 74 | Index(absl::Span<llvm::Value* const> multidim, const Shape& shape, |
| 75 | llvm::Type* index_type); |
| 76 | |
| 77 | // Same as above, but only the dimensions of the shape without layout is |
| 78 | // passed. The layout is assumed to be the default (descending |
| 79 | // minor-to-major) layout. |
| 80 | Index(absl::Span<llvm::Value* const> multidim, |
| 81 | absl::Span<int64 const> dimensions, llvm::Type* index_type); |
| 82 | |
| 83 | // Returns an index that adds `addend` to the given `dim` of the object. |
| 84 | Index AddOffsetToDim(llvm::Value* addend, int64 dim, |
| 85 | llvm::IRBuilder<>* b) const { |
| 86 | Index with_offset = *this; |
| 87 | with_offset.linear_ = nullptr; |
| 88 | with_offset.multidim_[dim] = |
| 89 | b->CreateAdd(with_offset.multidim_[dim], addend); |
| 90 | return with_offset; |
| 91 | } |
| 92 | |
| 93 | const std::vector<llvm::Value*>& multidim() const { return multidim_; } |
| 94 | const std::vector<int64>& dims() const { return dims_; } |
| 95 | llvm::Value* linear() const { return linear_; } |
| 96 | |
| 97 | size_t size() const { return multidim().size(); } |
| 98 | |
| 99 | llvm::Value* operator[](size_t i) const { return multidim()[i]; } |
| 100 | |
| 101 | using const_iterator = std::vector<llvm::Value*>::const_iterator; |
no outgoing calls
no test coverage detected