1731def sqsPrologue = (lambda unrestricted order : Nat . (lambda unrestricted seq : Nat . (lambda unrestricted heads : Nat . (lambda unrestricted keyValueHeads : Nat .
1732 (lambda unrestricted rest : (family SM86Program) .
1733 (let unrestricted groups = (naturalDivideUnchecked heads keyValueHeads) in
1734 (let unrestricted headBytes = (naturalMultiply seq (naturalMultiply saHeadWidth saHalfBytes)) in
1735 (let unrestricted rowBytes = (naturalMultiply heads (naturalMultiply saHeadWidth saHalfBytes)) in
1736 (saS2R saTid (constructor SM86SpecialRegister SM86ThreadIdX)
1737 (saS2R saTileIndex (saTileAxis order)
1738 (saS2R saKeyValueHead (saKeyValueAxis order)
1739 (saS2R saHead (saGroupAxis order)
1740 (saMovImm saOne 1
1741 (saMovImmAfter saScratch 31 saWait5
1742 (saImad saHead saKeyValueHead groups saHead
1743 (saCountDownTile order (naturalDivideUnchecked seq saTile)
1744 (saAnd saLane saTid saScratch
1745 (saShr saWarp saTid 5
1746 (saShr saQuad saLane 2
1747 (saMovImm saScratch 3
1748 (saAnd saQuadLane saLane saScratch
1749 (saImad saRow saWarp 16 saQuad
1750 (saImad saRow saTileIndex saTile saRow
1751 -- Q's and dO's A fragments: head h's plane, the row, the quad lane's pair
1752 (saMovImm saScratch 0
1753 (saImad saScratch saRow (naturalMultiply saHeadWidth saHalfBytes) saScratch
1754 (saImad saScratch saQuadLane 4 saScratch
1755 (saImad saScratch saHead headBytes saScratch
1756 (saWide saPointer saScratch saOne (saArgument 0)
1757 (saFor 16 (lambda unrestricted i : Nat .
1758 (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
1759 (saLoad (saQReg kt j) saPointer
1760 (naturalAdd (naturalMultiply kt 32)
1761 (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
1762 (naturalMultiply (naturalDivideUnchecked j 2) 16)))
1763 saSB1 saWaitNone))))
1764 (saWide sqOutputGradientLoadPointer saScratch saOne (saArgument 4)
1765 (saFor 16 (lambda unrestricted i : Nat .
1766 (let unrestricted kt = (naturalDivideUnchecked i 4) in (let unrestricted j = (naturalModuloUnchecked i 4) in
1767 (saLoad (naturalAdd (sqDOut kt) j) sqOutputGradientLoadPointer
1768 (naturalAdd (naturalMultiply kt 32)
1769 (naturalAdd (naturalMultiply (naturalModuloUnchecked j 2) (naturalMultiply 8 (naturalMultiply saHeadWidth saHalfBytes)))
1770 (naturalMultiply (naturalDivideUnchecked j 2) 16)))
1771 saSB1 saWaitNone))))
1772 -- dQ's offset (binary32 rows), kept in saKeyPointer's partner until the
1773 -- epilogue: saScratch becomes O's (merged plane) for the row terms
1774 (saMovImm saKeyOffset 0
1775 (saImad saKeyOffset saRow (naturalMultiply saHeadWidth 4) saKeyOffset
1776 (saImad saKeyOffset saQuadLane 8 saKeyOffset
1777 (saImad saKeyOffset saHead (naturalMultiply seq (naturalMultiply saHeadWidth 4)) saKeyOffset
1778 (saMovImm saScratch 0
1779 (saImad saScratch saRow rowBytes saScratch
1780 (saImad saScratch saHead (naturalMultiply saHeadWidth saHalfBytes) saScratch
1781 (saImad saScratch saQuadLane 4 saScratch
1782 -- the row terms' offset: 4 (h seq + row)
1783 (saImad saRow saHead seq saRow
1784 (saMovImm saValueOffset 0
1785 (saImad saRow saRow 4 saValueOffset
1786 (sqRowTerms seq heads
1787 -- dQ's offset into saScratch; K's and V's from key 0, K^T's
1788 (saAddImm saScratch saKeyOffset 0
1789 (saMovImm saKeyOffset 0
1790 (saImad saKeyOffset saQuad (naturalMultiply saHeadWidth saHalfBytes) saKeyOffset
1791 (saImad saKeyOffset saQuadLane 4 saKeyOffset
1792 (saImad saKeyOffset saKeyValueHead headBytes saKeyOffset
1793 (saMovImm sqKTOffset 0
1794 (saImad sqKTOffset saQuad (naturalMultiply seq saHalfBytes) sqKTOffset
1795 (saImad sqKTOffset saQuadLane 4 sqKTOffset
1796 (saImad sqKTOffset saKeyValueHead headBytes sqKTOffset
1797 (saAddImm saCount saTileIndex 1
1798 (saImad saMaskBase saWarp 16 saQuad
1799 (saImad saMaskBase saQuadLane 4294967294 saMaskBase
1800 (saAddImm saMaskBase saMaskBase (naturalSaturatingSubtract 4294967296 (naturalSaturatingSubtract 4096 64))
1801 (saMovConst saScale (saArgument 9)
1802 (saFor 32 (lambda unrestricted i : Nat . (saMovImm (naturalAdd 68 i) 0))
1803 (sqsSetup headBytes rest)))))))))))))))))))))))))))))))))))))))))))))))))))))))))))The compiler supplied declaration spans and resolved links from this source snapshot. This page does not assert that this file belongs to a checked closure.