#include <poplar/Vertex.hpp>

class [[poplar::constraint("region(*nonexecutableDummy) != region(*plans)"),
        poplar::constraint("elem(*source) != elem(*plans)"),
        poplar::constraint("elem(*received) != elem(*plans)")]]
MultiTileMatMul : public poplar::SupervisorVertex {
public:
  poplar::InOut<poplar::Vector<unsigned>> plans;
  poplar::Input<poplar::Vector<unsigned>> roles;
  poplar::Input<poplar::Vector<unsigned>> slots;
  poplar::Input<poplar::Vector<
      unsigned, poplar::VectorLayout::SPAN, 4, true>>
      nonexecutableDummy;
  poplar::Input<poplar::Vector<int>> source;
  poplar::InOut<poplar::Vector<int>> received;
  poplar::Input<unsigned> isOutput;
  poplar::Output<int> result;

  __attribute__((target("supervisor"))) bool compute() {
#if defined(__IPU__) && defined(__POPC__)
    asm volatile("setzi $m0, 1\n\tput 0xa6, $m0" ::: "$m0", "memory");
#endif

    for (unsigned pass = 0; pass < roles.size(); ++pass) {
      unsigned *code = &plans[pass * 9];
      if (roles[pass] == 1) {
        const unsigned sourceWord = reinterpret_cast<unsigned>(&source[0]) >> 2;
        for (unsigned i = 0; i < 9; ++i) {
          if ((code[i] & 0xf8000000u) == 0x78000000u) {
            code[i] = (code[i] & ~0x001ffff8u) |
                      ((sourceWord << 3) & 0x001ffff8u);
            break;
          }
        }
#if defined(__IPU__) && defined(__POPC__)
        asm volatile(
            "setzi $m10, 1f\n\t"
            "br %[entry]\n\t"
            "1:\n\t"
            :
            : [entry] "r"(code)
            : "$m0", "$m10", "memory");
#endif
      } else if (roles[pass] == 2) {
        int *destination = &received[slots[pass]];
        const unsigned windowWord =
            (reinterpret_cast<unsigned>(destination) - 0x50000u) >> 2;
        for (unsigned i = 0; i < 9; ++i) {
          if ((code[i] & 0x00014000u) == 0x00014000u &&
              (code[i] & 0xfc000000u) == 0x60000000u) {
            code[i] = (code[i] & ~0x1fffu) | (windowWord & 0x1fffu);
            break;
          }
        }
#if defined(__IPU__) && defined(__POPC__)
        asm volatile(
            "setzi $m10, 1f\n\t"
            "br %[entry]\n\t"
            "1:\n\t"
            :
            : [entry] "r"(code)
            : "$m0", "$m10", "memory");
#endif
      } else {
#if defined(__IPU__) && defined(__POPC__)
        asm volatile(
            ".long 0x40c00000\n\t"
            ".long 0x41800001\n\t"
            :
            :
            : "memory");
#endif
      }
    }

    int sum = 0;
    if (isOutput != 0) {
      for (unsigned k = 0; k < 4; ++k)
        sum += received[k] * received[4 + k];
    }
    *result = sum;
    return true;
  }
};