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

TempArrayType := (child, y, x, index) ->
    When( IsBound(child.a.t_in), child.a.t_in[index], let(
          X := Flat([x])[1], # without type unification everything has same type
          When(X.t.t=TComplex or Flat([y])[1].t.t=TComplex, TComplex, X.t.t)
));

TempArray := (y, x, child) -> let(
    cols := Flat([Cols(child)]),
    StripList( List( [1..Length(cols)], i ->
        TempVec( TArray( TempArrayType(child, y, x, i), cols[i]))
)));

# HACKed in. This should just be rolled into TempArray or otherwise be handled cleanly.
TempArraySeq := function(y, x, child)
    local dims, P;
    if ObjId(child)=BB then
      dims := child.child(1).dims();
      P    := child.child(1).P;
    else
      dims := child.dims();
      P    := child.P;
    fi;

    return(
        StripList(List(
          Flat([ dims[2]/P ]), z ->
            TempVec(
                #TArray(TempArrayType(child, Flat([y])[1], Flat([x])[1]), z))
                TArray(TempArrayType(child, y, x, 1), z))
            )));
end;

#F _CodeSums(<sums-spl>, <y-output-var>, <x-input-var>)
#F
#F Generates unoptimize loop code for <sums-spl>
#F
_CodeSums := function(sums, y, x)
    local code;
    code := sums.code(y,x);
    code.dimensions := sums.dimensions;
    code.root := sums;
    return code;
end;

#CodeSums := _CodeSums;

_CodeSumsAcc := function(sums, y, x)
    local code;
    code := _CodeSums(sums, y, x);
    code := SubstTopDownNR(code, [assign, [@(0,nth), @(1, var, e->Same(e,y)), ...], @(3)],
                     e -> assign_acc(@(0).val, @(3).val));
    return code;
end;

# =======================================================================================
# Code generation templates
# [ Data, Blk, Blk1, Prm, R, W, SUM, ISum, Compose ]
# =======================================================================================
Formula.code := (self,y,x) >> self.child(1).code(y, x);

BB.code := (self,y,x) >> MarkForUnrolling(_CodeSums(self.child(1), y, x));

Buf.code := (self,y,x) >> _CodeSums(self.child(1), y, x);

COND.code := (self,y,x) >>
    IF(self.cond,
    _CodeSums(self.child(1), y, x),
    _CodeSums(self.child(2), y, x));

Data.code := meth(self, y, x)
    local val, t;
    if self.inline then
    # we set live_out to prevent copy propagation
    # self.var.setAttr("live_out")
    return decl(self.var, chain(
        assign(self.var, self.value.tolist()),
        _CodeSums(self.child(1), y, x)));
    else
    val := self.value;
    if IsBound(val.tolist) then
        if IsBound(val.range) and val.range() <> false then
        t := TArray(val.range(), val.domain());
        val := t.value(val.tolist());
        else
        val := V(val.tolist());
        fi;
    else val := Cond(IsSPL(val), V(MatSPL(val)), val);
    fi;
    return data(self.var, val, _CodeSums(self.child(1), y, x));
    fi;
end;

Diag.code := meth(self, y, x)
    local i, elt;
    i := Ind();
    elt := self.element.lambda().at(i);
    return loop(i, self.element.domain(), assign(nth(y,i), elt * nth(x,i)));
end;

RCDiag.code := meth(self, y, x)
    local i, elt, re, im;
    i := Ind();
    elt := self.element.lambda();
    re := elt.at(2*i);
    im := elt.at(2*i+1);
    return loop(i, self.element.domain()/2, chain(
	    assign(nth(y,2*i),   re * nth(x,2*i) - im * nth(x,2*i+1)),
	    assign(nth(y,2*i+1), im * nth(x,2*i) + re * nth(x,2*i+1))));
end;

ColVec.code := meth(self, y, x)
    local i, datavar, diag, func;
    i := Ind();
    func := self.element;
    if IsBound(self.inline) and self.inline then
    return loop(i, Rows(self), assign(nth(y,i), mul(func.lambda().at(i), nth(x,0))));
    else
    diag := V(func.lambda().tolist());
    datavar := Dat(diag.t);
    return
    data(datavar, diag,
        loop(i, Rows(self), assign(nth(y,i), mul(nth(datavar,i), nth(x,0)))));
    fi;
end;

RowVec.code := meth(self, y, x)
    local i, datavar, diag, func, t;
    i := Ind();
    t := TempVar(x.t.t);
    func := self.element;
    if IsBound(self.inline) and self.inline then
    return chain(
        assign(t,0),
        loop(i, Cols(self), assign(t, add(t, mul(func.lambda().at(i), nth(x,i))))),
        assign(nth(y,0), t));
    else
    diag := V(func.lambda().tolist());
    datavar := Dat(diag.t);
    return
    data(datavar, diag,
        chain(
        assign(t,0),
        loop(i, Cols(self), assign(t, add(t, mul(nth(datavar,i), nth(x,i))))),
        assign(nth(y,0), t)));
    fi;
end;

Scale.code := (self, y, x) >> let(i := Ind(),
    chain(
	_CodeSums(self.child(1), y, x),
	loop(i, Rows(self),
	    assign(nth(y,i), mul(self.scalar, nth(y,i))))));

I.code := (self, y, x) >> let(i := Ind(),
    loop(i, Rows(self), assign(nth(y,i), nth(x,i))));

Blk2code := (blk, y, x) -> let(b := V(blk.element).v,
    t0 := TempVar(x.t.t), t1 := TempVar(x.t.t),
    decl([t0, t1], chain(
	    assign(t0, mul(b[1].v[1], nth(x,0)) + mul(b[1].v[2], nth(x,1))),
	    assign(t1, mul(b[2].v[1], nth(x,0)) + mul(b[2].v[2], nth(x,1))),
	    assign(nth(y,0), t0),
	    assign(nth(y,1), t1)
	)));

Blk4code := (blk, y, x) -> let(b := V(blk.element).v,
    t0 := TempVar(x.t.t), t1 := TempVar(x.t.t),
    t2 := TempVar(x.t.t), t3 := TempVar(x.t.t),

    x0 := TempVar(x.t.t), x1 := TempVar(x.t.t),
    x2 := TempVar(x.t.t), x3 := TempVar(x.t.t),
    decl([x0,x1,x2,x3, t0, t1, t2, t3], chain(
	    assign(x0, nth(x,0)),
	    assign(x1, nth(x,1)),
	    assign(x2, nth(x,2)),
	    assign(x3, nth(x,3)),

	    assign(t0, mul(b[1].v[1], x0) + mul(b[1].v[2], x1) +
		       mul(b[1].v[3], x2) + mul(b[1].v[4], x3)),
	    assign(t1, mul(b[2].v[1], x0) + mul(b[2].v[2], x1) +
		       mul(b[2].v[3], x2) + mul(b[2].v[4], x3)),
	    assign(t2, mul(b[3].v[1], x0) + mul(b[3].v[2], x1) +
		       mul(b[3].v[3], x2) + mul(b[3].v[4], x3)),
	    assign(t3, mul(b[4].v[1], x0) + mul(b[4].v[2], x1) +
		       mul(b[4].v[3], x2) + mul(b[4].v[4], x3)),
	    assign(nth(y,0), t0),
	    assign(nth(y,1), t1),
	    assign(nth(y,2), t2),
	    assign(nth(y,3), t3)
    )));

Blk.code := meth(self, y, x)
    local i, j, t, datavar, matval;
    if Rows(self)=2 and Cols(self)=2 then return Blk2code(self, y, x); fi;
    if Rows(self)=4 and Cols(self)=4 then return Blk4code(self, y, x); fi;
    j := Ind();
    i := Ind();
    t := TempVar(x.t.t);
    matval := V(self.element);
    datavar := Dat(matval.t);
    return
    data(datavar, matval,
    loop(j, Rows(self),
        decl(t,
        chain(
            assign(t, 0),
            loop(i, Cols(self),
            assign(t, add(t, mul(nth(nth(datavar, j), i), nth(x, i))))),
            assign(nth(y,j), t)))));
end;

toeplitz.code := meth(self, y, x)
    local i, j, t, datavar, matval, dataload;
    j := Ind();
    i := Ind();
    t := TempVar(x.t.t);
    matval := V(self.obj.element);
    datavar := Dat(matval.t);
    dataload := TempVec(TArray(TDouble, Rows(self)*Cols(self)));
    return
    data(datavar, matval, decl(dataload, chain(
    loop(j, Rows(self),
        loop(i, Cols(self),
        assign(nth(dataload, j*Cols(self)+i), nth(nth(datavar,j),i)))),

    loop(j, Rows(self),
        decl(t,
        chain(
            assign(t, 0),
            loop(i, Cols(self),
            assign(t, add(t, mul(nth(dataload, j*Cols(self)+i), nth(x, i))))),
            assign(nth(y,j), t)))))));
end;

Blk1.code := (self, y, x) >>
   assign(nth(y,0), mul(toExpArg(self.element), nth(x,0)));

BlkConj.code := (self, y, x) >>
   assign(nth(y,0), conj(nth(x,0)));

Prm.code := (self, y, x) >> let(i := Ind(), func := self.func.lambda(),
    loop(i, Rows(self),
    assign(nth(y, i),
           nth(x, func.at(i)))));

Z.code := (self, y, x) >> let(i := Ind(), ii := Ind(), z := self.params[2],
    chain(
    loop(i, Rows(self)-z,
        assign(nth(y, i), nth(x, i+z))),
    loop(i, z,
        assign(nth(y, i+(Rows(self)-z)), nth(x, i)))));

O.code := (self, y, x) >> let(i := Ind(),
    loop(i, self.params[1],
        assign(nth(y, i), V(0))));

Gath.code := meth(self, y, x)
    local i, func, rfunc, ix;
    i := Ind(); func := self.func.lambda();

    if IsBound(self.func.rlambda) then
	rfunc := self.func.rlambda();
	ix := var.fresh_t("ix", TInt);
	return chain(
	    assign(ix, func.at(0)),
	    assign(nth(y, 0), nth(x, ix)),
	    loop(i, self.func.domain()-1,
		chain(assign(ix, rfunc.at(ix)),
		      assign(nth(y, i+1), nth(x, ix)))));
    else
	return loop(i, self.func.domain(), assign(nth(y,i), nth(x, func.at(i))));
    fi;
end;

Scat.code := meth(self, y, x)
    local i, func, rfunc, ix;
    i := Ind(); func := self.func.lambda();
    if IsBound(self.func.rlambda) then
	rfunc := self.func.rlambda();
	ix := var.fresh_t("ix", TInt);
	return chain(
	    assign(ix, func.at(0)),
	    assign(nth(y, ix), nth(x, 0)),
	    loop(i, self.func.domain()-1,
		chain(assign(ix, rfunc.at(ix)),
		      assign(nth(y, ix), nth(x, i+1)))));
    else
	return loop(i, self.func.domain(), assign(nth(y,func.at(i)), nth(x, i)));
    fi;
end;

SUM.code := (self, y, x) >>
       chain(List(self.children(), c -> _CodeSums(c, y, x)));

SUMAcc.code := (self, y, x) >>
   let(ii := Ind(),
       When(not Same(ObjId(self.child(1)), Gath),  # NOTE: come up with a general condition
       chain(
           loop(ii, Rows(self), assign(nth(y, ii), V(0))),
           List(self.children(), c -> _CodeSumsAcc(c, y, x))),
       chain(
           _CodeSums(self.child(1), y, x),
           List(Drop(self.children(), 1), c -> _CodeSumsAcc(c, y, x)))));

ISum.code := (self, y, x) >> let(myloop := When(IsSymbolic(self.domain), loopn, loop),
   myloop(self.var, self.domain,
        _CodeSums(self.child(1), y, x)));

ISumAcc.code := (self, y, x) >>
   let(ii := Ind(),
       chain(
       loop(ii, Rows(self),
           assign(nth(y, ii), V(0))),
       loop(self.var,
            self.domain,
        _CodeSumsAcc(self.child(1), y, x))));

Compose.code := meth(self,y,x)
    local ch, numch, vecs, allow, i;
#VIENNA filtering DMPGath, DMPScat out here because they do not generate code
#    ch := self.children();
    ch := Filtered(self.children(), i-> not i.name in ["DMPGath","DMPScat"]);
    numch := Length(ch);
    vecs := [y];
    allow := (x<>y);
    for i in [1..numch-1] do
        if allow and ObjId(ch[i])=Inplace then vecs[i+1] := vecs[i];
	else vecs[i+1] := TempVec(TArray(TempArrayType(y, x), Cols(ch[i])));
	fi;
    od;
    vecs[numch+1] := x;
    for i in Reversed([1..numch]) do
        if allow and ObjId(ch[i])=Inplace then vecs[i] := vecs[i+1];
	fi;
    od;

#Print(vecs, "\n");
    if vecs[1] = vecs[numch+1] then # everything was inplace
	vecs[1] := y;
    fi;
    if vecs[1] = vecs[numch+1] then # everything was inplace
	vecs[numch+1] := x;
    fi;
#Print(vecs, "\n");
    vecs := Reversed(vecs);
    ch := Reversed(ch);
    return decl( Difference(vecs{[2..Length(vecs)-1]}, [x,y]),
	chain( List([1..numch], i -> When(vecs[i+1]=vecs[i],
		    compiler._IPCodeSums(ch[i], vecs[i]),
		    _CodeSums(ch[i], vecs[i+1], vecs[i])))));
end;

ICompose.code := (self, y, x) >>

#DoneSpiralHelp - How to make this work for any datatype (var t)
#Look into Compose.code for tempvec

  chain([let(
    t :=  Dat1d(x.t.t, Rows(self)),
    newind := Ind(self.domain),
    decl([t], chain(
       When(IsOddInt(self.domain),
       SubstVars(Copy(_CodeSums(self.child(1), y, x)), tab((self.var.id) := (V(0)))),
          skip()
       ),
       When(IsEvenInt(self.domain),
          chain(
          SubstVars(Copy(_CodeSums(self.child(1), t, x)), tab((self.var.id) := (V(0)))),
          SubstVars(Copy(_CodeSums(self.child(1), y, t)), tab((self.var.id) := (V(1))))
          ),
          skip()
       ),
      loop(newind, (self.domain/2),
        chain([
        SubstVars(Copy(_CodeSums(self.child(1), x, y)), tab((self.var.id) := (2*newind)+2)),
        SubstVars(Copy(_CodeSums(self.child(1), y, x)), tab((self.var.id) := (2*newind)+3))
        ])
      )
    ))
  )]);


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

Class(LegacySumsGen, HierarchicalVisitor, rec(

   __call__ := meth(arg)

        Constraint(Length(arg) >= 2);

        #BWD: I believe LegacySumsGen.visit only expects Formula objects
        if not PatternMatch(arg[2], Formula, empty_cx()) then
            arg[2] := Formula(arg[2]);
        fi;

        return ApplyFunc(arg[1].visit, arg{[2..Length(arg)]});

   end,

	Formula := meth(self, o, opts)
		local spl, sums;
		spl := _children(o)[1];

    	if IsBound(spl.memo_sums) then
    		sums := spl.memo_sums;
    	else
    	sums := spl.sums();
    		spl._sums := sums;
    	fi;
    	sums.root := spl;
    	#sums.sums := self >> self;
    	return sums;
	end,
));

Class(DefaultSumsGen, HierarchicalVisitor, rec(

    __call__ := meth(arg)
        local res;
        res := ApplyFunc(arg[1].visit, arg{[2..Length(arg)]});
        trace_log.addExpansion(ObjId(arg[2]), arg[2], res, var);
        return res;
   end,

    _directSumSums := (self, obj, roverlap, coverlap, opts) >>
        let(cspans  :=  BaseOverlap.spans(roverlap, List(obj.children(), Cols)),
            rspans  :=  BaseOverlap.spans(coverlap, List(obj.children(), Rows)),
            nblocks  :=  obj.numChildren(),
            C  :=  Cols(obj),
            R  :=  Rows(obj),
            OP  :=  When(coverlap > 0, SUMAcc, SUM),
            OP( Map([1..nblocks],
                 i -> Compose( #W(N, bksize_cols, fTensor(fBase(nblocks, ind), fId(bksize)))
                      WID(R, 1 + rspans[i][2] - rspans[i][1], rspans[i][1]-1),
                      self(obj.child(i), opts),
                      RID(C, 1 + cspans[i][2] - cspans[i][1], cspans[i][1]-1) )))),

    Formula     := (self, o, opts) >> self(_children(o)[1], opts),
    NonTerminal := (self, o, opts) >> o,
    PermClass   := (self, o, opts) >> Prm(o),

    Rot := (self, o, opts) >> self(ExpandRotationsSPL(o), opts),
    Sym := (self, o, opts) >> self(o.getObj(), opts),

    O  := (self, o, opts) >> o,
    L  := (self, o, opts) >> Prm(o),
    F  := (self, o, opts) >> self(o.getObj(), opts),
    I  := (self, o, opts) >> Prm(fId(o.params[1])),

    RC := (self, o, opts) >> CopyFields(o, rec(_children := [self(o.child(1), opts)])),

    Perm := (self, o, opts) >> let(
	N := o.size,
	plist_r := Permuted([1..N], o.element^-1) - 1, # we want 0-based perms
	#plist_w := Permuted([1..N], o.element) - 1,    # we want 0-based perms
	Gath(FList(N, plist_r).setRange(N))
    ),

    Diag := (self, o, opts) >> o,
    RCDiag := (self, o, opts) >> o,
    Mat := (self, o, opts) >> Blk(o.element),
    Sparse := (self, o, opts) >> Blk(MatSparseSPL(o)),

    B := (self, o, opts) >> Prm(fB(o.size, o.l)),
    B2 := (self, o, opts) >> Prm(o),
    Scale := (self, o, opts) >> Inherit(o, rec(_children := [self(o.child(1), opts)])),

    Conjugate := (self, o, opts) >> let(
	O := o.child(1), perm := o.child(2), Tperm := perm.transpose(),
	self(Compose(Tperm, O, perm), opts)
    ),

    Compose := (self, o, opts) >> Cond(o.isPermutation(),
	Prm(ApplyFunc(fCompose, List(Reversed(o.children()), c -> self(c, opts).func))),
	Compose(Map(o.children(), c -> self(c, opts)))
    ),

    Cross := (self, o, opts) >> Cond(o.isPermutation(),
	Prm(ApplyFunc(fCross, List(o.children(), c -> self(c, opts).func))),
	Cross(Map(o.children(), c -> self(c, opts)))
    ),

    ICompose := (self, o, opts) >> CopyFields(o, rec(
        _children := List(o.children(), c -> self(c, opts))
    )),

    DirectSum := (self, o, opts) >> Cond(o.isPermutation(),
	Prm(ApplyFunc(fDirsum, List(o.children(), c->self(c, opts).func))),
        self._directSumSums(o, 0, 0, opts)
    ),

    DelayedDirectSum := (self, o, opts) >> CopyFields(o, rec(
        _children := List(o.children(), c -> self(c, opts))
    )),

    RowDirectSum := (self, o, opts) >> self._directSumSums(o, o.overlap, 0, opts),
    ColDirectSum := (self, o, opts) >> self._directSumSums(o, 0, o.overlap, opts),

    RowTensor := (self, o, opts) >> let(
        A := o.child(1),
	i := Ind(o.isize),
	ISum(i, i.range,
	    Scat(fTensor(fBase(o.isize, i), fId(Rows(A)))) *
	    SumsSPL(A, opts) *
	    Gath(fAdd(Cols(o), Cols(A), i * (Cols(A) - o.overlap))))
    ),

    ColTensor := (self, o, opts) >> let(
        A := o.child(1),
	i := Ind(o.isize),
	ISumAcc(i, i.range,
	    Scat(fAdd(Rows(o), Rows(A), i * (Rows(A) - o.overlap))) *
	    SumsSPL(A, opts) *
	    Gath(fTensor(fBase(o.isize, i), fId(Cols(A)))))
    ),

    Tensor := meth(self, o, opts)
    	local ch, col_prods, row_prods, col_prods_rev, row_prods_rev, i, j1, j2, j1b, j2b, bkcols, bkrows, prod, term;

    	if o.isPermutation() then
	    return Prm(ApplyFunc(o.fTensor, List(o.children(), c->self(c, opts).func)));
	elif ForAll(o.children(), x->ObjId(x)=Diag) then
	    return Diag(ApplyFunc(diagTensor, List(o.children(), c->c.element)));
	fi;

	ch := o.children();
	col_prods := ScanL(ch, (x,y)->x*Cols(y), 1);
	col_prods_rev := Drop(ScanR(ch, (x,y)->x*Cols(y), 1), 1);

	row_prods := ScanL(ch, (x,y)->x*Rows(y), 1);
	row_prods_rev := Drop(ScanR(ch, (x,y)->x*Rows(y), 1), 1);

	prod := [];
	for i in [1..Length(ch)] do
	    if not IsIdentitySPL(ch[i]) then
		bkcols := Cols(ch[i]);
		bkrows := Rows(ch[i]);
		j1 := Ind(col_prods[i]);
		j2 := Ind(row_prods_rev[i]);
		j1b := When(j1.range = 1, 0, fBase(j1));
		j2b := When(j2.range = 1, 0, fBase(j2));

		term := Scat(o.scatTensor(i)(Filtered([j1b, fId(bkrows), j2b], x->x<>0))) *
		        self(ch[i], opts) *
			Gath(o.gathTensor(i)(Filtered([j1b, fId(bkcols), j2b], x->x<>0)));

		if j2.range <> 1 then
		    term := ISum(j2, j2.range, term);
		fi;

		if j1.range <> 1 then
		    term := ISum(j1, j1.range, term);
		fi;

		Add(prod, term);
	    fi;
	od;
	return Compose(prod);
    end,

    IDirSum := (self, o, opts) >> let(
	bkcols := Cols(o.child(1)),
	bkrows := Rows(o.child(1)),
	nblocks := o.domain,
	cols := Cols(o),
	rows := Rows(o),

	ISum(o.var, o.domain,
	    Compose(Scat(fTensor(fBase(nblocks, o.var), fId(bkrows))),
		    self(o.child(1), opts),
		    Gath(fTensor(fBase(nblocks, o.var), fId(bkcols)))))
    ),

    HStack := (self, o, opts) >> let(
	ch   := List(o.children(), e -> self(e, opts)),
	l    := Length(ch),
	it0  := Scat(fId(Rows(o))) * ch[1] * Gath(H(Cols(o), Cols(ch[1]), 0, 1)),
	ofs  := ScanL(ch, (p,x)->p+Cols(x), 0),
	its  := List([2..l], i ->
	           ScatAcc(fId(Rows(o))) * ch[i] * Gath(H(Cols(o), Cols(ch[i]), ofs[i], 1))),
	SUM( [it0] :: its )
    ),

    VStack := (self, o, opts) >> let(
	ch   := List(o.children(), e -> self(e, opts)),
	l    := Length(ch),
	ofs  := ScanL(ch, (p,x)->p+Rows(x), 0),
	SUM(List([1..l], i ->
	        Scat(H(Rows(o), Rows(ch[i]), ofs[i], 1)) * ch[i] * Gath(fId(Cols(o)))))
    ),

    HStack1 := (self, o, opts) >> let(
    	ch   := List(o.children(), e -> self(e, opts)),
    	l    := Length(ch),
    	it0  := Scat(fId(Rows(o))) * ch[1] * Gath(H(Cols(o), Cols(ch[1]), 0, 1)),
    	ofs  := ScanL(ch, (p,x)->p+Cols(x), 0),
    	its  := List([2..l], i ->
    	           Scat(fId(Rows(o))) * ch[i] * Gath(H(Cols(o), Cols(ch[i]), ofs[i], 1))),
    	SUM( [it0] :: its )
    ),

    VStack := (self, o, opts) >> let(
	ch   := List(o.children(), e -> self(e, opts)),
	l    := Length(ch),
	ofs  := ScanL(ch, (p,x)->p+Rows(x), 0),
	SUM(List([1..l], i ->
	        Scat(H(Rows(o), Rows(ch[i]), ofs[i], 1)) * ch[i] * Gath(fId(Cols(o)))))
    ),

    IterHStack := (self, o, opts) >> let(
	bkcols := Cols(o.child(1)),
	bkrows := Rows(o.child(1)),
	nblocks := o.domain,
	cols := Cols(o), rows := Rows(o),
	j := Ind(nblocks-1),
	ch := self(o.child(1), opts),

	SUM(
	    Scat(fId(bkrows)) *
	    SubstVars(Copy(ch), rec((o.var.id) := V(0))) *
	    Gath(fTensor(fBase(nblocks, 0), fId(bkcols))),

	    ISum(j, j.range,
		ScatAcc(fId(bkrows)) *
		SubstVars(Copy(ch), rec((o.var.id) := j+1)) *
		Gath(fTensor(fBase(nblocks, j+1), fId(bkcols)))
	    )
	)
    ),

    IterHStack1 := (self, o, opts) >> let(
    	bkcols := Cols(o.child(1)),
    	bkrows := Rows(o.child(1)),
    	nblocks := o.domain,
    	cols := Cols(o), rows := Rows(o),
    	j := Ind(nblocks),
    	ch := self(o.child(1), opts),

	    ISum(j, j.range,
		Scat(fId(bkrows)) *
		SubstVars(Copy(ch), rec((o.var.id) := j)) *
		Gath(fTensor(fBase(nblocks, j), fId(bkcols)))
    	)
    ),

    DPWrapper := (self, o, opts) >> self(o.child(1), opts),

    IterVStack := (self, o, opts) >> let(
	bkcols := Cols(o.child(1)),
	bkrows := Rows(o.child(1)),
	nblocks := o.domain,
	cols := Cols(o), rows := Rows(o),
	ISum(o.var, o.domain,
	    Scat(fTensor(fBase(nblocks, o.var), fId(bkrows))) *
	    self(o.child(1), opts) *
	    Gath(fId(bkcols)))
    ),

    ISum := (self, o, opts) >> CopyFields(o, rec(_children := [self(o._children[1], opts)])),

    RecursStep := ~.ISum,
    VPerm      := (self, o, opts) >> o,
    BlockVPerm := (self, o, opts) >> o,
    ExpDiag    := (self, o, opts) >> o,

    VTensor    := (self, o, opts) >> ObjId(o)(self(o.child(1), opts), o.vlen),

    SumsBase := meth(self, o, opts)
       local children, i, res;
       res := Copy(o);
       ## MRT, aug08
       # IsSPL -> IsSPL and IsFunction
       # joint SPL/Sigma objects like fId were improperly triggering the
       # SPL condition, resulting in illegal rewrites like:
       #   Scat(fId(2)) -> Scat(Prm(fId(2))

       # YSV: This is still buggy, because VTensor(L) --> VTensor(L)
       children := Map(res.rChildren(), c -> Cond(IsSPL(c) and not IsFunction(c), self(c, opts), c));
       for i in [1..Length(children)] do
           res.rSetChild(i, children[i]);
       od;
       return res;
    end,

     Split := (self, o, opts) >> o,
     Glue := (self, o, opts) >> o,
     VIxL := (self, o, opts) >> o,

));


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

Import(fftx.platforms.cuda);
Import(fftx.codegen);
Import(simt);

Class(HockneyMLC_SIMTRules, RuleSet);

# Set of rules for applying Tile and Fission SIMT flags
RewriteRules(HockneyMLC_SIMTRules, rec(
	#Enable synchronization when different thread clusters are mapped to innermost sums
	# SUM(A*[S,IS], B*[S,IS]) => SUM(A, B)*SUM([S,IS],[S,IS])
    collect := Rule([SIMTSUM, [@(1, Compose),
                                @(2).cond(e -> not ObjId(e) in [Scat]), @(3).cond(s -> IsBound(s.simt_dim))
                            ],
                            [@(6, Compose),
                                @(7).cond(e -> not ObjId(e) in [Scat]), @(8).cond(s -> IsBound(s.simt_dim))
                            ]
                    ], e -> let(dsum1 := @(3).val.dimensions[1], dsum2 := @(8).val.dimensions[1],
                                Compose(    SIMTSUM(Copy(e.simt_dim),   @(2).val*Gath(H(dsum1+dsum2, dsum1, 0, 1)),
                                                                    @(7).val*Gath(H(dsum1+dsum2, dsum2, dsum1, 1)) ),
                                            SIMTSUM(Copy(e.simt_dim),   Scat(H(dsum1+dsum2, dsum1, 0, 1))*@(3).val,
                                                                    Scat(H(dsum1+dsum2, dsum2, dsum1, 1))*@(8).val )
                                    )
                            ),
                        "CollectCompose"
    ),

    tile    := Rule(@(1, SIMTISum, e->ObjId(e.simt_dim)=ASIMTTileFlag ),
                    e -> let(d1 := e.domain/e.simt_dim.tsize(),
                             v1 := var.fresh(e.var.id{[1]}, e.var.t, d1),
                             d2 := e.simt_dim.tsize(),
                             v2 := var.fresh(e.var.id{[1]}, e.var.t, d2),
                             newb := SubstVars(e.children()[1], rec( (e.var.id) := v1*d2 + v2 ) ),
                             simt_dlst := e.simt_dim.simt_dim(),
                             SIMTISum(Copy(simt_dlst.simt_dim(1)), v1, d1,
                                SIMTISum(Copy(simt_dlst.simt_dim(2)), v2, d2, newb )
                                )
                        ),
                    "tile_SIMTISum"
    ),

    fission_SIMTISum := Rule([@(1, SIMTISum, e->ObjId(e.simt_dim)=ASIMTFissionFlag ),
                            @(2, Compose, e-> not ForAny(e.children(), cch -> ObjId(cch) in [Gath, Scat]) )
                        ], function(e)
                            local inner_comp_list, ch, cch, cchlen, pos, simt_dlst, b, newb, v, sl, gl;

                            ch := @(2).val;
                            inner_comp_list := [];
                            cchlen :=  Length(ch.children());
                            for cch in ch.children() do
                                b := Copy( cch );
                                v := e.var.clone();
                                newb := SubstVars(b, rec( (e.var.id) := v ));
                                sl := newb.dimensions[1];
                                gl := newb.dimensions[2];
                                pos := Length(inner_comp_list);
                                simt_dlst := e.simt_dim.simt_dim();
                                Add(inner_comp_list,
                                    Cond(
                                        pos = 0,        SIMTISum(Copy(simt_dlst.simt_dim(pos+1)), v, e.domain, newb*Gath(H(e.domain*gl, gl, gl*v, 1)) ),
                                        pos = cchlen-1, SIMTISum(Copy(simt_dlst.simt_dim(pos+1)), v, e.domain, Scat(H(e.domain*sl, sl, sl*v, 1))*newb ),
                                        SIMTISum(Copy(simt_dlst.simt_dim(pos+1)), v, e.domain, Scat(H(e.domain*sl, sl, sl*v, 1))*newb*Gath(H(e.domain*gl, gl, gl*v, 1)) )
                                    )
                                );
                            od;
                            e := Compose( inner_comp_list );
                            return e;
                        end,
                        "fission_SIMTISum"
    ),

    # If outer loop is mapped to cluster (warp) dimension and it contains other block-parallel-loops then loop must be fissioned
    # to enable synchronization.
    xy_fission := Rule([@(1, SIMTISum, e->ObjId(e.simt_dim)=ASIMTBlockDimX ),
                            @(2, Compose, e->   not ForAny(e.children(), cch -> ObjId(cch) in [Gath, Scat]) and
                                                ForAny(List(Collect(e, @(3,[SIMTISum, SIMTSUM]) ), s -> s.simt_dim.simt_sync()),
                                                        s -> s > @(1).val.simt_dim.simt_sync())
                            )
                        ], function(e)
                            local   inner_comp_list, sync_list,
                                    ch, cch, cchlen, b, newb, v, sl, gl;

                            inner_comp_list := Collect(e.children(),
                                                    @@(4, Compose,
                                                        (expr, cx) -> IsBound(cx.SIMTSUM) and cx.SIMTSUM <> []
                                                                    or
                                                                    IsBound(cx.SIMTISum) and cx.SIMTISum <> []
                                                                    or
                                                                    IsBound(cx.Compose) and cx.Compose <> []
                                                    )
                                                );
                            sync_list := List(Collect(inner_comp_list, @(5,[SIMTISum, SIMTSUM]) ), s -> s.simt_dim.simt_sync());
                            if ForAny(sync_list, s -> s > e.simt_dim.simt_sync()) then
                                Error("Cannot properly syncronize within following structure: ", e);
                            fi;

                            ch := @(2).val;
                            inner_comp_list := [];
                            cchlen :=  Length(ch.children());
                            for cch in ch.children() do
                                b := Copy( cch );
                                v := e.var.clone();
                                newb := SubstVars(b, rec( (e.var.id) := v ));
                                sl := newb.dimensions[1];
                                gl := newb.dimensions[2];
                                Add(inner_comp_list,
                                    Cond(
                                        Length(inner_comp_list) = 0,        SIMTISum(Copy(e.simt_dim), v, e.domain, newb*Gath(H(e.domain*gl, gl, gl*v, 1)) ),
                                        Length(inner_comp_list) = cchlen-1, SIMTISum(Copy(e.simt_dim), v, e.domain, Scat(H(e.domain*sl, sl, sl*v, 1))*newb ),
                                        SIMTISum(Copy(e.simt_dim), v, e.domain, Scat(H(e.domain*sl, sl, sl*v, 1))*newb*Gath(H(e.domain*gl, gl, gl*v, 1)) )
                                    )
                                );
                            od;
                            e := Compose( inner_comp_list );
                            return e;
                        end,
                        "xy_fission_SIMTISum"
    ),
));

HockneyMLC_fftxGen := function(t, opts)
    local sb, cb, rtb, symbl, tt, tb,
    nd3, nfreq, n2, ns3, nfreq, ns2,
    tag_prdft_pd_loop1, rtp_prdft_pd_loop1, tag_dft_pd_loop1, rtp_dft_pd_loop1, tag_prdft_pd_loop2, rtp_prdft_pd_loop2, tag_dft_pd_loop2, rtp_dft_pd_loop2,
    tag_prdft_rec1, rtp_prdft_rec1, tag_prdft_rec2, rtp_prdft_rec2, tag_prdft_rec_z, rtp_prdft_rec_z, tag_dft_rec_y, rtp_dft_rec_y, BX, tag_idft_rec_x, rtp_idft_rec_x,
    tag_dft_rec_x, rtp_dft_rec_x, tag_conv_x, rtp_conv_x, tag_pridft_rec_z, rtp_pridft_rec_z, tag_idft_rec_y, rtp_idft_rec_y, TS_stg1, TS_stg2, TS_stg543, trstg5to1;

    # promote t
    tt := opts.preProcess(t);

    # parse parameters
    symbl := tt.params[2].params[1];
    tb := tt.params[1];
    nd3 := Length(tt.params[1].params[4][1]);
    nfreq := tt.params[1].params[1][1]/2+1;
    n2 := tt.params[1].params[1][1];
    ns3 := Length(tt.params[1].params[6][1]);
    ns2 := Length(tt.params[1].params[6][1]);
    opts.symbol := [symbl];

	if IsBound(tt.params[2].fname) then
	   opts.cudasubName := tt.params[2].fname;
	fi;

    # code gen script
    rtb := RuleTreeMid(tb, opts);

    ## Pruned Real DFT vvvvv
    # rec1: PRDFT_PD_loop presents 4 parallelizable objects
    tag_prdft_pd_loop1 := [ASIMTBlockDimY(32,[0,0],[0,31]), ASIMTLoopDim(), ASIMTBlockDimY(32, [1,6]), ASIMTLoopDim()];
    # rec1: PRDFT_PD_loop produces 0 subrules
    rtp_prdft_pd_loop1 := tag_rec(tag_prdft_pd_loop1, Ignore);

    # rec1: DFT_PD_loop presents 3 parallelizable objects
    tag_dft_pd_loop1 := [ASIMTBlockDimY(32,[7,7],[0,31]), ASIMTBlockDimY(32, [8,14]), ASIMTLoopDim()];
    # rec1: DFT_PD_loop produces 0 subrules
    rtp_dft_pd_loop1 := tag_rec(tag_dft_pd_loop1, Ignore);

    # rec2: PRDFT_PD_loop presents 4 parallelizable objects
    tag_prdft_pd_loop2 := [ASIMTBlockDimY(32,[15,15],[0,31]), ASIMTLoopDim(), ASIMTBlockDimY(32, [16,21]), ASIMTLoopDim()];
    # rec2: PRDFT_PD_loop produces 0 subrules
    rtp_prdft_pd_loop2 := tag_rec(tag_prdft_pd_loop2, Ignore);

    # rec2: DFT_PD_loop presents 3 parallelizable objects
    tag_dft_pd_loop2 := [ASIMTBlockDimY(32,[22,22],[0,31]), ASIMTBlockDimY(32, [23,29]), ASIMTLoopDim()];
    # rec2: DFT_PD_loop produces 0 subrules
    rtp_dft_pd_loop2 := tag_rec(tag_dft_pd_loop2, Ignore);

    tag_prdft_rec1 := [ASIMTBlockDimY(32), ASIMTLoopDim(), ASIMTBlockDimY(32), ASIMTBlockDimY(32, [0,3]), ASIMTBlockDimY(32, [4,12])];
    rtp_prdft_rec1 := tag_rec(tag_prdft_rec1, [rtp_prdft_pd_loop1, rtp_dft_pd_loop1] :: Replicate(2, Ignore));

    tag_prdft_rec2 := [ASIMTBlockDimY(32), ASIMTLoopDim(), ASIMTBlockDimY(32), ASIMTBlockDimY(32, [13,15]), ASIMTBlockDimY(32, [16,25])];
    rtp_prdft_rec2 := tag_rec(tag_prdft_rec2, [rtp_prdft_pd_loop2, rtp_dft_pd_loop2] :: Replicate(2, Ignore));

    #Top prdft rec rule presents 3 parallelizable objects
    tag_prdft_rec_z := [ASIMTBlockDimX(16,[0,0],[0,15]), ASIMTBlockDimX(16), ASIMTBlockDimY(32)];
    #Top prdft rec rule produces 4 subrules
    rtp_prdft_rec_z := tag_rec(tag_prdft_rec_z, [Ignore, Ignore, rtp_prdft_rec1, rtp_prdft_rec2]);

    ## Pruned Real DFT ^^^^^

    ## Pruned DFT vvvvv

    #Top dft rec rule presents 4 parallelizable objects
    tag_dft_rec_y := [ASIMTBlockDimY(13, [0,9]), ASIMTBlockDimY(13), ASIMTBlockDimY(13, [0,6]), ASIMTBlockDimY(13, [7,12])];
    #Top dft rec rule produces 3 subrules
    rtp_dft_rec_y := tag_rec(tag_dft_rec_y, Replicate(3, Ignore));

    ## Pruned DFT ^^^^^

    ## Pruned Conv vvvvv
    BX := 16;
    #Top idft rec rule presents 4 parallelizable objects
    tag_idft_rec_x := [ASIMTBlockDimX(BX), ASIMTBlockDimX(BX, [0,7]), ASIMTBlockDimX(BX, [8,12]), ASIMTBlockDimX(BX, [0,9])];
    #Top idft rec rule produces 3 subrules
    rtp_idft_rec_x := tag_rec(tag_idft_rec_x, Replicate(3, Ignore));

    #Top dft rec rule presents 4 parallelizable objects
    tag_dft_rec_x := [ASIMTBlockDimX(BX, [0,9]), ASIMTBlockDimX(BX), ASIMTBlockDimX(BX, [0,6]), ASIMTBlockDimX(BX, [7,12])];
    #Top dft rec rule produces 3 subrules
    rtp_dft_rec_x := tag_rec(tag_dft_rec_x, Replicate(3, Ignore));

    #Top IOPrunedConv rule presents no parallelizable objects
    tag_conv_x := Ignore;
    #Top IOPrunedConv rule produces 2 subrules
    rtp_conv_x := tag_rec(tag_conv_x, [rtp_idft_rec_x, rtp_dft_rec_x]);

    ## Pruned IDFT YZ vvvvv

    BX := 32;
    #Z-dim pridft rec rule presents 3 parallelizable objects
    tag_pridft_rec_z := [ASIMTBlockDimX(BX, [0,1]), ASIMTBlockDimX(BX, [0,0], [0,31]), ASIMTBlockDimX(BX)];
    #Z-dim pridft rec rule produces 3 subrules
    rtp_pridft_rec_z := tag_rec(tag_pridft_rec_z, Replicate(3, Ignore));

    #Y-dim idft rec rule presents 4 parallelizable objects
    tag_idft_rec_y := [ASIMTBlockDimX(BX), ASIMTBlockDimX(BX, [0,7]), ASIMTBlockDimX(BX, [8,12]), ASIMTBlockDimX(BX, [0,9])];
    #Y-dim idft rec rule produces 3 subrules
    rtp_idft_rec_y := tag_rec(tag_idft_rec_y, Replicate(3, Ignore));

    ## Pruned IDFT YZ ^^^^^

    TS_stg1 := 11;
    TS_stg2 := 11;
    TS_stg543 := 4;

    # Define TagTree
    trstg5to1 := tag_rec(
                        # Taggable operators (Tensors, Sums, etc.) as they appear in Collect( IOPrunedMDRConv.apply(...), @([Tag. ops]) )
                        [ ASIMTLoopDim(), ASIMTKernelFlag( ASIMTGridDimX(nd3) ), ASIMTBlockDimY(6), ASIMTBlockDimY(6), # Stages 7-6: 4 tag operators
                            ASIMTLoopDim(), ASIMTLoopDim(), ASIMTTileFlag(
                                                            ASIMTListDim(
                                                                ASIMTKernelFlag(ASIMTGridDimX(nfreq*n2/TS_stg543)),
                                                                ASIMTBlockDimY(TS_stg543)
                                                            ), TS_stg543), ASIMTLoopDim(), ASIMTLoopDim(), # Stages 5-4-3: 5 tag operators
                            ASIMTKernelFlag(ASIMTGridDimX(ns3)), ASIMTTileFlag(
                                                                ASIMTListDim(
                                                                    ASIMTGridDimY(nfreq/TS_stg2),
                                                                    ASIMTBlockDimX(11) # << 11 OK because perfectly nests Inner loop mapped to Y (so block sync is used at compose boundaries)
                                                                ), TS_stg2), # Stage 2: 2 tag operators
                            ASIMTKernelFlag( ASIMTGridDimX(ns2) ),
                                    ASIMTTileFlag(
                                            ASIMTListDim(
                                                ASIMTGridDimY(ns3/TS_stg1),
                                                ASIMTFissionFlag( ASIMTListDim(ASIMTBlockDimY(32), ASIMTBlockDimX(16)) )
                                            ), TS_stg1), ASIMTLoopDim() ], # Stage 1: 3 tag operators
                        # Children as they appear in IOPrunedMDRConv.children(...)
                        [ rtp_prdft_rec_z, rtp_dft_rec_y, rtp_conv_x, rtp_idft_rec_y, rtp_pridft_rec_z, Ignore ]);

    tmp_tag_rt(rtb, trstg5to1, opts); # RuleTree + TagTree = TaggedRuleTree

    sb := SumsRuleTree(rtb, opts);
    cb := CodeSums(sb, opts);
    cb.transform := t;

    return cb;
end;

Class(HockneyMlcCUDADeviceOpts, TitanVDefaults, rec(
    operations := rec(Print := s -> Print("<Hockney MLC CUDA Device options record>")),
    tagIt := (self, t) >> t.withTags(self.tags)
));

hockneyMlcCUDADeviceOpts := function(arg)
    local opts;

    opts := Copy(HockneyMlcCUDADeviceOpts);
    opts.includes := [];

    opts.breakdownRules.PRDFT := List([ PRDFT1_Base2, CopyFields(PRDFT_PD, rec(maxSize := 7)), PRDFT_PD_loop], _noT);
    opts.breakdownRules.IPRDFT := List([ IPRDFT1_Base1, IPRDFT1_Base2, IPRDFT_PD_loop, CopyFields(IPRDFT_PD, rec(maxSize := 7))], _noT);
    opts.breakdownRules.IPRDFT2 := List([ IPRDFT2_Base1, IPRDFT2_Base2, IPRDFT2_CT], _noT);
    opts.breakdownRules.PRDFT3 := [ ];
    opts.breakdownRules.URDFT := List([ URDFT1_Base1, URDFT1_Base2, URDFT1_Base4, URDFT1_CT ], _noT);
    opts.breakdownRules.DFT := List([ DFT_Base, CopyFields(DFT_PD, rec(maxSize := 7)),  DFT_PD_loop], _noT);
    opts.breakdownRules.PrunedPRDFT := List([ CopyFields(PrunedPRDFT_base, rec(maxSize := 5)), PrunedPRDFT_CT_rec_block ], _noT);
    opts.breakdownRules.PrunedIPRDFT := List([ CopyFields(PrunedIPRDFT_base, rec(maxSize := 5)), PrunedIPRDFT_CT_rec_block ], _noT);
    opts.breakdownRules.PrunedDFT := List([ PrunedDFT_base, PrunedDFT_CT_rec_block ], _noT);
    opts.breakdownRules.IOPrunedMDRConv := List([ IOPrunedMDRConv_3D_2trip_zyx_freqdata ], _noT);
    opts.breakdownRules.MDRConv := [ MDRConv_3D_2trip_zyx_freqdata ]; # [MDRConv_Base]

    wrap_rule_apply(opts.breakdownRules.PRDFT, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.IPRDFT, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.IPRDFT2, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.URDFT, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.DFT, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.PrunedPRDFT, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.PrunedIPRDFT, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.PrunedDFT, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.PrunedIDFT, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.TTensorI{[2..3]}, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.IOPrunedConv, wrap_apply);
    wrap_rule_apply(opts.breakdownRules.IOPrunedMDRConv, wrap_apply);

    opts.codegen._Formula := Copy(opts.codegen.Formula);
    opts.codegen.Formula := meth ( self, o, y, x, opts )
        local prog, prog1, s;
        s := Copy(o);
        s := fixUpSigmaSPL(s, opts);

        opts.("postCompileStrategy") := function(icode, opts)
            local icode1;
            icode1 := skip();
            # This reduces the amount of constant matrices but might prevent to store
            # some of them in constant mem. E.g., two identical matrices but only one has broadcast-type access.
            # while icode1 <> icode do
            #     icode1 := icode;
            #     icode := fixReplicatedData(icode, opts);
            # od;
            icode := scalarizeAccumulator(icode, opts);
            return icode;
        end;

        prog := opts.codegen._Formula(s, y, x, opts);

        return prog;
    end;

    opts.formulaStrategies.postProcess := opts.formulaStrategies.postProcess
        :: opts.formulaStrategies.sigmaSpl :: [MergedRuleSet(HockneyMLC_SIMTRules, RulesSums)];

    opts.fftxGen := (self, t) >> HockneyMLC_fftxGen(t, self);
    opts.tags := [];

    opts.preProcess := (self, t) >> let(t1 := RulesFFTXPromoteNT(Copy(t)), RulesFFTXPromoteNT_Cleanup(t1));
    opts.prettyPrint := meth(self, c)
        local name;
        name := "fftx_generated";
        if (IsBound(c.transform) and IsBound(c.transform.params[2].fname)) then
            name := c.transform.params[2].fname;
        fi;
        PrintCode(name, c, self);
    end;

    return opts;
end;

confHockneyMlcCUDADevice := rec(
    defaultName := "confHockneyMlcCUDADevice",
    defaultOpts := (arg) >> rec(useHockneyMlcCUDADevice := true),
    confHandler := hockneyMlcCUDADeviceOpts
);

fftx.FFTXGlobals.registerConf(confHockneyMlcCUDADevice);


# Concatenate arrays
Concatenation([1, 2, 3], [4, 5, 6], [7, 8, 9]);
# [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ]

# Append to a variable
a := [1, 2, 3];
Append(a, [4, 5, 6);
Append(a, [7, 8, 9]);
a;
# [ 1, 2, 3, 4, 5, 6, 7, 8, 9 ]


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

VarFrontier := (e, depth) -> Cond(
    IsBound(e.isExpComposite) and e.isExpComposite,
        Union(List(VarArgsExp(e), x->VarFrontier(x, depth))),
    IsVar(e),
        When(depth=0,
            Set([e]),
            Union([e], Union(List(PredLoc(e), x->VarFrontier(x, depth-1))))),
    Set([])
);

IsXRef := exp -> Intersection([var.table.X,var.table.X1,var.table.XY1], VarFrontier(exp, -1)) <> [];

IsYRef := exp -> Intersection([var.table.Y,var.table.Y1,var.table.XY1], VarFrontier(exp, -1)) <> [];

CollectLoads := cmd -> Collect(cmd, [assign, @, @(1, [deref, nth], IsXRef)]);

CollectStores := cmd -> Collect(cmd, [assign, @(1, [deref, nth], IsYRef), @]);


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

#F IsNonTerminal ( <obj> )
#F   returns true, if <obj> is a non-terminal
#F
IsNonTerminal := N -> IsRec(N) and IsBound(N.isNonTerminal) and N.isNonTerminal = true;

Class(NonTerminalOps, SPLOps, rec(
));

Declare(NonTerminal);

# ==========================================================================
# NonTerminal
# ==========================================================================
Class(NonTerminal, BaseMat, rec(
    #F create getter methods
    #F    ex: NonTerminal.WithParams("param1_getter_name", "param2_getter_name")
    WithParams := arg >> CopyFields(arg[1],
	RecList( ConcatList( [2..Length(arg)],
	    i -> [arg[i], DetachFunc(Subst(self >> self.params[$(i-1)]))] ))),

    isNonTerminal := true, # for nonterm.g and IsNonTerminal() function
    _short_print := false,

    #--------- Transformation rules support --------------------------------
    from_rChildren := (self, rch) >> let(
        t := ApplyFunc(ObjId(self), DropLast(rch, 1)),
        When(Last(rch), t.transpose(), t)
    ),

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

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

    #-----------------------------------------------------------------------
    canonizeParams := meth(self)
        local A, nump;
        nump := Length(self.params);
    if nump = 0 then return; fi;
        for A in self.abbrevs do
            if NumArgs(A) = -1 or NumArgs(A) = nump then
                self.params := ApplyFunc(A, self.params);
                return;
            fi;
        od;
        Error("Nonterminal ", self.name, " can take ", List(self.abbrevs, NumArgs),
              "  arguments, but not ", nump);
    end,

    new := meth(arg)
        local result, self, params;
        self := arg[1];
        params := arg{[2..Length(arg)]};
        result := SPL(WithBases(self, rec(params := params, transposed := false )));
        result.canonizeParams();
        result.dimensions := result.dims();
        return result;
    end,
    #-----------------------------------------------------------------------
    __call__ := ~.new,
    #-----------------------------------------------------------------------
    isTerminal := False,
    #-----------------------------------------------------------------------
    isPermutation := False,
    #-----------------------------------------------------------------------
    transpose := self >>
        Inherit(self, rec(transposed := not self.transposed,
                  dimensions := Reversed(self.dimensions))),
    setTransposed := (self, v) >> When( self.transposed <> v,
        self.transpose(), self),
    #-----------------------------------------------------------------------
    # .conjTranspose() is needed for RC(T).transpose() == RC(T.conjTranspose())
    # The inert form will only work is non-terminal is not used inside RC,
    # or used inside RC, but not transposed
    conjTranspose := self >> InertConjTranspose(self),
    isInertConjTranspose := self >> self.conjTranspose = NonTerminal.conjTranspose,

    #-----------------------------------------------------------------------
    print := (self, i, is) >>
	Cond(
            not IsBound(self.params), Print(self.name),
	    Print(self._print(self.params, i, is),
		  When(self.transposed, ".transpose()"))),
    #-----------------------------------------------------------------------
    toAMat := self >> self.terminate().toAMat(),
    #-----------------------------------------------------------------------
    export := arg >> Error("Can't export a non-terminal"),
    #-----------------------------------------------------------------------
    arithmeticCost := arg >> Error("Can't compute arithmetic cost of a non-terminal"),

    transposeSymmetric := True,
    isSymmetric := False,

    area := self >> Product(self.dims())
));


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

# ================================================================
# Diagonals
# ================================================================
Class(RulesDiag, RuleSet);

LeftPull  := [Diag, DiagCpxSplit, RCDiag];
RightPull := [Diag, DiagCpxSplit, RCDiag];
PullCont  := [SUM, BB, Buf, Inplace, Grp];
PullContL  := [SUM, BB, Buf, Inplace, Grp, NoDiagPullinRight];
PullContR  := [SUM, BB, Buf, Inplace, Grp, NoDiagPullinLeft];

RewriteRules(RulesDiag, rec(
 DiagPullInRight := ARule( Compose,
       [ @(1, LeftPull), @(2, PullContL) ],
  e -> [ ApplyFunc(ObjId(@2.val), List(@2.val.children(), c -> @1.val * c)) ]),

 DiagPullInLeft  := ARule( Compose,
       [ @(1, PullContR), @(2, RightPull) ],
  e -> [ ApplyFunc(ObjId(@1.val), List(@1.val.children(), c -> c * @2.val)) ]),

 DiagISumRight := ARule( Compose,
       [ @(1, LeftPull), @(2, ISum, canReorder) ],
  e -> [ ISum(@2.val.var, @2.val.domain, @1.val * @2.val.child(1)).attrs(@(2).val) ]),

 DiagISumLeft := ARule( Compose,
       [ @(1, ISum, canReorder), @(2, RightPull) ],
  e -> [ ISum(@1.val.var, @1.val.domain, @1.val.child(1) * @2.val).attrs(@(1).val) ]),

 DiagIComposeRight := ARule( Compose,
       [ @(1, LeftPull), @(2, ICompose) ],
  e -> [ ICompose(@2.val.var, @2.val.domain, @1.val * @2.val.child(1)).attrs(@(2).val) ]),

 DiagIComposeLeft := ARule( Compose,
       [ @(1, ICompose), @(2, RightPull) ],
  e -> [ ICompose(@1.val.var, @1.val.domain, @1.val.child(1) * @2.val).attrs(@(1).val) ]),

 DiagDataRight := ARule( Compose,
       [ @(1, LeftPull), @(2, Data, canReorder) ],
  e -> [ Data(@2.val.var, @2.val.value, @1.val * @2.val.child(1)).attrs(@(2).val) ]),

 DiagDataLeft := ARule( Compose,
       [ @(1, Data, canReorder), @(2, RightPull) ],
  e -> [ Data(@1.val.var, @1.val.value, @1.val.child(1) * @2.val).attrs(@(1).val) ]),

 DiagRecursStepRight := ARule( Compose,
       [ @(1, LeftPull), @(2, RecursStep) ],
  e -> [ RecursStep(@2.val.yofs, @2.val.xofs, @1.val * @2.val.child(1)) ]),

 DiagRecursStepLeft := ARule( Compose,
       [ @(1, RecursStep), @(2, RightPull) ],
  e -> [ RecursStep(@1.val.yofs, @1.val.xofs, @1.val.child(1) * @2.val) ])
));

Class(RulesDiagStandalone, RuleSet);
RewriteRules(RulesDiagStandalone, rec(
 # Gath * Diag
 CommuteGathDiag := ARule( Compose,
       [ @(1, Gath), @(2, Diag) ], # o 1-> 2->
  e -> [ Diag(fCompose(@2.val.element, @1.val.func)).attrs(@(2).val), @1.val ]),

 # Diag * Scat
 CommuteDiagScat := ARule( Compose,
       [ @(1, Diag), @(2, Scat) ], # <-1 <-2 o
  e -> [ @2.val, Diag(fCompose(@1.val.element, @2.val.func)).attrs(@(1).val) ]),

 # Gath * DiagCpxSplit
 CommuteGathDiagCpxSplit := ARule( Compose,
       [ @(1, Gath), @(2, DiagCpxSplit) ], # o 1-> 2->
  e -> [ DiagCpxSplit(fCompose(@2.val.element, fTensor(@1.val.func, fId(2)))).attrs(@(2).val), @1.val ]),

 # DiagCpxSplit * Scat
 CommuteDiagCpxSplitScat := ARule( Compose,
       [ @(1, DiagCpxSplit), @(2, Scat) ], # <-1 <-2 o
  e -> [ @2.val, DiagCpxSplit(fCompose(@1.val.element, fTensor(@2.val.func, fId(2)))).attrs(@(1).val) ]),

 # Gath * RCDiag
 CommuteGathRCDiag := ARule( Compose,
       [ [@(1, Gath), [@(0,fTensor), ..., [fId,@(2).cond(IsEvenInt)]]],
      @(4, RCDiag) ],
  e -> [ RCDiag(fCompose(@(4).val.element, @(0).val), @(4).val.post),
         @(1).val ]),

 # RCDiag * Scat
 CommuteRCDiagScat := ARule( Compose,
       [ @(4, RCDiag),
     [@(1, Scat), [@(0,fTensor), ..., [fId,@(2).cond(IsEvenInt)]]] ],
  e -> [ @(1).val,
         RCDiag(fCompose(@(4).val.element, @(0).val), @(4).val.post) ]),

 # Gath * IDirSum
 CommuteGathIDirSum := ARule( Compose,
       [ [@(1, Gath), [@(3,fTensor), ..., [fId,@(2)]]],
     [@(0, [PushL, PushLR]), @(4, IDirSum, e -> (@(2).val mod Rows(e.child(1))) = 0)] ],
  e -> let(ch := DropLast(@(3).val.children(), 1),
           ratio := @(2).val / Rows(@(4).val.child(1)),
           f := When(ratio=1, fTensor(ch), fTensor(Concatenation(ch, [fId(ratio)]))),
           j := @(4).val.var,
           jj := Ind(f.domain()),

           [ ObjId(@(0).val)(IDirSum(jj,      #SubstVars(@(4).val.child(1), rec((j.id) := f.at(jj)))),
                   Data(j, f.at(jj), SubstTopDown(@(4).val.child(1), @(5,fBase,e->e.params[2]=j),
                                       e -> fCompose(f, fBase(jj)))))),
             @(1).val ])
 ),

 # IDirSum * Scat
 CommuteIDirSumScat := ARule( Compose,
       [ [@(0,[PushR, PushLR]), @(4, IDirSum)],
         [@(1, Scat), [@(3,fTensor), ..., [fId,@(2).cond(e -> (e mod Cols(@(4).val.child(1))) = 0)]]]],

  e -> let(ch := DropLast(@(3).val.children(), 1),
           ratio := @(2).val / Cols(@(4).val.child(1)),
           f := When(ratio=1, fTensor(ch), fTensor(Concatenation(ch, [fId(ratio)]))),
           j := @(4).val.var,
           jj := Ind(f.domain()),

           [ @(1).val,
             ObjId(@(0).val)(IDirSum(jj,      #SubstVars(@(4).val.child(1), rec((j.id) := f.at(jj)))),
                   Data(j, f.at(jj), SubstTopDown(@(4).val.child(1), @(5,fBase,e->e.params[2]=j),
                                       e -> fCompose(f, fBase(jj)))))) ])
  ),

));

#--------------------------------------------------
#   new stuff for II

Class(RulesDiagII, RuleSet);
RewriteRules(RulesDiagII, rec(
    fCompose_diagXX := Rule([fCompose, @(1, [diagAdd, diagMul]), @(2)],
        e -> ApplyFunc(ObjId(@(1).val), List(@(1).val.children(), c->(fCompose(c, @(2).val ))))),

    simp1 := Rule([fCompose, @(1, II, e->e.params[2]=0), [fTensor, @(2, fId), @(3, fBase, e->@(1).val.params[3] = e.range() )] ],
        e -> II(@(2).val.domain(), 0, 1)),

    fBase_cond := Rule([fCompose, @(1, II), @(2, fBase)],
        e -> let(ii := @(1).val.params,
            diagCond(ivElOf(ivII(ii[1], ii[2], ii[3]), @(2).val.params[2]), fConst(1,1), fConst(1,0)))),

    tensor_cond := ARule(diagTensor, [@(1), @(2, diagCond)],
        e -> [diagCond(@(2).val.child(1), diagTensor(@(1).val, @(2).val.child(2)), diagTensor(@(1).val, @(2).val.child(3)))]),

    ivElOf_fold_val := Rule([ivElOf, @(1, ivII), @(2, Value)],
        e-> When(@2.val.v >= @1.val.params[2] and @2.val.v < @1.val.params[3], V(1), V(0))),

    diagCond_fold := Rule([diagCond, @(1, Value), @(2), @(3)],
        e-> When(@1.val.v = 1, @(2).val, @(3).val)),

    diagAdd_combine1 := ARule(diagAdd, [@(1, II), @(2, II, e->@(1).val.params[3]=e.params[2])],
        e -> [II(@(1).val.params[1], @(1).val.params[2], @(2).val.params[3])]),

    diagAdd_combine2 := ARule(diagAdd, [@(1, II), @(2, II, e->@(1).val.params[2]=e.params[3])],
        e -> [II(@(1).val.params[1], @(2).val.params[2], @(1).val.params[3])]),

    II_to_fConst := Rule(@(1, II, e->e.params[1]=e.params[3] and e.params[2]=0), e->fConst(@(1).val.params[1], 1)),

    diagMul0r := ARule(diagMul, [@(1), @(2, fConst, e->e.params[2] = 0)], e -> [@(2).val]),
    diagMul0l := ARule(diagMul, [@(1, fConst, e->e.params[2] = 0), @(2)], e -> [@(1).val]),

    diagAdd0r := ARule(diagAdd, [@(1), @(2, fConst, e->e.params[2] = 0)], e -> [@(1).val]),
    diagAdd0l := ARule(diagAdd, [@(1, fConst, e->e.params[2] = 0), @(2)], e -> [@(2).val]),

    Diag_fConst_to_I := Rule([@(1, Diag), @(2, fConst, e->e.params[2] = 1)],
        e->Prm(fId(@(2).val.params[1]))),

    adjust_ivII := Rule ([ivElOf, @(1, ivII), [@(2, add), @(3, var), @(4, Value)]],
        e -> let(iv := @(1).val.params, shift := @(4).val.v,
            ivElOf(ivII(iv[1], Maximum(0, iv[2]-shift), Maximum(0, iv[3]-shift)), @(3).val))),

    ivElOf_fold_var1 := Rule([ivElOf, @(1, ivII, e->e.params[2]=0),
        @(2, var, e->IsBound(e.range) and e.range <= @(1).val.params[3])],
        e-> V(1)),

    ivElOf_fold_var2 := Rule([ivElOf, @(1, ivII, e->e.params[3]=0), @(2, var)],
        e-> V(0)),
));

RulesDiagAll := MergedRuleSet(RulesDiag, RulesDiagStandalone);


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

#F standard cache params structure, it takes
#F
#F blksize (block size in bytes)
#F nsets   (number of sets)
#F assoc   (associativity eg. 2-way,4-way,etc)
#F rpolicy  (replacement policy string eg, "FIFO","LRU", etc)
#F
#F and computes a hidden value, csize (cache size)

Declare(CacheDesc);

Class(CacheDesc, rec(
    __call__ := (self, blksize, nsets, assoc, rpolicy) >>
        Checked(Is2Power(blksize),
            Is2Power(nsets),
            Is2Power(assoc),
            IsString(rpolicy),
            WithBases(self, rec(
                blksize := blksize,
                nsets := nsets,
                assoc := assoc,
                rpolicy := rpolicy,
                csize := blksize * nsets * assoc,
            ))
        ),

    operations := rec(
        Print := self >> When(ObjId(self) = CacheDesc,
            Print("CacheDesc(",
                PrintCS([
                    self.blksize,
                    self.nsets,
                    self.assoc,
                    Concat("\"", self.rpolicy, "\"")
                ]), ")"
            ),
            Print("CacheDesc(<blksize>, <nsets>, <assoc>, <rpolicy>)")
        )
    ),
));

# Cache tag. This is a generic tag, which means it takes any arguments.
# However, the correct arguments are:
#
# arg1: positive integer symbolizing cache level
# arg2: a CacheSpec()

Class(ACache, AGenericTag);

# Inplace tag. Takes one argument, maximum BB size.
#
# use this tag to specify that a transform should be computed inplace
Class(AInplace, AGenericTag);

# Right expansion tag
#
# one argument, max left kernel size
#
Class(AExpRight, AGenericTag);

# AMem tag.
#
# params are:
# 1: size
# 2: packet/block size
# 3: per use defined "fit" function
#
# what fit-function's parameters are is TBD

Class(AMem, AGenericTag, rec(
	__call__ := meth(arg)
		local self, params;
		self := arg[1];
		params := Drop(arg, 1);
		return WithBases(self,
			rec(
				params := params,
				size := (self) >> self.params[1],
				blocks := (self) >> self.params[2],
				assoc := (self) >> self.params[3],
				operations := RewritableObjectOps
			)
		);
	end,

));

#
# Basic block tag.
#
# param1 is K, size of the basic block
# param2 is mu, the block/row/whatever size
#
# we assume K >= mu
#
Class(ABB, AGenericTag);


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

NewRulesFor(GT, rec(

# Why must we do parallelization and streaming in one big rule? Because they
# must both be done together. If done separately, we'll get two sets of outer
# perms (LxI)s, and we don't want that. In addition, the parallel breakdown
# rule has to tell a special mbuf rule how many iterations to multibuffer for,
# and a) this is messy, and b) this is trivial to handle in one single rule,
# and therefore best done within a single rule.

#F GT_ParStreamCore
#F Params: (p, u). p=# of procs, u=minimum packet size
    GT_ParStreamCore := rec(
        kmin          := 16, #128  # Minimum kernel size. DMA packet size = 4*k bytes for single precision real
        kmax          := 2048,  # Based on Max packet size for single precision (2048*4*2 = 16 Kbytes) NOTE: HARDCODED!

        requiredFirstTag := ParStream,

        # Variable naming scheme: either (In x Am) or (Am x In)

        applicable := (self, t) >> let(
            rank := Length(t.params[4]),     # Eg.: (Am x In)
            g    := t.params[2],
            s    := t.params[3],

            m    := Rows(t.params[1]),       # (splsize)
            n    := t.params[4][1],          # Available iterations (n)
            p    := t.firstTag().params[1],  # procs
            u    := t.firstTag().params[2],  # Specified minimum packet size

            b    := Maximum((m*n)/(self.kmax*p), 4), # Number of multibuffered iterations
            r    := n/(b*p),                 # Remaining iterations.

            # Conditions:
            rank = 1
            and PatternMatch(t, [GT, @(1), @(2,XChain), @(3,XChain), ...], empty_cx())
            and r >= 1       # If not, we have too few iterations for mbuf+par
            and When(g=GTPar and s=GTPar, r*m >= u,  r >= u)  # Respect specified packet size
        ),

        children := (self, t) >> let(
            g    := t.params[2],
            s    := t.params[3],

            m    := Rows(t.params[1]),
            n    := t.params[4][1],
            p    := t.firstTag().params[1],  # procs

            b    := Maximum((m*n)/(self.kmax*p), 4), # Number of multibuffered iterations
            r    := n/(b*p),                         # Remaining iterations.

            pkSizeG := When(g=GTPar, r*m, r),
            pkSizeS := When(s=GTPar, r*m, r),

            tags   := Drop(t.getTags(), 1), # Eat up tag

            When(r > 1,
              [ [ GT(t.params[1], g, s, [r]).withTags(tags),
                InfoNt(b, pkSizeG, pkSizeS) ] ],
              [ [ t.params[1].withTags(tags), InfoNt(b, pkSizeG, pkSizeS) ] ]
            )
        ),

        apply := (self, t, C, Nonterms) >> let(
            dft := Nonterms[1],
            g   := t.params[2],
            s   := t.params[3],
            n   := t.params[4][1],
            m   := Rows(t.params[1]),
            mc  := Cols(t.params[1]),
            p   := t.firstTag().params[1],
            b   := Nonterms[2].params[1],
            pkSizeG := Nonterms[2].params[2],
            pkSizeS := Nonterms[2].params[3],
            gg      := When(g.params[1]=[0,1], XChain([0,1,2]), XChain([1,2,0])),
            ss      := When(s.params[1]=[0,1], XChain([0,1,2]), XChain([1,2,0])),

            i      := var("spuid", TInt, p),
            j      := Ind(b),

            sp     := When(s=GTPar,
                        ScatDist(p, n*m/p, p, i),
                        ScatSend(ss.part(1, i, m, [p, 1]), n/p, p, i)),

            gp     := When(g=GTPar,
                        GathDist(p, n*m/p, p, i),
                        GathRecv(gg.part(1, i, mc, [p, 1]), n/p, p, i)),

            sm     := When(s=GTPar,
                        ScatMem(s.part(1, j, 1, [b]), pkSizeS),
                        ScatMem(s.part(1, j, m,  [b]), pkSizeS)),

            gm     := When(g=GTPar,
                        GathMem(g.part(1, j, 1, [b]), pkSizeG),
                        GathMem(g.part(1, j, mc, [b]), pkSizeG)),

            DistSum(p, i, p,
                sp * MultiBufISum(j, b, sm, C[1], gm) * gp
            )

        ) #apply (let)
    ) #GT_ParStreamCore
));

NewRulesFor(GT, rec(
#F GT_StreamParChip
#F Params: (p, um, uc=u/p). p=# of procs, um=min pksize (mem), uc=min pksize (on-chip)
    GT_StreamParChip := rec(
        kmin          := 16, #128  # Minimum kernel size. DMA packet size = 4*k bytes for single precision real
        kmax          := 2048,  # Based on Max packet size for single precision (2048*4*2 = 16 Kbytes) NOTE: HARDCODED!

        requiredFirstTag := ParStream,

        # Variable naming scheme: either (In x Am) or (Am x In)
        applicable := (self, t) >> let(
            rank := Length(t.params[4]),     # Eg.: (Am x In)
            g    := t.params[2],
            s    := t.params[3],

            m    := Rows(t.params[1]),       # (splsize)
            n    := t.params[4][1],          # Available iterations (n)
            p    := t.firstTag().params[1],  # procs
            u    := t.firstTag().params[2],  # Specified minimum packet size (mem)

            b    := Maximum((m*n)/(self.kmax*p), 4), # Number of multibuffered iterations
            r    := n/b,                     # Remaining iterations.

            # Conditions:
            rank = 1
            and PatternMatch(t, [GT, @(1), @(2,XChain), @(3,XChain), ...], empty_cx())
            and r >= 1       # If not, we have too few iterations for mbuf+par
            and When(g=GTPar and s=GTPar, r*m >= u,  r >= u)  # Respect specified packet size
            #and not GT_ParStreamCore.applicable(t) # For now, prioritize the GT_ParStreamCore rule
        ),

        children := (self, t) >> let(

            g    := t.params[2],
            s    := t.params[3],

            m    := Rows(t.params[1]),
            n    := t.params[4][1],
            p    := t.firstTag().params[1],  # procs
            u    := t.firstTag().params[2],  # Specified minimum packet size (mem)
            uc   := u/p,                     # Minimum packet size (on-chip)

            b    := Maximum((m*n)/(self.kmax*p), 4), # Number of multibuffered iterations
            r    := n/b,                     # Remaining iterations.

            tags   := Drop(t.getTags(), 1), # Eat up tag

            # Add a parallel tag to child
            When(r > 1,
               [[ GT(t.params[1], g, s, [r]).setTags(Concatenation([ParCell(p, uc)], tags)) ]],
               [[ t.params[1].setTags(Concatenation([ParCell(p, uc)], tags)) ]]
            )
        ),

        apply := (self, t, C, Nonterms) >> let(
            dft     := Nonterms[1],
            g       := t.params[2],
            s       := t.params[3],

            m    := Rows(t.params[1]),
            n    := t.params[4][1],
            p    := t.firstTag().params[1],  # procs
            u    := t.firstTag().params[2],  # Specified minimum packet size (mem)

            b    := Maximum((m*n)/(self.kmax*p), 4), # Number of multibuffered iterations
            r    := n/b,                             # Remaining iterations.

            pkSizeS := When(s=GTPar, r*m, r),
            pkSizeG := When(g=GTPar, r*m, r),

            i      := Ind(b),
            sm     := When(s=GTPar,
                        ScatMem(s.part(1, i, 1, [b]), pkSizeS),
                        ScatMem(s.part(1, i, Rows(dft)/pkSizeS, [b/pkSizeS]), pkSizeS)),
            gm     := When(g=GTPar,
                        GathMem(g.part(1, i, 1, [b]), pkSizeG),
                        GathMem(g.part(1, i, Cols(dft)/pkSizeG, [b/pkSizeG]), pkSizeG)),

            MultiBufISum(i, b,
                sm,
                C[1],
                gm
            )

        ) #let/apply
    ), #GT rule

));

# NOTE: Works if the specified u is the max u. Else, fails.
# NOTE: Should convert this rule so that u is automatically picked.
NewRulesFor(GT, rec(
#F This is the ParStream version of the MBufCell_spec vector recursion rule
#F Automatically choses k,n, subject to a specified minimum
    GT_DFT_ParStreamCore_VecRecur := rec(
        kmin          := 16, #128  # Minimum kernel size. DMA packet size = 4*k bytes for single precision real
        kmax          := 2048,  # Maximum kernel size on a single SPE

        requiredFirstTag := ParStream,

        # Variable naming scheme: either (In x Am) or (Am x In)

        #applicable := false, # THIS RULE IS INCORRECT!

        applicable := (self, t) >> let(
            rank := Length(t.params[4]),     # Eg.: (Am x In)
            g    := t.params[2],
            s    := t.params[3],

            k    := t.params[4][1],          # Available iterations
            mn   := Rows(t.params[1]),       # (splsize)
            p    := t.firstTag().params[1],  # procs
            u    := t.firstTag().params[2],  # Specified minimum packet size

            mns  := Filtered(DivisorPairs(mn),
                       i->let(m:=i[1], n:=i[2],
                          n >= u                # n is a packet size in second factor
                          and m >= 4*p          # multibuffer and parallelize around m iterations in second factor
                          and m*n >= self.kmin  # First factor kernel size is big enough
                          and m*n <= self.kmax  # First  factor fits on SPEs
                          and k*n <= self.kmax) # Second factor fits on SPEs
                    ),
#           Error("BP--applicable"),
            rank = 1
#           and not GT_ParStreamCore.applicable(t)  # Regular rule is not applicable
            and g=GTVec and s=GTPar             # Only consider (IxD)L
            and PatternMatch(t, [GT, DFT, XChain, XChain, @, @, @], empty_cx()) # Only consider GTs with DFT for a kernel
            and k >= 4*p                  # First loop is multibuffered and parallelized around k
            and k >= u                    # k is also the one of the packet sizes (2nd factor)
            and Length(mns) >= 1          # We can break the DFT based on all our constraints
        ),

        children := (self, t) >> let(
            g    := t.params[2],
            s    := t.params[3],
            k    := t.params[4][1],          # Available iterations (n)
            mn   := Rows(t.params[1]),       # (splsize)
            p    := t.firstTag().params[1],  # procs
            u    := t.firstTag().params[2],  # Specified minimum packet size

            tags := Drop(t.getTags(), 1),
            #oldtags := Drop(t.getTags(), 1),
            #tags := Concatenation([ParStream(p, u)], oldtags),

            mns  := Filtered(DivisorPairs(mn),
                       i->let(m:=i[1], n:=i[2],
                          n >= u                # n is a packet size in second factor
                          and m >= 4*p          # multibuffer around m iterations in second factor
                          and m*n >= self.kmin  # First factor kernel size is big enough
                          and m*n <= self.kmax  # First  factor fits on SPEs
                          and k*n <= self.kmax) # Second factor fits on SPEs
                    ),

            Map2(mns, (m,n)-> [
              GT(DFT(m,1), GTVec, GTVec, [n]).withTags(tags),
              GT(DFT(n,1), GTVec, GTPar, [k]).withTags(tags),
              InfoNt(m, n)
              ])
        ),

        apply := (self, t, C, Nonterms) >> let(
            k    := t.params[4][1],
            m    := Nonterms[3].params[1],
            n    := Nonterms[3].params[2],
            p    := t.firstTag().params[1],  # procs

            j1    := Ind(k/p),
            j2    := Ind(m/p),
            pvar := var("spuid", TInt, p),
            pkSize1 := Rows(C[1]),

            #Error("BP----"),

            # WORKS:
            #MultiBufISum(j1, k,
            #      ScatMem(GTPar.part(1, j1, 1, [k]), pkSize1),
            #      C[1] * Diag(fPrecompute(Tw1(m*n, n, 1))),
            #      GathMem(GTPar.part(1, j1, 1, [k]), pkSize1)
            #)
            #*
            #MultiBufISum(j2, m,
            #     ScatMem(fTensor(fId(k), fBase(j2)), n),
            #     C[2],
            #     GathMem(fTensor(fId(n), fBase(j2)), k)
            #)

            perms := L(k*m, m),
            ixfs  := fTensor(fBase(pvar), fId(m*k/p)),

            permg := L(m*n, m),
            ixfg  := fTensor(fBase(pvar), fId(m*n/p)),

            DistSum(p, pvar, p,
                ScatDist(p, k*n*m/p, p, pvar) *
                MultiBufISum(j1, k/p,
                  ScatMem(GTPar.part(1, j1, 1, [k]), pkSize1),
                  C[1] * Diag(fPrecompute(Tw1(m*n, n, 1))),
                  GathMem(GTPar.part(1, j1, 1, [k]), pkSize1)
                ) * GathDist(p, k*n*m/p, p, pvar)
            ) *
            DistSum(p, pvar, p,
               ScatSend(fCompose(perms, ixfs), n, p, pvar) *

               MultiBufISum(j2, m/p,
                 ScatMem(fTensor(fBase(j2), fId(k)), n),
                  C[2],
                 GathMem(fTensor(fBase(j2), fId(n)), k)
               ) *

               GathRecv(fCompose(permg, ixfg), k, p, pvar)
            )

        ) #let/apply
    )
));

NewRulesFor(GT, rec(
#F Vector recursion using streaming to entire chip
#F Automatically choses k,n, subject to a specified minimum pksize
    GT_DFT_StreamParChip_VecRecur := rec(
        kmin          := 128,  # Minimum kernel size. DMA packet size = 4*k bytes for single precision real
        kmax          := 2048,  # Maximum kernel size on a single SPE

        requiredFirstTag := ParStream,

        # Variable naming scheme: either (In x Am) or (Am x In)

        applicable := (self, t) >> let(
            rank := Length(t.params[4]),     # Eg.: (Am x In)
            g    := t.params[2],
            s    := t.params[3],

            k    := t.params[4][1],          # Available iterations
            mn   := Rows(t.params[1]),       # (splsize)
            p    := t.firstTag().params[1],  # procs
            u    := t.firstTag().params[2],  # Specified minimum packet size

            mns  := Filtered(DivisorPairs(mn),
                       i->let(m:=i[1], n:=i[2],
                          n >= u                  # n is a packet size in second factor
                          and m >= 4              # multibuffer around m iterations in second factor
                          and m*n >= self.kmin    # First factor kernel size is big enough
                          and m*n <= self.kmax*p  # First  factor fits on SPEs
                          and k*n <= self.kmax*p) # Second factor fits on SPEs
                    ),
            rank = 1
#          and not GT_ParStreamCore.applicable(t)  # None of the 3 regular rules is applicable: GT_StreamParChip and GT_ParStreamCore should also not be applicable
#           and not GT_DFT_ParStreamCore_VecRecur.applicable(t)
           and not GT_StreamParChip.applicable(t)
            and g=GTVec and s=GTPar             # Only consider (IxD)L
            and PatternMatch(t, [GT, DFT, XChain, XChain, @, @, @], empty_cx()) # Only consider GTs with DFT for a kernel
            and k >= 4                    # First loop is multibuffered around k
            and k >= u                    # k is also the one of the packet sizes (2nd factor)
            and Length(mns) >= 1          # We can break the DFT based on all our constraints
        ),

        children := (self, t) >> let(
            g    := t.params[2],
            s    := t.params[3],
            k    := t.params[4][1],          # Available iterations (n)
            mn   := Rows(t.params[1]),       # (splsize)
            p    := t.firstTag().params[1],  # procs
            u    := t.firstTag().params[2],  # Specified minimum packet size (mem)
            uc   := u/p,                     # Minimum packet size (on-chip)

            oldtags := Drop(t.getTags(), 1),
            tags := Concatenation([ParCell(p, uc)], oldtags),

            mns  := Filtered(DivisorPairs(mn),
                       i->let(m:=i[1], n:=i[2],
                          n >= u                  # n is a packet size in second factor
                          and m >= 4              # multibuffer around m iterations in second factor
                          and m*n >= self.kmin    # First factor kernel size is big enough
                          and m*n <= self.kmax*p  # First  factor fits on SPEs
                          and k*n <= self.kmax*p) # Second factor fits on SPEs
                    ),

            Map2(mns, (m,n)-> [
              GT(DFT(m,1), GTVec, GTVec, [n]).withTags(tags),
              GT(DFT(n,1), GTVec, GTPar, [k]).withTags(tags),
              InfoNt(m, n)
              ])
        ),

        apply := (self, t, C, Nonterms) >> let(
            k    := t.params[4][1],
            m    := Nonterms[3].params[1],
            n    := Nonterms[3].params[2],

            i    := Ind(k),
            j    := Ind(m),
            pkSize1 := Rows(C[1]),

            MultiBufISum(i, k,
                ScatMem(GTPar.part(1, i, 1, [k]), pkSize1),
                C[1] * Diag(fPrecompute(Tw1(m*n, n, 1))),
                GathMem(GTPar.part(1, i, 1, [k]), pkSize1)
            )
            *
            MultiBufISum(j, m,
                ScatMem(fTensor(fId(k), fBase(j)), n),
                C[2],
                GathMem(fTensor(fId(n), fBase(j)), k)
            )

        ) #let/apply
    )
));

#NOTES
#-----------------------------------------------------------------------------
# Mon 11 Jan 2010 12:56:58 AM EST

# Outline for cleanup:

# We need to separate and disentangle all these rules: vecrecur should be
# independent of ParStreamCore or StreamParChip. ParStreamCore and
# StreamParChip should simply split the parallelism and let the lower level
# mbuf and par rules handle the rest. SigmaSPL rewrite rules should either
# /naturally/ be able to understand whether something is ParStreamCore or
# StreamParChip and act accordingly, or, if this is not possible naturally,
# this should be marked up somehow at the breakdown-rule level. Ugly ways for
# SigmaSPL rules to 'try' to figure out what's going on is as bad as Clippy.


# Test of GAP potential for Si system

units		metal
boundary	p p p

atom_style	atomic

read_data	data_gap

pair_style	quip
pair_coeff	* * gap_example.xml "Potential xml_label=GAP_2015_2_20_0_10_54_35_765" 14

neighbor	0.3 bin
neigh_modify	delay 10

fix		1 all nve
thermo		10
timestep	0.001

#dump		1 all custom 10 dump.gap id fx fy fz

run		40


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

# RuleStrategy := [
#     RulesSums, RulesFuncSimp, RulesGMonTensorPullOut, RulesFuncSimp,
#     RulesGMonToSGMon, RulesFuncSimp,
#     RulesSGMonTensorPullIn, RulesSums, RulesFuncSimp
# ];

StandardSumsRules := MergedRuleSet(
    RulesSums, RulesFuncSimp, RulesDiag, RulesDiagStandalone, RulesStrengthReduce, RulesRC, RulesII, OLRules
);

StandardSumsRulesNoRC := MergedRuleSet(
    RulesSums, RulesFuncSimp, RulesDiag, RulesDiagStandalone, RulesStrengthReduce, RulesII, OLRules
);

HfuncSumsRules := MergedRuleSet(
    StandardSumsRules, RulesHfunc
);

LibStrategy := [ StandardSumsRules, HfuncSumsRules ];

# HasRM := e -> ObjId(e) in [Gath, Scat] and e.func.range() > 16 and
#               Collect(e.func, @(1,RM,e->e.range()>16)) <> [];

# PrecomputeRM := s -> SubstBottomUp(s, @.cond(HasRM), e->ObjId(e)(GenerateData(e.func)));

#F RecompileStrategies(<opts>)
#F
RecompileStrategies := function(opts)
    local fld, lst, rset;
    for fld in UserRecFields(opts.formulaStrategies) do
        lst := opts.formulaStrategies.(fld);
        if not IsList(lst) then lst := [lst]; fi;
        for rset in lst do
            if IsRec(rset) and IsBound(rset.compile) then rset.compile(); fi;
        od;
    od;
end;


local reps,tori,labels,I,II,ords,blocks;

# Permutation from standard to gap
# / 1 2 3 4 \
# \ 2 4 3 1 /
# Position(pr,Permuted([2,3,4,2],(2,4,1)))=24;

#......Ross Jordan Block Sizes........Donna & Ross nilpotents........(in GAP notation)
reps:=[[[2,1],[4,1],[3,1],[1,1]],    # F_4
       [[2,1],[4,1],[7,1],[5,1]],    # F_4(a_1) [[2,1],[4,1],[10,1],[1,1]],
       #[[6,1],[10,1],[1,1],[5,1]],   # F_4(a_2) [[9,1],[1,1],[10,1],[4,1]],
       # F_4(a_2): s_11s_9*[4,3,1,24]=[4,5,9,10] also good .. has 2 semisimple elements?
       [[1,1],[4,1],[9,1],[10,1]],   # F_4(a_2) good for all p\neq 2 <C_3A_1
       [[4,1],[3,1],[1,1]],          # C_3      [[1,1],[3,1],[4,1]],
       [[2,1],[4,1],[3,1]],          # B_3      the same
       [[4,1],[6,1],[10,1],[18,1]],  # F_4(a_3) [[4,1],[13,1],[11,1],[12,1]],
       [[4,1],[1,1],[10,1]],         # C_3(a_1) [[1,1],[10,1],[4,1]],
       [[2,1],[3,1],[1,1]],          # sA_2A_1  [[3,1],[1,1],[2,1]],
       [[4,1],[3,1]],                # B_2      the same
       [[2,1],[4,1],[1,1]],          # A_2sA_1  the same
       [[3,1],[1,1]],                # sA_2     the same
       [[2,1],[4,1]],                # A_2      the same
#       [[2,1],[3,1]],                # A_1sA_1  [[2,1],[1,1]],
#       [[3,1]],                      # sA_1     [[1,1]],
#       [[2,1]]                       # A_1      [[2,1]]
       [[2,1],[1,1]],                # A_1sA_1  [[2,1],[1,1]],
       [[1,1]],                      # sA_1     [[1,1]],
       [[2,1]]                       # A_1      [[2,1]]
       ];

I:=[[],     # F_4
    [3],    # F_4(a_1)
    [3,4],  # F_4(a_2)
    [],     # C_3
    [],     # B_3
    [2,3,1],# F_4(a_3)
    [3],    # C_3(a_1)
    [],     # sA_2A_1
    [],     # B_2
    [],     # A_2sA_1
    [],     # sA_2
    [],     # A_2
    [],     # A_1sA_1
    [],     # sA_1
    []      # A_1
    ];

#
# The parabolic where the representative is distinguished
#
II:=[[1,2,3,4],   # F_4
     [1,2,3,4],   # F_4(a_1)
     [1,2,3,4],   # F_4(a_2)
     [4,3,1],     # C_3
     [2,4,3],     # B_3
     [1,2,3,4],   # F_4(a_3)
     [4,3,1],     # C_3(a_1)
     [3,1,2],     # sA_2A_1
     [4,3],       # B_2
     [2,4,1],     # A_2sA_1
     [3,1],       # sA_2
     [2,4],       # A_2
     [2,1],       # A_1sA_1
     [1],         # sA_1
     [2]          # A_1
     ];

labels:=["F_4",
         "F_4(a_1)",
         "F_4(a_2)",
         "C_3",
         "B_3",
         "F_4(a_3)",
         "C_3(a_1)",
         "\\tilde A_2A_1",
         "B_2",
         "A_2\\tilde A_1",
         "\\tilde A_2",
         "A_2",
         "A_1\\tilde A_1",
         "\\tilde A_1",
         "A_1"
         ];

tori:=[
       [2,2,2,2],  # F_4
       [2,2,0,2],  # F_4(a_1)
       [2,0,0,2],  # F_4(a_2)
       [2,-9,2,2], # C_3
       [-6,2,2,2], # B_3
       [0,0,0,2],  # F_4(a_3)
       [2,-5,0,2], # C_3(a_1)
       [2,2,2,-5], # sA_2A_1
       [-3,-4,2,2],# B_2
       [2,2,-3,2], # A_2sA_1
       [2,0,2,-4], # sA_2
       [0,2,-2,2], # A_2
       [2,2,-1,-1],# A_1sA_1
       [2,0,-1,0], # sA_1
       [0,2,0,-1]  # A_1
    ];

#blocks:=[];
blocks:=[[[19, 1], [15, 1], [9, 2]],                      # F_4
         [[9, 5], [7, 1]],                                # F_4(a_1)
         [[9, 3], [7, 1],[6, 2], [3, 2]],                 # F_4(a_2)
         [[9, 2], [8, 2], [6, 2], [3, 1], [1, 3]],        # C_3
         [[9, 2], [7, 4], [3, 1], [1, 3]],                # B_3
         [[7, 1], [6, 2], [5, 3], [3, 6]],                # F_4(a_3)
         [[7, 1], [6, 2], [5, 1], [4, 4], [3, 3], [1, 3]],# C_3(a_1)
         [[3, 16], [2, 2]],                               # sA_2A_1
         [[7, 1], [5, 4], [4, 4], [3, 1], [1, 6]],        # B_2
         [[3, 16], [2, 2]],                               # A_2sA_1
         [[3, 15], [1, 7]],                               # sA_2
         [[3, 15], [1, 7]],                               # A_2
         [[3, 10], [2, 8], [1, 6]],                       # A_1sA_1
         [[3, 7], [2, 8], [1, 15]],                       # sA_1
         [[3, 1], [2, 14], [1, 21]]                       # A_1
         ];

#ords:=[];
ords:=[ 27, 9, 9, 9, 9, 9, 9, 3, 9, 3, 3, 3, 3, 3, 3 ];

return [reps,tori,labels,I,II,ords,blocks];


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

_SVCT_THRSHOLD := 2^13;

NewRulesFor(GT, rec(
    # Vectorize AxI: A x I_n -> (A x I_n/v) x I_v
    GT_Vec_AxI := rec(
        forTransposition := false,
        requiredFirstTag := [AVecReg, AVecRegCx],
        applicable := t -> t.rank()=1 and t.params[2]=GTVec and t.params[3]=GTVec and
                           IsInt(t.params[4][1] / t.firstTag().v),
        children := t -> let(
        r := t.params[4][1] / t.firstTag().v,
        isa := t.firstTag().isa,
        spl := t.params[1], #.setWrap(VWrap(isa)),
        tags := Drop(t.getTags(), 1),
        [[ GT(spl, GTVec, GTVec, [r]).withTags(tags).setWrap(VWrap(isa)) ]]),

        apply := (t, C, Nonterms) -> VTensor(C[1], t.firstTag().v),
    ),

    GT_Vec_IxA := rec(
        forTransposition := false,
        requiredFirstTag := [AVecReg, AVecRegCx],
        applicable := t -> t.rank()=1 and t.params[2]=GTPar and t.params[3]=GTPar and
                           IsInt(t.params[4][1] / t.firstTag().v),
        children := t -> let(
            tags := t.getTags(),
            v := t.firstTag().v,
#            r := t.params[4][1] / v,
            spl := t.params[1],

            [[  TL(Rows(spl)*v, v, 1, 1).withTags(tags).setWrap(VWrapId),
                spl.withTags(Drop(tags, 1)).setWrap(VWrap(tags[1].isa)),
                TL(Cols(spl)*v, Cols(spl), 1, 1).withTags(tags).setWrap(VWrapId)
            ]]),

        apply := (t, C, Nonterms) ->
            let(P := t.params, v := t.firstTag().v, r := P[4][1] / v,
                A := C[1] * VTensor(C[2], v) * C[3], When(r=1, A, Tensor(I(r), A)))
    ),

    # Push vector tag down to preserve locality: vec(I x A) -> (I x vec(A))
    GT_Vec_IxA_Push := rec(
        forTransposition := false,
        requiredFirstTag := [AVecReg, AVecRegCx],
        applicable := t -> t.rank()=1 and t.params[2]=GTPar and t.params[3]=GTPar,
        children := t -> [[ t.params[1].withTags(t.getTags()) ]],
        apply := (t, C, Nonterms) -> Tensor(I(t.params[4][1]), C[1])
    ),

    # Split off L:  vectorizes (I x A), (I x A) L, or L (I x A)
    #
    GT_Vec_SplitL := rec(
        minSize := _SVCT_THRSHOLD,
        forTransposition := false,
        requiredFirstTag := [AVecReg, AVecRegCx],

        applicable := (self, t) >> t.rank()=1 and Maximum(t.dims()) > self.minSize and
            (t.params{[2,3]} in [[GTPar, GTPar], [GTVec, GTPar], [GTPar, GTVec]]),

        children := t -> let(tags := t.getTags(), spl := t.params[1], n := t.params[4][1],
            g := t.params[2], s := t.params[3],
            [[ GT(spl, GTVec, GTVec, [n]).withTags(tags) ]]),

        apply := (t, C, Nonterms) -> let(spl := t.params[1], n := t.params[4][1], g := t.params[2], s := t.params[3],
                        c := Concatenation(
                            When(s=GTPar, [ L(Rows(spl)*n, n) ], []),
                            [ C[1] ],
                            When(g=GTPar, [ L(Cols(spl)*n, Cols(spl)) ], [])), Product(c))
    ),

    GT_Vec_IxA_L := rec(
        maxSize := _SVCT_THRSHOLD,
        forTransposition := false,
        requiredFirstTag := [AVecReg, AVecRegCx],

        applicable := (self, t) >> t.rank()=1
            and t.params[2]=GTVec
            and t.params[3]=GTPar
            and Maximum(t.dims()) <= self.maxSize
            and IsInt(t.params[4][1] / t.firstTag().v),

        children := t -> let(tags := t.getTags(),
            spl := t.params[1], v:=tags[1].v,
            isa:=tags[1].isa, spl := t.params[1], d:=spl.dims(),

            [[ TL(d[1]*v, v).withTags(tags).setWrap(VWrapId),
               spl.withTags(Drop(tags,1)).setWrap(VWrap(isa)) ]]
        ),

        apply := (t, C, Nonterms) -> let(spl := t.params[1],
            d:=spl.dims(), n := t.params[4][1], tag := t.firstTag(),
            v := tag.v, m := n / v,

            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)
            ))
        ),
    ),

    GT_Vec_L_IxA := rec(
        maxSize := _SVCT_THRSHOLD,
        forTransposition := false,
        requiredFirstTag := [AVecReg, AVecRegCx],

        applicable := (self, t) >> t.rank()=1 and t.params[2]=GTPar and t.params[3]=GTVec and Maximum(t.dims()) <= self.maxSize and
                           IsInt(t.params[4][1] / t.firstTag().v),

        children := t -> let(tags := t.getTags(), spl := t.params[1], v:=tags[1].v,
            isa:=tags[1].isa, spl := t.params[1], d:=spl.dims(),
            [[ spl.withTags(Drop(tags,1)).setWrap(VWrap(isa)),
               TL(d[2]*v, d[2]).withTags(tags).setWrap(VWrapId) ]] ),

        apply := (t, C, Nonterms) -> let(spl := t.params[1], d:=spl.dims(), n := t.params[4][1], tag := t.firstTag(), v := tag.v, m := n / v,
            NoDiagPullinRight(When(m=1, VTensor(C[1], v) * C[2], VTensor(L(m*d[2], d[2]), v) * Tensor(I(m), VTensor(C[1], v) * C[2]))))
    )
));