#include <algorithm>

#include "AlignLoads.h"
#include "Bounds.h"
#include "HexagonAlignment.h"
#include "IRMutator.h"
#include "IROperator.h"
#include "ModulusRemainder.h"
#include "Scope.h"
#include "Simplify.h"
using std::vector;

namespace Halide {
namespace Internal {

namespace {

// This mutator attempts to rewrite unaligned or strided loads to
// sequences of aligned loads by loading aligned vectors that cover
// the original unaligned load, and then slicing or shuffling the
// intended vector out of the aligned vector.
class AlignLoads : public IRMutator {
public:
 AlignLoads(int alignment, int min_bytes)
 : alignment_analyzer(alignment), required_alignment(alignment), min_bytes_to_align(min_bytes) {
 }

private:
 HexagonAlignmentAnalyzer alignment_analyzer;

 // Loads and stores should ideally be aligned to the vector width in bytes.
 int required_alignment;

 // Minimum size of load to align.
 int min_bytes_to_align;

 using IRMutator::visit;

 // Rewrite a load to have a new index, updating the type if necessary.
 Expr make_load(const Load *load, const Expr &index, ModulusRemainder alignment) {
 internal_assert(is_const_one(load->predicate)) << "Load should not be predicated.\n";
 return mutate(Load::make(load->type.with_lanes(index.type().lanes()), load->name,
 index, load->image, load->param,
 const_true(index.type().lanes()),
 alignment));
 }

 Expr visit(const Load *op) override {
 if (!is_const_one(op->predicate)) {
 // TODO(psuriana): Do nothing to predicated loads for now.
 return IRMutator::visit(op);
 }

 if (!op->type.is_vector()) {
 // Nothing to do for scalar loads.
 return IRMutator::visit(op);
 }

 if (op->image.defined()) {
 // We can't reason about the alignment of external images.
 return IRMutator::visit(op);
 }

 if (required_alignment % op->type.bytes() != 0) {
 return IRMutator::visit(op);
 }

 if (op->type.bytes() * op->type.lanes() <= min_bytes_to_align) {
 // These can probably be treated as scalars instead.
 return IRMutator::visit(op);
 }

 Expr index = mutate(op->index);
 const Ramp *ramp = index.as<Ramp>();
 const int64_t *const_stride = ramp ? as_const_int(ramp->stride) : nullptr;
 if (!ramp || !const_stride) {
 // We can't handle indirect loads, or loads with
 // non-constant strides.
 return IRMutator::visit(op);
 }
 if (!(*const_stride == 1 || *const_stride == 2 || *const_stride == 3 || *const_stride == 4)) {
 // Handle ramps with stride 1, 2, 3 or 4 only.
 return IRMutator::visit(op);
 }

 int64_t aligned_offset = 0;
 bool is_aligned =
 alignment_analyzer.is_aligned(op, &aligned_offset);
 // We know the alignment_analyzer has been able to reason about alignment
 // if the following is true.
 bool known_alignment = is_aligned || (!is_aligned && aligned_offset != 0);
 int lanes = ramp->lanes;
 int native_lanes = required_alignment / op->type.bytes();
 int stride = static_cast<int>(*const_stride);
 if (stride != 1) {
 internal_assert(stride >= 0);

 // If we know the offset of this strided load is smaller
 // than the stride, we can just make the load aligned now
 // without requiring more vectors from the dense
 // load. This makes loads like f(2*x + 1) into an aligned
 // load of double length, with a single shuffle.
 int shift = known_alignment && aligned_offset < stride ? aligned_offset : 0;

 // Load a dense vector covering all of the addresses in the load.
 Expr dense_base = simplify(ramp->base - shift);
 ModulusRemainder alignment = op->alignment - shift;
 Expr dense_index = Ramp::make(dense_base, 1, lanes * stride);
 Expr dense = make_load(op, dense_index, alignment);

 // Shuffle the dense load.
 return Shuffle::make_slice(dense, shift, stride, lanes);
 }

 // We now have a dense vector load to deal with.
 internal_assert(stride == 1);
 if (lanes < native_lanes) {
 // This load is smaller than a native vector. Load a
 // native vector.
 Expr ramp_base = ramp->base;
 ModulusRemainder alignment = op->alignment;
 int slice_offset = 0;

 // If load is smaller than a native vector and can fully fit inside of it and offset is known,
 // we can simply offset the native load and slice.
 if (!is_aligned && aligned_offset != 0 && Int(32).can_represent(aligned_offset) && (aligned_offset + lanes <= native_lanes)) {
 ramp_base = simplify(ramp_base - (int)aligned_offset);
 alignment = alignment - aligned_offset;
 slice_offset = aligned_offset;
 }

 Expr native_load = make_load(op, Ramp::make(ramp_base, 1, native_lanes), alignment);

 // Slice the native load.
 return Shuffle::make_slice(native_load, slice_offset, 1, lanes);
 }

 if (lanes > native_lanes) {
 // This load is larger than a native vector. Load native
 // vectors, and concatenate the results.
 vector<Expr> slices;
 for (int i = 0; i < lanes; i += native_lanes) {
 int slice_lanes = std::min(native_lanes, lanes - i);
 Expr slice_base = simplify(ramp->base + i);
 ModulusRemainder alignment = op->alignment + i;
 slices.push_back(make_load(op, Ramp::make(slice_base, 1, slice_lanes), alignment));
 }
 return Shuffle::make_concat(slices);
 }

 if (!is_aligned && aligned_offset != 0 && Int(32).can_represent(aligned_offset)) {
 // We know the offset of this load from an aligned
 // address. Rewrite this is an aligned load of two
 // native vectors, followed by a shuffle.
 Expr aligned_base = simplify(ramp->base - (int)aligned_offset);
 ModulusRemainder alignment = op->alignment - (int)aligned_offset;
 Expr aligned_load = make_load(op, Ramp::make(aligned_base, 1, lanes * 2), alignment);

 return Shuffle::make_slice(aligned_load, (int)aligned_offset, 1, lanes);
 }

 return IRMutator::visit(op);
 }
};

} // namespace

Stmt align_loads(const Stmt &s, int alignment, int min_bytes_to_align) {
 return AlignLoads(alignment, min_bytes_to_align).mutate(s);
}

} // namespace Internal
} // namespace Halide