#include <poplar/Vertex.hpp>

class NoBlobSetupRecv : public poplar::SupervisorVertex {
public:
  poplar::Output<poplar::Vector<unsigned>> planBuf;
  poplar::Input<unsigned> count;

  __attribute__((target("supervisor"))) bool compute() { return true; }
};

class NoBlobSetupSend : public poplar::SupervisorVertex {
public:
  poplar::InOut<poplar::Vector<unsigned>> planBuf;
  poplar::Input<poplar::Vector<unsigned>> recvTile;
  poplar::Input<unsigned> count;

  __attribute__((target("supervisor"))) bool compute() { return true; }
};

class [[poplar::constraint("region(*nonexecutableDummy) != region(*planBuf)"),
        poplar::constraint("elem(*data) != elem(*planBuf)")]]
NoBlobSend : public poplar::SupervisorVertex {
public:
  poplar::InOut<poplar::Vector<unsigned>> planBuf;
  poplar::Input<poplar::Vector<
      unsigned, poplar::VectorLayout::SPAN, 4, true>>
      nonexecutableDummy;
  poplar::Input<int> elementSelector;
  poplar::Input<poplar::Vector<int>> data;

  __attribute__((target("supervisor"))) bool compute() {
#if defined(__IPU__) && defined(__POPC__)
    asm volatile(
        "add $m11, $m11, -8\n\t"
        "setzi $m6, 0\n\t"
        "setzi $m6, 0\n\t"
        "setzi $m6, 0\n\t"
        "setzi $m6, 0\n\t"
        "st32 $m10, $m11, $m15, 1\n\t"
        "st32 $m7, $m11, $m15, 0\n\t"
        "ld32 $m1, $m0, $m15, 4\n\t"
        "ld32 $m2, $m0, $m15, 5\n\t"
        "ld32 $m1, $m1, $m15, 0\n\t"
        "shl $m1, $m1, 2\n\t"
        "add $m1, $m2, $m1\n\t"
        "get $m7, 0xa7\n\t"
        "put 0xa7, $m1\n\t"
        "ld32 $m0, $m0, $m15, 0\n\t"
        "setzi $m10, 1f\n\t"
        "br $m0\n\t"
        "1:\n\t"
        "put 0xa7, $m7\n\t"
        "setzi $m0, 1\n\t"
        "ld32 $m7, $m11, $m15, 0\n\t"
        "ld32 $m10, $m11, $m15, 1\n\t"
        "add $m11, $m11, 8\n\t"
        "br $m10\n\t"
        :
        :
        : "memory");
    __builtin_unreachable();
#endif

    return true;
  }
};

class [[poplar::constraint("region(*nonexecutableDummy) != region(*planBuf)"),
        poplar::constraint("elem(*result) != elem(*planBuf)")]]
NoBlobRecv : public poplar::SupervisorVertex {
public:
  poplar::InOut<poplar::Vector<unsigned>> planBuf;
  poplar::Input<poplar::Vector<
      unsigned, poplar::VectorLayout::SPAN, 4, true>>
      nonexecutableDummy;
  poplar::Input<int> tileSelector;
  poplar::Input<poplar::Vector<unsigned short>> logicalToPhysicalTile;
  poplar::Output<poplar::Vector<int>> result;

  __attribute__((target("supervisor"))) bool compute() {
    unsigned *plan = &planBuf[0];
    unsigned patchIndex = plan[0];
    unsigned *code = plan + 1;
    unsigned physicalTile = logicalToPhysicalTile[*tileSelector];
    code[patchIndex] ^= physicalTile;

    int *dest = &result[0];
    unsigned saved = 0;

#if defined(__IPU__) && defined(__POPC__)
    asm volatile(
        "get %[saved], 0xa4\n\t"
        "put 0xa4, %[dest]\n\t"
        "setzi $m10, 1f\n\t"
        "br %[entry]\n\t"
        "1:\n\t"
        "put 0xa4, %[saved]\n\t"
        : [saved] "=&r"(saved)
        : [dest] "r"(dest), [entry] "r"(code)
        : "$m0", "$m10", "memory");
#endif

    code[patchIndex] &= ~0x1fffu;
    return true;
  }
};

class NoBlobNonParticipation : public poplar::SupervisorVertex {
public:
  __attribute__((target("supervisor"))) bool compute() {
#if defined(__IPU__) && defined(__POPC__)
    asm volatile(
        ".long 0x40c00000\n\t"
        ".long 0x41800001\n\t"
        :
        :
        : "memory");
#endif
    return true;
  }
};
