From 247c81dfb3e196d8b1f715aededcee3f28853564 Mon Sep 17 00:00:00 2001 From: cqchu Date: Wed, 3 Jan 2024 12:16:55 +0800 Subject: [PATCH] doc(xla): modify xla doc --- source/user-guide/model-development/jit/xla.rst | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/source/user-guide/model-development/jit/xla.rst b/source/user-guide/model-development/jit/xla.rst index e27030257d9..ba44eb6096c 100644 --- a/source/user-guide/model-development/jit/xla.rst +++ b/source/user-guide/model-development/jit/xla.rst @@ -93,6 +93,23 @@ XLA executable 加速训练过程。 print (xla_fused_softmax(inp)) # run in xla +如果我们想看到 mge 和 xla 优化的一些中间 IR 表示，可以通过设置环境变量 MGE_VERBOSE_XLA_IR 来打印相关结果。MGE_VERBOSE_XLA_IR 为 1 时，会打印 mge trace 出来的图 IR，MGE_VERBOSE_XLA_IR 为 2 时，会打印xla 的 hlo 图结构，在 MGE_VERBOSE_XLA_IR 为 3 时会打印 xla 编译优化后的图结构。如果我们 export MGE_VERBOSE_XLA_IR=1 后再执行上述代码，则可以看到： + +.. code-block:: python + + please_realize_func_name_system_1( + 0%:<256x1000x1000,f32> + ) { + 1%:<256x1000x1000,f32> = io_mark_var(0%:<256x1000x1000,f32>) + 2%:<256x1000x1,f32> = ReduceMAX(1%:<256x1000x1000,f32>) + 3%:<256x1000x1000,f32> = SUB(1%:<256x1000x1000,f32>, 2%:<256x1000x1,f32>) + 4%:<256x1000x1000,f32> = EXP(3%:<256x1000x1000,f32>) + 5%:<256x1000x1,f32> = ReduceSUM(4%:<256x1000x1000,f32>) + 6%:<256x1000x1000,f32> = TRUE_DIV(4%:<256x1000x1000,f32>, 5%:<256x1000x1,f32>) + 7%:<256x1000x1000,f32> = io_mark_var(6%:<256x1000x1000,f32>) + return 1 7%:<256x1000x1000,f32> + } + 当模型训练迭代（Iteration）完全静态的情况下， 您也可以使用 jit.xla_trace 装饰器将训练迭代全部交由XLA执行。 需要将 optimizer， module 作为train_func 参数传入，同时 train_func 中需包含包含模型前向、 反向 、 参数更新等代码，