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

Import(paradigms.vector.breakdown, paradigms.vector.rewrite, paradigms.vector.sigmaspl);

SwitchRulesName([IxLxI_kmn_n, IxLxI_kmn_km, IxLxI_IxLxI_up, IxLxI_IxLxI_down], false);
SwitchRulesName([L_base, L_nv_n_vec, L_nv_v_vec, L_mn_m_vec, IxLxI_vtensor], true);

# FF: if this is turned off, DP tries to measure TL and complex code using CR(.) dies. Why did I need it??
#TL.doNotMeasure := false;

FixDataType := function(c, opts)
    local vars, v, srec;

    # Vectors of vectors must be flattened
#    c := SubstTopDown(c, @(1, TVect, e->ObjId(e.t)=TVect), e-> TVect(@(1).val.t.t, @(1).val.size*@(1).val.t.size));

    if IsBound(opts.vector.isa.needScalarVarFix) and opts.vector.isa.needScalarVarFix then
        vars := Collect(c, @(1, var, e->e.t=TVect(TReal, 2)));
        srec := rec();

        for v in vars do
            srec.(v.id) := opts.vector.isa.scalarVar();
        od;
        SubstVars(c, srec);
#        c := DeclareVars(c);
    fi;
    return c;
end;

# Populate this with default parameter values
VectorDefaults := rec(
    expandConstants := false,
    dontRuntimeLocalize := false,
);

DEFAULT_UNROLL := 55;   # largest prime: 13, 4-way -> need 40 and 52 be unrolled
UNROLL_LO := 16;
UNROLL_HI := 64;
VERBOSITY := 0;
FAILTOL := false;
HASH := false;

_help := function()
    PrintLine("\nopts := SIMDGlobals.getOpts(<arch>, flags);" );
    PrintLine("dpopts := SIMDGlobals.getDPOpts(<arch>, opts); # for correct DP options");
    PrintLine("\nsupported flags:");
    PrintLine("scalar DFT configuartion: CT, PFA, PFA_maxSize, PD, Rader, RealRader, Rader_maxSize, Bluestein, Bluestein_maxSize, minCost, splitRadix, PRDFT, URDFT");
    PrintLine("DFT format: interleavedComplex");
    PrintLine("tspl DFT configuartion: tsplCT, tsplPFA, tsplRader, tsplBluestein, bluesteinMinPrime, bluesteinExtraSizes, raderAvoidSizes");
    PrintLine("vectorization: svct, splitL, oddSizes, realVect, cplxVect, useConj, stdTTensor, flipIxA, pushTag");
    PrintLine("options: globalUnrolling, useArea, verify, verifyDP, verifyTolerance, verbosity, faultTolerant, mode, language, highPerf, propagateNth, useDeref\n");
end;

Class(SIMDGlobals, rec(
    svct := True,
    DPOpts := rec(
        globalUnrolling := true,
        globalUnrollingMin := 128,
        globalUnrollingMax := 128,
        defaultUnrolling := 128, # NOTE: THIS AFFECTS Rader.maxSize and other things. BAD IDEA.
        verbosity := 0,
        faultTolerant := false
    ),

    getDPOpts := meth(arg)
        local self, isa, opts;
        if Length(arg) = 1 then
            _help();
            return;
        fi;
        self := arg[1];
        isa := arg[2];
        opts := When(Length(arg) >=3, arg[3], rec());

        return rec(
            globalUnrolling := When(IsBound(opts.globalUnrolling) or (IsBound(opts.globalUnrollingMin) and IsBound(opts.globalUnrollingMax)), true, self.DPOpts.globalUnrolling),
            globalUnrollingMin := Cond(IsBound(opts.globalUnrolling), opts.globalUnrolling, IsBound(opts.globalUnrollingMin), opts.globalUnrollingMin, self.DPOpts.globalUnrollingMin * isa.v),
            globalUnrollingMax := Cond(IsBound(opts.globalUnrolling), opts.globalUnrolling, IsBound(opts.globalUnrollingMax), opts.globalUnrollingMax, self.DPOpts.globalUnrollingMax * isa.v),
            verbosity := When(IsBound(opts.verbosity), opts.verbosity, self.DPOpts.verbosity),
            timeBaseCases := true
        );
    end,

#------------------------------------------
    getOpts := meth(arg)

         local self, isa, argrec, opts, profile_rec, arg3, bsizes, unr;

        if Length(arg) = 1 then
            _help();
            return;
        fi;
        self := arg[1];
        isa := arg[2];

        #   some defaults
        argrec := rec(
            #   stride permutation configuration
            stride := false,

            #   DFT configuration
            #   scalar algorithms
            splitRadix := false,
            SplitRadix_maxSize := SIMDGlobals.DPOpts.defaultUnrolling/2,
            CT  := true,
            CT_forcePrimeFactor := false,
            PFA := true,
            PFA_maxSize := SIMDGlobals.DPOpts.defaultUnrolling/2,
            PD  := true,
            Rader := false,
            RealRader := true,
            Rader_maxSize := SIMDGlobals.DPOpts.defaultUnrolling/2,
            Rader_minSize := 3,
            Bluestein := false,
            Bluestein_maxSize := false,
            minCost := false,
            Mincost_maxSize := SIMDGlobals.DPOpts.defaultUnrolling/2,
            PRDFT := false,
            PRDFT_PF_maxSize := SIMDGlobals.DPOpts.defaultUnrolling,
            URDFT := false,
            URDFT_maxRadix := SIMDGlobals.DPOpts.defaultUnrolling/2,
            #   vector algorithms -- preset for small sizes unrolled code.
            svct := true,
            splitL := false,    #   turn on for large sizes, as SVCT's payoff diminishes and complicates matters
            useConj := false,   #   not supported right now
            oddSizes := false,   #   only for unrolled code. turn off splitL as DP takes a wrong turn for small sizes
            realVect := true,   #   use real-based vectorization
            cplxVect := false,   #   use complex-based vectorization
            pushTag := true,
            flipIxA := true,
            splitComplexTPrm := false, # used in SplitComplex DFTs and in DCTs
            TRCDiag_VRCLR := false,
            #   tSPL breakdown rules
            tsplBluestein := true,
            bluesteinMinPrime := 997,   # maybe level-based threshold? basically no automatic Bluestein applicability for now
            bluesteinExtraSizes := [ 23, 46, 47, 49, 59, 67, 79, 83, 94, 103, 106, 107, 115 ],  #   for these sizes I needed Bluestein so far
            raderAvoidSizes := [ 47, 59, 107 ], # avoid Rader for these sizes as codesize explodes for unrolled code
            tsplRader := true,
            tsplPFA := true,
            stdTTensor := true,
            tsplCT := true,
            tsplVBase := true,
            tsplCxVBase := true,
            tsplCT_oddvloop := false,
            loopOddSizes := false,
            TRCbyDef := false,
            #   other defaults
            includeMath   := true,
            faultTolerant := SIMDGlobals.DPOpts.faultTolerant,
            language      := Cond(LocalConfig.cpuinfo.default_lang <> "", LocalConfig.cpuinfo.default_lang, SpiralDefaults.language),
            verify        := false,
            verifyDP := false,
            verifyTolerance := 1E-3,
            verbosity := 0,
            interleavedComplex := true,
            highPerf := true,
            SIMD := LocalConfig.cpuinfo.SIMDname,
            fracbits := When(IsBound(isa.fracbits), isa.fracbits, false),
            useArea := false,
            processIntTables := false,
            propagateNth := false,
            useDeref := true,
            globalUnrolling := SIMDGlobals.DPOpts.defaultUnrolling,

            fma := false,
            cxfma := false,
            splitVDiag := true,
            finalSReduce := false,
            fixUnalignedLoadStore := false
        );

        if Length(arg) = 3 then
            arg3 := arg[3];
            if IsBound(arg3.__name__)  then Unbind(arg3.__name__); fi;
            if IsBound(arg3.__doc__)   then Unbind(arg3.__doc__); fi;
            if IsBound(arg3.__bases__) then Unbind(arg3.__bases__); fi;
            argrec := CopyFields(argrec, arg3);
        fi;

     profile_rec :=
        When(IsBound(SpiralDefaults.profile),
           rec(
              profile := SpiralDefaults.profile,
	          measureFunction := SpiralDefaults.profile.meas
           ),
           rec()
     );

     bsizes := Filtered(argrec.bluesteinExtraSizes, i -> not IsInt(i/isa.v^2));

     opts := CopyFields(When(argrec.highPerf, SpiralDefaults.highPerf(), SpiralDefaults), isa.getOpts(), profile_rec, rec(
         baseHashes := When(argrec.svct or argrec.flipIxA, [SIMD_ISA_DB.getHash()], []),
         breakdownRules := rec(
             TTwiddle := [TTwiddle_Tw1],
             DFT := Concat([DFT_Base],
                When(argrec.realVect and argrec.tsplVBase,   [DFT_tSPL_VBase], []),
                When(argrec.cplxVect and argrec.tsplCxVBase, [DFT_tSPL_CxVBase, DFT_tSPL_CxVBase2], []),
                When(argrec.splitRadix, [CopyFields(DFT_SplitRadix, rec(maxSize := argrec.SplitRadix_maxSize)) ], []),
                When(argrec.PD,         [DFT_PD], []),

                When(argrec.Bluestein,  [CopyFields(DFT_Bluestein, rec(
		            maxSize := argrec.Bluestein_maxSize, switch := true))], []),

                When(argrec.CT,         [CopyFields(DFT_CT, rec(forcePrimeFactor := argrec.CT_forcePrimeFactor))],
		            When(argrec.URDFT and not argrec.minCost, [CopyFields(DFT_CT, rec(maxSize := 4))], [])),

                When(argrec.PFA,        [CopyFields(DFT_GoodThomas, rec(maxSize := argrec.PFA_maxSize))], []),
                When(argrec.Rader,      [CopyFields(DFT_Rader,      rec(minSize := argrec.Rader_minSize, maxSize := argrec.Rader_maxSize))], []),
                When(argrec.RealRader,  [CopyFields(DFT_RealRader,  rec(maxSize := argrec.Rader_maxSize))], []),
                When(argrec.minCost,    [CopyFields(DFT_CT_Mincost, rec(maxSize := argrec.Mincost_maxSize))], []),
                When(argrec.tsplCT,     [DFT_tSPL_CT], []),
                When(argrec.tsplCT_oddvloop,[DFT_tSPL_CT_oddvloop], []),
                When(argrec.tsplBluestein, [ CopyFields(DFT_tSPL_Bluestein, rec(
                    applicableSizes := i -> (i in bsizes) or (not IsInt(i/isa.v^2)
			                    and ForAny(Factors(i), j -> j >= argrec.bluesteinMinPrime)),
                    minRoundup := isa.v^2,
                    customFilter := DetachFunc(Subst(i -> IsInt(i/$(isa.v^2)))))) ], []),

                When(argrec.tsplPFA,    [DFT_tSPL_GoodThomas], []),

                When(argrec.tsplRader,  [CopyFields(DFT_tSPL_Rader, rec(
                    useSymmetricAlgorithm := true,
                    avoidSizes := argrec.raderAvoidSizes)) ], []),

                When(argrec.PRDFT,      [DFT_PRDFT], []),
                When(argrec.URDFT,      [CopyFields(spiral.sym.DFT_URDFT_Decomp, rec(
		    forTransposition := true,
		    maxRadix := argrec.URDFT_maxRadix))], [])
             ),
             DFT3 := [ DFT3_Base, DFT3_CT ],
             TDCT2 := [ DCT2_DCT4_tSPL ],
             TDCT3 := [ DCT3_DCT4_tSPL ],
             TDCT4 := [ DCT4_CT_tSPL ],
             TDST2 := [ DST2_DST4_tSPL ],
             TDST3 := [ DST3_DST4_tSPL ],
             TDST4 := [ DST4_CT_tSPL ],
             TMDCT := [TMDCT_DCT4_tSPL],
             TIMDCT   := [TIMDCT_TMDCT_tSPL],
             TRDFT    := [TRDFT_By_Def, TRDFT_By_Def_tr, TRDFT_DFT_NR_tSPL_New, TRDFT_CT_tSPL_New],
             TRDFT2D  := [TRDFT2D_ColRow_tSPL],
             TIRDFT2D := [TIRDFT2D_RowCol_tSPL],
             TRConv2D := [TRConv2D_TRDFT2D_tSPL],
             TDHT := [DHT_DFT_tSPL],
             TS := [ TS_vect ],
             TConjEven := [TConjEven_vec, TConjEven_vec_tr ],
             TXMatDHT := [TXMatDHT_vec],
             WHT := [ WHT_Base, WHT_BinSplit, WHT_tSPL_BinSplit, WHT_tSPL_Base ],
             MDDFT := [ MDDFT_Base, MDDFT_RowCol, MDDFT_tSPL_RowCol ],
             PrunedDFT := [ PrunedDFT_base, PrunedDFT_DFT, PrunedDFT_CT, PrunedDFT_tSPL_CT ],
             IOPrunedDFT := [
		 IOPrunedDFT_tSPL_CT, IOPrunedDFT_base, IOPrunedDFT__PrunedDFT,
		 IOPrunedDFT__PrunedDFT_T, IOPrunedDFT__Gath_PrunedDFT, IOPrunedDFT__PrunedDFT_T_Scat, IOPrunedDFT_CT ],
             InterpolateDFT := [ InterpolateDFT_tSPL_PrunedDFT ],
             Downsample := [ Downsample_base, Downsample_tag ],
             InterpolateSegmentDFT := [ InterpolateSegmentDFT_PrunedDFT, InterpolateSegmentDFT_tSPL_PrunedDFT ],

             TTensor := Concat(When(argrec.stdTTensor, [ AxI_IxB, IxB_AxI], []),
                        When(argrec.tsplPFA, [splitL_BxI__L_AxI, AxI_L__BxI_splitL], [L_BxI__L_AxI, AxI_L__BxI_L ])),

             TTensorI := Concat([ IxA_base, AxI_base, IxA_L_base, L_IxA_base, AxI_vec ],
                            When(argrec.pushTag, [ IxA_vec_push ], []),
                            When(argrec.flipIxA, [ IxA_vec ], []),
                            When(argrec.svct, [ IxA_L_vec, L_IxA_vec], []),
                            When(argrec.splitL, [ IxA_L_split_vec, L_IxA_split_vec, IxA_split_vec ], []),
                            When(argrec.oddSizes, [ AxI_svec ], []),
                            When(argrec.loopOddSizes, [ TTensorI_oddvloop ], []),
                            When(argrec.oddSizes and argrec.svct, [ IxA_L_svec ], []),
                            When(argrec.useConj, [ IxA_conj_vec ], [])
                            ),
             TL := Concat(When(argrec.cplxVect, [L_cx_real], []),
                        When(argrec.stride, [L_GV1_vtensor], []),#[L_base_vec, L_mn_m_vec, IxLxI_vtensor],
                        When(argrec.svct or argrec.flipIxA, [ L_nv_n_vec, L_nv_v_vec, L_mn_m_vec, IxLxI_vtensor ], [])
#                        When(argrec.splitL,[ L_base_vec, IxLxI_vtensor ], [])
                    ),
             TTensorInd := [ dsA_base_vec_push, L_dsA_L_base_vec, L_dsA_L_vec, L_dsA_L_base ],

             TRC := Concat(
                When(argrec.realVect, [TRC_vect], []),
                When(argrec.cplxVect, [TRC_cplx, TRC_cplx_v2], []),
		        When(argrec.TRCbyDef, [CopyFields(TRC_By_Def, rec(maxSize := 2*isa.v))], []),
                When(not(argrec.realVect) and argrec.cplxVect, [TRC_cplxvect], [])
             ),
	     TMat    := [ TMat_Base, TMat_Vec],
             TDiag   := [ TDiag_tag ],
             TRCDiag := When(argrec.TRCDiag_VRCLR, [TRCDiag_VRCLR], [ TRCDiag_tag ]),
             TGath   := [ TGath_base ],
             TScat   := [ TScat_base ],
             TRDiag  := [ TRDiag_Vec ],
             TPrm    := When(argrec.splitComplexTPrm, [TPrm_format], [ TPrm_IJ_Vec, TPrm_IP_Base1, TPrm_J, TPrm_Jv ]),
             TGrp    := [ TGrp_tag ],
             TCompose   := [ TCompose_tag ],
             TICompose  := [ TICompose_tag ],
             TRaderMid  := [ Pad_vec ],
             TDirectSum := [ A_dirsum_B_delayed ],

             TConj   := [ TConj_perm ],
             PRDFT   := [ PRDFT1_Base1, PRDFT1_Base2, PRDFT1_CT,
		 CopyFields(PRDFT1_PF, rec(maxSize := argrec.PRDFT_PF_maxSize)), PRDFT_PD, PRDFT_Rader],
             IPRDFT  := [ IPRDFT1_Base1, IPRDFT1_Base2, IPRDFT1_CT, IPRDFT_PD, IPRDFT_Rader],
             IPRDFT2 := [ IPRDFT2_Base1, IPRDFT2_Base2, IPRDFT2_CT],
             PRDFT3  := [ PRDFT3_Base1, PRDFT3_Base2, PRDFT3_CT, PRDFT3_OddToPRDFT1],
             URDFT   := [ URDFT1_Base1, URDFT1_Base2, URDFT1_Base4, CopyFields(URDFT1_CT, rec(maxRadix := argrec.URDFT_maxRadix)) ],
             GT     := [ CopyFields(GT_Base, rec(maxSize:=false)), GT_NthLoop, GT_Vec_AxI ],
             InfoNt := [ Info_Base],
             Filt   := [ spiral.transforms.filtering.Filt_Base, spiral.transforms.filtering.Filt_Blocking ],

             # old DCT/DST rules without DFT/RDFT termination
             DCT2   := [ DCT2_DCT2and4],
             DCT3   := [ DCT3_Base2, DCT3_Base3, DCT3_Base5],
             DCT4   := [ DCT4_Base2, DCT4_Base3, DCT4_DCT2andDST2, DCT4_DST4andDST2, DCT4_DCT2, DCT4_DCT2t],
             DST2   := [ DST2_Base2, DST2_toDCT2],
             DST4   := [ DST4_Base, DST4_toDCT4]
         ),
         formulaStrategies := rec(
             sigmaSpl := VectorStrategySum,
             postProcess := Concatenation(
                VectorStrategySum,
                VectorStrategyTerm,
                VectorStrategySum,
                VectorStrategySum,
                VectorStrategyTerm2,
                VectorStrategySum,
                VectorStrategySum,
                VectorStrategyRC,
                VectorStrategyRC,
                [RulesVRCTerm, s -> SubstBottomUp(s, VIxL, e->e.implement(isa))], # NOTE: WHY do I need to have these???
                VectorStrategySum,
                VectorStrategyRC,
                VectorStrategySum,
                [ BlockSumsOpts,
                  (s, opts) -> Process_fPrecompute(s, opts)
                   ],
                fix_fAdd
             ),
             preRC := [MergedRuleSet(StandardSumsRules, RulesPropagate, JoinDirectSums),
                       MergedRuleSet(StandardSumsRules, RulesPropagate, StretchRaderMid),
                       MergedRuleSet(StandardSumsRules, RulesPropagate, RulesSplitComplex),
                       TerminateDirectSums],
             rc := []
         ),

         includes := Concatenation(
            When(argrec.includeMath = false, "", ["<math.h>"]),
            When(IsBound(isa.includes), isa.includes(), [])
         ),

         generateInitFunc := true,

         XType := TPtr(isa.t.base_t()),
         YType := TPtr(isa.t.base_t()),

     # NOTE: use layering w/ CopyFields instead of When/IsBound
     language := argrec.language,
     faultTolerant := argrec.faultTolerant,
     verify := argrec.verify,
     verifyDP := argrec.verifyDP,
     verifyTolerance := argrec.verifyTolerance,
     verbosity := argrec.verbosity,
     interleavedComplex := argrec.interleavedComplex,
     useDeref := argrec.useDeref,
     propagateNth := argrec.propagateNth,
     unparser := When(IsBound(isa.unparser), isa.unparser, CMacroUnparserProg),
     codegen := VectorCodegen,
     compileStrategy := Concatenation(
                            Cond( argrec.cxfma, IndicesCS_CXFMA,
                                  argrec.fma,   IndicesCS_FMA,
                                  # else
                                  SpiralDefaults.compileStrategy),
                            When( isa.isFixedPoint, [ c -> FixedPointCode(c, isa.bits, argrec.fracbits) ] ,[]),
                            [(c,opts) -> opts.vector.isa.fixProblems(c, opts)],
                            [HashConsts, FixDataType]),

     simpIndicesInside := SpiralDefaults.simpIndicesInside :: isa.simpIndicesInside,

     vector := CopyFields(VectorDefaults, rec(
            vlen := isa.v,
            isa := isa,
            conf := argrec,
            SIMD := argrec.SIMD
        )),
        tags := isa.getTags(),
        cxtags := isa.getTagsCx()
     ));

     if IsBound(isa.countrec) then opts.countrec := isa.countrec; fi;

     if IsBound(isa.useDeref) then opts.useDeref := isa.useDeref; fi;
     if IsBound(isa.codegenStrat) then opts.codegenStrat := isa.codegenStrat; fi;
     if IsBound(isa.compileStrategy) then opts.compileStrategy := isa.compileStrategy(); fi;

     if IsBound(isa.declareConstants) then opts.declareConstants := isa.declareConstants; fi;
     opts.scalarDataModifier := "const";
     if IsBound(isa.expandVectorConstants) then opts.expandVectorConstants := isa.expandVectorConstants; fi;

     if IsBound(isa.backendConfig) then
        opts.profile := isa.backendConfig.profile;
        opts.measureFunction := isa.backendConfig.measureFunction;
     fi;

     # vector has VRC to deal with complex code so whe turn off implicit RC business in SumsRuleTreeXXX()
     opts.generateComplexCode := true;

     if IsBound(isa.arrayBufModifier) then
        opts.arrayBufModifier :=isa.arrayBufModifier;
	 elif IsBound(isa.alignmentBytes) then
	     opts.arrayBufModifier := Concat("static ", LocalConfig.compilerinfo.alignmentSpecifier(isa.alignmentBytes));
     else
         opts.arrayBufModifier := Concat("static ", LocalConfig.compilerinfo.alignmentSpecifier());
     fi;
     if IsBound(isa.arrayDataModifier) then
        opts.arrayDataModifier :=isa.arrayDataModifier;
	 elif IsBound(isa.alignmentBytes) then
	     opts.arrayDataModifier := Concat("static ", LocalConfig.compilerinfo.alignmentSpecifier(isa.alignmentBytes));
     else
         opts.arrayDataModifier := Concat("static ", LocalConfig.compilerinfo.alignmentSpecifier());
     fi;

     if IsBound(argrec.globalUnrolling) then
        # complex vectorization requires smaller global unrolling
        opts.globalUnrolling := argrec.globalUnrolling * isa.v; # /When(argrec.cplxVect, 2, 1); #NOTE: bad problem when wrapping complex guys...
     fi;
     if argrec.useArea then
        opts.globalUnrolling := 2 * opts.globalUnrolling; # * Log2Int(opts.globalUnrolling);
        opts.markBlock := MarkBlocksAreaSums;
     fi;

     if argrec.splitVDiag then
        opts.hack_vRC_VDiag := "split";
     else
        opts.hack_vRC_VDiag := "compact";
     fi;
     opts.finalSReduce := argrec.finalSReduce;
     opts.fixUnalignedLoadStore := argrec.fixUnalignedLoadStore;

#     # use Fred's magic flags to make the compiler behave -- hopefully
#     opts.propagateNth:=false;
#     opts.useDeref := true;
#     opts.doScalarReplacement:=false;
#     if argrec.safeMode then
#         opts.useDeref:=false;
#         opts.doScalarReplacement:=false;
#     fi;

      if argrec.processIntTables then
          Append(opts.formulaStrategies.preRC, [compiler.MergeIntData]);
      fi;

     opts.operations := rec(Print := (s) -> Print("<Spiral SIMD options>"));

     if IsBound(argrec.measureFinal) then opts.measureFinal := argrec.measureFinal; fi;

     return opts;

    end
));


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

Class(cu_allocate_managed, assign);

cudaMemcpyHostToHost     := "cudaMemcpyHostToHost";
cudaMemcpyHostToDevice   := "cudaMemcpyHostToDevice";
cudaMemcpyDeviceToHost   := "cudaMemcpyDeviceToHost";
cudaMemcpyDeviceToDevice := "cudaMemcpyDeviceToDevice";

Class(cu_check_errors, call);

Class(T_Class, T_Type, rec(

    fields := rec(),

    updateParams := meth(self)
        Constraint(IsString(self.params[1]));
        Constraint(When(Length(self.params)>1, IsRec(self.params[2]), true));
    end,
));

Class(Dim3, T_Class, rec(

    __call__ := (self) >> let(fields := rec(x := TUInt, y := TUInt, z := TUInt), Inherited("Dim3", fields)),

    )
);

var.("fresh_t_obj") := meth(self,id,t,params)
    local v, f;
    if IsType(t) and IsRec(params) and ForAll(RecFields(params), f->f in RecFields(t.params[2])) then
        v := var.fresh_t(id, t);
        for f in Filtered(RecFields(params), e -> IsType(t.params[2].(e))) do
            v.(f) := var.fresh_t(f, t.params[2].(f));
            v.(f).setValue(params.(f));
        od;
        return v;
    else
        Error("<t> must be a type or an integer that represents an interval");
    fi;
end;

Class(T_CudaEvent, T_Type);
Class(cu_event_create, call);
Class(cu_event_record, call);
Class(cu_event_destroy, call);
Class(cu_event_synchronize, call);
Class(cu_device_synchronize, call);
Class(cu_event_elapsed_time, call);

Class(cu_allocate, assign, rec(
   isAssign := true,
   __call__ := (self, loc, type, size) >> WithBases(self,
       rec(operations := CmdOps,
       loc  := toAssignTarget(loc),
       type  := type,
       size := size)
       ),

   rChildren := self >> [self.loc, self.type, self.size],
   rSetChild := rSetChildFields("loc", "type", "size"),

   print := (self,i,si) >> let(name := Cond(IsBound(self.isCompute) and self.isCompute,
                                            gap.colors.DarkYellow(self.__name__),
                                            IsBound(self.isLoad) and self.isLoad,
                                            gap.colors.DarkRed(self.__name__),
                                            IsBound(self.isStore) and self.isStore,
                                            gap.colors.DarkGreen(self.__name__),
                                            self.__name__),
                                Print(name, "(", self.loc, ", ", self.type, ", ", self.size, ")"))
));

Class(cu_memcpy, assign, rec(
   isAssign := true,
   __call__ := (self, loc, exp, size, kind) >> WithBases(self,
       rec(operations := CmdOps,
       loc  := toAssignTarget(loc),
       exp  := toExpArg(exp),
       size := size,
       kind := kind)),

   rChildren := self >> [self.loc, self.exp, self.size, self.kind],
   rSetChild := rSetChildFields("loc", "exp", "size", "kind"),

   print := (self,i,si) >> let(name := Cond(IsBound(self.isCompute) and self.isCompute,
                                            gap.colors.DarkYellow(self.__name__),
                                            IsBound(self.isLoad) and self.isLoad,
                                            gap.colors.DarkRed(self.__name__),
                                            IsBound(self.isStore) and self.isStore,
                                            gap.colors.DarkGreen(self.__name__),
                                            self.__name__),
                                Print(name, "(", self.loc, ", ", self.exp, ", ", self.size, ", ", self.kind, ")"))
));

Class(cu_memcpy_to_sym, assign, rec(
   isAssign := true,
   __call__ := (self, loc, exp, size) >> WithBases(self,
       rec(operations := CmdOps,
       loc  := toAssignTarget(loc),
       exp  := toExpArg(exp),
       size := size)),

   rChildren := self >> [self.loc, self.exp, self.size],
   rSetChild := rSetChildFields("loc", "exp", "size"),

   print := (self,i,si) >> let(name := Cond(IsBound(self.isCompute) and self.isCompute,
                                            gap.colors.DarkYellow(self.__name__),
                                            IsBound(self.isLoad) and self.isLoad,
                                            gap.colors.DarkRed(self.__name__),
                                            IsBound(self.isStore) and self.isStore,
                                            gap.colors.DarkGreen(self.__name__),
                                            self.__name__),
                                Print(name, "(", self.loc, ", ", self.exp, ", ", self.size, ")"))
));

Class(cu_free, call);

Class(cu_call, call, rec(
    __call__ := arg >> let(
        len := Length(arg),
        self := arg[1],
        f := arg[2],
        dim_grid := Checked(IsList(arg[3]) or IsVar(arg[3]), arg[3]),
        dim_block := Checked(IsList(arg[4]) or IsVar(arg[4]), arg[4]),
        args := arg{[5..len]},
        WithBases(self, rec(
            func := f,
            dim_grid := dim_grid,
            dim_block := dim_block,
            operations := CmdOps,
            args := List(args, toExpArg)))
        ),

    rChildren := self >> [self.func, self.dim_grid, self.dim_block]::self.args,

    rSetChild := meth(self, n, newChild)
        local len;
        len := Length(self.args);
        if n > len+3 then Error(); fi;
        if n = 1 then
            self.func := newChild;
        elif n = 2 then
            self.dim_grid := newChild;
        elif n = 3 then
            self.dim_block := newChild;
        elif n > 3 and n <= len+3 then
            self.args[n-3] := newChild;
        fi;
    end,

    print := (self,i,si) >> Print(self.func, "<<<", self.dim_grid, ",", self.dim_block, ">>>", "(", PrintCS(self.args), ")")

));

ext: gi

max_stars_repo_name: mikefranusich/spiral-package-simt

max_issues_repo_name: mikefranusich/spiral-package-simt

max_forks_repo_name: mikefranusich/spiral-package-simt


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

domain := f -> f.domain(); pdomain := p -> p.val.domain();
range  := f -> f.range();  prange := p -> p.val.range();
is_perm := f -> f.isPermutation();

Class(fId, PermClass, rec(
    domain := self >> self.params[1],
    range  := self >> self.params[1],
    def    := size -> Checked(IsPosInt0Sym(size), rec(size := size)),
    lambda := self >> let(i := Ind(self.params[1]), Lambda(i,i)),
    transpose := self >> self,
    isIdentity := True
));

Class(fBase, FuncClass, rec(
    abbrevs := [ var -> Checked(IsVar(var) or ObjId(var)=ind, [var.range, var]) ],
    def    := (N, pos) -> rec(),
    domain := self >> 1,
    range  := self >> self.params[1],
    print  := (self, i, is) >> Print(self.name, "(",
        When(ObjId(self.params[2]) in [var, ind], Print(self.params[2]), PrintCS(self.params)), ")"),
    lambda := self >> let(i := Ind(1), Lambda(i, self.params[2]))
));

# used for I_n dirsum J_n
# below represent an alternator <I_n | J_n>_x o fBase(j)
# x = (pos+par) mod 2
Class(fXbase, FuncClass, rec(
    def := (n,j,par,pos) -> rec(),
    domain := self >> 1,
    range := self >> self.params[1],
    lambda := self >> let(
	n:=self.params[1], j:=self.params[2], par:=self.params[3],
	pos:=self.params[4], i := Ind(1),
	Lambda(i, j + imod(pos+par, 2)*(n - 1 - 2*j))
    )
));

# used for I_n dirsum (I_1 dirsum J_n-1)
# below represent an alternator <I_n | I_1 dirsum J_n_1>_x o fBase(j)
# x = (pos+par) mod 2
Class(fYbase, FuncClass, rec(
    def := (n,j,par,pos) -> rec(),
    lambda := self >> let(
	n := self.params[1], j := self.params[2], par := self.params[3],
	pos := self.params[4], i := Ind(1),
	Lambda(i, cond(eq(j, 0), 0, j + imod(pos+par, 2)*(n - 2*j)))
    )
));

Class(fDot, FuncClass, rec(
    def    := (v, func) -> Chain(v.setRange(func.domain()), rec()),
    range  := self >> self.params[2].range(),
    domain := self >> self.params[2].domain(),
    lambda := self >> let(
        v    := self.params[1],
        func := self.params[2],
        i    := Ind(func.domain()),
        Lambda(i, SubstVars(func.at(i), tab((v.id) := i))))
));

Class(fPlaceholder, FuncClass, rec(
    def := func -> rec(),
    range := self >> self.params[1].range(),
    domain := self >> self.params[1].domain(),
    lambda := self >> self.params[1].lambda(),
    at := (self, n) >> self.params[1].at(n)
));

#F fAdd is the "add constant" index-mapping function
#F fAdd(N, n, k): II_n -> II_N, i->i+k
#F References: "Formal Loop Merging for Signal Transforms"
Class(fAdd, FuncClass, rec(
    def := (N, n, val) -> rec(),
    range  := self >> self.params[1],
    domain := self >> self.params[2],
    lambda := self >> let(i := Ind(self.params[2]), Lambda(i, i + self.params[3]))
));

#F fSub is the "subtract constant" index-mapping function
#F fSub(N, n, k): II_n -> II_N, i->i-k
#F References: "fAdd" function
Class(fSub, FuncClass, rec(
    def := (N, n, val) -> rec(N := N, n := n),
    lambda := self >> let(i := Ind(self.params[2]), Lambda(i, i - self.params[3]))
));

Function.hash := false;

Class(FuncClassOper, Function, rec(
    _perm := true,
    _children := [],

    child := (self, n) >> self._children[n],
    children := self >> self._children,
    numChildren := self >> Length(self._children),

    from_rChildren := (self, rch) >> ApplyFunc(ObjId(self), rch),
    rChildren := self >> self._children,
    rSetChild := meth(self, n, what)
        self._children[n] := what;
    end,
    free := self >> Union(List(self.rChildren(), FreeVars)),

    skipOneChild := false,

    advdims   := self >> [ [[ self.range() ]], [[ self.domain() ]] ],
    advrange  := self >> [[ self.range() ]],
    advdomain := self >> [[ self.domain() ]],
    # dimensionality of range (1-d, 2-d, etc)
    advrangeDim := self >> Length(self.advrange()[1]),
    # dimensionality of domain (1-d, 2-d, etc)
    advdomainDim := self >> Length(self.advdomain()[1]),

    __call__ := meth(arg)
        local self, children, lkup, res, h;
        self := arg[1];
        children := Flat(Drop(arg, 1));
	Constraint(ForAll(children, c -> IsFunction(c) or IsFuncExp(c)));
        if self.skipOneChild and Length(children)=1 then return children[1]; fi;

        h := self.hash;
        if h<>false then
            lkup := h.objLookup(self, children);
            if lkup[1]<>false then return lkup[1]; fi;
        fi;
        res := WithBases(self, rec(operations := RewritableObjectOps, _children := children));
        if h<>false then return h.objAdd(res, lkup[2]);
        else return res;
        fi;
    end,

    print := self >> Print(self.name, "(", PrintCS(self._children), ")"),
    range := self >> Error("not implemented"),
    domain := self >> Error("not implemented"),

    rightBinary := self >> FoldR1(self._children, (p,x) -> let(base:=self.__bases__[1], base(x, p))),
    leftBinary  := self >> FoldL1(self._children, (p,x) -> let(base:=self.__bases__[1], base(p, x))),

    equals := (self, o) >> ObjId(o) = ObjId(self) and self.rChildren() = o.rChildren(),

    lessThan := (self, o) >> Cond(
        ObjId(self) <> ObjId(o), ObjId(self) < ObjId(o),
        [ ObjId(self), self.rChildren() ] < [ ObjId(o), o.rChildren() ]
    ),
));

FuncClass.tensor_op := (v,range,domain,fl,gl) ->
    Lambda(v, range * fl.at(idiv(v, domain)) + gl.at(imod(v, domain)));

FuncClassOper.tensor_op := FuncClass.tensor_op;

Class(fTensorBase, SumsBase, FuncClassOper, rec(
    domain := self >> Product(self._children, x->x.domain()),

    range := self >> let(cht := List(self._children, x->x.range()),
        Cond(ForAll(cht, x->not IsType(x)), Product(cht),
             ForAll(cht, x->not IsType(x) or x=TInt), TInt,
             Error("self.range() can't handle this combination of self._children ranges"))),

    skipOneChild := true,

    split_op := self >> let(
        c := self.children(), v := Ind(self.domain()),
        c1dom := c[1].domain(), c2dom := c[2].domain(),
        Checked(Length(c)=2,
            [ Cond(c1dom=1, 0, c2dom=1, v, idiv(v, c2dom)),
              Cond(c1dom=1, v, c2dom=1, 0, imod(v, c2dom)),
              v ])),

    lambda := self >> Cond(
	self.numChildren() > 2,
	    self.rightBinary().lambda(),
        let(split := self.split_op(),
            f := self.child(1).lambda(),
            g := self.child(2).lambda(),
            jv := DropLast(When(Length(f.vars) > Length(g.vars), f.vars, g.vars), 1),
            Lambda(Concatenation(jv, [split[3]]),
                   self.combine_op(jv, split, f, g)))),

    isIdentity := self >> ForAll(self._children, IsIdentity),
));

Class(fTensor, fTensorBase, rec(
    combine_op := (self, jv, split, f, g) >>
        self.child(2).range() * f.relaxed_at(jv, split[1]) + g.relaxed_at(jv, split[2]),

    transpose := self >> self.__bases__[1](List(self.children(), c->c.transpose()))
));

Class(diagTensor, fTensorBase, rec(
    print := FuncClassOper.print,
    range := self >> UnifyTypes(List(self.children(), x->x.range())),
    combine_op := (self, jv, split, f, g) >> f.relaxed_at(jv, split[1]) * g.relaxed_at(jv, split[2])
));

Class(gammaTensor, fTensorBase, rec(
    isCyclic := true,
    print := FuncClassOper.print,
    combine_op := (self, jv, split, f, g) >> let(
        r := self.child(1).range(), s := self.child(2).range(),
        imod(s * no_mod(f.relaxed_at(jv,split[1])) + r * no_mod(g.relaxed_at(jv,split[2])), self.range()))
));

_rankedLambdaCompose := function(l1, l2)
    local jv, v, vars;
    l1.expr := l1.expr.eval();
    l2.expr := l2.expr.eval();

    vars := When(Length(l1.vars) > Length(l2.vars), l1.vars, l2.vars);
    v := Last(vars);
    jv := vars{[1..Length(vars)-1]};

    return Lambda(vars, l1.relaxed_at(jv, l2.relaxed_at(jv, v)));
end;

Class(fCompose, FuncClassOper, rec(
    domain := self >> Last(self._children).domain(),
    range := self >> self._children[1].range(),
    advdomain := self >> Last(self._children).advdomain(),
    advrange := self >> self._children[1].advrange(),
    skipOneChild := true,

    lambda := self >>
        FoldL1(List(self.children(), z->z.lambda()), _rankedLambdaCompose),

    transpose := self >> self.__bases__[1](List(Reversed(self.children()), c->c.transpose())),

    isIdentity := self >> ForAll(self._children, IsIdentity),
));

Class(fDirsum, FuncClassOper, rec(
    domain := self >> Sum(self._children, x->x.domain()),
    range  := self >> Sum(self._children, x->x.range()),
    skipOneChild := true,

    lambda := meth(self)
        local dom_spans, row_spans, i, conds, v;
        # spans uses 1-based offsets
        dom_spans := BaseOverlap.spans(0, List(self.children(), x->x.domain())) - 1;
        row_spans := BaseOverlap.spans(0, List(self.children(), x->x.range())) - 1;
        conds := [];
        v := Ind(self.domain());
        for i in [1..self.numChildren()] do
            Add(conds, leq(dom_spans[i][1], v, dom_spans[i][2]));
            Add(conds, row_spans[i][1] + self.child(i).lambda().at(v - dom_spans[i][1]));
        od;
        return Lambda(v, ApplyFunc(cond, conds));
    end,

    transpose := self >> self.__bases__[1](List(self.children(), c->c.transpose()))
));

Class(fStack, FuncClassOper, rec(
    domain := self >> Sum(self._children, x->x.domain()),
    # NB: use lowercase max that handles symbolics
    range := self >> ApplyFunc(max, List(self._children, x->x.range())),
    subdomainsDivisibleBy := (self, n) >> ForAll(self._children, x -> x.domain() mod n = 0),
    skipOneChild := true,

    lambda := meth(self)
        local dom_spans, i, conds, v;
        # spans uses 1-based offsets
        dom_spans := BaseOverlap.spans(0, List(self.children(), x->x.domain())) - 1;
        conds := [];
        v := Ind(self.domain());
        for i in [1..self.numChildren()] do
            Add(conds, leq(dom_spans[i][1], v, dom_spans[i][2]));
            Add(conds, self.child(i).lambda().at(v - dom_spans[i][1]));
        od;
	# Else branch needed below (=errExp), it ensures that cond(..).ev() does not crash
	# with nested conds. For example:
	#   cond(c1, cond(...., errExp), b)
	# if c1 is false, our inner cond can be invalid, because then b is taken instead
	Add(conds, errExp(conds[2].t));
        return Lambda(v, ApplyFunc(cond, conds));
    end
));

Class(diagDirsum, fStack, rec(
    range := self >> UnifyTypes(List(self.children(), x->x.range()))
));

# Function for Pease twiddle access. Maps n->n to index into
#   Diag(TC(n,2,0)) based upon iteration l.
#   fAccPease(n, r, l)
Class(fAccPease, FuncClass, rec(
   def := (n, r, l) -> rec(N:=n, n:=n),
   lambda := self >> let(
       N := self.params[1], k := Ind(N),
       R := self.params[2], l := self.params[3],

       Lambda(k, let(
           bits:=Log2Int(N)-Log2Int(R),
	       rCyclicShift(
               bin_and(
                   k,
                   concat(
                       bin_shl(
                           2^(bits)-1,
                           l*Log2Int(R),
                           bits
                       ),
		               R-1,
                       Log2Int(R)
                   )
               ),
               Log2Int(R),
               Log2Int(N)
           )
	   )))
));

#F fPad(<diag_func>)
#F
#F   Index function with 0 range and non zero domain.
#F   Returns result of <diag_func> diagonal function wrapped into funcExp().
#F   Relies on rewriting nth(<loc>, funcExp(A)) -> A.

Class(fPad, FuncClass, rec(
    def       := (f) -> rec(),

    domain    := self >> self.params[1].domain(),
    range     := self >> 0,

    lambda    := self >> let(
                    i := Ind(self.domain()),
                    l := self.params[1].lambda(),
                    Lambda(i, funcExp(l.at(i)))),

    transpose := self >> Error("fPad is not transposable, needs inert form"),
));

#F fInsert(<map>, <value>)
#F
#F Index function with padding.
#F
#F   Ex: Gath(fInsert([1, 1, 0, 1, 0], 3))*X ==
#F
#F   [ [ 1, 0, 0 ],
#F     [ 0, 1, 0 ],
#F     [ 0, 0, 0 ],  * X + [ 0, 0, 3, 0, 3 ]^T
#F     [ 0, 0, 1 ],
#F     [ 0, 0, 0 ] ]
#F
#F

Class(fInsert, FuncClass, rec(
    def := (map, val) -> rec(),
    domain := self >> Length(self.params[1]),
    range  := self >> Sum(self.params[1]),
    lambda := self >> let(
        i := Ind(self.domain()),
        map := self.params[1],
        v   := funcExp(self.params[2]),
        ctr := Counter(0),
        Lambda(i, ApplyFunc( cond,
            ConcatList([1..Length(map)], j ->[eq(i, j-1), Cond(map[j]<>0, ctr.next(), v)])
            :: [errExp(TInt)]))),
    transpose := self >> Error("fInsert is not transposable, needs inert form"),
));