Merge Dev into Main

.gitignore

@@ -1,3 +1,6 @@
 .idea/*
+.vscode/*
 */__pycache__/*
-cmake-build-debug/*
+__pycache__/*
+cmake-build-debug/*
+run/.tmp/*

apps/compression/asg.py

@@ -0,0 +1,29 @@
+from core.asg import *
+from apps import compression
+
+
+def truncate(data, nbits):
+    return Encoder('truncate', data, nbits)
+
+class Encoder(Tensor):
+    def __init__(self, op_type, *operators):
+        self.op_type = op_type
+        self.operators = []
+        for opr in operators:
+            self.operators.append(opr)
+            if isinstance(opr, ASTNode):
+                opr.ref_by.append(self)
+
+        if type(self.operators[1]) == int:
+            self.operators[1] = Const(self.operators[1], 'int')
+        assert len(self.operators[0].ref_size) == 1
+        ref_size = [self.operators[1]]
+        fix_size = []
+
+        name = f'{op_type}_' + '_'.join([op.name if hasattr(op, 'name') else '' for op in self.operators])
+        super().__init__(name, ref_size, 'int', fix_size)
+
+    def _gen_ir(self):
+        return compression.asg2ir.gen_ir(self)
+
+

apps/compression/asg2ir.py

@@ -0,0 +1,35 @@
+import helpers
+from core.asg import *
+from core.ir import *
+from apps.compression.asg import *
+
+
+def gen_ir(node):
+    assert isinstance(node, ASTNode)
+    if node.eval or len(node.decl) > 0 or (type(node) == TensorOp and len(node.compute) > 0):
+        return node
+
+    if type(node) == Encoder:
+        if node.op_type == 'truncate':
+            node.operators[0]._gen_ir()
+            node.operators[1]._gen_ir()
+            size = helpers.get_ir_of_size(node._size())
+            node.eval = Ndarray(node.dtype, size)
+            node.decl.append(Decl(node.eval))
+            src = '''unsigned m = 0x0000FFFF;
+    for (int j = 16; j != 0; j = j >> 1, m = m ^ (m << j)) {
+        for (int k = 0; k < 32; k = (k + j + 1) & ~j) {
+            unsigned t = (DATA[k] ^ (DATA[k+j] >> j)) & m;
+            DATA[k] = DATA[k] ^ t;
+            DATA[k+j] = DATA[k+j] ^ (t << j);
+        }
+    }'''
+            node.compute = [Code(src, {'DATA': node.operators[0].eval})]
+
+            output_loop = Loop(0, node.operators[1].eval, 1, [])
+            assign = Assignment(Indexing(node.eval, output_loop.iterate), Indexing(node.operators[0].eval, Expr(Expr(Literal(32, 'int'), node.operators[1].eval, '-'), output_loop.iterate, '+')))
+            output_loop.body.append(assign)
+            node.compute.append(output_loop)
+
+
+    return node

apps/compression/test.py

@@ -0,0 +1,80 @@
+from core.asg import *
+import codegen
+import torch
+import run
+
+
+def compression():
+    nblocks = 50
+    input = Tensor('input', (nblocks, 32), dtype='float')
+    quant_res = (input * 1000).round()
+    lorenzo_res = quant_res[:][0:32] - quant_res[:][-1:31]
+    encode_nbits = lorenzo_res.abs().max(axis=1).nbits()
+    ofs = encode_nbits.prefix_sum(inclusive=False)
+    compressed_res = apply(lambda x, y: truncate(x, y), (lorenzo_res, encode_nbits), out_ofs=ofs)
+    res = compressed_res
+    code = codegen.cpu.print_cpp(res._gen_ir())
+
+    print(code)
+
+    code = '#include "../../compression/util.h"\n' + code
+
+    data = torch.rand(nblocks, 32)
+
+    d = run.cpu.compile_and_run(code, data)
+    # d = run.cpu.run(data)
+    print(d)
+
+
+def invert_lorenzo():
+    nblocks = 50
+    input = Tensor('input', (nblocks, 32), dtype='float')
+    quant_res = (input * 1000).round()
+    lorenzo_res = quant_res[:][0:32] - quant_res[:][-1:31]
+
+    recover_res = Tensor('recovered', input._size(), dtype='float')
+    recover_res = setval(recover_res, recover_res[:][-1:31] + lorenzo_res[:][0:32])
+    # recover_res = lorenzo_res.prefix_sum(axis=1)
+    res = recover_res
+
+    code = codegen.cpu.print_cpp(res._gen_ir())
+    code = '#include "../../compression/util.h"\n' + code
+
+    data = torch.rand(nblocks, 32)
+
+    print((data * 1000).round())
+
+    d = run.cpu.compile_and_run(code, data)
+    print(d)
+
+def regression():
+    nblocks = 50
+    input = Tensor('input', (nblocks, 32), dtype='float')
+    quant_res = (input * 1000).round()
+    res = apply(lambda x: einsum('i,i->', x, Const(slice(1, 33, 1), 'slice')), (quant_res, ))
+
+    code = codegen.cpu.print_cpp(res._gen_ir())
+    print(code)
+
+
+def regression2():
+    nblocks = 50
+    input = Tensor('input', (nblocks, 6, 6, 6), dtype='float')
+    quant_res = (input * 1000).round()
+    tmp1 = apply(lambda x: einsum('ijk,i->', x, Const(slice(1, 33, 1), 'slice')), (quant_res, ))
+    tmp2 = apply(lambda x: einsum('ijk,j->', x, Const(slice(1, 33, 1), 'slice')), (quant_res, ))
+    tmp3 = apply(lambda x: einsum('ijk,k->', x, Const(slice(1, 33, 1), 'slice')), (quant_res, ))
+    tmp4 = apply(lambda x: einsum('ijk,->', x, None), (quant_res, ))
+
+    coefficient = (tmp1 + tmp2 + tmp3) / tmp4
+
+    code = codegen.cpu.print_cpp(coefficient._gen_ir())
+    print(code)
+
+
+
+if __name__ == "__main__":
+    # compression()
+    # invert_lorenzo()
+    # regression()
+    regression2()

apps/compression/util.h

@@ -0,0 +1,22 @@
+//
+// Created by Jiang, Peng on 11/10/23.
+//
+
+#ifndef UTIL_H
+#define UTIL_H
+
+#include <math.h>
+#include <stdlib.h>
+
+template <class T=int>
+int nbits(T n) {
+    for (int i = 30; i>=0; i--) {
+        if (n & (1 << i)) {
+            return i + 1;
+        }
+    }
+    return 0;
+}
+
+
+#endif //UTIL_H