677 lines
24 KiB
C++
677 lines
24 KiB
C++
#include "triton/ir/builder.h"
|
|
#include <functional>
|
|
#include <iostream>
|
|
#include <pybind11/pybind11.h>
|
|
|
|
namespace ir = triton::ir;
|
|
namespace py = pybind11;
|
|
|
|
static const std::string _builder_doc = R"pbdoc(
|
|
:param builder: IR builder to generate code into, optional, set automatically when called inside a @triton.jit function
|
|
:type builder: triton.ir.builder
|
|
)pbdoc";
|
|
|
|
#define VA_ARGS(...) , ##__VA_ARGS__
|
|
#define DEF_FUNC(MOD, PY_NAME, C_FUNC, ...) \
|
|
MOD.def(PY_NAME, C_FUNC, (C_FUNC##_docstr + _builder_doc).c_str(), \
|
|
ret::reference VA_ARGS(__VA_ARGS__), "builder"_a)
|
|
|
|
void throw_not_implemented(std::string key) {
|
|
throw std::runtime_error("Encountered unimplemented code path in `" + key + "`. This is likely a bug on our side.");
|
|
}
|
|
|
|
void throw_not_int_or_float(std::string key) {
|
|
throw std::runtime_error("`" + key + "` only supported for integer and floating point types.");
|
|
}
|
|
|
|
enum type_code {
|
|
_bool,
|
|
int8,
|
|
int16,
|
|
int32,
|
|
int64,
|
|
float16,
|
|
float32,
|
|
float64
|
|
};
|
|
|
|
ir::type *make_ir(type_code ty, ir::builder *builder) {
|
|
switch (ty) {
|
|
case float16:
|
|
return builder->get_half_ty();
|
|
case float32:
|
|
return builder->get_float_ty();
|
|
default:
|
|
throw_not_implemented("make_ir");
|
|
}
|
|
}
|
|
|
|
type_code from_ir(ir::type *ty) {
|
|
if (ty->is_half_ty())
|
|
return float16;
|
|
if (ty->is_float_ty())
|
|
return float32;
|
|
throw_not_implemented("from_ir");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.cast / triton.ir.value.to
|
|
----------------------------------------------*/
|
|
std::string cast_docstr = R"pbdoc(
|
|
Tries to cast a block to a new data type.
|
|
|
|
:param input: The input block.
|
|
:type input: triton.ir.value
|
|
)pbdoc";
|
|
|
|
ir::value *cast(ir::value *input, type_code _dtype, ir::builder *builder) {
|
|
ir::type *src_ty = input->get_type();
|
|
ir::type *dst_ty = make_ir(_dtype, builder);
|
|
if (src_ty->is_block_ty())
|
|
dst_ty = ir::block_type::get(dst_ty, input->get_type()->get_block_shapes());
|
|
ir::type *src_sca_ty = src_ty->get_scalar_ty();
|
|
ir::type *dst_sca_ty = dst_ty->get_scalar_ty();
|
|
// FP Truncation
|
|
bool truncate_fp = src_sca_ty->is_floating_point_ty() &&
|
|
dst_sca_ty->is_floating_point_ty() &&
|
|
src_sca_ty->get_fp_mantissa_width() > dst_sca_ty->get_fp_mantissa_width();
|
|
if (truncate_fp)
|
|
return builder->create_fp_trunc(input, dst_ty);
|
|
// FP Extension
|
|
bool ext_fp = src_sca_ty->is_floating_point_ty() &&
|
|
dst_sca_ty->is_floating_point_ty() &&
|
|
src_sca_ty->get_fp_mantissa_width() < dst_sca_ty->get_fp_mantissa_width();
|
|
if (ext_fp)
|
|
return builder->create_fp_ext(input, dst_ty);
|
|
// Int cast
|
|
if (src_sca_ty->is_integer_ty() && dst_sca_ty->is_integer_ty() &&
|
|
src_sca_ty->get_integer_bitwidth() != dst_sca_ty->get_integer_bitwidth())
|
|
return builder->create_int_cast(input, dst_ty, true);
|
|
// Float -> Int
|
|
if (src_sca_ty->is_floating_point_ty() && dst_sca_ty->is_integer_ty())
|
|
return builder->create_fp_to_si(input, dst_ty);
|
|
// int -> Float
|
|
if (src_sca_ty->is_integer_ty() && dst_sca_ty->is_floating_point_ty())
|
|
return builder->create_si_to_fp(input, dst_ty);
|
|
// Ptr -> Ptr
|
|
if (src_sca_ty->is_pointer_ty() && dst_sca_ty->is_pointer_ty())
|
|
return builder->create_cast(ir::BitCast, input, dst_ty);
|
|
// * -> Bool
|
|
if (dst_sca_ty->is_bool_ty()) {
|
|
if (src_sca_ty->is_pointer_ty())
|
|
input = cast(input, int64, builder);
|
|
ir::value *other = builder->get_int64(0);
|
|
if (src_ty->is_bool_ty())
|
|
other = builder->create_splat(other, src_ty->get_block_shapes());
|
|
return builder->create_icmpNE(input, other);
|
|
}
|
|
throw_not_implemented("cast");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.broadcast_check
|
|
----------------------------------------------*/
|
|
std::string try_broadcast_docstr = R"pbdoc(
|
|
Tries to broadcast two blocks to a common compatible shape.
|
|
|
|
:param input: The first input block.
|
|
:type input: triton.ir.value
|
|
:param other: The second input block.
|
|
:type other: triton.ir.value
|
|
)pbdoc";
|
|
|
|
std::tuple<ir::value *, ir::value *> try_broadcast(ir::value *lhs, ir::value *rhs, ir::builder *builder) {
|
|
ir::type *lhs_ty = lhs->get_type();
|
|
ir::type *rhs_ty = rhs->get_type();
|
|
// make_shape_compatible(block, scalar)
|
|
if (lhs_ty->is_block_ty() && !rhs_ty->is_block_ty())
|
|
rhs = builder->create_splat(rhs, lhs_ty->get_block_shapes());
|
|
// make_shape_compatible(scalar, block)
|
|
else if (!lhs_ty->is_block_ty() && rhs_ty->is_block_ty())
|
|
lhs = builder->create_splat(lhs, rhs_ty->get_block_shapes());
|
|
// make_shape_compatible(block, block)
|
|
else if (lhs_ty->is_block_ty() && rhs_ty->is_block_ty()) {
|
|
auto lhs_shape = lhs_ty->get_block_shapes();
|
|
auto rhs_shape = rhs_ty->get_block_shapes();
|
|
if (lhs_shape.size() != rhs_shape.size())
|
|
throw std::runtime_error("Cannot make_shape_compatible: blocks must have the same rank");
|
|
ir::type::block_shapes_t ret_shape;
|
|
for (size_t i = 0; i < lhs_shape.size(); ++i) {
|
|
unsigned left = lhs_shape[i];
|
|
unsigned right = rhs_shape[i];
|
|
if (left == 1)
|
|
ret_shape.push_back(right);
|
|
else if (right == 1)
|
|
ret_shape.push_back(left);
|
|
else if (left == right)
|
|
ret_shape.push_back(left);
|
|
else
|
|
throw std::runtime_error("Cannot make_shape_compatible: incompatible dimensions at index " + std::to_string(i) +
|
|
": " + std::to_string(left) + " and " + std::to_string(right));
|
|
}
|
|
if (lhs_shape != ret_shape)
|
|
lhs = builder->create_broadcast(lhs, ret_shape);
|
|
if (rhs_shape != ret_shape)
|
|
rhs = builder->create_broadcast(rhs, ret_shape);
|
|
}
|
|
return std::make_tuple(lhs, rhs);
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.broadcast_to
|
|
----------------------------------------------*/
|
|
std::string broadcast_to_docstr = R"pbdoc(
|
|
Tries to broadcast a block to a new shape.
|
|
|
|
:param input: The input block.
|
|
:type input: triton.value
|
|
:param shape: The new shape.
|
|
:type shape: tuple of int
|
|
)pbdoc";
|
|
|
|
ir::value *broadcast_to(ir::value *input, const ir::type::block_shapes_t &shape, ir::builder *builder) {
|
|
if (!input->get_type()->is_block_ty())
|
|
return builder->create_splat(input, shape);
|
|
auto src_shape = input->get_type()->get_block_shapes();
|
|
if (src_shape.size() != shape.size())
|
|
throw std::runtime_error("Cannot broadcast");
|
|
return builder->create_broadcast(input, shape);
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.load
|
|
----------------------------------------------*/
|
|
std::string load_docstr = R"pbdoc(
|
|
Return a block of data whose values are, elementwise, loaded from memory at location defined by `pointer`.
|
|
|
|
:param pointer: Pointer to the data to be loaded.
|
|
:type pointer: Block of triton.pointer
|
|
:param mask: if mask[idx] is false, do not load the data at `pointer[idx]`.
|
|
:type mask: Block of triton.bool, optional
|
|
:param other: if mask[idx] is false, return other[idx] instead of 'pointer[idx]`
|
|
:type other: Block of triton.value, optional
|
|
)pbdoc";
|
|
|
|
ir::value *load(ir::value *pointer, std::optional<ir::value *> _mask, std::optional<ir::value *> _other, ir::builder *builder) {
|
|
if (!_mask.has_value() && !_other.has_value())
|
|
return builder->create_load(pointer);
|
|
if (!_mask.has_value())
|
|
throw std::runtime_error("`other` cannot be provided without `mask`");
|
|
ir::value *mask = _mask.value();
|
|
ir::type *elt_ty = pointer->get_type()->get_scalar_ty()->get_pointer_element_ty();
|
|
auto shape = pointer->get_type()->get_block_shapes();
|
|
ir::value *other = _other.has_value() ? _other.value() : ir::undef_value::get(elt_ty);
|
|
other = cast(other, from_ir(elt_ty), builder);
|
|
other = broadcast_to(other, shape, builder);
|
|
mask = broadcast_to(mask, shape, builder);
|
|
return builder->create_masked_load(pointer, mask, other);
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.store
|
|
----------------------------------------------*/
|
|
std::string store_docstr = R"pbdoc(
|
|
Stores `value` block of elements in memory, element-wise, at the memory locations specified by `pointer`.
|
|
|
|
:param pointer: The memory locations where the elements of `value` are stored.
|
|
:type pointer: Block of triton.pointer
|
|
:param value: The block of elements to be stored.
|
|
:type value: Block of triton.value
|
|
:param mask: If mask[idx] is false, do not store `value[idx]` at `pointer[idx]`.
|
|
:type mask: Block of triton.bool, optional
|
|
)pbdoc";
|
|
ir::value *store(ir::value *ptr, ir::value *val, std::optional<ir::value *> _mask, ir::builder *builder) {
|
|
if (!_mask.has_value())
|
|
return builder->create_store(ptr, val);
|
|
ir::value *mask = _mask.value();
|
|
return builder->create_masked_store(ptr, val, mask);
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.dot
|
|
----------------------------------------------*/
|
|
std::string dot_docstr = R"pbdoc(
|
|
Returns the matrix product of two blocks.
|
|
The two blocks must be two dimensionals and have compatible inner dimensions.
|
|
|
|
:param input: The first block to be multiplied.
|
|
:type input: 2D block of scalar-type in {`float16`, `float32`}
|
|
:param other: The second block to be multiplied.
|
|
:type other: 2D block of scalar-type in {`float16`, `float32`}
|
|
)pbdoc";
|
|
ir::value *dot(ir::value *lhs, ir::value *rhs, ir::builder *builder) {
|
|
ir::value *_0 = builder->get_float32(0);
|
|
unsigned M = lhs->get_type()->get_block_shapes()[0];
|
|
unsigned N = rhs->get_type()->get_block_shapes()[1];
|
|
_0 = builder->create_splat(_0, {M, N});
|
|
return builder->create_dot(lhs, rhs, _0);
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.where
|
|
----------------------------------------------*/
|
|
std::string where_docstr = R"pbdoc(
|
|
Returns a block of elements from either `x` or `y`, depending on `condition`.
|
|
Note that `x` and `y` are always evaluated regardless of the value of `condition`.
|
|
If you want to avoid unintented memory operations, use the `mask` arguments in `triton.load` and `triton.store` instead.
|
|
|
|
:param condition: When True (nonzero), yield x, otherwise yield y.
|
|
:type condition: Block of triton.bool
|
|
:param x: values selected at indices where condition is True.
|
|
:param y: values selected at indices where condition is False.
|
|
)pbdoc";
|
|
ir::value *where(ir::value *condition, ir::value *x, ir::value *y, ir::builder *builder) {
|
|
return builder->create_select(condition, x, y);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.arange
|
|
----------------------------------------------*/
|
|
std::string arange_docstr = R"pbdoc(
|
|
Returns contiguous values within the open interval [start, end).
|
|
|
|
:param start: Start of the interval.
|
|
:type start: int
|
|
:param stop: End of the interval.
|
|
:type stop: int
|
|
)pbdoc";
|
|
ir::value *arange(int start, int end, ir::builder *builder) {
|
|
return builder->get_range(start, end);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.program_id
|
|
----------------------------------------------*/
|
|
std::string program_id_docstr = R"pbdoc(
|
|
Returns the id of the current program instance along the given `axis`.
|
|
Triton uses an SPMD model in which different @triton.jit functions run in parallel with different `program_id`s.
|
|
|
|
:param axis: The axis of the 3D launch grid. Has to be either 0, 1 or 2.
|
|
:type axis: int
|
|
)pbdoc";
|
|
ir::value *program_id(int axis, ir::builder *builder) {
|
|
return builder->create_get_program_id(axis);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.num_programs
|
|
----------------------------------------------*/
|
|
std::string num_programs_docstr = R"pbdoc(
|
|
Returns the number of program instances launched along the given `axis`.
|
|
|
|
:param axis: The axis of the 3D launch grid. Has to be either 0, 1 or 2.
|
|
:type axis: int
|
|
)pbdoc";
|
|
ir::value *num_programs(int axis, ir::builder *builder) {
|
|
return builder->create_get_num_programs(axis);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.zeros
|
|
----------------------------------------------*/
|
|
std::string zeros_docstr = R"pbdoc(
|
|
Returns a block filled with the scalar value 0 and the given shape.
|
|
|
|
:param shape: Shape of the new array, e.g., (8, 16) or (8, )
|
|
:type shape: tuple of ints
|
|
:param dtype: Data-type of the new array, e.g., tl.float16
|
|
:type dtype: triton.ir.dtype
|
|
)pbdoc";
|
|
ir::value *zeros(ir::type::block_shapes_t shape, type_code _dtype, ir::builder *builder) {
|
|
ir::type *dtype = make_ir(_dtype, builder);
|
|
ir::value *_0 = ir::constant::get_null_value(dtype);
|
|
return builder->create_splat(_0, shape);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.exp
|
|
----------------------------------------------*/
|
|
std::string _exp_docstr = R"pbdoc(
|
|
Returns the element-wise exponential of `input`.
|
|
)pbdoc";
|
|
ir::value *_exp(ir::value *input, ir::builder *builder) {
|
|
return builder->create_exp(input);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.log
|
|
----------------------------------------------*/
|
|
std::string _log_docstr = R"pbdoc(
|
|
Returns the element-wise natural logarithm of `input`.
|
|
)pbdoc";
|
|
ir::value *_log(ir::value *input, ir::builder *builder) {
|
|
return builder->create_log(input);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.sqrt
|
|
----------------------------------------------*/
|
|
std::string sqrt_docstr = R"pbdoc(
|
|
Returns the element-wise square root of `input`.
|
|
)pbdoc";
|
|
ir::value *sqrt(ir::value *input, ir::builder *builder) {
|
|
return builder->create_sqrt(input);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.min
|
|
----------------------------------------------*/
|
|
ir::value *reduce_impl(ir::value *input, unsigned int axis, ir::builder *builder, const std::string &name,
|
|
ir::reduce_inst::op_t FLOAT_OP, ir::reduce_inst::op_t INT_OP) {
|
|
ir::type *scalar_ty = input->get_type()->get_scalar_ty();
|
|
if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_reduce(input, FLOAT_OP, axis);
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_reduce(input, INT_OP, axis);
|
|
else
|
|
throw_not_int_or_float(name);
|
|
}
|
|
|
|
std::string min_docstr = R"pbdoc(
|
|
Returns the minimum value of `input`.
|
|
)pbdoc";
|
|
ir::value *min(ir::value *input, unsigned int axis, ir::builder *builder) {
|
|
return reduce_impl(input, axis, builder, "min", ir::reduce_inst::FMIN, ir::reduce_inst::MIN);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.max
|
|
----------------------------------------------*/
|
|
std::string max_docstr = R"pbdoc(
|
|
Returns the maximum value of `input`.
|
|
)pbdoc";
|
|
ir::value *max(ir::value *input, unsigned int axis, ir::builder *builder) {
|
|
return reduce_impl(input, axis, builder, "max", ir::reduce_inst::FMAX, ir::reduce_inst::MAX);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.sum
|
|
----------------------------------------------*/
|
|
std::string sum_docstr = R"pbdoc(
|
|
Returns the sum of `input`.
|
|
)pbdoc";
|
|
ir::value *sum(ir::value *input, unsigned int axis, ir::builder *builder) {
|
|
return reduce_impl(input, axis, builder, "sum", ir::reduce_inst::FADD, ir::reduce_inst::ADD);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.atomic_cas
|
|
----------------------------------------------*/
|
|
std::string atomic_cas_docstr = R"pbdoc(
|
|
Atomic compare-and-swap.
|
|
)pbdoc";
|
|
ir::value *atomic_cas(ir::value *ptr, ir::value *cmp, ir::value *val, ir::builder *builder) {
|
|
return builder->create_atomic_cas(ptr, cmp, val);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
definition of triton.atomic_xchg
|
|
----------------------------------------------*/
|
|
std::string atomic_xchg_docstr = R"pbdoc(
|
|
Atomic exchange.
|
|
)pbdoc";
|
|
ir::value *atomic_xchg(ir::value *ptr, ir::value *val, ir::builder *builder) {
|
|
return builder->create_atomic_exch(ptr, val);
|
|
};
|
|
|
|
/*----------------------------------------------
|
|
debug barrier
|
|
----------------------------------------------*/
|
|
std::string debug_barrier_docstr = R"pbdoc(
|
|
Temporary hacky fixup for when the compiler forgets to insert sync barriers
|
|
)pbdoc";
|
|
ir::value *debug_barrier(ir::builder *builder) {
|
|
return builder->create_barrier();
|
|
}
|
|
|
|
#define DEF_BINARY_OP(MOD, PY_NAME, C_FUNC, ...) \
|
|
MOD.def(PY_NAME, binary_op(C_FUNC), (C_FUNC##_docstr + _builder_doc).c_str(), \
|
|
ret::reference VA_ARGS(__VA_ARGS__), "builder"_a)
|
|
|
|
template <class FN>
|
|
std::function<ir::value *(ir::value *, ir::value *, ir::builder *builder)>
|
|
binary_op(const FN &fn) {
|
|
auto ret = [&fn](ir::value *self, ir::value *other, ir::builder *builder) {
|
|
//std::tie(self, other) = try_broadcast(self, other, builder);
|
|
return fn(self, other, builder);
|
|
};
|
|
return ret;
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self + other
|
|
----------------------------------------------*/
|
|
std::string add_docstr = R"pbdoc(
|
|
Returns self + other, element-wise.
|
|
)pbdoc";
|
|
ir::value *add(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
ir::type *scalar_ty = self->get_type()->get_scalar_ty();
|
|
// ptr + offset
|
|
if (scalar_ty->is_pointer_ty())
|
|
return builder->create_gep(self, {other});
|
|
// float + float
|
|
else if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_fadd(self, other);
|
|
// int + int
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_add(self, other);
|
|
throw_not_implemented("add");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self - other
|
|
----------------------------------------------*/
|
|
std::string sub_docstr = R"pbdoc(
|
|
Returns self - other, element-wise.
|
|
)pbdoc";
|
|
ir::value *sub(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
ir::type *scalar_ty = self->get_type()->get_scalar_ty();
|
|
// ptr + offset
|
|
if (scalar_ty->is_pointer_ty())
|
|
return builder->create_gep(self, {other});
|
|
// float + float
|
|
if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_fsub(self, other);
|
|
// int + int
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_sub(self, other);
|
|
throw_not_implemented("sub");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self * other
|
|
----------------------------------------------*/
|
|
std::string mul_docstr = R"pbdoc(
|
|
Returns self * other, element-wise.
|
|
)pbdoc";
|
|
ir::value *mul(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
ir::type *scalar_ty = self->get_type()->get_scalar_ty();
|
|
// float * float
|
|
if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_fmul(self, other);
|
|
// int * int
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_mul(self, other);
|
|
throw_not_implemented("mul");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self > other
|
|
----------------------------------------------*/
|
|
std::string greater_than_docstr = R"pbdoc(
|
|
Returns self > other, element-wise.
|
|
)pbdoc";
|
|
ir::value *greater_than(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
ir::type *scalar_ty = self->get_type()->get_scalar_ty();
|
|
// float > float
|
|
if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_fcmpOGT(self, other);
|
|
// int > int
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_icmpSGT(self, other);
|
|
throw_not_implemented("greater_than");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self >= other
|
|
----------------------------------------------*/
|
|
std::string greater_equal_docstr = R"pbdoc(
|
|
Returns self >= other, element-wise.
|
|
)pbdoc";
|
|
ir::value *greater_equal(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
ir::type *scalar_ty = self->get_type()->get_scalar_ty();
|
|
// float >= float
|
|
if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_fcmpOGE(self, other);
|
|
// int >= int
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_icmpSGE(self, other);
|
|
throw_not_implemented("greater_equal");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self < other
|
|
----------------------------------------------*/
|
|
std::string less_than_docstr = R"pbdoc(
|
|
Returns self < other, element-wise.
|
|
)pbdoc";
|
|
ir::value *less_than(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
ir::type *scalar_ty = self->get_type()->get_scalar_ty();
|
|
// float < float
|
|
if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_fcmpOLT(self, other);
|
|
// int < int
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_icmpSLT(self, other);
|
|
throw_not_implemented("less_than");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self <= other
|
|
----------------------------------------------*/
|
|
std::string less_equal_docstr = R"pbdoc(
|
|
Returns self <= other, element-wise.
|
|
)pbdoc";
|
|
ir::value *less_equal(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
ir::type *scalar_ty = self->get_type()->get_scalar_ty();
|
|
// float < float
|
|
if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_fcmpOLE(self, other);
|
|
// int < int
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_icmpSLE(self, other);
|
|
throw_not_implemented("less_equal");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self == other
|
|
----------------------------------------------*/
|
|
std::string equal_docstr = R"pbdoc(
|
|
Returns self == other, element-wise.
|
|
)pbdoc";
|
|
ir::value *equal(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
ir::type *scalar_ty = self->get_type()->get_scalar_ty();
|
|
// float == float
|
|
if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_fcmpOEQ(self, other);
|
|
// int == int
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_icmpEQ(self, other);
|
|
throw_not_implemented("equal");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self / other
|
|
----------------------------------------------*/
|
|
std::string _div_docstr = R"pbdoc(
|
|
Returns self / other, element-wise.
|
|
)pbdoc";
|
|
ir::value *_div(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
ir::type *scalar_ty = self->get_type()->get_scalar_ty();
|
|
// float / float
|
|
if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_fdiv(self, other);
|
|
// int / int
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_sdiv(self, other);
|
|
throw_not_implemented("div");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self % other
|
|
----------------------------------------------*/
|
|
std::string mod_docstr = R"pbdoc(
|
|
Returns self % other, element-wise.
|
|
)pbdoc";
|
|
ir::value *mod(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
ir::type *scalar_ty = self->get_type()->get_scalar_ty();
|
|
// float % int
|
|
if (scalar_ty->is_floating_point_ty())
|
|
return builder->create_frem(self, other);
|
|
// int % int
|
|
else if (scalar_ty->is_integer_ty())
|
|
return builder->create_srem(self, other);
|
|
throw_not_implemented("mod");
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self & other
|
|
----------------------------------------------*/
|
|
std::string _and_docstr = R"pbdoc(
|
|
Returns self & other, element-wise.
|
|
)pbdoc";
|
|
ir::value *_and(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
return builder->create_and(self, other);
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of minimum(self, other)
|
|
----------------------------------------------*/
|
|
std::string minimum_docstr = R"pbdoc(
|
|
Returns element-wise minimum of self and other
|
|
)pbdoc";
|
|
ir::value *minimum(ir::value *self, ir::value *other, ir::builder *builder) {
|
|
return where(less_than(self, other, builder), self, other, builder);
|
|
}
|
|
|
|
/*----------------------------------------------
|
|
definition of self[slices]
|
|
----------------------------------------------*/
|
|
|
|
enum slice_mode_t {
|
|
NEWAXIS,
|
|
ALL
|
|
};
|
|
|
|
std::string subscript_docstr = R"pbdoc(
|
|
returns self[slices].
|
|
|
|
:param slices: The slices to subscript with.
|
|
:type slices: List of `None` or `:` slices.
|
|
)pbdoc";
|
|
ir::value *subscript(ir::value *self, std::vector<py::object> slices, ir::builder *builder) {
|
|
std::vector<slice_mode_t> modes;
|
|
for (py::object slice : slices) {
|
|
py::object none = py::none();
|
|
py::object all = py::make_tuple(none, none, none);
|
|
if (slice.is(none))
|
|
modes.push_back(NEWAXIS);
|
|
else if (all.attr("__eq__")(slice))
|
|
modes.push_back(ALL);
|
|
else
|
|
throw std::runtime_error("slice must be None or (None, None, None)");
|
|
}
|
|
|
|
ir::type::block_shapes_t shape;
|
|
size_t curr = 0;
|
|
for (slice_mode_t mode : modes) {
|
|
if (mode == NEWAXIS)
|
|
shape.push_back(1);
|
|
else {
|
|
assert(mode == ALL);
|
|
shape.push_back(self->get_type()->get_block_shapes()[curr++]);
|
|
}
|
|
}
|
|
return builder->create_reshape(self, shape);
|
|
}
|