# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

# determine ldepth by object type
_ldepth := o -> let(
    x := When(ObjId(o)=tcast, o.args[2], o),

    # cond functions like a switch statement.
    Cond(
        IsBound(x.ldepth),
            x.ldepth,
        not IsVar(x),
            2^20,
        let(
            f := DefFrontier(x, 5),
            l := Length(f),
            last := Last(f),

            Cond(
                l = 1,
                    0,
                Length(last)=1 and IsLoopIndex(last[1]),
                    last[1].ldepth,
                2^20
            )
        )
    )
);

# added a high value for constants so that they sort at one end
# when the variables are sorted by ldepth
_ldepth2 := (o, params_first) -> let(
    x := When(ObjId(o)=tcast, o.args[2], o),

    # cond functions like a switch statement.
    Cond(
        IsBound(x.ldepth),
            x.ldepth,
        IsValue(x),
            (2^20)+1,
        ObjId(x) in [param, in_param], 100,
#            (2^20) + When(params_first, -2-(InternalHash(x.id) mod 1024),
#                                        +2+(InternalHash(x.id) mod 1024)),

        not IsVar(x),
            2^20,
        let(
            f := DefFrontier(x, 5),
            l := Length(f),
            last := Last(f),

            Cond(
                l = 1,
                    0,
                Length(last)=1 and IsLoopIndex(last[1]),
                    last[1].ldepth,
                2^21
            )
        )
    )
);

Class(BinSplit, rec(

    canSplit := (self, exp) >> ((IsExpCommand(exp) and not IsBound(exp.isExpComposite)) or
        (IsBound(exp.isExpComposite) and exp.isExpComposite
         and ForAny(exp.rChildren(), e->IsBound(e.isExpComposite))))
         and not (exp _is self._splitDisabled),

    expSplit := meth(self, exp)
        local c, v, composites, res, ch, args, i, prev, newexp;
        res := [];
        if self.canSplit(exp) then
            ch := exp.rChildren();
            composites := SplitBy([1..Length(ch)], i->IsBound(ch[i].isExpComposite))[1];
            for c in composites do
                v := var.fresh_t("a", ch[c].t);
                Append(res, self.assignSplit(assign(v, ch[c])));
                exp.rSetChild(c, v);
            od;
        fi;

        # check for non-binary add/mult
	# NOTE: use a list of assoc expr? use Exp property!!
        if (ObjId(exp) in [add, mul, xor, bin_xor, bin_and]) and Length(exp.args) > 2 then
            args := ShallowCopy(exp.args);
            for v in args do v.ldepth := _ldepth2(v, ObjId(exp)=mul); od;
            Sort(args, (a,b) -> a.ldepth < b.ldepth);

            prev := args[1];
            for i in [2..Length(args)-1] do
	        newexp := ObjId(exp)(prev, args[i]);
                v := var.fresh_t("b", newexp.t);
                Add(res, assign(v, newexp));
                prev := v;
            od;
            exp.args := [prev, Last(args)];
        fi;
	return [res, exp];
    end,

    assignSplit := meth(self, cmd)
        local loc, exp, cmds_loc, cmds_exp;
        [loc, exp] := [cmd.loc, cmd.exp];
	[cmds_exp, exp] := self.expSplit(exp);
	[cmds_loc, loc] := self.expSplit(loc);

        if ObjId(exp) in [div,mul,add,xor, stickyNeg] then
            # sets kdepth correctly for stuff like a+(b+c)
            loc.ldepth:=Minimum(List(exp.args,x->_ldepth2(x, false)));
            # Sets ldepth correctly for stuff like a2= 2*i3 => a2.ldepth=_ldepth(i3)
            if (loc.ldepth>=2^18) then
                loc.ldepth:=Minimum0(Filtered(List(exp.args,x->_ldepth2(x, false)), x->x<2^18));
            fi;
        fi;
        return cmds_exp :: cmds_loc :: [cmd];
    end,

    expCommandSplit := (self, cmd) >> self.expSplit(cmd),
#    expCommandSplit := meth(self, cmd)
#        local a, args, res, cmds;
#	res := [];
#	for a in cmd.args do
#	    [cmds, a] := self.expSplit(a);
#	    Append(res, cmds);
#	od;
#        return res :: [cmd];
#    end,

    # BinSplit(code[, opts])

    __call__ := meth(arg)
        local self, code, opts;
        [self, code, opts] := [arg[1], arg[2], Cond(Length(arg)>2, arg[3], rec())];
        if IsBound(opts.binsplitOnIndices) and not opts.binsplitOnIndices then
            self := CopyFields(self, rec( _splitDisabled := opts.simpIndicesInside ));
        else
            self := CopyFields(self, rec( _splitDisabled := [] ));
        fi;
        return self.apply(code);
    end,

    apply := meth(self, code)
        local chains, blk, i, cmd, aux, exp, split;
        chains := Collect(code, chain);
	for blk in chains do
	    for i in [1..Length(blk.cmds)] do
	        cmd := blk.cmds[i];
	        if IsAssign(cmd) then
                    # cmd could have parts aliased with different command
                    # since we do inplace modifications of cmd inside assignSplit, we copy it
                    split := self.assignSplit(Copy(cmd));
                    blk.cmds[i] := When(Length(split)=1, split[1], chain(split));
		elif IsExpCommand(cmd) then
                    split := self.expCommandSplit(Copy(cmd));
                    blk.cmds[i] := When(Length(split)=1, split[1], chain(split));
                fi;
            od;
        od;
        return FlattenCode(code);
    end,
));


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

# ==========================================================================
# SPLSymbol(<name>, <params>)  - obsolete function for creating spl symbols
#
# Note: in this revision all symbol names can be used as functions.
#       SPLSymbol("F", 2) is now F(2)
#
# ==========================================================================
SPLSymbol := function(s, params)
    local sym, perm;
    if not IsList(params) then params := [params]; fi;
    sym  := IsBound(spl.(s));
    if not sym then
	Error("Symbol '", s, "' does not exist. See Dir(spl).");
    else
	return ApplyFunc(spl.(s), params);
    fi;
end;

#F IsSPLSym(<x>) - returns true for parametrized symbols (F, I, ...)
#F
IsSPLSym := x -> IsRec(x) and IsBound(x._sym) and x._sym;

# ==========================================================================
# Sym - base class for SPL symbols
# ==========================================================================
Class(Sym, BaseMat, rec(
    _sym := true,
    transposed := false,
    type := "Sym",

    new := (self, s, obj) >>
        SPL(Inherit(self, rec( symbol := s,
                   obj := obj,
                   dimensions := obj.dims()))),

   # -------- Transformation rules support ---------------------------------
    rChildren := self >> When(IsList(self.params), self.params, [self.params]),
    rSetChild := meth(self, n, newChild)
        if IsList(self.params) then
            self.params[n] := newChild;
        else
            if n<>1 then Error("<n> must be 1"); fi;
            self.params := newChild;
        fi;
        # self.canonizeParams(); ??
        self.dimensions := self.dims();
    end,

    from_rChildren := (self, rch) >> let(
	res := ApplyFunc(ObjId(self), rch).appendAobj(self),
	When(self.transposed, res.transpose(), res)),
    # ----------------------------------------------------------------------

    checkParams := meth(self, params)
        local nargs, nump;
	nargs := NumArgs(self.def);
	nump := When(IsList(params), Length(params), 1);
        if nargs <> -1 and nargs <> nump then
	    Error("Symbol '", self.name, "' needs ", NumArgs(self.def), " parameters: ",
		ParamsFunc(self.def), "\n");
	fi;
    end,

    canonizeParams := meth(self, params)
        local A, nump;
	nump := Length(params);
        if IsBound(self.abbrevs) then
	    for A in self.abbrevs do
                if NumArgs(A) = -1 or NumArgs(A) = nump then
		    return ApplyFunc(A, params);
		fi;
	    od;
	    return params;
	else
	    return params;
	fi;
    end,

    fromDef := meth(arg)
        local result, self, params, h, lkup;
	self := arg[1];
	params := arg{[2..Length(arg)]};
	params := self.canonizeParams(params);
        self.checkParams(params);
	params := When(IsList(params), params, [params]);

	h := self.hash;
	if h<>false then
	    lkup := h.objLookup(self, params);
	    if lkup[1] <> false then return lkup[1]; fi;
	fi;

	result := SPL(WithBases(self, rec(params := params, transposed := false)));
	result.obj := ApplyFunc(result.def, params);
	result.dimensions := Dimensions(result.obj);

	if h<>false then return h.objAdd(result, lkup[2]);
	else return result;
	fi;
    end,

    __call__ := ~.fromDef,

    getObj := self >> When(self.transposed, self.obj.transpose(), self.obj),

    dims          := self >> self.getObj().dims(),
    isPermutation := self >> self.obj.isPermutation(),
    isReal        := self >> self.obj.isReal(),
    toAMat        := self >> self.getObj().toAMat(),
    free          := self >> Union(List(self.params, FreeVars)),

    #-----------------------------------------------------------------------
    transpose := self >>
        Inherit(self, rec(transposed := not self.transposed,
                  dimensions := Reversed(self.dimensions))),
    #-----------------------------------------------------------------------
    equals := meth(self, other)
        return ObjId(other) = ObjId(self)
           and other.params = self.params
           and self.transposed = other.transposed
	   and self.a = other.a;
    end,
    #-----------------------------------------------------------------------
    print := (self, i, is) >> let(
        params := When(IsList(self.params), self.params, [self.params]),
        Print(self.__name__, "(", PrintCS(params), ")",
              When(self.transposed, ".transpose()", ""),
              self.printA())),
    #-----------------------------------------------------------------------
    arithmeticCost := (self, costMul, costAddMul) >>
        self.getObj().arithmeticCost(costMul, costAddMul)
));

#F SymFunc(<name>, <func>)
#F      Convert a function (which returns SPL) to SPL symbol.
#F      Resulting symbol is automatically added to spiral.spl.symbols namespace.
SymFunc := function(name, func)
    local sym;
    sym := rec(notSupported := true, def := func);
    sym.name := name;
    sym.operations := ClassOps;
    sym.__bases__ := [ Sym ];
    CantCopy(sym);
    return sym;
end;


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

Class(NEONUnparser, CMacroUnparserProg, rec(

    T_Real := (self, t, vars, i, is) >> Cond(
	t.params[1] = 32, Print("float32_t ", self.infix(vars, ", ", i+is)),
	Inherited(t, vars, i, is)),

    # composite instructions, normally they should go away before unparser
    vrev_neon  := (self, o, i, is) >> self.printf("vrev64q_f32($1)", o.args),
    vrev_half  := (self, o, i, is) >> self.printf("vrev64_f32($1)", o.args),

#    vunpacklo2_neon    := (self, o, i, is) >> self.printf("vzipq_f32(vtrnq_f32($1, $2).val[0], vtrnq_f32($1, $2).val[1]).val[0]", [o.args[1], o.args[2]]),
#    vunpackhi2_neon    := (self, o, i, is) >> self.printf("vzipq_f32(vtrnq_f32($1, $2).val[0], vtrnq_f32($1, $2).val[1]).val[1]", [o.args[1], o.args[2]]),

    vunpacklo_half    := (self, o, i, is) >> self.printf("vtrn_f32($1,$2).val[0]", o.args),
    vunpackhi_half    := (self, o, i, is) >> self.printf("vtrn_f32($1,$2).val[1]", o.args),

    vunpacklolo2_neon := (self, o, i, is) >> self.printf(
	"vcombine_f32(vget_low_f32($1),vget_low_f32($2))", o.args),
    vunpacklohi2_neon := (self, o, i, is) >> self.printf(
	"vcombine_f32(vget_low_f32($1),vget_high_f32($2))", o.args),
    vunpackhilo2_neon := (self, o, i, is) >> self.printf(
	"vcombine_f32(vget_high_f32($1),vget_low_f32($2))", o.args),
    vunpackhihi2_neon := (self, o, i, is) >> self.printf(
	"vcombine_f32(vget_high_f32($1),vget_high_f32($2))", o.args),

    # native instructions
    vuzpq_32f  := (self, o, i, is) >> self.printf("vuzpq_f32($1, $2)", o.args),
    vzipq_32f  := (self, o, i, is) >> self.printf("vzipq_f32($1, $2)", o.args),
    vtrnq_32f  := (self, o, i, is) >> self.printf("vtrnq_f32($1, $2)", o.args),
    vtrnq_half := (self, o, i, is) >> self.printf("vtrn_f32($1, $2)", o.args),

    vextract_neon_4x32f  := (self, o, i, is) >> self.printf(
	"$1.val[$2]", [o.args[1], o.args[2].p[1]]),
    vextract_half        := (self, o, i, is) >> self.printf(
	"$1.val[$2]", [o.args[1], o.args[2].p[1]]),

    vext_half  := (self, o, i, is) >> self.prefix("vext_f32", o.args),
    vext_neon  := (self, o, i, is) >> self.prefix("vextq_f32", o.args),

    fdiv := (self, o, i, is) >> self.printf("(((float)$1) / $2)", o.args),
    div  := (self, o, i, is) >> self.printf("($1 / $2)", o.args),
    idiv := (self, o, i, is) >> self.printf("($1 / $2)", o.args),
    imod := (self, o, i, is) >> self.printf("($1 % $2)", o.args),

    Value := (self, o, i, is) >> let(zero := "0", Cond(
	    (o.t = TVect(T_Real(32), 4) or o.t = TVect(TReal, 4))  and Length(Set(o.v)) = 1,
	        Print("vdupq_n_f32(", self(o.v[1], i, is), ")"),
	    (o.t = TVect(T_Real(32), 2) or o.t = TVect(TReal, 2))  and Length(Set(o.v)) = 1,
	        Print("vdup_n_f32(", self(o.v[1], i, is), ")"),
	    o.t = T_Real(32) or o.t = TReal,
	        When(o.v < 0, Print("(", o.v, ")"), Print(o.v)),
	    IsIntT(o.t) or IsUIntT(o.t),
	        When(o.v < 0, Print("(", o.v, ")"), Print(o.v)),
	    o.t = TVect(T_Real(32), 4) or o.t = TVect(TReal, 4),
	        Print("(float32x4_t){",o.v[1].v,", ",o.v[2].v,", ",o.v[3].v,", ",o.v[4].v,"}"),
	    o.t = TVect(T_Real(32), 2) or o.t = TVect(TReal, 2),
	        Print("(float32x2_t){",o.v[1].v,", ",o.v[2].v,"}"),
	    o.t = TString,
	        Print(o.v),
	    Inherited(o, i, is))),

    nth := (self, o, i, is) >> self.printf("$1[$2]", [o.loc, o.idx]),

    TReal := (self, t, vars, i, is) >> Print("float ", self.infix(vars, ", ")),

    TVect := (self, t, vars, i, is) >> Print(
        Cond( t = NEON.t,           "float32x4_t ",
              t = NEON_HALF.t,      "float32x2_t ",
              t = TVect(TReal, 4), "float32x4_t ",
              t = TVect(T_Real(32), 2), "float32x2_t ",
              t = TVect(TReal, 2), "float32x2_t ",
              t = TVect(NEON.t, 2), "float32x4x2_t ",
              t = TVect(NEON_HALF.t, 2), "float32x2x2_t ",
              Error("unexpected vector type")),
        self.infix(vars, ", ")),

    vcombine_neon   := (self, o, i, is) >> self.prefix("vcombine_f32", o.args),
    vload_half_neon := (self, o, i, is) >> self.printf("vld1_f32((float32_t*)$1)", o.args),

    vload1_neon  := (self, o, i, is) >> self.prefix("vld1q_lane_f32",
	[o.args[1], o.args[2], o.args[3].p-1]),
    vload1_half  := (self, o, i, is) >> self.prefix("vld1_lane_f32",
	[o.args[1], o.args[2], o.args[3].p-1]),
    vstore1_neon  := (self, o, i, is) >> Print(Blanks(i),
	self.prefix("vst1q_lane_f32", [o.args[1], o.args[2], o.args[3]-1]), ";\n"),
    vstore1_half  := (self, o, i, is) >> Print(Blanks(i),
	self.prefix("vst1_lane_f32", [o.args[1], o.args[2], o.args[3]-1]), ";\n"),

    vstore2lo_neon  := (self, o, i, is) >> Print(Blanks(i), self.printf(
	    "vst1_f32($1,vget_low_f32($2));\n", o.args)),
    vstore2hi_neon  := (self, o, i, is) >> Print(Blanks(i), self.printf(
	    "vst1_f32($1,vget_high_f32($2));\n", o.args)),

    data := (self, o, i, is) >> Cond(
	o.var.t=TVect(T_Real(32), 4) or o.var.t=TVect(TReal, 4),
	    Print(Blanks(i), self.printf(
              "const float32x4_t $1 = {(float32_t)$2,(float32_t)$3,(float32_t)$4,(float32_t)$5}",
	      [o.var] :: List(o.value.v, x -> x.v)), ";\n", self(o.cmd, i, is)),

	o.var.t=TVect(T_Real(32), 2) or o.var.t=TVect(TReal, 2),
	    Print(Blanks(i), self.printf(
	      "const float32x2_t $1 = {(float32_t)$2, (float32_t)$3}",
	      [o.var] :: List(o.value.v, x->x.v)), ";\n", self(o.cmd, i, is)),

	o.var.t = T_Real(32) or o.var.t = TReal,
	    Print(Blanks(i), self.printf(
	      "float32_t $1 = (float32_t)$2", [o.var, o.value.v]), ";\n", self(o.cmd, i, is)),

        Inherited(o, i, is)),

    vpack  := (self, o, i, is) >> Cond(
	o.t=TVect(T_Real(32), 4) or o.t=TVect(TReal, 4),
	self.printf("(float32x4_t){(float32_t)$1, (float32_t)$2, (float32_t)$3, (float32_t)$4}", o.args),
	o.t=TVect(T_Real(32), 2) or o.t=TVect(TReal, 2),
	self.printf("(float32x2_t){(float32_t)$1, (float32_t)$2}", o.args),

	Inherited(o, i, is)),

    vdup := (self, o, i, is) >> Cond(
	o.t=TVect(T_Real(32), 4) or o.t=TVect(TReal, 4),
	    self.printf("vdupq_n_f32($1)", o.args),
	o.t=TVect(T_Real(32), 2) or o.t=TVect(TReal, 2),
	    self.printf("vdup_n_f32($1)", o.args),
	Inherited(o, i, is)),

    vdup_lane_neon := (self, o, i, is) >> Cond(
	o.args[2] = 1 or o.args[2] = 2,
	    self.printf("vdupq_lane_f32(vget_low_f32($1), $2)", [o.args[1], o.args[2]-1]),
	o.args[2] = 3 or o.args[2] = 4,
	    self.printf("vdupq_lane_f32(vget_high_f32($1), $2)", [o.args[1], o.args[2]-3])
    ),

    vdup_lane_half := (self, o, i, is) >> self.prefix("vdup_lane_f32",
	[o.args[1], o.args[2]-1]),

    vswapcx_neon := (self, o, i, is) >> self.printf("vrev64q_f32($1)", o.args),

    vswapcx_half := (self, o, i, is) >> self.printf("vrev64_f32($1)", o.args),

    mul := (self, o, i, is) >> Cond(ObjId(o.t) = TVect and o.t.size=4,
	CondPat( o,
	    [mul, @TVect, @TVect],
	      self.printf("vmulq_f32($1, $2)", o.args),
	    [mul, @TReal, @TVect],
	      self.printf("vmulq_f32(vdupq_n_f32($1), $2)", o.args),
	    [mul, @TReal, @TReal],
	      self.printf("($1 * $2)", o.args),
	    Inherited(o, i, is)),

	CondPat( o,
	    [mul, @TVect, @TVect],
	      self.printf("vmul_f32($1, $2)", o.args),
	    [mul, @TReal, @TVect],
	      self.printf("vmul_f32(vdup_n_f32($1), $2)", o.args),
	    [mul, @TReal, @TReal],
	      self.printf("($1 * $2)", o.args),
	    Inherited(o, i, is))),

    mul  := (self, o, i, is) >> CondPat( o,
	[mul, @(1).cond(x -> x=TVect(T_Real(32), 4) or x=TVect(TReal,4)),
	      @(2).cond(x -> x=TVect(T_Real(32), 4) or x=TVect(TReal,4))],
	    self.printf("vmulq_f32($1, $2)", o.args),
	[mul, @(1).cond(x -> x=TVect(T_Real(32), 2) or x=TVect(TReal,2)),
	      @(2).cond(x -> x=TVect(T_Real(32), 2) or x=TVect(TReal, 2))],
	    self.printf("vmul_f32($1, $2)", o.args),
	[mul, @TReal, @(1).cond(x -> x=TVect(T_Real(32), 4) or x=TVect(TReal,4))],
	    self.printf("vmulq_f32(vdupq_n_f32($1), $2)", o.args),
	[mul, @TReal, @(1).cond(x -> x=TVect(T_Real(32), 2) or x=TVect(TReal,2))],
	    self.printf("vmul_f32(vdup_n_f32($1), $2)", o.args),
	[mul, @TReal, @TReal],
	    self.printf("($1 * $2)", o.args),

        Inherited(o.ooo, i, is)),

    add  := (self, o, i, is) >> Cond(
	not IsVecT(o.t),
	  self.pinfix(o.args, " + "),
	Length(o.args)=2 and o.t.size=4,
	  self.printf("vaddq_f32($1, $2)", o.args),
	Length(o.args)=2 and o.t.size=2, self.printf("vadd_f32($1, $2)", o.args),
	let(rem := ApplyFunc(add, Drop(o.args, 1)),
	    self.printf(Cond(o.t.size=4,"vaddq_f32($1,$2)","vadd_f32($1,$2)"), [o.args[1], rem]))),

    sub  := (self, o, i, is) >> Cond(
	not IsVecT(o.t),                 self.pinfix(o.args, " - "),
	Length(o.args)=2 and o.t.size=4, self.printf("vsubq_f32($1, $2)", o.args),
	Length(o.args)=2 and o.t.size=2, self.printf("vsub_f32($1, $2)", o.args),
	let(rem := ApplyFunc(add, Drop(o.args, 1)),
	    self.printf(Cond(o.t.size=4,"vsubq_f32($1,$2)","vsub_f32($1,$2)"),
		[o.args[1], rem]))),

		neg  := (self, o, i, is) >> Cond(
		    IsVecT(o.t),
		        self.printf(Cond(o.t.size=4,"vnegq_f32($1)","vneg_f32($1)"), o.args),
		    o.t=TReal,
		        self.printf("(-$1)", o.args),
		    Inherited(o, i, is)),

		fma  := (self, o, i, is) >> CondPat(o,
		    [fma, @TVect, @TVect, @TVect],
		       self.printf(Cond(o.t.size=4,"vmlaq_f32($1, $2, $3)","vmla_f32($1, $2, $3)"), o.args),
		    [fma, @TVect, @TReal, @TVect],
		       self.printf(Cond(o.t.size=4,"vmlaq_f32($1, vdupq_n_f32($2), $3)","vmla_f32($1, vdup_n_f32($2), $3)"), o.args),
		    [fma, @TReal, @TReal, @TReal],
		       self.printf("($1 + ($2 * $3))", o.args),
		    Inherited(o, i, is)),

		fms  := (self, o, i, is) >> CondPat(o,
		    [fms, @TVect, @TVect, @TVect],
		       self.printf(Cond(o.t.size=4,"vmlsq_f32($1, $2, $3)","vmls_f32($1, $2, $3)"), o.args),
		    [fms, @TVect, @TReal, @TVect],
		       self.printf(Cond(o.t.size=4,"vmlsq_f32($1, vdupq_n_f32($2), $3)","vmls_f32($1, vdup_n_f32($2), $3)"), o.args),
		    [fms, @TReal, @TReal, @TReal],
		       self.printf("($1 - ($2 * $3)", o.args),
		    Inherited(o, i, is)),

		nfma  := (self, o, i, is) >> CondPat(o,
		    [nfma, @TVect, @TVect, @TVect],
		       self.printf(Cond(o.t.size=4,"vmlaq_f32(vnegq_f32($1), $2, $3)","vmla_f32(vneg_f32($1), $2, $3)"), o.args),
		    [nfma, @TVect, @TReal, @TVect],
		       self.printf(Cond(o.t.size=4,"vmlaq_f32(vnegq_f32($1), vdupq_n_f32($2), $3)","vmla_f32(vneg_f32($1), vdup_n_f32($2), $3)"), o.args),
		    [nfma, @TReal, @TReal, @TReal],
		       self.printf("(($2 * $3) - $1)", o.args),
		    Inherited(o, i, is)),

		_const := (self,o) >> Cond(
		    o.t = T_Real(32) or o.t = TReal,
		        ["float32_t", [o.v]],
		    IsVecT(o.t), Cond(
			o.t.t = T_Real(32) and o.t.size=4,
			    ["float32x4_t", List(o.v, x->x.v)],
			o.t.t = T_Real(32) and o.t.size=2,
			    ["float32x2_t", List(o.v, x->x.v)],
			Inherited(o)
		    ),
#				o.t = TString,                     ["", [o.v]],
		    Inherited(o)
		)
));


# Copyright 2018-2019, Carnegie Mellon University
# See LICENSE for details

Class(TMovingWindow, Tagged_tSPL_Container, rec(
  abbrevs := [(size, op) -> [size, op] ],
  dims := self >> self.params[2].dims(),
  transpose := self >> self,
  isReal := True
));

hf_size := var.fresh_t("t", TInt);

Class(HF, Tagged_tSPL_Container, rec(
  abbrevs := [(op, it) -> [op, it]],
  dims := self >> [self.params[1].dims()[1], hf_size],
  transpose := self >> self,
  isReal := True
));

Class(TBinOp, Tagged_tSPL_Container, rec(
    abbrevs :=  [(x, op, y)->[x, op, y] ],
    dims := self >> [self.params[3].dims()[1], self.params[3].dims()[2]*2 ],
    transpose := self >> self, #CopyFields(self, rec(transposed := not self.transposed)),
   isReal := True
));

NewRulesFor(TBinOp, rec(
  TBinOp_base := rec(
    forTransposition := false,
    applicable := (self, nt) >> true,
    children := nt -> [ nt.params ],
    apply := (t, c, nt) ->
     let(
       OLCompose(
         Reduction(t.dims()[1], op, 0, False)
       )
     )
  )
));

Class(TForAllInd, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (ind, exp) -> Checked(IsVar(ind), IsSPL(exp), [ind, exp]) ],
    dims := self >> [self.params[2].dims()[1]*self.params[1].range, self.params[2].dims()[2]],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
   isReal := True
));

Class(TComputeDist, Tagged_tSPL_Container, rec(
  abbrevs := [n -> Checked(IsInt(n), [n]) ],
  transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
  dims := self >> [1, self.params[1]],
  isReal := True
));

Class(TGreater, Tagged_tSPL_Container, rec(
    binop :=  (a,b)->_unwrap(gt(a,b)),
    lamdaop := (a,b) -> gt(a,b),
    abbrevs :=  [ (t1, t2) -> Checked(IsSPL(t1) and IsSPL(t2), [t1, t2]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [self.params[1].dims()[1], self.params[1].dims()[2] + self.params[2].dims()[2]],
    isReal := True,
    toOperator := self >> (vec -> let(a := vec{[1..self.params[1].dims()[2]]}, b := vec{[self.params[1].dims()[2]+1..self.dims()[2]]},
        o1 := self.params[1].toOperator(), o2 := self.params[2].toOperator(),
        List(Transposed([o1(a), o2(b)]), i->self.binop(i[1], i[2]))))
));

Class(TLessEqual, Tagged_tSPL_Container, rec(
    binop :=  (a,b)->_unwrap(leq(a,b)),
    lamdaop := (a,b) -> gt(b,a),
    abbrevs :=  [ (t1, t2) -> Checked(IsSPL(t1) and IsSPL(t2), [t1, t2]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [self.params[1].dims()[1], self.params[1].dims()[2] + self.params[2].dims()[2]],
    isReal := True,
    toOperator := self >> (vec -> let(a := vec{[1..self.params[1].dims()[2]]}, b := vec{[self.params[1].dims()[2]+1..self.dims()[2]]},
        o1 := self.params[1].toOperator(), o2 := self.params[2].toOperator(),
        List(Transposed([o1(a), o2(b)]), i->self.binop(i[1], i[2]))))
));

Class(TGreaterEqual, Tagged_tSPL_Container, rec(
    binop :=  (a,b)->_unwrap(geq(a,b)),
    lamdaop := (a,b) -> geq(a,b),
    abbrevs :=  [ (t1, t2) -> Checked(IsSPL(t1) and IsSPL(t2), [t1, t2]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [self.params[1].dims()[1], self.params[1].dims()[2] + self.params[2].dims()[2]],
    isReal := True,
    toOperator := self >> (vec -> let(a := vec{[1..self.params[1].dims()[2]]}, b := vec{[self.params[1].dims()[2]+1..self.dims()[2]]},
        o1 := self.params[1].toOperator(), o2 := self.params[2].toOperator(),
        List(Transposed([o1(a), o2(b)]), i->self.binop(i[1], i[2]))))
));

Class(TEqual, Tagged_tSPL_Container, rec(
    binop :=  (a,b)->_unwrap(eq(a,b)),
    lamdaop := (a,b) -> eq(a,b),
    abbrevs :=  [ (t1, t2) -> Checked(IsSPL(t1) and IsSPL(t2), [t1, t2]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [self.params[1].dims()[1], self.params[1].dims()[2] + self.params[2].dims()[2]],
    isReal := True,
    toOperator := self >> (vec -> let(a := vec{[1..self.params[1].dims()[2]]}, b := vec{[self.params[1].dims()[2]+1..self.dims()[2]]},
        o1 := self.params[1].toOperator(), o2 := self.params[2].toOperator(),
        List(Transposed([o1(a), o2(b)]), i->self.binop(i[1], i[2]))))
));

Class(TLess, Tagged_tSPL_Container, rec(
    binop :=  (a,b)->_unwrap(lt(a,b)),
    lamdaop := (a,b) -> geq(b,a),
    abbrevs :=  [ (t1, t2) -> Checked(IsSPL(t1) and IsSPL(t2), [t1, t2]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [self.params[1].dims()[1], self.params[1].dims()[2] + self.params[2].dims()[2]],
    isReal := True,
    toOperator := self >> (vec -> let(a := vec{[1..self.params[1].dims()[2]]}, b := vec{[self.params[1].dims()[2]+1..self.dims()[2]]},
        o1 := self.params[1].toOperator(), o2 := self.params[2].toOperator(),
        List(Transposed([o1(a), o2(b)]), i->self.binop(i[1], i[2]))))
));

Class(TEvalPolynomial, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (n, s) -> Checked(IsInt(n), [n, s]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [1, 1],
    isReal := True,
    toOperator := self >> (vec -> [_unwrapVec(self.params[2]) * List([0..self.params[1]], i->vec[1]^i)])
));

Class(TDistance, Tagged_tSPL_Container, rec(
    abbrevs :=  [ s -> Checked(IsSPL(s), [s]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [1, 2*self.params[1].dims()[2]],
    isReal := True,
    toOperator := self >> (vec -> let(n:=Length(vec)/2, self.params[1].toOperator()(vec{[1..n]} - vec{[n+1..2*n]})))
));

Class(TPNorm, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (n,p) -> Checked(IsInt(n) and IsInt(p), [n,p]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [1, self.params[1]],
    isReal := True,
    toOperator := self >> (vec -> [Double(Sum(List(vec, i->i^self.params[2])))^(1/self.params[2])])
));

Class(TInfinityNorm, Tagged_tSPL_Container, rec(
    abbrevs :=  [ n -> Checked(IsInt(n), [n]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [1, self.params[1]],
    isReal := True,
    toOperator := self >> (vec -> [Maximum(List(vec, AbsFloat))])
));

Class(TEvalLinSys, Tagged_tSPL_Container, rec(
    abbrevs :=  [ s -> Checked(IsSPL(s), [s]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [1, self.params[1].dims()[2]],
    isReal := True,
));

Class(TMatVecProd, Tagged_tSPL_Container, rec(
    abbrevs :=  [ s -> Checked(IsSPL(s) or IsExp(s), [s]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> self.params[1].dims(),
    isReal := True
));

Class(TInsidePoly, Tagged_tSPL_Container, rec(
    abbrevs :=  [ s -> Checked(IsSPL(s), [s]) ],
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    dims := self >> [1, self.params[1].dims()[2]],
    isReal := True,
    toOperator := self >> (vec -> ForAll(_unwrapMat(self.params[1].mat) * vec - _unwrapVec(self.params[1].rhs), i->i<0)),
));

Class(TForAnyInd, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (ind, exp) -> Checked(IsVar(ind), IsSPL(exp), [ind, exp]) ],
    dims := self >> self.params[2].dims(),
    transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    isReal := True,
    toOperator := self >> (vec -> [ForAny(List([0..self.params[1].range-1], j->RulesStrengthReduce(SubstVars(Copy(self.params[2]), rec((self.params[1].id) := V(j)))).toOperator()), cnd->cnd(vec))])
));

Class(TConstant, Tagged_tSPL_Container, rec(
   abbrevs := [x -> [x]],
   dims := self>>[1, 1],
   transpose := self >>  CopyFields(self, rec(transposed := not self.transposed)),
   isReal := True
));

if IsBound(AllClasses.TCond) then
    Unbind(AllClasses.TCond);
fi;

Class(TCond, Tagged_tSPL_Container, rec(
	abbrevs := [(test, body1, body2) -> Checked(IsSPL(test), IsSPL(body1), IsSPL(body2), [test, body1, body2])],
	dims := self >> self.params[1].dims(),
	transpose := self >> CopyFields(self, rec(transposed := not self.transposed)),
    isReal := True
));


# MPI Baseline options

ImportAll(fftx);
ImportAll(simt);

Declare(MPIOpts);
Declare(MPIGPUDeviceOpts);

Class(MPIDefaultConf, LocalConfig, rec(
    getOpts := meth(self, t)
                   local _opts;
                   _opts := CopyFields(MPIOpts, rec(tags := self.tags));
                   _opts.breakdownRules.TTensorI := [IxA_base];
                   _opts.breakdownRules.TRC := [CopyFields(TRC_tag, rec(applicable := True))];
                   _opts.breakdownRules.TPrm := [TPrm_MPI];
                   return _opts;
               end,

    operations := rec(Print := s -> Print("<MPI Default Configuration>")),
    tags := []
));

Class(MPIGPUDeviceConf, fftx.platforms.cuda.FFTXCUDADeviceOpts, rec(
    getOpts := meth(self, t)
                   local opts;

                   opts := batchFftCUDADeviceOpts();
                   opts.operations := rec(Print := s -> Print("<MPI Distributed FFT/CUDA Device options>"));

                   opts.tags := self.tags::opts.tags;
                   opts.breakdownRules.TTensorI := [IxA_MPI, IxA_base]::opts.breakdownRules.TTensorI;
                   opts.breakdownRules.TTensorII := [batch_cuFFT_cuberot_3D];
                   opts.breakdownRules.TRC := [CopyFields(TRC_tag, rec(applicable := True))];
                   opts.breakdownRules.TPrm := [TPrm_MPI];

                   opts.breakdownRules.MDDFT := When(Length(self.copts) > 0 and IsRec(self.copts[1]) and IsBound(self.copts[1].useCUFFT) and self.copts[1].useCUFFT,
                       [ MDDFT_tSPL_pencil_cuFFT_MPI ], [ MDDFT_tSPL_pencil ]);

                   opts.includes := [ "\"include/fftx_mpi_int.h\"" ];
                   opts.preProcess := (self, t) >> RulesMPIPromoteNT(t);

                   opts.preProcess := t -> ApplyStrategy(t,
                    [ RulesMPIPromoteNT,
                      RulesFFTXPromoteNT,
                      RulesFFTXPromoteNT_Cleanup ],
                       BUA, opts);

                   opts.fftxGen := MPIGPUDeviceOpts.fftxGen;
                   opts._sumsRuleTree := opts.sumsRuleTree;
                   opts.sumsRuleTree := MPIGPUDeviceOpts.sumsRuleTree;

                   opts.codegen.MPIRCPrm := MPICUDACodeGenMixin.MPIRCPrm;
                   opts.unparser.mpi_rcperm := MPICUDAUnparserMixin.mpi_rcperm;
                   opts.sumsgen.MPITensor := MPISumsgenMixin.MPITensor;

                   opts.dynamicDeviceMemory := true;

                   return opts;
               end,

    operations := rec(Print := s -> Print("<MPI GPU Device Configuration>")),
    tags := []
));

Class(MPIGPUDeviceOpts, SpiralDefaults, rec(
    tags := [],
    includes := [ "\"include/fftx_mpi_int.h\"" ],
    tagIt := (self, t) >> t.withTags(self.tags),
    operations := rec(Print := s -> Print("<MPI default options>")),

    preProcess := (self, t) >> RulesMPIPromoteNT(t),

    search := (self, t) >> RandomRuleTree(t, self),

    sumsRuleTree := meth(self, rt)
        local ss;
        ss := self._sumsRuleTree(rt);
        ss := ApplyStrategy(ss, FullLocalizeMPI, BUA, self);
        return ss;
    end,

    fftxGen := meth(self, tt)
        local _tt, rt, ss, c, plist, plist1, tf, tf2;

        _tt := self.preProcess(tt);
        rt := self.search(_tt);
        ss := self.sumsRuleTree(rt);

        self.symbol := When(IsBound(rt.node.params[2].params), rt.node.params[2].params, []);

        c := CodeSums(ss, self);
        c := FixUpMPICUDAPerm(c);
        c := fixReplicatedData(c, self);

        plist := Collect(c, @(1,func, e->e.id="transform"))[1].params;
        plist1 := _orderedUniquify(plist);
        if Length(plist1) = 2 then
            plist1 := plist1 :: self.symbol;
        fi;
        c := SubstTopDown(c, @(1).cond(e->IsList(e) and e = plist), e->plist1);

        if self.symbol <> [] then
            c := SubstTopDown(c, @(1,func, e->e.id="transform"),
                e->SubstVars(e, rec((self.symbol[1].id) := var("_"::self.symbol[1].id, TPtr(TReal)))));

            c := SubstTopDown(c, @(1,specifiers_func),
                e->SubstVars(e, rec((self.symbol[1].id) := var("__"::self.symbol[1].id, TPtr(TReal)))));

            c := SubstTopDown(c, @(1, decl, e-> self.symbol[1] in e.vars),
                e-> decl(Filtered(e.vars, j->j <> self.symbol[1]), e.cmd));
        fi;

        c.ruletree := rt;
        c.ss := ss;
        return c;
    end
));

buildMPIGPUDeviceConf := function(arg)
    local conf;

    conf := CopyFields(MPIGPUDeviceConf,
        rec(
            tags := arg{[1]},
            copts := Drop(arg,1),
            gpuconf := fftx.platforms.cuda.FFTXCUDADeviceDefaultConf
        ));

    return conf;
end;

Class(MPIOpts, SpiralDefaults, rec(
    tags := [],
    includes := [ "\"include/fftx_mpi_int.h\"" ],
    tagIt := (self, t) >> t.withTags(self.tags),
        operations := rec(Print := s -> Print("<MPI default options>")),

    preProcess := (self, t) >> RulesMPIPromoteNT(t),

    search := (self, t) >> RandomRuleTree(t, self),
));

Class(MPIGlobals, rec(
    confMPI := (arg) >> CopyFields(MPIDefaultConf, rec(tags := [arg[2]])),
    defaultConf := arg >> ApplyFunc(arg[1].confMPI, Drop(arg, 1)),
    confMPIGPUDevice := arg >> ApplyFunc(buildMPIGPUDeviceConf, Drop(arg, 1))
));

spiral.LocalConfig.mpi := MPIGlobals;


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

NewRulesFor(TS, rec(
    TS_vect := rec(
        forTransposition := false,
        applicable := t -> t.isTag(1, AVecReg) and IsInt(t.params[1] / t.firstTag().v),
        apply := (t, C, Nonterms) -> let(
            v := t.firstTag().v,
            N := t.params[1],
            VS(N, v)
        )
    )
));


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

#######################################################################################
#   AltiVec 4-way 32-bit float instructions

# shuffle operations
Class(vbinop_4x32f_av, vbinop_av, rec(
    v := self >> 4,
    #ctype := self >> "float",
)); #Class added for AltiVec

Class(vunpacklo_4x32f_av, vbinop_4x32f_av, rec(
    altivec := "vec_mergeh",
    semantic := (in1, in2, p) -> unpacklo(in1, in2, 4, 1)
)); #Class added for AltiVec

Class(vunpackhi_4x32f_av, vbinop_4x32f_av, rec(
    altivec := "vec_mergel",
    semantic := (in1, in2, p) -> unpackhi(in1, in2, 4, 1)
)); #Class added for AltiVec

#   binary instructions
vunpacklo_4x32f.altivec := "vec_mergel";
vunpackhi_4x32f.altivec := "vec_mergeh";

Class(vperm_4x32f, vbinop_4x32f_av, rec(
    altivec := "vec_perm",
    semantic := (in1, in2, p) -> vpermop(in1, in2, p, 4),
    params := self >> sparams(4,8),
    permparams := aperm
));

#   unary instructions
Class(vuperm_4x32f, vunbinop_av, rec(
    binop := vperm_4x32f
));

#######################################################################################
#   AltiVec 8-way 16-bit integer instructions

#   binary instructions
vunpacklo_8x16i.altivec := "vec_mergel";
vunpackhi_8x16i.altivec := "vec_mergeh";

Class(vperm_8x16i, VecExp_8.binary(), rec(
    altivec := "vec_perm",
    semantic := (in1, in2, p) -> vpermop(in1, in2, p, 8),
    params := self >> sparams(8,16),
    permparams := aperm
));

#   unary instructions
Class(vuperm_8x16i, VecExp_8.unaryFromBinop(vperm_8x16i));

#######################################################################################
#   AltiVec 16-way 8-bit integer instructions

#   binary instructions
vunpacklo_16x8i.altivec := "vec_mergel";
vunpackhi_16x8i.altivec := "vec_mergeh";

Class(vperm_16x8i, VecExp_8.binary(), rec(
    altivec := "vec_perm",
    semantic := (in1, in2, p) -> vpermop(in1, in2, p, 16),
    params := self >> sparams(16,32),
    permparams := aperm
));

#   unary instructions
Class(vuperm_16x8i, VecExp_8.unaryFromBinop(vperm_16x8i));


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

_computeExpType := function(e)
    local t;
    if IsBound(e.computeType) then t := e.computeType();
#    elif Length(e.args)=0 then t := TUnknown;
    else t := UnifyTypes(List(e.args, x->x.t));
    fi;
    e.t := t;
    return e;
end;

PropagateTypes := c -> SubstBottomUp(c,
    @.cond(e->IsExp(e) and not IsBound(e.t)),
    _computeExpType);


##  Copyright (c) 2018-2021, Carnegie Mellon University
##  See LICENSE for details

# from FFTX C++ unparser
Class(TFCall, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (nt, cconf) -> Checked(IsSPL(nt), [nt, cconf]) ],
    transpose := self >> ObjId(self)(self.params[1].transpose(), self.params[2]).withTags(self.getTags())
));

Class(TDeviceCall, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (nt, cconf) -> Checked(IsSPL(nt), [nt, cconf]) ],
    transpose := self >> ObjId(self)(self.params[1].transpose(), self.params[2]).withTags(self.getTags())
));

Class(TDecl, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (nt, vlist) -> Checked(IsSPL(nt), [nt, vlist]) ],
    transpose := self >> ObjId(self)(self.params[1].transpose(), self.params[2]).withTags(self.getTags())
));

Class(TDAGNode, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (nt, ylist, xlist) -> Checked(IsSPL(nt), [nt, When(IsList(ylist), ylist, [ylist]), When(IsList(xlist), xlist, [xlist])]) ],
    transpose := self >> ObjId(self)(self.params[1].transpose(), self.params[2], self.params[3]).withTags(self.getTags())
));

Class(TDAG, TCompose, rec(
    terminate := self >> Error("Not yet implemented."),

    from_rChildren := (self, rch) >> let(
        len := Length(rch),
        transposed := rch[len-1],
        tags := rch[len],
        t := ApplyFunc(ObjId(self), [rch{[1..len-2]}]),
        tt := When(transposed, t.transpose(), t),
        attrTakeA(tt.withTags(tags), self)
    ),

    rChildren := self >>
        Concatenation(self.params[1], [self.transposed, self.tags]),

    rSetChild := meth(self, n, newChild)
        local l;
        l := Length(self.params[1]);
        if n <= l then
            self.params[1][n] := newChild;
        elif n = l+1 then
            self.transposed := newChild;
        elif n = l+2 then
            self.tags := newChild;
        else Error("<n> must be in [1..", l+2, "]");
        fi;
        # self.canonizeParams(); ??
        self.dimensions := self.dims();
    end,

));

## the dims here are ordered as in the MDDFT case: e.g., for 3D: odims[1] = x, odims[2] = y, odims[3] = z
#Class(TResample, Tagged_tSPL_Container, rec(
#    abbrevs :=  [ (odims, idims) -> [ odims, idims, [] ],
#                  (odims, idims, shifts) -> [ odims, idims, shifts ],],
#    dims := self >> [Product(self.params[1]), Product(self.params[2])],
#    isReal := True,
#    terminate := self >> Error("Not yet implemented."),
#    transpose := self >> ObjId(self)(self.params[2], self.params[1], self.params[3]).withTags(self.getTags())
#));

Declare(TSparseMat);
Class(TSparseMat, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (dims, entries) -> [ dims, entries ] ],
    dims := self >> self.params[1],
    terminate := self >> Error("Not yet implemented."),
    transpose := self >> TSparseMat(Reversed(self.params[1]), [x->Error("Not yet done")])
));

Declare(TIterVStack);

Class(TIterHStack, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (nt, idx) -> Checked(IsSPL(nt), IsVar(idx),
	    [nt, idx]) ],

    dims := self >> self.params[1].dims(){[1]} :: self.params[1].dims(){[2]} * self.params[2].range,

    terminate := self >> IterHStack(self.params[2], self.params[2].range, self.params[1].terminate()),

    transpose := self >> TIterVStack(self.params[1].transpose(), self.params[2])
          .withTags(self.getTags()),

    isReal := self >> self.params[1].isReal(),

    normalizedArithCost := self >>
        self.params[1].normalizedArithCost() * self.params[2].range,

#    doNotMeasure := true,

    HashId := self >> let(
    	p := self.params,
    	h := When(IsBound(p[1].HashId), p[1].HashId(), p[1]),
            [h, p[2].range] :: When(IsBound(self.tags), self.tags, [])
    )
));

Class(TIterVStack, Tagged_tSPL_Container, rec(
    abbrevs :=  [ (nt, idx) -> Checked(IsSPL(nt), IsVar(idx),
	    [nt, idx]) ],

    dims := self >> (self.params[1].dims(){[1]} * self.params[2].range) :: self.params[1].dims(){[2]},

    terminate := self >> IterVStack(self.params[2], self.params[2].range, self.params[1].terminate()),

    transpose := self >> TIterHStack(self.params[1].transpose(), self.params[2])
          .withTags(self.getTags()),

    isReal := self >> self.params[1].isReal(),

    normalizedArithCost := self >>
        self.params[1].normalizedArithCost() * self.params[2].range,

#    doNotMeasure := true,

    HashId := self >> let(
    	p := self.params,
    	h := When(IsBound(p[1].HashId), p[1].HashId(), p[1]),
            [h, p[2].range] :: When(IsBound(self.tags), self.tags, [])
    )
));

Class(TNoDiagPullin, Tagged_tSPL_Container, rec(
    abbrevs :=  [ s -> Checked(IsSPL(s), [s]) ],
    HashId := self >> let(
    	p := self.params[1],
    	h := When(IsBound(p.HashId), p.HashId(), p),
            [ self.__name__, h ] :: When(IsBound(self.tags), self.tags, [])
        )
));

Class(TNoDiagPullinRight, Tagged_tSPL_Container, rec(
    abbrevs :=  [ s -> Checked(IsSPL(s), [s]) ],
    HashId := self >> let(
    	p := self.params[1],
    	h := When(IsBound(p.HashId), p.HashId(), p),
            [ self.__name__, h ] :: When(IsBound(self.tags), self.tags, [])
        )
));

Class(TNoDiagPullinLeft, Tagged_tSPL_Container, rec(
    abbrevs :=  [ s -> Checked(IsSPL(s), [s]) ],
    HashId := self >> let(
    	p := self.params[1],
    	h := When(IsBound(p.HashId), p.HashId(), p),
            [ self.__name__, h ] :: When(IsBound(self.tags), self.tags, [])
        )
));

Class(TNoPullRight, Tagged_tSPL_Container, rec(
    abbrevs :=  [ s -> Checked(IsSPL(s), [s]) ],
    HashId := self >> let(
    	p := self.params[1],
    	h := When(IsBound(p.HashId), p.HashId(), p),
            [ self.__name__, h ] :: When(IsBound(self.tags), self.tags, [])
        )
));

Class(TNoPullLeft, Tagged_tSPL_Container, rec(
    abbrevs :=  [ s -> Checked(IsSPL(s), [s]) ],
    HashId := self >> let(
    	p := self.params[1],
    	h := When(IsBound(p.HashId), p.HashId(), p),
            [ self.__name__, h ] :: When(IsBound(self.tags), self.tags, [])
        )
));

#F   tSPL ColMajor -> RowMajor transformation
Class(TColMajor, Tagged_tSPL_Container, rec(
    _short_print := true,
    abbrevs :=  [ (A) -> Checked(IsNonTerminal(A) or IsSPL(A), [A]) ],
    dims := self >> 2*self.params[1].dims(),
    terminate := self >> Mat(MatSPL(RC(self.params[1]))),

    transpose := self >> ObjId(self)(
	self.params[1].conjTranspose()).withTags(self.getTags()),

    conjTranspose := self >> self.transpose(),

    isReal := self >> true,

    # Do not use doNotMeasure, this will prevent TRC_By_Def from ever being found!
    doNotMeasure := false,
    normalizedArithCost := self >> self.params[1].normalizedArithCost(),

    HashId := self >> let(
	h := [ When(IsBound(self.params[1].HashId), self.params[1].HashId(),
		    self.params[1]) ],
        When(IsBound(self.tags), Concatenation(h, self.tags), h))
));


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

########################################################################
#   rules for A x I, I x A, (A x I)L, (I x A)L
Declare(AxI_vec);

NewRulesFor(TTensorI, rec(
#   vectorization cases

#   I_n x A_rxs -> I_n/v x L^rv_v(A_rxs x I_v)L^sv_s
    IxA_vec_push := rec(
        info := "IxA vec",
        forTransposition := false,
        applicable := nt -> IsParPar(nt.params) and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx)),
        children := nt -> let(
            pv := nt.firstTag(),
            v := pv.v,
            d := nt.params[1].dims(),
            krnl := nt.params[1].withTags(nt.getTags()),
            [[ krnl ]]
	    # YSV: below is the original piece of code and comment. This is no longer works
	    #      and does not seem to be necessary. Above works quite well. Delete the below
	    #      as soon as everything is tested.
            #D  When(krnl.isReal(), krnl.setWrap(VWrapId), krnl.setWrap(VWrapTRC(pv.isa))) ]]
        ),  #D  thats bad but required to keep wrapping working :(
        apply := (nt, c, cnt) -> Tensor(I(nt.params[2]), c[1])

    ),

#   A x I_n -> (A x I_n/v) x I_v
    AxI_vec := rec(
        info := "A x I_n -> (A x I_n/v) x I_v",
        forTransposition := false,
        applicable := nt -> nt.hasTags() and IsVecVec(nt.params) and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx)) and IsInt(nt.params[2]/nt.firstTag().v),
        children := nt -> let(
            r := nt.params[2] / nt.firstTag().v,
            isa := nt.firstTag().isa,
            [[ When(r = 1,
                When(nt.numTags() = 1,
                    nt.params[1].setWrap(VWrap(isa)),
                    nt.params[1].setWrap(Drop(nt.getTags(), 1)).setWrap(VWrap(isa))
                ),
                TTensorI(nt.params[1].setWrap(VWrap(isa)), r, AVec, AVec).withTags(Drop(nt.getTags(), 1))
            )]]
        ),
        apply := (nt, c, cnt) -> VTensor(c[1], nt.firstTag().v)
    ),
#   I_n x A_rxs -> I_n/v x L^rv_v(A_rxs x I_v)L^sv_s
    IxA_vec := rec(
        info := "IxA vec",
        forTransposition := false,
        applicable := nt -> IsParPar(nt.params) and nt.hasTags() and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx)) and IsInt(nt.params[2]/nt.firstTag().v),
        children := nt -> let(
            pv := nt.getTags(),
            v := pv[1].v,
            isa := pv[1].isa,
            d := nt.params[1].dims(),
            [[
                TL(d[1]*v, v, 1, 1).withTags(pv).setWrap(VWrapId),
                When(Length(pv)=1, nt.params[1].setWrap(VWrap(isa)), nt.params[1].setpv(Drop(pv, 1)).setWrap(VWrap(isa))),
                TL(d[2]*v, d[2], 1, 1).withTags(pv).setWrap(VWrapId)
            ]]
        ),
        apply := (nt, c, cnt) -> let(
            l := nt.params[2] / nt.firstTag().v,
            A := c[1] * VTensor(c[2], nt.firstTag().v) * c[3],
            NoDiagPullin(When(l=1, A , Tensor(I(l), A)))
        )
    ),
#   (I_n x A_rxs)L^nr_n -> (I_n/v x L^rv_v(A_rxs x I_v))(L^ns/v_n/v x I_v)
    IxA_L_vec := rec(
        info := "(IxA)L vec",
        forTransposition := false,
        maxSize := -1,
        applicable := (self, nt) >>
            ObjId(nt.params[1]) <> TTensorI
            and IsParVec(nt.params)
            and nt.hasTags()
            and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx))
            and IsInt(nt.params[2] / nt.firstTag().v)
            and (
                self.maxSize = -1
                or (
                    Rows(nt.params[1]) <= self.maxSize
                    and Cols(nt.params[1]) <= self.maxSize
                )
            ),
        children := nt -> let(
            pv := nt.getTags(),
            v := pv[1].v,
            isa := pv[1].isa,
            d := nt.params[1].dims(),
            [[
                TL(d[1]*v, v).withTags(pv).setWrap(VWrapId),
                When(Length(pv)=1,
                    nt.params[1].setWrap(VWrap(isa)),
                    nt.params[1].withTags(Drop(pv, 1)).setWrap(VWrap(isa))
                )
            ]]
        ),
        apply := (nt, c, cnt) -> let(
            v := nt.firstTag().v,
            m := nt.params[2] / v,
            d := nt.params[1].dims(),
            NoDiagPullinLeft(
                When(m = 1,
                    c[1] * VTensor(c[2], v),
                    Tensor(I(m), c[1] * VTensor(c[2], v)) * VTensor(L(m*d[2], m), v)
                )
            )
        ),
        switch := false

    ),
#   L^nr_r(I_n x A_rxs) -> (L^nr/v_r/v x I_v)(I_n/v x (A_rxs x I_v)L^sv_s)
    L_IxA_vec := rec(
        info := "L(IxA) vec",
        forTransposition := false,
        maxSize := -1,
        applicable := (self, nt) >>
            ObjId(nt.params[1]) <> TTensorI
            and IsVecPar(nt.params)
            and nt.hasTags()
            and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx))
            and IsInt(nt.params[2] / nt.firstTag().v)
            and (
                self.maxSize = -1
                or (
                    Rows(nt.params[1]) <= self.maxSize
                    and Cols(nt.params[1]) <= self.maxSize
                )
            ),
        children := nt -> let(
            pv := nt.getTags(),
            v := pv[1].v,
            isa := pv[1].isa,
            d := nt.params[1].dims(),
            [[
                When(nt.numTags() = 1,
                    nt.params[1].setWrap(VWrap(isa)),
                    nt.params[1].withTags(Drop(pv, 1)).setWrap(VWrap(isa))
                ),
                TL(d[2]*v, d[2]).withTags(pv).setWrap(VWrapId)
            ]]
        ),
        apply := (nt, c, cnt) -> let(
            v := nt.firstTag().v,
            m := nt.params[2] / v,
            d := nt.params[1].dims(),
            NoDiagPullinRight(
                When(m = 1,
                    VTensor(c[1], v) * c[2],
                    VTensor(L(m*d[1], d[1]), v) * Tensor(I(m), VTensor(c[1], v) * c[2])
                )
            )
        ),

        switch := false
    ),
############################################################################
#   L splitting variant
#   (I_n x A_rxs) L^ns_n
    IxA_L_split_vec := rec(
        info := "split (I_n x A_rxs) L^ns_n",
        forTransposition := false,
        applicable := nt -> IsParVec(nt.params) and nt.hasTags() and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx)),
        children := nt -> [[ TTensorI(nt.params[1], nt.params[2], AVec, AVec).withTags(nt.getTags()) ]],
        apply := (nt, c, cnt) -> let(d := nt.params[1].dims()[1],
            L(d * nt.params[2], nt.params[2]) * c[1]
        ),
        switch := false
    ),
#   L^nr_r (I_n x A_rxs)
    L_IxA_split_vec := rec(
        info := "split L^nr_r (I_n x A_rxs)",
        forTransposition := false,
        applicable := nt -> IsVecPar(nt.params) and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx)),
        children := nt -> [[ TTensorI(nt.params[1], nt.params[2], AVec, AVec).withTags(nt.getTags()) ]],
        apply := (nt, c, cnt) -> let(d := nt.params[1].dims()[2],
            c[1] * L(d * nt.params[2], d)
        ),
        switch := false
    ),
    IxA_split_vec := rec(
        info := "split (I_n x A_rxs)",
        forTransposition := false,
        applicable := nt -> IsParPar(nt.params) and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx)),
        children := nt -> [[ TTensorI(nt.params[1], nt.params[2], AVec, AVec).withTags(nt.getTags()) ]],
        apply := (nt, c, cnt) -> let(d := nt.params[1].dims(),
            L(d[1] * nt.params[2], nt.params[2]) * c[1] * L(d[2] * nt.params[2], d[2])
        ),
        switch := false
    ),
    IxA_conj_vec := rec(
        info := "split (I_n x A_rxs)",
        forTransposition := false,
        applicable := nt -> IsParPar(nt.params) and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx)),
        children := nt -> [[ TTensorI(nt.params[1], nt.params[2], AVec, AVec).withTags(nt.getTags()) ]],
        apply := (nt, c, cnt) -> let(d := nt.params[1].dims(),
            ConjLR(c[1], L(d[1] * nt.params[2], nt.params[2]), L(d[2] * nt.params[2], nt.params[2]))
        ),
        switch := false
    ),

############################################################################
#   for now, these rules are for unrolled cod only
#   non-vref TTensorI
#   A x I_n, v\nmid n
    AxI_svec := rec(
        info := "A x I_n for v\nmid n",
        forTransposition := false,
        applicable := nt ->
            IsVecVec(nt.params)
            and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx))
            #and let( d := nt.params[1].dims(), d[1] = d[2] )
            and not IsInt(nt.params[2] / nt.firstTag().v),

        children := nt -> let(
            v := nt.firstTag().v,
            r := v * QuoInt(nt.params[2] + v - 1, v),
            [[ TTensorI(nt.params[1], r, AVec, AVec).withTags(nt.getTags()) ]]
        ),
        apply := (nt, c, cnt) -> let(
            d := nt.params[1].dims(),
            v := nt.firstTag().v,
            IxVScat_pc(d[1], nt.params[2], nt.params[2], 0, v)
            * c[1]
            * IxVGath_pc(d[2], nt.params[2], nt.params[2], 0, v)
        ),
    ),
#   (I_m x A_n )L^mn_m, v\nmid m
    IxA_L_svec := rec(
        info := "(I_m x A_n )L^mn_m, v\nmid m",
        forTransposition := false,
        applicable := nt -> IsParVec(nt.params) and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx))
            and let(
                v := nt.firstTag().v,
                d := nt.params[1].dims(),
                d[1] = d[2]
                and (
                    not IsInt(nt.params[2] / v)
                    or (not IsInt(d[1] / v))
                )
            ),
        children := nt -> let(
            v := nt.firstTag().v,
            d := nt.params[1].dims(),
            r := v * QuoInt(nt.params[2]+v-1, v),
            s := v * QuoInt(d[1]+v-1, v),
            [[
                TL(r*s, r, 1, 1).withTags(nt.getTags()),
                TTensorI(nt.params[1], r, AVec, AVec).withTags(nt.getTags())
            ]]
        ),
        apply := (nt, C, cnt) -> let(
            d:=nt.params[1].dims()[1],
            v := nt.firstTag().v,
            c1c := Cols(C[1])/v, c1r := Rows(C[1])/v,
            c2c := Cols(C[2])/v, c2r := Rows(C[2])/v,
            C0 := IxVScat_pc(nt.params[2],d,d,0,v),
            c0c := Cols(C0)/v,
            C0
#           * VGath(fAdd(c1r,c0c,0), v)
            * VGath_zero(c1r,c0c, v)
#           * VTensor(HStack(I(c0c), O(c0c,c1r-c0c)),v)
            * C[1]
            * VScat_zero(c1c,c2r,v)
#           * VTensor(VStack(I(c2r), O(c1c-c2r,c2r)),v)
            * C[2]
            * IxVGath_pc(d,nt.params[2],nt.params[2],0,v)
        )
    ),
    TTensorI_oddvloop := rec(
        switch  := false,
        maxSize := false,
        filter  := e->true,
        applicable := nt ->
            IsVecVec(nt.params)
            and (nt.isTag(1,AVecReg) or nt.isTag(1,AVecRegCx))
            and not IsInt(nt.params[2] / nt.firstTag().v),

        children := nt -> let(
            v := nt.firstTag().v,
           [[ TTensorI(nt.params[1], v, AVec, AVec).withTags(nt.getTags()) ]]
        ),

        apply := (t, C, Nonterms) -> let(
  	        v := t.firstTag().v,
            kernel := C[1],
            m := Rows(t.params[1]),
            n := Cols(t.params[1]),
            k := t.params[2],
            kd := _rounddown(k,v).v,
            kr := k-kd,
            i := Ind(kd/v),
            #---------------------
            SUM(
                ISum(i,
                    IxVScat_pc(n, k, v, v*i, v) *
                    kernel *
                    IxVGath_pc(m, k, v, v*i, v)
                ),
                IxVScat_pc(n, k, k-kd, kd, v) *
                kernel *
                IxVGath_pc(m, k, k-kd, kd, v)
            )
        )
    )
));


#
# walrus: Computational Methods for Finitely Generated Monoids and Groups
#

## Places
InstallGlobalFunction(NewPlace,
function(loc, c, colour)
    local p;
    p := Objectify(PregroupPlaceType, [loc,c,colour]);
    Add(loc![3], p);
    return p;
end);

InstallMethod(Location
             , "for a pregroup place"
             , [IsPregroupPlaceRep],
             p -> p![1]);

InstallMethod(Relator
             , "for a pregroup place"
             , [IsPregroupPlace],
             p -> Relator(p![1]));

InstallMethod(Letter
             , "for a pregroup place"
             , [IsPregroupPlaceRep],
             p -> p![2]);

InstallMethod(Colour
             , "for a pregroup place"
             , [IsPregroupPlaceRep],
             p -> p![3]);

InstallMethod(NextPlaces
             , "for a pregroup place"
             , [IsPregroupPlaceRep],
function(p)
    if not IsBound(p![6]) then
        p![6] := Places(NextLocation(Location(p)));
    fi;
    return p![6];
end);

AddOrUpdate := function(map, key, val)
    local tmp;

    tmp := map[key];
    if tmp = fail or val < tmp then
        map[key] := val;
    fi;
end;

InstallGlobalFunction(OneStepRedCase,
function(P)
    local Q, Ql
          , b, c, d, x, y
          , Lp
          , v, v1, v2
          , res
          , xi1, xi2, binv, l, onv, pres, vg, xi;

    pres := PregroupPresentationOf(Relator(P));
    res := HashMap(1024);
    vg := VertexGraph(pres);

    c := Letter(P);

    for Q in NextPlaces(P) do
        # same relator, one position up
        Ql := Location(Q);
        b := InLetter(Ql); binv := PregroupInverse(b);
        d := OutLetter(Ql);
        x := Letter(Q);
        v := VertexFor(vg, [b, d, 0]);
        onv := OutNeighboursOfVertex(vg, v);

        for y in IntermultMap(binv) do
            v1 := VertexFor(vg, [y, binv, 1]);
            xi1 := Blob(pres, y, binv, c);
            for v2 in onv do
                #T this check is new, and assuming the paper
                #T is accurate correct
                if DigraphVertexLabel(vg, v2)[1][2] = x then
                    xi2 := Vertex(pres, v1, v, v2);

                    AddOrUpdate(res, [ __ID(Q), 1 ], xi1 + xi2);
                fi;
            od;
        od;
    od;
    # Note this list is not necessarily dense atm.
    return res;
end);

InstallGlobalFunction(OneStepGreenCase,
function(P)
 local L, L2, b, c, loc, pls, is_consoledge, v, v1, v2, xi1, xi2,
       R, R2, P2, P2s, P2T, i, j, next, res, len, n, l, pres, vg,
       P2_loc, P2_inletter, P2_outletter, P2_outletterinv, P2_letter;

    pres := PregroupPresentationOf(Relator(P));
    res := HashMap(1024);
    vg := VertexGraph(pres);

    L := Location(P);
    b := OutLetter(L);
    c := Letter(P);

    # Every place that is instanciated with location that is in an instantiation of
    # a place P'.
    # We're interested in consolidated edges between Relator(P) and Relator(L2)
    # from the locations that P and pls are at.
    #
    # Do we have to check that Relator(P) and Relator(pls) don't entirely cancel?
    for L2 in Locations(pres) do
        R2 := Relator(L2);

        # L2 instantiates place on R2, have to test consolidated
        #    edges between R and R2 starting from L/L2 on R/R2
        #    respectively
        if (InLetter(L2) = PregroupInverse(b))
            and (OutLetter(L2) = c) then

            # We compute all consolidated edge places,
            # we could do this incrementally, but I don't
            # currently see a use in doing so
            P2s := ConsolidatedEdgePlaces(L, L2);

            for P2T in P2s do
                P2 := P2T[1];  # Place reachable on R1 by consolidated edge

                P2_loc := Location(P2);
                P2_inletter := InLetter(P2_loc);
                P2_outletter := OutLetter(P2_loc);
                P2_outletterinv := PregroupInverse(P2_outletter);
                P2_letter := Letter(P2);

                # P2T[2] location on R2 reachable by the edge
                len := P2T[3]; # length of consolidated edge
                v1 := VertexFor(vg, [ InLetter(P2T[2]), OutLetter(P2T[2]), 0 ]);
                v := VertexFor(vg, [ P2_inletter, P2_outletter, 0 ]);
                for v2 in OutNeighboursOfVertex(vg, v) do
                    if DigraphVertexLabel(vg, v2)[1][1] = P2_outletterinv and
                       P2_letter = DigraphVertexLabel(vg, v2)[1][2] then

                        xi1 := Vertex(pres, v1, v, v2);
                        if Colour(P2) = "green" then
                            AddOrUpdate(res, [ __ID(P2), len ], xi1);

                        elif Colour(P2) = "red" then
                            next := OneStepRedCase(P2);
                            # testhack
                            for n in Keys(next) do
                                AddOrUpdate( res, [ n[1], len + 1 ]
                                           , xi1 + next[n] );
                            od;
                        else
                            Error("Invalid colour");
                        fi;
                    else

                    fi;
                od;
            od;
        fi;
    od;
    return res;
end);

InstallMethod(OneStepReachablePlaces
             , "for a pregroup place"
             , [IsPregroupPlaceRep],
function(p)
    if not IsBound(p![7]) then
        if Colour(p) = "red" then
            p![7] := OneStepRedCase(p);
        elif Colour(p) = "green" then
            p![7] := OneStepGreenCase(p);
        else
            Error("Invalid colour for place ", p, "\n");
        fi;
    fi;
    return p![7];
end);

InstallMethod(__ID
             , "for a pregroup place"
             , [IsPregroupPlace],
             p -> p![5]);

InstallMethod(ViewString
             , "for a pregroup place"
             , [IsPregroupPlaceRep],
function(p)
    return STRINGIFY("(", ViewString(p![1]),
                     ",", ViewString(p![2]),
                     ",", ViewString(p![3]),
                     ")");
end);

InstallMethod(\=
             , "for two pregroup places"
             , [IsPregroupPlaceRep, IsPregroupPlaceRep],
function(p1, p2)
    return (p1![1] = p2![1])
           and (p1![2] = p2![2])
           and (p1![3] = p2![3]);
end);


_diceIt := (k,d) -> let(fct := Factors(k), l := Length(fct), lx:= Replicate(d-1, Int(l/d)), ll := l-Sum(lx), prt := lx::[ll],
    pfx := [0]::List([1..d], i -> Sum(prt{[1..i]})), prtl := List([1..d], i->Product(fct{[pfx[i]+1..pfx[i+1]]})),
    Reversed(prtl));

NewRulesFor(MDDFT, rec(
    MDDFT_tSPL_pencil := rec(

        applicable := (self, t) >> Length(t.params[1]) > 1 and let(tg := t.getTags()[1], IsBound(tg.isParMPI) and tg.isParMPI),

        children := nt -> [let(tags := nt.getTags(), dd := Length(tags), peelof := Product(nt.getTags()[1].params[1]),
                              pencils := List([1..Length(nt.params)], d -> TMap(DFT(nt.params[1][d], nt.params[2]),
                                  [peelof]::_diceIt(Product(nt.params[1]{[1..d-1]}::nt.params[1]{[d+1..Length(nt.params[1])]})/peelof, dd-1), APar, APar).withTags(nt.getTags())),
                              perms := List([1..Length(nt.params)], d -> TPrm(fCubeTranspose(Product(nt.params[1]), d, nt.params[1]))),
                              DropLast(Flat(Zip2(pencils, perms)), 1)
                           )],

        apply := (t, C, Nonterms) -> Compose(C),
        switch := true
    ),
    MDDFT_tSPL_pencil_cuFFT_MPI := rec(

        applicable := (self, t) >> Length(t.params[1]) > 1 and let(tg := t.getTags()[1], IsBound(tg.isParMPI) and tg.isParMPI),

        children := nt -> [[ TCompose([
                                TTensorII(DFT(nt.params[1][3], nt.params[2]), nt.params[1]{[1,2]}, ACubeRot_ZXY, ACubeRot_XYZ),
                                TPrm(fCubeTranspose(Product(nt.params[1]), "FFTX_MPI_3D_CUFFT_STAGE1", nt.params[1])),
                                TTensorII(DFT(nt.params[1][2], nt.params[2]), nt.params[1]{[1,3]}, ACubeRot_ZXY, ACubeRot_XYZ),
                                TPrm(fCubeTranspose(Product(nt.params[1]), "FFTX_MPI_3D_CUFFT_STAGE2", nt.params[1])),
                                TTensorII(DFT(nt.params[1][1], nt.params[2]), nt.params[1]{[2,3]}, ACubeRot_ZXY, ACubeRot_XYZ)
                            ]).withTags(nt.getTags()) ]],

        apply := (t, C, Nonterms) -> C[1],
        switch := true
    )
));

NewRulesFor(TPrm, rec(
    TPrm_MPI := rec(
        forTransposition := false,
        children := nt -> [[  ]],
        apply := (nt, c, cnt) -> NoPull(MPIPrm(nt.params[1]))
    )
));

NewRulesFor(TTensorI, rec(
    IxA_MPI := rec(
        info := "IxA base",
        forTransposition := false,
        applicable := nt -> nt.hasTags() and let(tg := nt.getTags()[1], IsBound(tg.isParMPI) and tg.isParMPI) and IsParPar(nt.params),
        children := nt -> [[ nt.params[1].withTags(Drop(nt.getTags(), 1)) ]],
        apply := (nt, c, cnt) ->  MPITensor(c[1], nt.params[2])
    )
));

NewRulesFor(TTensorII, rec(
    batch_cuFFT_cuberot_3D := rec(
        info := "IxA base",
        forTransposition := false,
        applicable := nt -> nt.hasTags() and let(tg := nt.getTags()[1], IsBound(tg.isParMPI) and tg.isParMPI) and
            nt.params[3] = ACubeRot_ZXY and nt.params[4] = ACubeRot_XYZ and
            ObjId(nt.params[1]) = DFT,
        children := nt -> [[ ]],
        apply := (nt, c, cnt) ->  let(K := nt.params[1].params[1], k := nt.params[1].params[2], M := nt.params[2][1], N := nt.params[2][2],
                                      pg := nt.getTags()[1].params[1], p1 := pg[1], p2 := pg[2],
            MPITensor(
                CUFFTCall(
#                    Tensor(L(N*M/p1, M/p1), I(K/p2)) * Tensor(I(N), L(M*K/(p1*p2), M/p1)) * Tensor(I(p1*p2), DFT(K, k)), # this formula is only an approximation
                    TTensorII(nt.params[1], List([1..2], i-> nt.params[2][i]/nt.getTags()[1].params[1][i]), ACubeRot_ZXY, ACubeRot_XYZ),
                    rec(K := K, M := M, N := N, p1 := p1, p2 := p2, k := k)),
                Product(pg)))
    )
));


# Copyright (c) 2018-2021, Carnegie Mellon University
# See LICENSE for details

########################################
# These functions work on formula level
########################################

MarkBlock := function(sums) sums.isBlock := true; return sums; end;

IsMarkedBlock := o -> IsRec(o) and IsBound(o.isBlock) and o.isBlock;

DoNotMark := s -> (IsBound(s.doNotMarkBB) and s.doNotMarkBB = true) or
                  (IsIterative(s) and IsSymbolic(s.domain));

IsBlockTransitive := o -> IsBound(o.isBlockTransitive) and o.isBlockTransitive;

Compose.isBlockTransitive := true;
SUM.isBlockTransitive := true;
SUMAcc.isBlockTransitive := true;
Data.isBlockTransitive := true;

markBlocks := function(bk, sums, criterion_func)
    local doNotMark, c;
    if not IsSPL(sums) then return sums; fi;

    #Error(">>>>>>>------", ObjId(sums), "------>>>>>>>>>");
    doNotMark := false;

    for c in sums.rChildren() do
        markBlocks(bk, c, criterion_func);
        # if any of our children cannot be marked BB, we cannot be marked BB
        if IsSPL(c) and DoNotMark(c) then
          doNotMark := true;

          sums.doNotMarkBB := true; # To communicate this up to parent.
          # This is a slight HACK. Maybe we should mark these with a different
          # label (like sums.childIsNoMarkBB) and check for this, so we don't confuse
          # these with preset .doNotMarkBB

          #Error("Not marking: ", sums.__name__, " because of child: ", c.__name__);
        fi;
    od;

    if DoNotMark(sums) then
      doNotMark := true;
      #Error("Not marking ", sums.__name__, " because of doNotMarkBB");
    fi;

    #Error("<<<<<<<------", ObjId(sums), "------<<<<<<<<<");

    if not doNotMark then
        if IsBlockTransitive(ObjId(sums)) then
            if ForAll(sums.children(), IsMarkedBlock) then
               MarkBlock(sums);
            fi;
        elif criterion_func(sums, bk) then
            MarkBlock(sums);
        fi;
    fi;

    return sums;

end;

standard_criterion_func := (sums, bk) -> not IsBound(sums._perm)
    and not (IsIterative(sums) and IsSymbolic(sums.domain))
    and not ForAll(Flat(sums.dims()), IsSymbolic)
    and     ForAny(Flat(sums.dims()), d -> d<=bk);

MarkBlocksDimSums := (bk,sums) -> SubstBottomUp(sums, @,
    e -> Cond(not IsSPL(e) or (IsBound(e.doNotMarkBB) and e.doNotMarkBB), e,
          not IsBound(e._perm) and not (IsIterative(e) and IsSymbolic(e.domain)) and
             (((not IsSymbolic(Cols(e))) and Cols(e)<=bk) or # finish this
              ((not IsSymbolic(Rows(e))) and Rows(e)<=bk)),    MarkBlock(e),
          IsBlockTransitive(ObjId(e)) and ForAll(e.children(), IsMarkedBlock), MarkBlock(e),
          e));

MarkBlocksAreaSums := (bk,sums) -> SubstBottomUpRules(sums, rec(
    blockAreaSums := Rule(@,
        e -> Cond(not IsSPL(e) or not IsBound(e.area),    e,
              not IsSymbolic(e.area()) and e.area() <= bk,  MarkBlock(e),
              IsBlockTransitive(ObjId(e)) and ForAll(e.children(), IsMarkedBlock), MarkBlock(e),
              e))));

MarkBlocksOps := (bk,sums) -> SubstBottomUp(sums, @,
    e -> Cond(not IsSPL(e), e,
          not IsBound(e._perm) and e.numops()<=bk,    MarkBlock(e),e));

##
#MarkBlocksSums := MarkBlocksDimSums;
MarkBlocksSums := (bk, sums) -> markBlocks(bk, sums, standard_criterion_func);
##

Class(PostProcessBBs, RuleSet);
RewriteRules(PostProcessBBs, rec(
    mergeBBsinsideCompose := ARule(Compose, [@(1, BB), @(2, BB)], e -> [BB(Compose(@(1).val.rChildren()[1], @(2).val.rChildren()[1]))]),
    kickOutNestedBBs := Rule( @@(1, BB, (x, cx) -> cx.isInside(BB)), e -> e.child(1)),
));

#F BlockSums(<bksize>, <sums>)
#F     Uses MarkBlocksSums with given block size <bksize> to partition the
#F     Sigma-SPL formula <sums> into unrolled blocks.
#F
#F     MarkBlocksSums = MarkBlocksDimSums || MarkBlocksAreaSums.
#F     ..DimSums partitions by the smallest dimensions
#F     ..AreaSums partitions by the total dense area of a (potentially sparse) matrix
#F
BlockSums := function(bk, sums)
    sums := MarkBlocksSums(bk, sums);
    sums := SubstTopDownNR(sums, @.cond(IsMarkedBlock), e->When(ObjId(e)=BB, e, BB(e)));
    return sums;
end;

BlockSumsOpts := function(sums, opts)
    if IsBound(opts.markBlock) then
        sums := opts.markBlock(opts.globalUnrolling, sums);
    else
        sums := MarkBlocksSums(opts.globalUnrolling, sums);
    fi;
    sums := SubstTopDownNR(sums, @.cond(IsMarkedBlock), e->When(ObjId(e)=BB, e, BB(e)));
    return PostProcessBBs(sums);
end;

# Alternative implementation can use rules
#RuleSet(RulesBlockSums);
#Rule(@.cond(IsMarkedBlock), e -> BB(e));
#ARule(Compose, [..., @(1, BB), @(2, BB), ...], e -> BB(Compose(@(1).val, @(2).val));
#ARule(SUM,     [..., @(1, BB), @(2, BB), ...], e -> BB(SUM    (@(1).val, @(2).val)));
#ARule(SUMAcc,  [..., @(1, BB), @(2, BB), ...], e -> BB(SUMAcc (@(1).val, @(2).val)));

####################################
# These functions work on code level
####################################

#F MarkForUnrolling(<bbnum>, <code>, [<opts>])
MarkForUnrolling := (arg) -> Cond(
    # if has options and no SimFlexKernelFlag then go without kern()
    Length(arg)=3 and not IsBound(arg[3].SimFlexKernelFlag),
       unroll_cmd(arg[2]),
    # else
       kern(arg[1], unroll_cmd(arg[2])));

IsMarkedForUnrolling := code -> ObjId(code) = unroll_cmd;

#F BlockCollect(<code>)
#F
#F Returns a list of blocks marked for unrolling (field doUnroll:=true must be set)
#F
BlockCollect := c -> Collect(c, @.cond(IsMarkedForUnrolling));

Declare(_BlockUnroll);

LAST_BLOCKUNROLL := 0;
LAST_BLOCKUNROLLOPTS := 0;
#F BlockUnroll(<code>, <opts>)
#F
#F Generates unrolled code for marked blocks (field doUnroll:=true must be set)
#F
BlockUnroll := function(c, opts)
    LAST_BLOCKUNROLL := Copy(c);
    LAST_BLOCKUNROLLOPTS := opts;
    return _BlockUnroll(c, opts, 1);
end;

_BlockUnroll := function(c, opts, ldepth)
    local i, last, ch, newch, tocompile;

    if IsMarkedForUnrolling(c) then return Compile(c.cmds[1], opts); # strip unroll_cmd
    elif not IsCommand(c) then return c;
    else
    # a chain with some subcommands to be unrolled
    if ObjId(c) = chain and ForAny(c.cmds, IsMarkedForUnrolling) then
        newch := [];
        i := 1;
        # group each sequence of marked subcommands and unroll together
        # to facilitate copy propagation in-between
        while i <= Length(c.cmds) do
            ch := c.cmds[i];
           # recurse on blocks that aren't marked
            if not IsMarkedForUnrolling(ch) then
                Add(newch, _BlockUnroll(ch, opts, ldepth));
                # otherwise find maximal sequence of marked blocks and unroll them together
            else
                last := i;
                tocompile := [ch.cmds[1]];
                while last < Length(c.cmds) and IsMarkedForUnrolling(c.cmds[last+1]) do
                    last := last+1;
                    ch := c.cmds[last];
                    Add(tocompile, ch.cmds[1]); # strip unroll_cmd
                od;
                Add(newch, Compile(chain(tocompile), opts));
                i := last;
            fi;
            i := i+1;
        od;
        return chain(newch);
    # for an unmarked command (that is not a chain) simply recurse on subcommands
    else
        # mark loop variables with their loop nesting depth
        if ObjId(c) in [loop, loopn] then
            c.var.ldepth := ldepth;
            ldepth := ldepth + 1;
        fi;

        i := 1;
        for ch in c.rChildren() do
            c.rSetChild(i, _BlockUnroll(ch, opts, ldepth));
            i := i+1;
        od;
        return c;
    fi;
    fi;
end;

# enumerate the blocks. useful for when you want to turn blocks on and off at
# the code level.

EnumBlocks := function(sums, opts)
    local blocks, numbered, notnumb, max, i;
    blocks := Collect(sums, BB);

    numbered := Filtered(blocks, e -> IsBound(e.bbnum));
    notnumb := Filtered(blocks, e -> not IsBound(e.bbnum));

    if numbered <> [] then
        max := Maximum(List(numbered, e -> e.bbnum));
    else
        max := 0;
    fi;

    for i in notnumb do
        max := max + 1;
        i.bbnum := max;
    od;

    return sums;
end;