linalg snitch

import%20marimo%0A%0A__generated_with%20%3D%20%220.23.15%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20return%20(mo%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Aasync%20def%20_()%3A%0A%20%20%20%20%23%20Get%20the%20current%20notebook%20URL%2C%20drop%20the%20'blob'%20URL%20components%20that%20seem%20to%20be%20added%2C%0A%20%20%20%20%23%20and%20add%20the%20buildnumber%20that%20a%20makethedocs%20PR%20build%20seems%20to%20add.%20This%20allows%20to%20load%0A%20%20%20%20%23%20the%20wheel%20both%20locally%20and%20when%20deployed%20to%20makethedocs.%0A%20%20%20%20async%20def%20import_xdsl()%3A%0A%20%20%20%20%20%20%20%20import%20re%0A%20%20%20%20%20%20%20%20from%20urllib.parse%20import%20urlparse%0A%0A%20%20%20%20%20%20%20%20import%20micropip%0A%20%20%20%20%20%20%20%20from%20marimo%20import%20notebook_location%0A%0A%20%20%20%20%20%20%20%20url%20%3D%20str(notebook_location()).replace(%22blob%3A%22%2C%20%22%22)%0A%20%20%20%20%20%20%20%20print(f%22DEBUG%3A%20notebook%20url%20(full)%3A%20%7Burl%7D%22)%0A%0A%20%20%20%20%20%20%20%20url_parsed%20%3D%20urlparse(url)%0A%20%20%20%20%20%20%20%20scheme%20%3D%20url_parsed.scheme%0A%20%20%20%20%20%20%20%20netloc%20%3D%20url_parsed.netloc%0A%0A%20%20%20%20%20%20%20%20print(f%22DEBUG%3A%20notebook%20url%20(parsed)%3A%20%7Burl_parsed%7D%22)%0A%0A%20%20%20%20%20%20%20%20url%20%3D%20re.sub(%0A%20%20%20%20%20%20%20%20%20%20%20%20%22(%5B%5E%2F%5D)%2F(%5Ba-f0-9-%5D%2B-%5Ba-f0-9-%5D%2B-%5Ba-f0-9-%5D%2B-%5Ba-f0-9-%5D%2B)%22%2C%20%22%5C%5C1%2F%22%2C%20url%2C%20count%3D1%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20buildnumber%20%3D%20re.sub(%22.*--(%5B0-9%2B%5D%2B).*%22%2C%20%22%5C%5C1%22%2C%20url%2C%20count%3D1)%0A%0A%20%20%20%20%20%20%20%20new_url%20%3D%20scheme%20%2B%20%22%3A%2F%2F%22%20%2B%20netloc%0A%0A%20%20%20%20%20%20%20%20if%20buildnumber%20!%3D%20url%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20new_url%20%3D%20new_url%20%2B%20%22%2F%22%20%2B%20buildnumber%20%2B%20%22%2F%22%0A%0A%20%20%20%20%20%20%20%20print(f%22DEBUG%3A%20notebook%20url%20(trimmed)%3A%20%7Bnew_url%7D%22)%0A%0A%20%20%20%20%20%20%20%20await%20micropip.install(%22xdsl%20%40%20%22%20%2B%20new_url%20%2B%20%22%2Fxdsl-0.0.0-py3-none-any.whl%22)%0A%0A%20%20%20%20await%20import_xdsl()%0A%20%20%20%20from%20xdsl.utils%20import%20marimo%20as%20xmo%0A%0A%20%20%20%20return%20(xmo%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20%23%20Import%20all%20the%20necessary%20functionality%20from%20xDSL%20for%20this%20notebook%0A%20%20%20%20%23%20If%20you%20see%20an%20error%20about%20xdsl%20not%20being%20defined%20run%20this%20cell%20manually%0A%0A%20%20%20%20from%20xdsl.backend.riscv.lowering%20import%20(%0A%20%20%20%20%20%20%20%20convert_arith_to_riscv%2C%0A%20%20%20%20%20%20%20%20convert_func_to_riscv_func%2C%0A%20%20%20%20%20%20%20%20convert_memref_to_riscv%2C%0A%20%20%20%20%20%20%20%20convert_scf_to_riscv_scf%2C%0A%20%20%20%20)%0A%20%20%20%20from%20xdsl.backend.riscv.lowering.convert_riscv_scf_to_riscv_cf%20import%20(%0A%20%20%20%20%20%20%20%20ConvertRiscvScfToRiscvCfPass%2C%0A%20%20%20%20)%0A%20%20%20%20from%20xdsl.builder%20import%20ImplicitBuilder%0A%20%20%20%20from%20xdsl.context%20import%20Context%0A%20%20%20%20from%20xdsl.dialects%20import%20arith%2C%20func%2C%20linalg%0A%20%20%20%20from%20xdsl.dialects.builtin%20import%20AffineMap%2C%20AffineMapAttr%2C%20MemRefType%2C%20ModuleOp%2C%20f64%0A%20%20%20%20from%20xdsl.dialects.riscv%20import%20riscv_code%0A%20%20%20%20from%20xdsl.interpreters.utils.ptr%20import%20TypedPtr%0A%20%20%20%20from%20xdsl.ir%20import%20Block%2C%20Region%0A%20%20%20%20from%20xdsl.passes%20import%20PassPipeline%0A%20%20%20%20from%20xdsl.dialects%20import%20get_all_dialects%0A%20%20%20%20from%20xdsl.transforms%20import%20(%0A%20%20%20%20%20%20%20%20arith_add_fastmath%2C%0A%20%20%20%20%20%20%20%20convert_linalg_to_loops%2C%0A%20%20%20%20%20%20%20%20convert_linalg_to_memref_stream%2C%0A%20%20%20%20%20%20%20%20convert_riscv_scf_for_to_frep%2C%0A%20%20%20%20%20%20%20%20reconcile_unrealized_casts%2C%0A%20%20%20%20)%0A%20%20%20%20from%20xdsl.transforms.canonicalize%20import%20CanonicalizePass%0A%20%20%20%20from%20xdsl.transforms.riscv_allocate_registers%20import%20RISCVAllocateRegistersPass%0A%20%20%20%20from%20xdsl.transforms.riscv_lower_parallel_mov%20import%20RISCVLowerParallelMovPass%0A%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20AffineMap%2C%0A%20%20%20%20%20%20%20%20AffineMapAttr%2C%0A%20%20%20%20%20%20%20%20Block%2C%0A%20%20%20%20%20%20%20%20CanonicalizePass%2C%0A%20%20%20%20%20%20%20%20Context%2C%0A%20%20%20%20%20%20%20%20ConvertRiscvScfToRiscvCfPass%2C%0A%20%20%20%20%20%20%20%20ImplicitBuilder%2C%0A%20%20%20%20%20%20%20%20MemRefType%2C%0A%20%20%20%20%20%20%20%20ModuleOp%2C%0A%20%20%20%20%20%20%20%20PassPipeline%2C%0A%20%20%20%20%20%20%20%20RISCVAllocateRegistersPass%2C%0A%20%20%20%20%20%20%20%20RISCVLowerParallelMovPass%2C%0A%20%20%20%20%20%20%20%20Region%2C%0A%20%20%20%20%20%20%20%20TypedPtr%2C%0A%20%20%20%20%20%20%20%20arith%2C%0A%20%20%20%20%20%20%20%20arith_add_fastmath%2C%0A%20%20%20%20%20%20%20%20convert_arith_to_riscv%2C%0A%20%20%20%20%20%20%20%20convert_func_to_riscv_func%2C%0A%20%20%20%20%20%20%20%20convert_linalg_to_loops%2C%0A%20%20%20%20%20%20%20%20convert_linalg_to_memref_stream%2C%0A%20%20%20%20%20%20%20%20convert_memref_to_riscv%2C%0A%20%20%20%20%20%20%20%20convert_riscv_scf_for_to_frep%2C%0A%20%20%20%20%20%20%20%20convert_scf_to_riscv_scf%2C%0A%20%20%20%20%20%20%20%20f64%2C%0A%20%20%20%20%20%20%20%20func%2C%0A%20%20%20%20%20%20%20%20get_all_dialects%2C%0A%20%20%20%20%20%20%20%20linalg%2C%0A%20%20%20%20%20%20%20%20reconcile_unrealized_casts%2C%0A%20%20%20%20%20%20%20%20riscv_code%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20%23%20Compiling%20%60linalg%60%20to%20Snitch%0A%0A%20%20%20%20This%20notebook%20walks%20through%20compiling%20micro-kernels%20defined%20in%20%60linalg%60%20to%20RISC-V%20and%20RISC-V%20with%20extensions%20for%20%5BSnitch%5D(https%3A%2F%2Fpulp-platform.github.io%2Fsnitch%2F)%2C%20a%20neural%20network%20accelerator.%0A%0A%20%20%20%20_Toggle%20app%20view%20with%20%60%E2%8C%98%60%20%2B%20%60.%60%20or%20%60ctrl%60%20%2B%20%60.%60_%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20AffineMap%2C%0A%20%20%20%20AffineMapAttr%2C%0A%20%20%20%20Block%2C%0A%20%20%20%20ImplicitBuilder%2C%0A%20%20%20%20MemRefType%2C%0A%20%20%20%20ModuleOp%2C%0A%20%20%20%20Region%2C%0A%20%20%20%20a_shape%2C%0A%20%20%20%20arith%2C%0A%20%20%20%20b_shape%2C%0A%20%20%20%20c_shape%2C%0A%20%20%20%20f64%2C%0A%20%20%20%20func%2C%0A%20%20%20%20linalg%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20xmo%2C%0A)%3A%0A%20%20%20%20a_type%20%3D%20MemRefType(f64%2C%20a_shape)%0A%20%20%20%20b_type%20%3D%20MemRefType(f64%2C%20b_shape)%0A%20%20%20%20c_type%20%3D%20MemRefType(f64%2C%20c_shape)%0A%0A%20%20%20%20kernel_op%20%3D%20func.FuncOp(%22matmul%22%2C%20((a_type%2C%20b_type%2C%20c_type)%2C%20()))%0A%20%20%20%20with%20ImplicitBuilder(kernel_op.body)%20as%20(a%2C%20b%2C%20c)%3A%0A%20%20%20%20%20%20%20%20%23%20Add%20name%20hints%20to%20make%20it%20easier%20to%20track%20how%20values%20are%20lowered%0A%20%20%20%20%20%20%20%20a.name_hint%20%3D%20%22A%22%0A%20%20%20%20%20%20%20%20b.name_hint%20%3D%20%22B%22%0A%20%20%20%20%20%20%20%20c.name_hint%20%3D%20%22C%22%0A%20%20%20%20%20%20%20%20body%20%3D%20Region(Block(arg_types%20%3D%20(f64%2C%20f64%2C%20f64)))%0A%20%20%20%20%20%20%20%20linalg.ops.GenericOp(%0A%20%20%20%20%20%20%20%20%20%20%20%20inputs%3D(a%2C%20b)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20outputs%3D(c%2C)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20body%3Dbody%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20indexing_maps%3D(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20AffineMapAttr(AffineMap.from_callable(lambda%20m%2C%20n%2C%20k%3A%20(m%2C%20k)))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20AffineMapAttr(AffineMap.from_callable(lambda%20m%2C%20n%2C%20k%3A%20(k%2C%20n)))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20AffineMapAttr(AffineMap.from_callable(lambda%20m%2C%20n%2C%20k%3A%20(m%2C%20n)))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20iterator_types%3D(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20linalg.attrs.IteratorTypeAttr.parallel()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20linalg.attrs.IteratorTypeAttr.parallel()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20linalg.attrs.IteratorTypeAttr.reduction()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20with%20ImplicitBuilder(body)%20as%20(a_val%2C%20b_val%2C%20acc_old_val)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20prod_val%20%3D%20arith.MulfOp(a_val%2C%20b_val).result%0A%20%20%20%20%20%20%20%20%20%20%20%20acc_new_val%20%3D%20arith.AddfOp(acc_old_val%2C%20prod_val).result%0A%20%20%20%20%20%20%20%20%20%20%20%20linalg.ops.YieldOp(acc_new_val)%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%20Add%20more%20name%20hints%20to%20make%20it%20easier%20to%20track%20how%20values%20are%20lowered%0A%20%20%20%20%20%20%20%20%20%20%20%20a_val.name_hint%20%3D%20%22a%22%0A%20%20%20%20%20%20%20%20%20%20%20%20b_val.name_hint%20%3D%20%22b%22%0A%20%20%20%20%20%20%20%20%20%20%20%20acc_old_val.name_hint%20%3D%20%22acc_old%22%0A%20%20%20%20%20%20%20%20%20%20%20%20prod_val.name_hint%20%3D%20%22prod%22%0A%20%20%20%20%20%20%20%20%20%20%20%20acc_new_val.name_hint%20%3D%20%22acc_new%22%0A%20%20%20%20%20%20%20%20func.ReturnOp()%0A%0A%20%20%20%20linalg_module%20%3D%20ModuleOp((kernel_op%2C))%0A%0A%20%20%20%20mo.md(f%22%22%22%0A%0A%20%20%20%20Here%20is%20matrix%20multiplication%20defined%20in%20the%20%60linalg%60%20dialect%2C%20with%20the%20iteration%20space%20decoupled%20from%20the%20computation%3A%0A%0A%20%20%20%20%7Bxmo.module_html(linalg_module)%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%20(linalg_module%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20min_val%20%3D%201%0A%20%20%20%20max_val%20%3D%2010%0A%20%20%20%20m%20%3D%20mo.ui.slider(min_val%2C%20max_val%2C%20value%3D2%2C%20label%3D%22M%22)%0A%20%20%20%20n%20%3D%20mo.ui.slider(min_val%2C%20max_val%2C%20value%3D2%2C%20label%3D%22N%22)%0A%20%20%20%20k%20%3D%20mo.ui.slider(min_val%2C%20max_val%2C%20value%3D2%2C%20label%3D%22K%22)%0A%20%20%20%20return%20k%2C%20m%2C%20n%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(k%2C%20m%2C%20mo%2C%20n)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20We%20can%20parametrize%20the%20shapes%20of%20the%20matrices%20operated%20on%3A%0A%0A%20%20%20%20%7Bm%7D%7Bm.value%7D%0A%0A%20%20%20%20%7Bn%7D%7Bn.value%7D%0A%0A%20%20%20%20%7Bk%7D%7Bk.value%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(k%2C%20m%2C%20mo%2C%20n)%3A%0A%20%20%20%20a_shape%20%3D%20(m.value%2C%20k.value)%0A%20%20%20%20b_shape%20%3D%20(k.value%2C%20n.value)%0A%20%20%20%20c_shape%20%3D%20(m.value%2C%20n.value)%0A%0A%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20f%22%22%22%0A%0A%20%20%20%20%20%20%20%20%60%60%60%0A%20%20%20%20%20%20%20%20A%3A%20%7B'x'.join(str(dim)%20for%20dim%20in%20a_shape)%7Dxf64%20B%3A%20%7B'x'.join(str(dim)%20for%20dim%20in%20b_shape)%7Dxf64%20C%3A%20%7B'x'.join(str(dim)%20for%20dim%20in%20c_shape)%7Dxf64%0A%20%20%20%20%20%20%20%20%60%60%60%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20return%20a_shape%2C%20b_shape%2C%20c_shape%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20%23%23%23%20Compiling%20to%20RISC-V%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Context%2C%20get_all_dialects)%3A%0A%20%20%20%20linalg_ctx%20%3D%20Context()%0A%0A%20%20%20%20for%20dialect_name%2C%20dialect_factory%20in%20get_all_dialects().items()%3A%0A%20%20%20%20%20%20%20%20linalg_ctx.register_dialect(dialect_name%2C%20dialect_factory)%0A%20%20%20%20return%20(linalg_ctx%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20We%20can%20take%20this%20representation%2C%20and%20lower%20to%20RISC-V-specific%20dialects%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20PassPipeline%2C%0A%20%20%20%20convert_arith_to_riscv%2C%0A%20%20%20%20convert_func_to_riscv_func%2C%0A%20%20%20%20convert_linalg_to_loops%2C%0A%20%20%20%20convert_memref_to_riscv%2C%0A%20%20%20%20convert_scf_to_riscv_scf%2C%0A%20%20%20%20linalg_ctx%2C%0A%20%20%20%20linalg_module%2C%0A%20%20%20%20reconcile_unrealized_casts%2C%0A%20%20%20%20xmo%2C%0A)%3A%0A%20%20%20%20lower_to_riscv%20%3D%20PassPipeline(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20convert_linalg_to_loops.ConvertLinalgToLoopsPass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20convert_func_to_riscv_func.ConvertFuncToRiscvFuncPass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20convert_memref_to_riscv.ConvertMemRefToRiscvPass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20convert_arith_to_riscv.ConvertArithToRiscvPass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20convert_scf_to_riscv_scf.ConvertScfToRiscvPass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20reconcile_unrealized_casts.ReconcileUnrealizedCastsPass()%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20riscv_ctx%2C%20riscv_module%2C%20riscv_html%20%3D%20xmo.pipeline_html(%0A%20%20%20%20%20%20%20%20linalg_ctx%2C%20linalg_module%2C%20tuple((%22%22%2C%20p)%20for%20p%20in%20lower_to_riscv.passes)%0A%20%20%20%20)%0A%0A%20%20%20%20riscv_html%0A%20%20%20%20return%20riscv_ctx%2C%20riscv_module%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20%23%23%23%23%20Register%20allocation%0A%0A%20%20%20%20xDSL%20provides%20a%20register%20allocator%20for%20our%20RISC-V%20representation%2C%20that%20works%20on%20functions%20with%20structured%20control%20flow%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20CanonicalizePass%2C%0A%20%20%20%20PassPipeline%2C%0A%20%20%20%20RISCVAllocateRegistersPass%2C%0A%20%20%20%20RISCVLowerParallelMovPass%2C%0A%20%20%20%20riscv_ctx%2C%0A%20%20%20%20riscv_module%2C%0A%20%20%20%20xmo%2C%0A)%3A%0A%20%20%20%20allocate_registers%20%3D%20PassPipeline(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20RISCVAllocateRegistersPass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20CanonicalizePass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20RISCVLowerParallelMovPass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20CanonicalizePass()%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20regalloc_ctx%2C%20regalloc_module%2C%20regalloc_html%20%3D%20xmo.pipeline_html(%0A%20%20%20%20%20%20%20%20riscv_ctx%2C%20riscv_module%2C%20tuple((%22%22%2C%20p)%20for%20p%20in%20allocate_registers.passes)%2C%0A%20%20%20%20)%0A%0A%20%20%20%20regalloc_html%0A%20%20%20%20return%20regalloc_ctx%2C%20regalloc_module%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20CanonicalizePass%2C%0A%20%20%20%20ConvertRiscvScfToRiscvCfPass%2C%0A%20%20%20%20PassPipeline%2C%0A%20%20%20%20regalloc_ctx%2C%0A%20%20%20%20regalloc_module%2C%0A%20%20%20%20xmo%2C%0A)%3A%0A%20%20%20%20lower_to_asm%20%3D%20PassPipeline(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20ConvertRiscvScfToRiscvCfPass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20CanonicalizePass()%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20riscv_asm_ctx%2C%20riscv_asm_module%2C%20assembly_html%20%3D%20xmo.pipeline_html(%0A%20%20%20%20%20%20%20%20regalloc_ctx%2C%20regalloc_module%2C%20((%22%22%2C%20lower_to_asm)%2C)%2C%0A%20%20%20%20)%0A%0A%20%20%20%20assembly_html%0A%20%20%20%20return%20(riscv_asm_module%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20This%20representation%20of%20the%20program%20in%20xDSL%20corresponds%20~1%3A1%20to%20RISC-V%20assembly%2C%20and%20we%20can%20use%20a%20helper%20function%20to%20print%20that%20out.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo%2C%20riscv_asm_module%2C%20riscv_code%2C%20xmo)%3A%0A%20%20%20%20riscv_asm%20%3D%20riscv_code(riscv_asm_module)%0A%0A%20%20%20%20mo.md(f%22%22%22%5C%0A%20%20%20%20**RISC-V%20Assembly%3A**%0A%0A%20%20%20%20%7Bxmo.asm_html(riscv_asm)%7D%0A%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20%23%23%23%20Compiling%20to%20Snitch%0A%0A%20%20%20%20xDSL%20is%20also%20capable%20of%20targeting%20Snitch%2C%20and%20making%20use%20of%20its%20streaming%20registers%20and%20fixed-repetition%20loop.%20We%20use%20a%20different%20lowering%20flow%20from%20the%20linalg.generic%20representation%20to%20represent%20a%20high-level%2C%20structured%2C%20but%20Snitch-specific%20representation%20of%20the%20code%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20PassPipeline%2C%0A%20%20%20%20arith_add_fastmath%2C%0A%20%20%20%20convert_linalg_to_memref_stream%2C%0A%20%20%20%20convert_riscv_scf_for_to_frep%2C%0A%20%20%20%20linalg_ctx%2C%0A%20%20%20%20linalg_module%2C%0A%20%20%20%20xmo%2C%0A)%3A%0A%20%20%20%20from%20xdsl.transforms.test_lower_linalg_to_snitch%20import%20LOWER_MEMREF_STREAM_TO_SNITCH_STREAM_PASSES%2C%20OPTIMISE_MEMREF_STREAM_PASSES%0A%0A%20%20%20%20convert_linalg_to_snitch%20%3D%20PassPipeline(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20convert_linalg_to_memref_stream.ConvertLinalgToMemRefStreamPass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20arith_add_fastmath.AddArithFastMathFlagsPass()%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20*OPTIMISE_MEMREF_STREAM_PASSES%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20*LOWER_MEMREF_STREAM_TO_SNITCH_STREAM_PASSES%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20convert_riscv_scf_for_to_frep.ConvertRiscvScfForToFrepPass()%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20snitch_stream_ctx%2C%20snitch_stream_module%2C%20snitch_stream_html%20%3D%20xmo.pipeline_html(%0A%20%20%20%20%20%20%20%20linalg_ctx%2C%20linalg_module%2C%20tuple((%22%22%2C%20p)%20for%20p%20in%20convert_linalg_to_snitch.passes)%2C%0A%20%20%20%20)%0A%0A%20%20%20%20snitch_stream_html%0A%20%20%20%20return%20snitch_stream_ctx%2C%20snitch_stream_module%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20We%20can%20then%20lower%20this%20to%20assembly%20that%20includes%20assembly%20instructions%20from%20the%20Snitch-extended%20ISA%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(snitch_stream_ctx%2C%20snitch_stream_module%2C%20xmo)%3A%0A%20%20%20%20from%20xdsl.transforms.test_lower_linalg_to_snitch%20import%20LOWER_SNITCH_STREAM_TO_ASM_PASSES%0A%0A%20%20%20%20snitch_asm_ctx%2C%20snitch_asm_module%2C%20snitch_asm_html%20%3D%20xmo.pipeline_html(%0A%20%20%20%20%20%20%20%20snitch_stream_ctx%2C%20snitch_stream_module%2C%20tuple((%22%22%2C%20p)%20for%20p%20in%20LOWER_SNITCH_STREAM_TO_ASM_PASSES)%0A%20%20%20%20)%0A%0A%20%20%20%20snitch_asm_html%0A%20%20%20%20return%20(snitch_asm_module%2C)%0A%0A%0A%40app.cell%0Adef%20_(k%2C%20m%2C%20mo%2C%20n)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20We%20can%20see%20how%20changing%20our%20input%20sizes%20affects%20the%20assembly%20produced%3A%0A%0A%20%20%20%20%7Bm%7D%7Bm.value%7D%0A%0A%20%20%20%20%7Bn%7D%7Bn.value%7D%0A%0A%20%20%20%20%7Bk%7D%7Bk.value%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20riscv_code%2C%20snitch_asm_module%2C%20xmo)%3A%0A%20%20%20%20snitch_asm%20%3D%20riscv_code(snitch_asm_module)%0A%0A%20%20%20%20mo.md(f%22%22%22%5C%0A%20%20%20%20**Snitch%20Assembly%3A**%0A%0A%20%20%20%20%7Bxmo.asm_html(snitch_asm)%7D%0A%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20%23%23%23%20Interpreting%20the%20assembly%20using%20xDSL%0A%0A%20%20%20%20One%20of%20the%20useful%20features%20of%20xDSL%20is%20its%20interpreter.%20Here%20we've%20implemented%20all%20the%20necessary%20functions%20to%20interpret%20the%20code%20at%20a%20low%20level%2C%20to%20check%20that%20our%20compilation%20is%20correct.%20Here's%20the%20slider%20modifying%20the%20shape%20variable%20defined%20above%2C%20we%20can%20slide%20it%20to%20see%20the%20result%20of%20the%20code%20compiled%20with%20different%20input%20shapes%2C%20and%20interpreted%20at%20the%20RISC-V%20level.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(TypedPtr%2C%20a_shape%2C%20b_shape%2C%20c_shape%2C%20mo%2C%20riscv_ctx%2C%20riscv_module)%3A%0A%20%20%20%20from%20math%20import%20prod%0A%0A%20%20%20%20from%20xdsl.interpreter%20import%20Interpreter%2C%20OpCounter%0A%20%20%20%20from%20xdsl.interpreters%20import%20register_implementations%0A%20%20%20%20from%20xdsl.interpreters.shaped_array%20import%20ShapedArray%0A%0A%20%20%20%20a_len%20%3D%20prod(a_shape)%0A%20%20%20%20b_len%20%3D%20prod(b_shape)%0A%20%20%20%20c_len%20%3D%20prod(c_shape)%0A%0A%20%20%20%20a_shaped%20%3D%20ShapedArray(TypedPtr.new_float64(%5Bi%20%2B%201%20for%20i%20in%20range(a_len)%5D)%2C%20a_shape)%0A%20%20%20%20b_shaped%20%3D%20ShapedArray(TypedPtr.new_float64(%5B(i%20%2B%201)%20%2F%20100%20for%20i%20in%20range(b_len)%5D)%2C%20b_shape)%0A%20%20%20%20riscv_c_shaped%20%3D%20ShapedArray(TypedPtr.new_float64(%5B0.0%5D%20*%20c_len)%2C%20c_shape)%0A%0A%20%20%20%20riscv_op_counter%20%3D%20OpCounter()%0A%20%20%20%20riscv_interpreter%20%3D%20Interpreter(riscv_module%2C%20listeners%3D(riscv_op_counter%2C))%0A%0A%20%20%20%20register_implementations(riscv_interpreter%2C%20riscv_ctx)%0A%0A%20%20%20%20riscv_interpreter.call_op(%22matmul%22%2C%20(a_shaped.data_ptr.raw%2C%20b_shaped.data_ptr.raw%2C%20riscv_c_shaped.data_ptr.raw))%0A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20**RISC-V%20Results%3A**%0A%0A%20%20%20%20A%3A%20%7Ba_shaped%7D%0A%0A%20%20%20%20B%3A%20%7Bb_shaped%7D%0A%0A%20%20%20%20C%3A%20%7Briscv_c_shaped%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20Interpreter%2C%0A%20%20%20%20%20%20%20%20OpCounter%2C%0A%20%20%20%20%20%20%20%20ShapedArray%2C%0A%20%20%20%20%20%20%20%20a_shaped%2C%0A%20%20%20%20%20%20%20%20b_shaped%2C%0A%20%20%20%20%20%20%20%20c_len%2C%0A%20%20%20%20%20%20%20%20register_implementations%2C%0A%20%20%20%20%20%20%20%20riscv_c_shaped%2C%0A%20%20%20%20%20%20%20%20riscv_op_counter%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20Interpreter%2C%0A%20%20%20%20OpCounter%2C%0A%20%20%20%20ShapedArray%2C%0A%20%20%20%20TypedPtr%2C%0A%20%20%20%20a_shaped%2C%0A%20%20%20%20b_shaped%2C%0A%20%20%20%20c_len%2C%0A%20%20%20%20c_shape%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20register_implementations%2C%0A%20%20%20%20riscv_c_shaped%2C%0A%20%20%20%20snitch_stream_ctx%2C%0A%20%20%20%20snitch_stream_module%2C%0A)%3A%0A%20%20%20%20snitch_op_counter%20%3D%20OpCounter()%0A%20%20%20%20snitch_interpreter%20%3D%20Interpreter(%0A%20%20%20%20%20%20%20%20snitch_stream_module%2C%20listeners%3D(snitch_op_counter%2C)%0A%20%20%20%20)%0A%0A%20%20%20%20snitch_c_shaped%20%3D%20ShapedArray(TypedPtr.new_float64(%5B0.0%5D%20*%20c_len)%2C%20c_shape)%0A%0A%20%20%20%20register_implementations(snitch_interpreter%2C%20snitch_stream_ctx)%0A%0A%20%20%20%20snitch_interpreter.call_op(%0A%20%20%20%20%20%20%20%20%22matmul%22%2C%0A%20%20%20%20%20%20%20%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20a_shaped.data_ptr.raw%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20b_shaped.data_ptr.raw%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20snitch_c_shaped.data_ptr.raw%2C%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20)%0A%0A%20%20%20%20mo.md(f%22%22%22%0A%0A%20%20%20%20**Snitch%20Results%3A**%0A%0A%20%20%20%20A%3A%20%7Ba_shaped%7D%0A%0A%20%20%20%20B%3A%20%7Bb_shaped%7D%0A%0A%20%20%20%20C%3A%20%7Briscv_c_shaped%7D%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%20(snitch_op_counter%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(k%2C%20m%2C%20mo%2C%20n%2C%20riscv_op_counter%2C%20snitch_op_counter)%3A%0A%20%20%20%20rv_dict%20%3D%20dict(riscv_op_counter.ops)%0A%20%20%20%20sn_dict%20%3D%20dict(snitch_op_counter.ops)%0A%0A%20%20%20%20all_keys%20%3D%20sorted(set(rv_dict)%20%7C%20set(sn_dict))%0A%20%20%20%20max_len_key%20%3D%20max(len(k)%20for%20k%20in%20all_keys)%0A%20%20%20%20max_len_value%20%3D%20max(len(str(val))%20for%20val%20in%20(*rv_dict.values()%2C%20*sn_dict.values()))%0A%0A%20%20%20%20def%20format_row(key%3A%20str%2C%20*values%3A%20str)%3A%0A%20%20%20%20%20%20%20%20paddings%20%3D%20tuple(%22%20%22%20*%20(max_len_value%20-%20len(val))%20for%20val%20in%20values)%0A%20%20%20%20%20%20%20%20vals%20%3D%20%22%22.join(f%22%5Ct%7Bpadding%7D%7Bval%7D%22%20for%20padding%2C%20val%20in%20zip(paddings%2C%20values))%0A%20%20%20%20%20%20%20%20return%20f%22%7Bkey%7D%7B'%20'%20*%20(max_len_key%20-%20len(key))%7D%7Bvals%7D%5Cn%22%0A%0A%20%20%20%20rows%20%3D%20%22%20%22%20*%20max_len_key%20%2B%20%22%5Ctrv%5Ctsn%5Ctdiff%5Cn%22%0A%0A%20%20%20%20ZERO_VAL%20%3D%20%22.%22%0A%0A%20%20%20%20for%20key%20in%20all_keys%3A%0A%20%20%20%20%20%20%20%20rv_val%20%3D%20rv_dict.get(key%2C%200)%0A%20%20%20%20%20%20%20%20sn_val%20%3D%20sn_dict.get(key%2C%200)%0A%20%20%20%20%20%20%20%20diff_val%20%3D%20sn_val%20-%20rv_val%0A%0A%20%20%20%20%20%20%20%20rv_str%20%3D%20str(rv_val)%20if%20rv_val%20else%20ZERO_VAL%0A%20%20%20%20%20%20%20%20sn_str%20%3D%20str(sn_val)%20if%20sn_val%20else%20ZERO_VAL%0A%20%20%20%20%20%20%20%20diff_str%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20(f%22%2B%7Bdiff_val%7D%22%20if%20diff_val%20%3E%200%20else%20f%22%7Bdiff_val%7D%22)%20if%20diff_val%20else%20%22%3D%22%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20key%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20%22%20%22%20*%20(max_len_key%20-%20len(key))%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20%22%5Ct%22%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20%22%20%22%20*%20(max_len_value%20-%20len(rv_str))%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20rv_str%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20%22%5Ct%22%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20%22%20%22%20*%20(max_len_value%20-%20len(sn_str))%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20sn_str%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20%22%5Ct%22%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20%22%20%22%20*%20(max_len_value%20-%20len(diff_str))%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20diff_str%0A%20%20%20%20%20%20%20%20rows%20%2B%3D%20%22%5Cn%22%0A%0A%20%20%20%20rv_sum%20%3D%20sum(rv_dict.values())%0A%20%20%20%20sn_sum%20%3D%20sum(sn_dict.values())%0A%20%20%20%20total_diff%20%3D%20sn_sum%20-%20rv_sum%0A%0A%20%20%20%20rows%20%2B%3D%20format_row(%22total%22%2C%20str(rv_sum)%2C%20str(sn_sum)%2C%20str(total_diff))%0A%0A%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20f%22%22%22%0A%20%20%20%20The%20interpreter%20kept%20track%20of%20the%20number%20of%20times%20an%20operation%20was%20executed%2C%20which%20we%20can%20use%20as%20a%20proxy%20for%20performance.%0A%0A%20%20%20%20For%20example%2C%20we%20can%20see%20that%20one%20version%20has%20many%20more%20instructions%20executed%20overall%20(%7Brv_sum%7D%20vs%20%7Bsn_sum%7D)%2C%20and%20that%20one%20version%20uses%20explicit%20load%20and%20store%20instructions%2C%20while%20the%20other%20uses%20the%20streaming%20equivalents%3A%0A%0A%20%20%20%20%7Bm%7D%7Bm.value%7D%0A%0A%20%20%20%20%7Bn%7D%7Bn.value%7D%0A%0A%20%20%20%20%7Bk%7D%7Bk.value%7D%0A%0A%20%20%20%20%60%60%60%0A%20%20%20%20%7Brows%7D%0A%20%20%20%20%60%60%60%0A%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A