%23%20Copyright%202026%20Google%20LLC%0A%23%0A%23%20Licensed%20under%20the%20Apache%20License%2C%20Version%202.0%20(the%20%22License%22)%3B%0A%23%20you%20may%20not%20use%20this%20file%20except%20in%20compliance%20with%20the%20License.%0A%23%20You%20may%20obtain%20a%20copy%20of%20the%20License%20at%0A%23%0A%23%20%20%20%20%20http%3A%2F%2Fwww.apache.org%2Flicenses%2FLICENSE-2.0%0A%23%0A%23%20Unless%20required%20by%20applicable%20law%20or%20agreed%20to%20in%20writing%2C%20software%0A%23%20distributed%20under%20the%20License%20is%20distributed%20on%20an%20%22AS%20IS%22%20BASIS%2C%0A%23%20WITHOUT%20WARRANTIES%20OR%20CONDITIONS%20OF%20ANY%20KIND%2C%20either%20express%20or%20implied.%0A%23%20See%20the%20License%20for%20the%20specific%20language%20governing%20permissions%20and%0A%23%20limitations%20under%20the%20License.%0A%0A%23%20pylint%3A%20skip-file%0A%0Aimport%20marimo%0A%0A__generated_with%20%3D%20%220.23.9%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)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%3Cstyle%3E%0A%20%20%20%20%20%20.nav-link%20%7B%0A%20%20%20%20%20%20%20%20color%3A%20%231a73e8%3B%0A%20%20%20%20%20%20%20%20text-decoration%3A%20none%3B%0A%20%20%20%20%20%20%20%20font-weight%3A%20500%3B%0A%20%20%20%20%20%20%20%20margin%3A%200%202px%3B%0A%20%20%20%20%20%20%7D%0A%20%20%20%20%20%20.nav-link%3Ahover%20%7B%0A%20%20%20%20%20%20%20%20text-decoration%3A%20underline%3B%0A%20%20%20%20%20%20%7D%0A%20%20%20%20%20%20.nav-strong%20%7B%0A%20%20%20%20%20%20%20%20color%3A%20%23202124%3B%0A%20%20%20%20%20%20%20%20font-weight%3A%20600%3B%0A%20%20%20%20%20%20%20%20margin%3A%200%202px%3B%0A%20%20%20%20%20%20%7D%0A%20%20%20%20%3C%2Fstyle%3E%0A%20%20%20%20%3Cdiv%20style%3D'margin-bottom%3A%2024px%3B%20font-size%3A%2013px%3B%20font-family%3A%20-apple-system%2C%20BlinkMacSystemFont%2C%20%22Segoe%20UI%22%2C%20Roboto%2C%20sans-serif%3B%20text-align%3A%20center%3B%20border-bottom%3A%201px%20solid%20%23e0e0e0%3B%20padding-bottom%3A%2012px%3B%20line-height%3A%201.8%3B'%3E%0A%20%20%20%20%20%20%3Ca%20class%3D'nav-link'%20data-page%3D'index'%20href%3D'index.html'%3EHome%3C%2Fa%3E%20%7C%0A%20%20%20%20%20%20%3Ca%20class%3D'nav-link'%20data-page%3D'extending_modules_and_functions_via_composition'%20href%3D'extending_modules_and_functions_via_composition.html'%3EComposition%3C%2Fa%3E%20%7C%0A%20%20%20%20%20%20%3Ca%20class%3D'nav-link'%20data-page%3D'customizing_autograd_via_torch_autograd_function'%20href%3D'customizing_autograd_via_torch_autograd_function.html'%3ECustom%20Autograd%3C%2Fa%3E%20%7C%0A%20%20%20%20%20%20%3Ca%20class%3D'nav-link'%20data-page%3D'background_on_aten_ops'%20href%3D'background_on_aten_ops.html'%3EATen%20Ops%3C%2Fa%3E%20%7C%0A%20%20%20%20%20%20%3Ca%20class%3D'nav-link'%20data-page%3D'background_on_torch_compile'%20href%3D'background_on_torch_compile.html'%3Etorch.compile%3C%2Fa%3E%20%7C%0A%20%20%20%20%20%20%3Cspan%20class%3D'nav-strong'%3ECompile%20on%20TPU%3C%2Fspan%3E%20%7C%0A%20%20%20%20%20%20%3Ca%20class%3D'nav-link'%20data-page%3D'intro_custom_ops'%20href%3D'intro_custom_ops.html'%3EIntro%20to%20Custom%20Ops%3C%2Fa%3E%20%7C%0A%20%20%20%20%20%20%3Ca%20class%3D'nav-link'%20data-page%3D'pt_op_quant'%20href%3D'pt_op_quant.html'%3EPython%20Baseline%3C%2Fa%3E%20%7C%0A%20%20%20%20%20%20%3Ca%20class%3D'nav-link'%20data-page%3D'custom_ops_via_pallas'%20href%3D'custom_ops_via_pallas.html'%3EPallas%20TPU%20Kernel%3C%2Fa%3E%0A%20%20%20%20%3C%2Fdiv%3E%0A%20%20%20%20%3Cscript%3E%0A%20%20%20%20%20%20if%20(window.location.search%20%26%26%20window.location.search.indexOf('file%3D')%20!%3D%3D%20-1)%20%7B%0A%20%20%20%20%20%20%20%20document.querySelectorAll('.nav-link').forEach(link%20%3D%3E%20%7B%0A%20%20%20%20%20%20%20%20%20%20const%20page%20%3D%20link.getAttribute('data-page')%3B%0A%20%20%20%20%20%20%20%20%20%20link.href%20%3D%20'%3Ffile%3D'%20%2B%20page%20%2B%20'.py'%3B%0A%20%20%20%20%20%20%20%20%7D)%3B%0A%20%20%20%20%20%20%7D%0A%20%20%20%20%3C%2Fscript%3E%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Background%20on%20%60torch.compile%60%20on%20TPUs%20via%20XLA%0A%0A%20%20%20%20TorchTPU%20is%20the%20backend%20for%20PyTorch%20to%20run%20on%20Google%20TPUs.%20TorchTPU%20uses%20XLA%20to%0A%20%20%20%20compile%20PyTorch%20rather%20than%20Inductor.%20XLA%20stands%20for%20%22Accelerated%20Linear%0A%20%20%20%20Algebra%22%20and%20is%20an%20%5Bopen-source%20ML%20compiler%5D(https%3A%2F%2Fopenxla.org%2Fxla).%0A%0A%20%20%20%20The%20path%20from%20your%20Python%20code%20to%20assembly%20on%20a%20TPU%20is%20as%20follows%3A%0A%0A%20%20%20%201.%20%20Your%20Python%20code%20will%20require%20minor%20changes%20(e.g.%20setting%20%22TPU%22%20as%20the%0A%20%20%20%20%20%20%20%20device).%0A%20%20%20%202.%20%20PyTorch's%20dispatcher%20lowers%20your%20Python%20code%20to%20ATen%20and%20c10d%20(no%20change).%0A%20%20%20%203.%20%20TorchTPU%20lowers%20the%20ATen%20and%20c10d%20ops%20to%20**StableHLO**%2C%20a%20dialect%20of%0A%20%20%20%20%20%20%20%20%5BMLIR%5D(https%3A%2F%2Fmlir.llvm.org%2F)%20that%20is%20device%20agnostic.%0A%20%20%20%204.%20%20XLA%20lowers%20the%20StableHLO%20to%20**HLO**%20ops.%0A%20%20%20%205.%20%20XLA%20optimizes%20the%20HLO%20over%20several%20passes.%0A%20%20%20%206.%20%20XLA%20lowers%20the%20HLO%20ops%20to%20device-specific%20**LLO**%20(Low-Level%20Ops)%20operators.%0A%20%20%20%207.%20%20XLA%20optimizes%20the%20LLO%20over%20several%20passes%2C%20targeting%20the%20specific%20TPU%0A%20%20%20%20%20%20%20%20hardware.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%201.%20Setting%20up%20HLO%20%26%20LLO%20Dumps%0A%0A%20%20%20%20To%20inspect%20how%20XLA%20compiles%20a%20pointwise%20addition%20and%20ReLU%20operation%2C%20we%20can%20configure%20environment%20variables%20to%20dump%20the%20intermediate%20HLO%20ops%20and%20LLO%20ops.%0A%0A%20%20%20%20%23%23%23%23%20HLO%20Dump%20Configuration%3A%0A%20%20%20%20%60%60%60python%0A%20%20%20%20XLA_DUMP_TO%3A%20Final%5Bstr%5D%20%3D%20%22%2Ftmp%2Fxla_dump%22%0A%20%20%20%20os.environ%5B%22XLA_FLAGS%22%5D%20%3D%20f%22--xla_dump_to%3D%7BXLA_DUMP_TO%7D%20--xla_dump_hlo_as_text%22%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20%23%23%23%23%20LLO%20Dump%20Configuration%3A%0A%20%20%20%20%60%60%60python%0A%20%20%20%20LLO_DUMP_TO%3A%20Final%5Bstr%5D%20%3D%20%22%2Ftmp%2Fllo_dump%22%0A%20%20%20%20flags%20%3D%20%5B%20f%22--xla_jf_dump_to%3D%7BLLO_DUMP_TO%7D%22%2C%20%22--xla_jf_dump_llo_text%3Dtrue%22%20%5D%0A%20%20%20%20sys.argv.extend(flags)%0A%20%20%20%20os.environ%5B%22LIBTPU_INIT_ARGS%22%5D%20%3D%20%22%20%22.join(flags)%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20Let's%20run%20the%20compiled%20pointwise%20operation%20and%20dump%20the%20generated%20HLO%2FLLO%20logs.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20os%0A%20%20%20%20import%20sys%0A%20%20%20%20from%20typing%20import%20Final%0A%0A%20%20%20%20%23%20Setup%20compilation%20dumps%20directories%0A%20%20%20%20LLO_DUMP_TO%3A%20Final%5Bstr%5D%20%3D%20%22%2Ftmp%2Fllo_dump%22%0A%20%20%20%20flags%20%3D%20%5B%0A%20%20%20%20%20%20%20%20f%22--xla_jf_dump_to%3D%7BLLO_DUMP_TO%7D%22%2C%0A%20%20%20%20%20%20%20%20%22--xla_jf_dump_llo_text%3Dtrue%22%2C%0A%20%20%20%20%5D%0A%20%20%20%20sys.argv.extend(flags)%0A%20%20%20%20os.environ%5B%22LIBTPU_INIT_ARGS%22%5D%20%3D%20%22%20%22.join(flags)%0A%0A%20%20%20%20XLA_DUMP_TO%3A%20Final%5Bstr%5D%20%3D%20%22%2Ftmp%2Fxla_dump%22%0A%20%20%20%20os.environ%5B%22XLA_FLAGS%22%5D%20%3D%20f%22--xla_dump_to%3D%7BXLA_DUMP_TO%7D%20--xla_dump_hlo_as_text%22%0A%0A%20%20%20%20from%20absl%20import%20app%20as%20absl_app%0A%20%20%20%20import%20torch%0A%0A%20%20%20%20return%20LLO_DUMP_TO%2C%20XLA_DUMP_TO%2C%20torch%0A%0A%0A%40app.cell%0Adef%20_(torch)%3A%0A%20%20%20%20%40torch.compile(backend%3D%22tpu%22)%0A%20%20%20%20def%20fwd(x%2C%20y)%3A%0A%20%20%20%20%20%20return%20torch.nn.functional.relu(torch.add(x%2C%20y))%0A%0A%20%20%20%20return%20(fwd%2C)%0A%0A%0A%40app.function%0Adef%20dump_dir(directory)%3A%0A%20%20import%20os%0A%20%20entries%20%3D%20os.listdir(directory)%20if%20os.path.exists(directory)%20else%20%5B%5D%0A%20%20print(f%22%5Cn%3D%3D%3D%20Files%20in%20%7Bdirectory%7D%20(Total%3A%20%7Blen(entries)%7D)%20%3D%3D%3D%22)%0A%0A%20%20%23%20Only%20print%20the%20filenames%20to%20keep%20standard%20output%20fast%0A%20%20for%20f%20in%20sorted(entries)%3A%0A%20%20%20%20print(f)%0A%0A%20%20print(%22%5Cn%F0%9F%92%A1%20Tip%3A%20All%20compilation%20files%20are%20written%20to%20disk.%20You%20can%20read%20any%20file%20interactively%20using%3A%20open(os.path.join(directory%2C%20filename)).read()%22)%0A%0A%0A%40app.cell%0Adef%20_(LLO_DUMP_TO%3A%20%22Final%5Bstr%5D%22%2C%20XLA_DUMP_TO%3A%20%22Final%5Bstr%5D%22%2C%20fwd%2C%20torch)%3A%0A%20%20%20%20def%20main()%3A%0A%20%20%20%20%20%20tpu_available%20%3D%20False%0A%20%20%20%20%20%20try%3A%0A%20%20%20%20%20%20%20%20x%20%3D%20torch.randn(1024).to(%22tpu%22)%0A%20%20%20%20%20%20%20%20y%20%3D%20torch.randn(1024).to(%22tpu%22)%0A%20%20%20%20%20%20%20%20tpu_available%20%3D%20True%0A%20%20%20%20%20%20except%20Exception%3A%0A%20%20%20%20%20%20%20%20pass%0A%0A%20%20%20%20%20%20if%20not%20tpu_available%3A%0A%20%20%20%20%20%20%20%20print(%22TPU%20not%20available.%20Run%20this%20notebook%20on%20a%20TPU%20VM%20to%20generate%20live%20compilation%20dumps.%22)%0A%20%20%20%20%20%20%20%20return%200%0A%0A%20%20%20%20%20%20res%20%3D%20fwd(x%2C%20y)%0A%20%20%20%20%20%20_%20%3D%20res.cpu()%0A%0A%20%20%20%20%20%20dump_dir(XLA_DUMP_TO)%0A%20%20%20%20%20%20dump_dir(LLO_DUMP_TO)%0A%0A%20%20%20%20%20%20print(%22Success.%22)%0A%20%20%20%20%20%20return%200%0A%0A%20%20%20%20main()%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%202.%20Analyzing%20the%20HLO%20(Before%20Optimizations)%0A%0A%20%20%20%20The%20output%20section%20will%20list%20several%20generated%20files%20in%20%60%2Ftmp%2Fxla_dump%60.%0A%0A%20%20%20%20Look%20for%20the%20module%20file%20named%20%60module_XXXX.tt_jit_pointwise_fusion_xla_...before_optimizations.txt%60.%20(Note%20that%20exact%20filenames%20vary%20based%20on%20your%20specific%20environment%20and%20line%20numbering).%0A%0A%20%20%20%20Opening%20the%20%60before_optimizations.txt%60%20file%2C%20you%20will%20find%20the%20initial%20unoptimized%20HLO%20representation.%20Notice%20that%20it%20explicitly%20outlines%20an%20%60add%60%20and%20a%20%60maximum%60%20(which%20represents%20the%20ReLU)%3A%0A%0A%20%20%20%20%60%60%60hlo%0A%20%20%20%20ENTRY%20%25main.1%20(Arg_0.1%3A%20f32%5B1024%5D%2C%20Arg_1.1%3A%20f32%5B1024%5D)%20-%3E%20f32%5B1024%5D%20%7B%0A%20%20%20%20%20%20%25Arg_0.1%20%3D%20f32%5B1024%5D%7B0%7D%20parameter(0)%0A%20%20%20%20%20%20%25Arg_1.1%20%3D%20f32%5B1024%5D%7B0%7D%20parameter(1)%0A%20%20%20%20%20%20%25add.1%20%3D%20f32%5B1024%5D%7B0%7D%20add(%25Arg_0.1%2C%20%25Arg_1.1)%2C%20metadata%3D%7Bop_name%3D%22add%2Fadd%22%20stack_frame_id%3D40%7D%0A%20%20%20%20%20%20%25constant.1%20%3D%20f32%5B%5D%20constant(0)%0A%20%20%20%20%20%20%25relu.2%20%3D%20f32%5B1024%5D%7B0%7D%20broadcast(%25constant.1)%2C%20dimensions%3D%7B%7D%2C%20metadata%3D%7Bop_name%3D%22relu%2Frelu%22%20stack_frame_id%3D40%7D%0A%20%20%20%20%20%20ROOT%20%25relu.3%20%3D%20f32%5B1024%5D%7B0%7D%20maximum(%25add.1%2C%20%25relu.2)%2C%20metadata%3D%7Bop_name%3D%22relu%2Frelu%22%20stack_frame_id%3D40%7D%0A%20%20%20%20%7D%0A%20%20%20%20%60%60%60%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%203.%20Analyzing%20the%20HLO%20(After%20Codegen)%0A%0A%20%20%20%20The%20XLA%20compiler%20will%20then%20optimize%20this%20HLO.%20Looking%20inside%20the%20optimized%20codegen%20output%20file%20(%60*after_codegen.txt%60)%2C%20you%20will%20see%20that%20XLA%20has%20successfully%20fused%20the%20two%20operations%20together%20into%20a%20single%20block%20assigned%20to%20a%20loop%20fusion%20computation%20variable%3A%20%60add_maximum_fusion%60.%0A%0A%20%20%20%20%60%60%60hlo%0A%20%20%20%20%25fused_computation%20(param_0.2%3A%20f32%5B1024%5D%2C%20param_1.2%3A%20f32%5B1024%5D)%20-%3E%20f32%5B1024%5D%20%7B%0A%20%20%20%20%20%20%25param_0.2%20%3D%20f32%5B1024%5D%7B0%3AT(1024)%7D%20parameter(0)%0A%20%20%20%20%20%20%25param_1.2%20%3D%20f32%5B1024%5D%7B0%3AT(1024)%7D%20parameter(1)%0A%20%20%20%20%20%20%25add.0%20%3D%20f32%5B1024%5D%7B0%3AT(1024)%7D%20add(%25param_0.2%2C%20%25param_1.2)%2C%20metadata%3D%7Bop_name%3D%22add%2Fadd%22%20stack_frame_id%3D40%7D%0A%20%20%20%20%20%20%25constant.0%20%3D%20f32%5B%5D%7B%3AT(128)%7D%20constant(0)%0A%20%20%20%20%20%20%25relu.1%20%3D%20f32%5B1024%5D%7B0%3AT(1024)%7D%20broadcast(%25constant.0)%2C%20dimensions%3D%7B%7D%2C%20metadata%3D%7Bop_name%3D%22relu%2Frelu%22%20stack_frame_id%3D40%7D%0A%20%20%20%20%20%20ROOT%20%25relu.0%20%3D%20f32%5B1024%5D%7B0%3AT(1024)%7D%20maximum(%25add.0%2C%20%25relu.1)%2C%20metadata%3D%7Bop_name%3D%22relu%2Frelu%22%20stack_frame_id%3D40%7D%0A%20%20%20%20%7D%0A%0A%20%20%20%20ENTRY%20%25main.1%20(Arg_0.1%3A%20f32%5B1024%5D%2C%20Arg_1.1%3A%20f32%5B1024%5D)%20-%3E%20f32%5B1024%5D%20%7B%0A%20%20%20%20%20%20%25Arg_1.1%20%3D%20f32%5B1024%5D%7B0%3AT(1024)%7D%20parameter(1)%2C%20backend_config%3D%7B%22flag_configs%22%3A%5B%5D%2C%22scoped_memory_configs%22%3A%5B%7B%22memory_space%22%3A%221%22%2C%22offset%22%3A%220%22%2C%22size%22%3A%2216777216%22%7D%5D%2C%22used_scoped_memory_configs%22%3A%5B%5D%7D%0A%20%20%20%20%20%20%25Arg_0.1%20%3D%20f32%5B1024%5D%7B0%3AT(1024)%7D%20parameter(0)%2C%20backend_config%3D%7B%22flag_configs%22%3A%5B%5D%2C%22scoped_memory_configs%22%3A%5B%7B%22memory_space%22%3A%221%22%2C%22offset%22%3A%220%22%2C%22size%22%3A%2216777216%22%7D%5D%2C%22used_scoped_memory_configs%22%3A%5B%5D%7D%0A%20%20%20%20%20%20ROOT%20%25add_maximum_fusion%20%3D%20f32%5B1024%5D%7B0%3AT(1024)%7D%20fusion(%25Arg_0.1%2C%20%25Arg_1.1)%2C%20kind%3DkLoop%2C%20calls%3D%25fused_computation%2C%20metadata%3D%7Bop_name%3D%22relu%2Frelu%22%20stack_frame_id%3D40%7D%2C%20backend_config%3D%7B%22flag_configs%22%3A%5B%5D%2C%22window_config%22%3A%7B%22kernel_window_bounds%22%3A%5B%5D%2C%22output_window_bounds%22%3A%5B%221%22%5D%2C%22input_window_bounds%22%3A%5B%5D%2C%22estimated_cycles%22%3A%221840%22%2C%22iteration_bounds%22%3A%5B%221%22%5D%2C%22cost_model_type%22%3A%22COST_MODEL_TYPE_INVALID%22%2C%22ml_estimated_microseconds%22%3A0%2C%22is_mask%22%3Afalse%2C%22pad_output_on_minor_dim%22%3A%220%22%2C%22pad_input_on_minor_dim%22%3A%220%22%2C%22estimated_vmem_bytes%22%3A%220%22%2C%22estimated_bundle_count%22%3A%220%22%2C%22estimated_scoped_vmem_bytes%22%3A%220%22%7D%2C%22scoped_memory_configs%22%3A%5B%7B%22memory_space%22%3A%221%22%2C%22offset%22%3A%220%22%2C%22size%22%3A%2216777216%22%7D%5D%2C%22used_scoped_memory_configs%22%3A%5B%7B%22memory_space%22%3A%221%22%2C%22offset%22%3A%220%22%2C%22size%22%3A%2212288%22%7D%5D%2C%22retry_config%22%3A%7B%22retry_count%22%3A%220%22%7D%2C%22aliasing_operands%22%3A%7B%22lists%22%3A%5B%5D%7D%7D%0A%20%20%20%20%7D%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20The%20key%20indicator%20here%20is%20%60kind%3DkLoop%60%20and%20the%20reference%20%60calls%3D%25fused_computation%60%2C%20meaning%20both%20operators%20compile%20into%20a%20single%20memory-efficient%20loop.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%204.%20Analyzing%20the%20LLO%20assembly%0A%0A%20%20%20%20In%20the%20LLO%20dump%20folder%20%60%2Ftmp%2Fllo_dump%60%2C%20search%20for%20files%20containing%20the%20%60add_maximum_fusion%60%20token.%20Look%20for%20the%20file%20ending%20in%20%60add_maximum_fusion-01-original.txt%60%20to%20view%20the%20physical%20TPU%20assembly%20instructions.%0A%0A%20%20%20%20Without%20needing%20to%20be%20a%20TPU%20architecture%20expert%2C%20you%20can%20isolate%20the%20exact%20vector%20addition%20(%60vadd%60)%20and%20vector%20maximum%20(%60vmax%60)%20instructions%20executing%20in%20TPU%20registers%3A%0A%0A%20%20%20%20%60%60%60assembly%0A%20%20%20%20%3D%3D%3D%20File%3A%201777396851540540302-add_maximum_fusion-01-original.txt%20%3D%3D%3D%0A%20%20%20%20%2F%2F%20%5BEnable%20stack%20traces%20via%20-xla_jf_collect_llo_stack_trace%20or%20-xla_jf_debug_level%3D2%5D%0A%20%20%20%20%24region0%3A%20%23%7Badd_maximum_fusion%7D%0A%20%20%20%20%20%20%23allocation6%20%5Bshape%20%3D%20's32%5B1%5D%7B0%7D'%2C%20space%3Dsflag%2C%20size%20%3D%200x4%2C%20scoped%2C%20tag%20%3D%20'scoped%20memory%20for%20add_maximum_fusion'%5D%0A%20%20%20%20%20%20%25s0%20%3D%20inlined_call_operand.hbm%20%5Bshape%3A%20f32%5B1024%5D%2C%20index%3A%200%2C%20kind%3A%20input%2C%20shape%20index%3A%20%7B%7D%5D%20%2F*%20operand%200%20*%2F%0A%20%20%20%20%20%20%25s1%20%3D%20inlined_call_operand.hbm%20%5Bshape%3A%20f32%5B1024%5D%2C%20index%3A%201%2C%20kind%3A%20input%2C%20shape%20index%3A%20%7B%7D%5D%20%2F*%20operand%201%20*%2F%0A%20%20%20%20%20%20%25s2%20%3D%20inlined_call_operand.hbm%20%5Bshape%3A%20f32%5B1024%5D%2C%20index%3A%202%2C%20kind%3A%20output%2C%20shape%20index%3A%20%7B%7D%5D%20%2F*%20operand%202%20*%2F%0A%20%20%20%20%20%20%24region1%3A%20%23%7Badd_maximum_fusion%7D%20parent%3D0%0A%20%20%20%20%20%20%20%20%23allocation0%20%5Bshape%20%3D%20'u8%5B4096%5D%7B0%7D'%2C%20space%3Dvmem%2C%20size%20%3D%200x1000%2C%20scoped%2C%20tag%20%3D%20'operand%20span%20for%20operand%200'%5D%0A%20%20%20%20%20%20%20%20%23allocation1%20%5Bshape%20%3D%20's32%5B1%5D%7B0%7D'%2C%20space%3Dsflag%2C%20size%20%3D%200x4%2C%20scoped%2C%20tag%20%3D%20'scoped%20memory%20for%20add_maximum_fusion'%5D%0A%20%20%20%20%20%20%20%20%23allocation2%20%5Bshape%20%3D%20's32%5B1%5D%7B0%7D'%2C%20space%3Dsflag%2C%20size%20%3D%200x4%2C%20scoped%2C%20tag%20%3D%20'scoped%20memory%20for%20add_maximum_fusion'%5D%0A%20%20%20%20%20%20%20%20%23allocation3%20%5Bshape%20%3D%20'u8%5B4096%5D%7B0%7D'%2C%20space%3Dvmem%2C%20size%20%3D%200x1000%2C%20scoped%2C%20tag%20%3D%20'operand%20span%20for%20operand%201'%5D%0A%20%20%20%20%20%20%20%20%23allocation4%20%5Bshape%20%3D%20's32%5B1%5D%7B0%7D'%2C%20space%3Dsflag%2C%20size%20%3D%200x4%2C%20scoped%2C%20tag%20%3D%20'scoped%20memory%20for%20add_maximum_fusion'%5D%0A%20%20%20%20%20%20%20%20%23allocation5%20%5Bshape%20%3D%20'u8%5B4096%5D%7B0%7D'%2C%20space%3Dvmem%2C%20size%20%3D%200x1000%2C%20scoped%2C%20tag%20%3D%20'operand%20span%20for%20operand%202'%5D%0A%20%20%20%20%20%20%20%20%253%20%3D%20vsyncpa%20%5B%23allocation1%5D%2C%200%0A%20%20%20%20%20%20%20%20%254%20%3D%20vsyncpa%20%5B%23allocation4%5D%2C%200%0A%20%20%20%20%20%20%20%20%255%20%3D%20vsyncpa%20%5B%23allocation2%5D%2C%200%0A%20%20%20%20%20%20%20%20%257%20%3D%20vsyncadd%20%5B%23allocation1%5D%2C%200%0A%20%20%20%20%20%20%20%20%25s9%20%3D%20sshll.u32%20%25s0%2C%204%0A%20%20%20%20%20%20%20%20%25s10%20%3D%20int_to_ptr.hbm%20%5Bresolvable%3A%24true%5D%20%25s9%0A%20%20%20%20%20%20%20%20%25s11%20%3D%20sshll.u32%20%5B%23allocation0%5D%2C%204%0A%20%20%20%20%20%20%20%20%25s12%20%3D%20int_to_ptr.vmem%20%5Bresolvable%3A%24true%5D%20%25s11%0A%20%20%20%20%20%20%20%20%2514%20%3D%20dma.hbm_to_vmem%20%5Bthread%3A%240%5D%20%20%2F*hbm%3D*%2F%25s10%2C%20%2F*size_in_granules%3D*%2F128%2C%20%2F*vmem%3D*%2F%25s12%2C%20%2F*dst_syncflagno%3D*%2F%5B%23allocation1%5D%0A%20%20%20%20%20%20%20%20...%0A%20%20%20%20%20%20%20%20%2525%20%3D%20dma.done%20%5B%23allocation1%5D%2C%20128%20%2F*%20pipeline-emitter-dma-wait%20*%2F%0A%20%20%20%20%20%20%20%20%2527%20%3D%20dma.done%20%5B%23allocation4%5D%2C%20128%20%2F*%20pipeline-emitter-dma-wait%20*%2F%0A%20%20%20%20%20%20%20%20%25v28%20%3D%20vld%20%5Bvmem%3A%5B%23allocation0%5D%20sm%3A%240xff%5D%0A%20%20%20%20%20%20%20%20%25v29%20%3D%20vld%20%5Bvmem%3A%5B%23allocation3%5D%20sm%3A%240xff%5D%0A%20%20%20%20%20%20%20%20%2530%20%3D%20xla_tuple%20%25v28%2C%20%25v29%0A%20%20%20%20%20%20%20%20%2531%20%3D%20xla_tuple%20%2530%0A%20%20%20%20%20%20%20%20%25v32%20%3D%20vadd.f32%20%25v28%2C%20%25v29%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%3C--%20Vector%20Addition%0A%20%20%20%20%20%20%20%20%2533%20%3D%20xla_tuple%20%25v32%0A%20%20%20%20%20%20%20%20%2534%20%3D%20xla_tuple%20%25v32%2C%200.0%0A%20%20%20%20%20%20%20%20%2535%20%3D%20xla_tuple%20%2534%0A%20%20%20%20%20%20%20%20%25v36%20%3D%20vmax.f32%20%25v32%2C%200.0%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%3C--%20Vector%20Maximum%20(ReLU)%0A%20%20%20%20%20%20%20%20%2537%20%3D%20xla_tuple%20%25v36%0A%20%20%20%20%20%20%20%20%2538%20%3D%20vst%20%5Bvmem%3A%5B%23allocation5%5D%20sm%3A%240xff%5D%20%2F*vst_source%3D*%2F%25v36%0A%20%20%20%20%20%20%20%20%2540%20%3D%20vsyncadd%20%5B%23allocation2%5D%2C%200%0A%20%20%20%20%20%20%20%20...%0A%20%20%20%20%20%20%20%20%2549%20%3D%20dma.done%20%5B%23allocation2%5D%2C%20128%20%2F*%20pipeline-emitter-dma-wait%20*%2F%0A%20%20%20%20%20%20%20%20%2550%20%3D%20vsyncpa%20%5B%23allocation1%5D%2C%201%0A%20%20%20%20%20%20%20%20%2551%20%3D%20vsyncpa%20%5B%23allocation4%5D%2C%201%0A%20%20%20%20%20%20%20%20%2552%20%3D%20vsyncpa%20%5B%23allocation2%5D%2C%201%0A%20%20%20%20%60%60%60%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Conclusion%0A%0A%20%20%20%20The%20heart%20of%20creating%20highly%20optimized%20custom%20TPU%20kernels%20is%20gainful%20control%20of%20this%20final%20outputted%20LLO%2C%20regardless%20of%20whether%20you%20write%20kernels%20using%20StableHLO%2C%20JAX-level%20ops%2C%20or%20Pallas.%0A%0A%20%20%20%20In%20the%20next%20chapters%2C%20we%20will%20explore%20how%20to%20build%20these%20custom%20ops%20using%20PyTorch's%20custom%20op%20API%20and%20low-level%20TPU%20kernel%20DSLs.%0A%20%20%20%20%22%22%22)%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
55728e8362cb5c62589d5c70c66851ba