14_affine_fusion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 前言\n",
    "\n",
    "目前AI编译领域中的Triton语言正迅速成为新一代开发者和研究者的首选. 在这一背景下, 针对Triton语言后端的编译优化技术, 其重要性不言而喻. 并且业界普遍认为, 将Triton语言编译优化的关键在于将其降低到MLIR的Linalg方言上(寒武纪Triton-Linalg/微软Triton-Shared), 进而在MLIR的Affine表示上进行深入优化. \n",
    "\n",
    "在这一编译优化的链条中, 基于Affine的优化方法扮演着基础而关键的角色. 它不仅能够提升代码的性能, 还能确保优化过程的可扩展性和可维护性. 本文将详细介绍基于Affine的Loop Fusion优化技术的实现原理、步骤和效果, 以帮助读者理解和掌握这一关键技术, 进一步推动AI编译技术的发展.\n",
    "\n",
    "注意:\n",
    "1. 关于更多不同框架的编译表示表示可以查看: [Tensor Expression DSL总结](https://zhuanlan.zhihu.com/p/674882975)\n",
    "2. 本文所实现的Affine Fusion优化大量参考了mlir内置的`affine-loop-fusion`优化, 并且根据最早的提交记录, 它是由Pluto论文作者Uday亲自操刀实现的.\n",
    "3. 本文简化了许多场景, 主要以讲解思路为主.\n",
    "4. 此教程基于llvm 17.04运行, 部分功能需要对mlir进行修改, 在[本教程仓库](https://github.com/zhen8838/isl_learn)中提供了`bindings.patch`和`python.patch`. 对应的isl在[这里](https://github.com/zhen8838/pet)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 实现过程拆解"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "from mlir.ir import *\n",
    "from mlir.dialects.builtin import ModuleOp\n",
    "from mlir.dialects.arith import ConstantOp\n",
    "from mlir.dialects.func import FuncOp\n",
    "from mlir.dialects.affine import AffineForOp, AffineLoadOp, AffineStoreOp, AffineIfOp\n",
    "import isl\n",
    "from typing import List, Tuple, Dict, Set, Optional\n",
    "from dataclasses import dataclass\n",
    "from mlir_utility import IrVisitor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. 载入与分析IR\n",
    "\n",
    "首先解析mlir源文件:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "module {\n",
      "  func.func @main(%arg0: memref<8x128x384xf32>, %arg1: memref<8x384x512xf32>, %arg2: memref<8x128x512xf32>, %arg3: memref<8x512x64xf32>, %arg4: memref<8x128x64xf32>) {\n",
      "    affine.for %arg5 = 0 to 8 {\n",
      "      affine.for %arg6 = 0 to 128 {\n",
      "        affine.for %arg7 = 0 to 512 {\n",
      "          affine.for %arg8 = 0 to 384 {\n",
      "            %0 = affine.load %arg0[%arg5, %arg6, %arg8] : memref<8x128x384xf32>\n",
      "            %1 = affine.load %arg1[%arg5, %arg8, %arg7] : memref<8x384x512xf32>\n",
      "            %2 = affine.load %arg2[%arg5, %arg6, %arg7] : memref<8x128x512xf32>\n",
      "            %3 = arith.mulf %0, %1 : f32\n",
      "            %4 = arith.addf %2, %3 : f32\n",
      "            affine.store %4, %arg2[%arg5, %arg6, %arg7] : memref<8x128x512xf32>\n",
      "          }\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "    affine.for %arg5 = 0 to 8 {\n",
      "      affine.for %arg6 = 0 to 128 {\n",
      "        affine.for %arg7 = 0 to 64 {\n",
      "          affine.for %arg8 = 0 to 512 {\n",
      "            %0 = affine.load %arg2[%arg5, %arg6, %arg8] : memref<8x128x512xf32>\n",
      "            %1 = affine.load %arg3[%arg5, %arg8, %arg7] : memref<8x512x64xf32>\n",
      "            %2 = affine.load %arg4[%arg5, %arg6, %arg7] : memref<8x128x64xf32>\n",
      "            %3 = arith.mulf %0, %1 : f32\n",
      "            %4 = arith.addf %2, %3 : f32\n",
      "            affine.store %4, %arg4[%arg5, %arg6, %arg7] : memref<8x128x64xf32>\n",
      "          }\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "    return\n",
      "  }\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "ctx = Context()\n",
    "\n",
    "with open(\"test1.mlir\") as f:\n",
    "    mod = Module.parse(f.read(), ctx)\n",
    "print(mod)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Affine Fusion的目标是是fuse两个相邻的循环体, 那么第一步则是分析Module找到相邻的循环体: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PerfectLoopNest(IrVisitor):\n",
    "  def __init__(self) -> None:\n",
    "    super().__init__()\n",
    "    self.forOps: List[AffineForOp] = []\n",
    "    self.loadOps: List[AffineLoadOp] = []\n",
    "    self.storeOps: List[AffineStoreOp] = []\n",
    "    self.hasNonAffineRegionOp = False\n",
    "\n",
    "  def create(op: OpView) -> PerfectLoopNest:\n",
    "    obj = PerfectLoopNest()\n",
    "    obj.visit(op)\n",
    "    return obj\n",
    "\n",
    "  def runBeforeOperation(self, op: OpView) -> bool:\n",
    "    if isinstance(op, AffineForOp):\n",
    "      self.forOps.append(op)\n",
    "    elif len(op.regions) != 0 and not isinstance(op, AffineIfOp):\n",
    "      self.hasNonAffineRegionOp = True\n",
    "    elif (isinstance(op, AffineLoadOp)):\n",
    "      self.loadOps.append(op)\n",
    "    elif (isinstance(op, AffineStoreOp)):\n",
    "      self.storeOps.append(op)\n",
    "    return super().runBeforeOperation(op)\n",
    "\n",
    "\n",
    "class LoopNestPairCollector(IrVisitor):\n",
    "  state = False\n",
    "  srcLoopNest: PerfectLoopNest\n",
    "  dstLoopNest: PerfectLoopNest\n",
    "\n",
    "  def __init__(self) -> None:\n",
    "    super().__init__()\n",
    "    self.state = False\n",
    "\n",
    "  @staticmethod\n",
    "  def collect(obj) -> Optional[None | Tuple[PerfectLoopNest, PerfectLoopNest]]:\n",
    "    collector = LoopNestPairCollector()\n",
    "    collector.visit(mod)\n",
    "    assert collector.state\n",
    "    return (collector.srcLoopNest, collector.dstLoopNest)\n",
    "\n",
    "  def runBeforeBlock(self, block: Block) -> bool:\n",
    "    for i in range(len(block.operations) - 1, -1, -1):\n",
    "      if isinstance(block.operations[i], AffineForOp) and i > 0 and isinstance(block.operations[i - 1], AffineForOp):\n",
    "        producer = block.operations[i - 1]\n",
    "        consumer = block.operations[i]\n",
    "        self.srcLoopNest = PerfectLoopNest.create(producer)\n",
    "        self.dstLoopNest = PerfectLoopNest.create(consumer)\n",
    "        self.state = True\n",
    "        return False\n",
    "    return True\n",
    "\n",
    "srcLoopNest, dstLoopNest = LoopNestPairCollector.collect(mod)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 分析与提取多面体表示\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "获得了两个相邻的循环体后, 需要尝试对他们进行fuse, 为了保证fuse后程序的正确性. 需要先确定循环体之间是否存在依赖关系, 并找到在不破坏依赖关系的情况下src循环可以合并到dst循环体最深的位置, 也就是dst循环体中所有读写操作所处的循环外层. 开始查找循环体之间的依赖:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def GatherProducerConsumerMemrefs(src: PerfectLoopNest, dst: PerfectLoopNest) -> Set[Value]:\n",
    "  producerConsumerMemrefs = set()\n",
    "  for store in src.storeOps:\n",
    "    for load in dst.loadOps:\n",
    "      if store.memref == load.memref:\n",
    "        for use in store.memref.uses:\n",
    "          owner: OpView = use.owner\n",
    "          if owner == store:\n",
    "            continue\n",
    "          elif owner == load:\n",
    "            continue\n",
    "          elif owner.operation.parent == store.operation.parent:\n",
    "            continue\n",
    "          else:\n",
    "            break\n",
    "        producerConsumerMemrefs.add(load.memref)\n",
    "  return producerConsumerMemrefs\n",
    "\n",
    "\n",
    "producerConsumerMemrefs = GatherProducerConsumerMemrefs(srcLoopNest, dstLoopNest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这里我简化了问题, 假设只存在一对producer/consumer依赖, 并直接获取这个dependentOpPair:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def GatherDependentOpPairs(src: PerfectLoopNest, dst: PerfectLoopNest) -> List[Tuple[AffineStoreOp, AffineLoadOp]]:\n",
    "  producerConsumerMemrefs = []\n",
    "  for store in src.storeOps:\n",
    "    for load in dst.loadOps:\n",
    "      if store.memref == load.memref:\n",
    "        for use in store.memref.uses:\n",
    "          owner: OpView = use.owner\n",
    "          if owner == store:\n",
    "            continue\n",
    "          elif owner == load:\n",
    "            continue\n",
    "          elif owner.operation.parent == store.operation.parent:\n",
    "            continue\n",
    "          else:\n",
    "            break\n",
    "        producerConsumerMemrefs.append((store, load))\n",
    "  return producerConsumerMemrefs\n",
    "\n",
    "\n",
    "dependentOpPair: Tuple[AffineStoreOp, AffineLoadOp] = GatherDependentOpPairs(srcLoopNest, dstLoopNest)[0]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "基于mlir构造出访问关系的多面体表示, 在本文中使用的Integer Set Library库来操作多面体表示, 而在mlir中则是使用内置的[FPL](https://grosser.science/FPL/)库来操作多面体表示. 提取的过程中主要分析AffineForOp和对应AffineLoadOp来构建access relation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "srcAccessRel { [i0, i1, i2, i3] -> [o0, o1, o2] : o0 = i0 and o1 = i1 and o2 = i2 and 0 <= i0 <= 7 and 0 <= i1 <= 127 and 0 <= i2 <= 511 and 0 <= i3 <= 383 }\n",
      "dstAccessRel { [i0, i1, i2, i3] -> [o0, o1, o2] : o0 = i0 and o1 = i1 and o2 = i3 and 0 <= i0 <= 7 and 0 <= i1 <= 127 and 0 <= i2 <= 63 and 0 <= i3 <= 511 }\n"
     ]
    }
   ],
   "source": [
    "class MemRefAccess:\n",
    "  memref: Value\n",
    "  op: OpView\n",
    "  indices: List[Value]\n",
    "  isStore: False\n",
    "\n",
    "  def __init__(self, op: OpView) -> None:\n",
    "    if isinstance(op, AffineLoadOp):\n",
    "      self.isStore = False\n",
    "    elif isinstance(op, AffineStoreOp):\n",
    "      self.isStore = True\n",
    "    else:\n",
    "      raise NotImplementedError()\n",
    "    self.op = op\n",
    "    self.memref = op.memref\n",
    "    self.indices = op.indices\n",
    "\n",
    "srcMemAccess = MemRefAccess(dependentOpPair[0])\n",
    "dstMemAccess = MemRefAccess(dependentOpPair[1])\n",
    "\n",
    "def GetBound(attr: AffineMapAttr) -> int:\n",
    "  \"\"\" currently only support constant bound \"\"\"\n",
    "  map: AffineMap = attr.value  # note mlir doesn't export the get value in python bindings.\n",
    "  if len(map.results) != 1:\n",
    "    raise NotImplementedError()\n",
    "  elif AffineConstantExpr.isinstance(map.results[0]):\n",
    "    return AffineConstantExpr(map.results[0]).value\n",
    "  else:\n",
    "    raise NotImplementedError()\n",
    "\n",
    "\n",
    "def GetEqualDimConstraint(bmap: isl.basic_map, in_index: int, out_index: int) -> isl.constraint:\n",
    "  ls = isl.local_space.from_space(bmap.space())\n",
    "  c = isl.constraint.alloc_equality(ls)\n",
    "  c = c.set_coefficient_si(isl.ISL_DIM_TYPE.IN, in_index, -1)\n",
    "  c = c.set_coefficient_si(isl.ISL_DIM_TYPE.OUT, out_index, 1)\n",
    "  return c\n",
    "\n",
    "\n",
    "def GetInEqualDimConstraint(bmap: isl.basic_map, in_index: int, value: int, coeff: int) -> isl.constraint:\n",
    "  ls = isl.local_space.from_space(bmap.space())\n",
    "  c = isl.constraint.alloc_inequality(ls)\n",
    "  c = c.set_constant_si(value)\n",
    "  c = c.set_coefficient_si(isl.ISL_DIM_TYPE.IN, in_index, coeff)\n",
    "  return c\n",
    "\n",
    "\n",
    "def AddRangeConstraint(bmap: isl.basic_map, loops: List[AffineForOp], indices: List[Value], out_index: int) -> isl.basic_map:\n",
    "  value: Value = indices[out_index]\n",
    "  owner: Optional[Block | Operation] = value.owner\n",
    "  if isinstance(owner, Block):\n",
    "    op: OpView = owner.owner\n",
    "    if isinstance(op, AffineForOp):\n",
    "      in_index = loops.index(op)\n",
    "      bmap = bmap.add_constraint(GetEqualDimConstraint(bmap, in_index, out_index))\n",
    "  elif isinstance(owner, Operation):\n",
    "    raise NotImplementedError()\n",
    "  else:\n",
    "    raise ValueError()\n",
    "  return bmap\n",
    "\n",
    "\n",
    "def AddDomainConstraint(bmap: isl.basic_map, loops: List[AffineForOp], in_index: int) -> isl.basic_map:\n",
    "  loop = loops[in_index]\n",
    "  lower_bound = GetBound(loop.attributes['lower_bound'])\n",
    "  bmap = bmap.add_constraint(GetInEqualDimConstraint(bmap, in_index, lower_bound, 1))\n",
    "  upper_bound = GetBound(loop.attributes['upper_bound'])\n",
    "  bmap = bmap.add_constraint(GetInEqualDimConstraint(bmap, in_index, upper_bound - 1, -1))\n",
    "  return bmap\n",
    "\n",
    "def GetAffineForIVs(op: OpView) -> List[AffineForOp]:\n",
    "  currOp: Operation = op.operation.parent\n",
    "  loops: List[AffineForOp] = []\n",
    "  while (currOp is not None):\n",
    "    if isinstance(currOp.opview, AffineForOp):\n",
    "      loops.append(currOp.opview)\n",
    "    currOp = currOp.parent\n",
    "  return loops[::-1]\n",
    "\n",
    "def GetAccessRelation(this: MemRefAccess) -> isl.basic_map:\n",
    "  domain = GetAffineForIVs(this.op)\n",
    "  domainRank = len(domain)\n",
    "  rangeRank = len(this.indices)\n",
    "  space = isl.space.unit()\n",
    "  space = space.add_unnamed_tuple(domainRank)\n",
    "  space = space.add_unnamed_tuple(rangeRank)\n",
    "  bmap = isl.basic_map.universe(space)\n",
    "  for i in range(domainRank):\n",
    "    bmap = AddDomainConstraint(bmap, domain, i)\n",
    "  for i in range(rangeRank):\n",
    "    bmap = AddRangeConstraint(bmap, domain, this.indices, i)\n",
    "  return bmap\n",
    "\n",
    "\n",
    "srcAccessRel = GetAccessRelation(srcMemAccess)\n",
    "dstAccessRel = GetAccessRelation(dstMemAccess)\n",
    "print(\"srcAccessRel\", srcAccessRel)\n",
    "print(\"dstAccessRel\", dstAccessRel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在造得到了srcAccRel/dstAccRel的多面体映射, 他们代表的两个不同循环domain下对于同一个buffer读写映射, 可以他们可以得到循环domain之间的映射关系:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dstSrcDomainRel { [i0, i1, i2, i3] -> [o0, o1, o2, o3] : o0 = i0 and o1 = i1 and o2 = i3 and 0 <= i0 <= 7 and 0 <= i1 <= 127 and 0 <= i2 <= 63 and 0 <= i3 <= 511 and 0 <= o3 <= 383 }\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def GetDstSrcDomainRelation(srcMap: isl.basic_map, dstMap: isl.basic_map) -> isl.basic_map:\n",
    "  srcR = srcMap.reverse()  # buffer -> srcdomain\n",
    "  return dstMap.apply_range(srcR)  # dst domain -> src domain\n",
    "\n",
    "\n",
    "dstSrcDomainRel = GetDstSrcDomainRelation(srcAccessRel, dstAccessRel)\n",
    "print(\"dstSrcDomainRel\", dstSrcDomainRel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. 计算收益\n",
    "\n",
    "目前为止拿到了足够的信息, 可以开始准备进行Fuse, 但想象现在需要将src循环移动到dst循环内部位置, 这个位置只要不违反依赖的话应该存在很多种可能, 所以这里第一步是计算可以插入的循环体深度(为了简化问题, 这里也只考虑在每个循环开始时插入, 不考虑循环体内block中的位置)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "InnermostLoopDepth: 4\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def FilterOps(dst: PerfectLoopNest, depMemrefs: Set[Value]) -> List[OpView]:\n",
    "  dstMemrefOps: List[OpView] = []\n",
    "  for load in dst.loadOps:\n",
    "    if load.memref in depMemrefs:\n",
    "      dstMemrefOps.append(load)\n",
    "  for store in dst.storeOps:\n",
    "    if store.memref in depMemrefs:\n",
    "      dstMemrefOps.append(store)\n",
    "  return dstMemrefOps\n",
    "\n",
    "\n",
    "def GetInnermostCommonLoopDepth(ops: List[OpView]) -> int:\n",
    "  numOps = len(ops)\n",
    "  assert numOps > 0 and \"Expected at least one operation\"\n",
    "\n",
    "  loops: List[List[AffineForOp]] = [[] for _ in range(numOps)]\n",
    "  loopDepthLimit = 1 << 31\n",
    "  for i in range(numOps):\n",
    "    loops[i] = GetAffineForIVs(ops[i])\n",
    "    loopDepthLimit = min(loopDepthLimit, len(loops[i]))\n",
    "\n",
    "  loopDepth = 0\n",
    "  for d in range(loopDepthLimit):\n",
    "    for i in range(1, numOps):\n",
    "      if (loops[i - 1][d] != loops[i][d]):\n",
    "        return loopDepth\n",
    "    loopDepth += 1\n",
    "  return loopDepth\n",
    "\n",
    "\n",
    "dstMemrefOps = FilterOps(dstLoopNest, producerConsumerMemrefs)\n",
    "InnermostLoopDepth = GetInnermostCommonLoopDepth(dstMemrefOps)\n",
    "print(\"InnermostLoopDepth:\", InnermostLoopDepth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "得到的InnermostLoopDepth意味着现在有四个fuse的机会`[0:InnermostLoopDepth)`, 那么为了确定哪个机会最好, 我们还需要进一步的分析收益. 考虑fuse后意味着domain映射关系中只有0到dstLoopDepth的循环变量会参与到src循环体计算, 而src循环体中的一部分将会在dstLoopDepth位置进行计算, 这里定义ComputationSliceState类来记录srcLoop所需要被重复执行的循环相关信息. 并且假设需要插入的位置为dstLoopDepthTest = 0来进行后续步骤."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ComputationSliceState:\n",
    "  def __init__(self, srcLoops: PerfectLoopNest, dstLoops: PerfectLoopNest, domainRel: isl.basic_map, dstDepth: int) -> None:\n",
    "    self.srcLoops = srcLoops\n",
    "    self.dstLoops = dstLoops\n",
    "    self.dstDepth = dstDepth\n",
    "    self.sliceDomainRel: isl.basic_map = domainRel.project_out(\n",
    "        isl.ISL_DIM_TYPE.IN, dstDepth + 1, len(dstLoops.forOps) - (dstDepth + 1))\n",
    "\n",
    "  def GetSliceTripCountMap(self) -> Dict[Operation, int]:\n",
    "    sliceTripCountMap: Dict[Operation, int] = {}\n",
    "    rg_set = self.sliceDomainRel.domain().space().universe_set()\n",
    "    for i in range(self.dstDepth + 1):\n",
    "      rg_set = rg_set.lower_bound_si(isl.ISL_DIM_TYPE.SET, i, 0)\n",
    "      rg_set = rg_set.upper_bound_si(isl.ISL_DIM_TYPE.SET, i, 0)\n",
    "    rg = rg_set.apply(self.sliceDomainRel)\n",
    "\n",
    "    for i in range(rg.tuple_dim()):\n",
    "      max = rg.dim_max_val(i).num_si()\n",
    "      min = rg.dim_min_val(i).num_si()\n",
    "      sliceTripCountMap[self.srcLoops.forOps[i]] = max - min + 1\n",
    "    return sliceTripCountMap\n",
    "\n",
    "\n",
    "dstLoopDepthTest = 0\n",
    "sliceStateTest = ComputationSliceState(srcLoopNest, dstLoopNest, dstSrcDomainRel, dstLoopDepthTest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "下面实现一个简单的类来计算嵌套循环的ComputeCost作为评估收益的cost model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "srcLoopNestCost 1207959552\n",
      "dstLoopNestCost 201326592\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def GetConstantTripCount(forOp: AffineForOp) -> int:\n",
    "  lb = GetBound(forOp.attributes['lower_bound'])\n",
    "  ub = GetBound(forOp.attributes['upper_bound'])\n",
    "  return ub - lb\n",
    "\n",
    "\n",
    "class LoopNestStats(IrVisitor):\n",
    "  loopMap: Dict[AffineForOp, List[AffineForOp]]\n",
    "  opCountMap: Dict[AffineForOp, int]\n",
    "  tripCountMap: Dict[AffineForOp, int]\n",
    "\n",
    "  def __init__(self, forOp: AffineForOp) -> None:\n",
    "    super().__init__()\n",
    "    self.rootForOp = forOp\n",
    "    self.loopMap = {}\n",
    "    self.opCountMap = {}\n",
    "    self.tripCountMap = {}\n",
    "\n",
    "  @staticmethod\n",
    "  def collect(forOp: AffineForOp) -> LoopNestStats:\n",
    "    stats = LoopNestStats(forOp)\n",
    "    stats.visit(forOp)\n",
    "    return stats\n",
    "\n",
    "  def runBeforeOperation(self, op: OpView) -> bool:\n",
    "    if not isinstance(op, AffineForOp):\n",
    "      return True\n",
    "\n",
    "    childForOp: AffineForOp = op\n",
    "    if childForOp != self.rootForOp:\n",
    "      parentForOp: AffineForOp = op.operation.parent.opview\n",
    "      lst = self.loopMap.get(parentForOp)\n",
    "      if lst:\n",
    "        lst.append(childForOp)\n",
    "      else:\n",
    "        self.loopMap.setdefault(parentForOp, [childForOp])\n",
    "\n",
    "    count = 0\n",
    "    self.opCountMap[childForOp] = 0\n",
    "    for iop in childForOp.region.blocks[0]:\n",
    "      if not isinstance(iop, AffineIfOp) and not isinstance(iop, AffineForOp):\n",
    "        count += 1\n",
    "    self.opCountMap[childForOp] = count\n",
    "    self.tripCountMap[childForOp] = GetConstantTripCount(childForOp)\n",
    "    return True\n",
    "\n",
    "\n",
    "def GetLoopComputeCost(forOp: AffineForOp, stats: LoopNestStats, tripCountOverrideMap: Dict[AffineForOp, int] = None, computeCostMap: Dict[AffineForOp, int] = None):\n",
    "  opCount = stats.opCountMap[forOp] - 1\n",
    "  if stats.loopMap.get(forOp) is not None:\n",
    "    for childForOp in stats.loopMap[forOp]:\n",
    "      opCount += GetLoopComputeCost(childForOp, stats, tripCountOverrideMap,\n",
    "                                    computeCostMap)\n",
    "  if computeCostMap is not None and computeCostMap.get(forOp) is not None:\n",
    "    opCount += computeCostMap[forOp]\n",
    "  tripCount = stats.tripCountMap[forOp]\n",
    "  if tripCountOverrideMap is not None and tripCountOverrideMap.get(forOp) is not None:\n",
    "    tripCount = tripCountOverrideMap[forOp]\n",
    "  return tripCount * opCount\n",
    "\n",
    "\n",
    "srcLoopStats = LoopNestStats.collect(srcLoopNest.forOps[0])\n",
    "srcLoopNestCost = GetLoopComputeCost(srcLoopNest.forOps[0], srcLoopStats)\n",
    "print(\"srcLoopNestCost\", srcLoopNestCost)\n",
    "dstLoopStats = LoopNestStats.collect(dstLoopNest.forOps[0])\n",
    "dstLoopNestCost = GetLoopComputeCost(dstLoopNest.forOps[0], dstLoopStats)\n",
    "print(\"dstLoopNestCost\", dstLoopNestCost)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "计算fuse到dstLoopDepth时的额外增加cost:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1409286144\n",
      "additional compute fraction: 0.0 %\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def GetFusedLoopComputeCost(srcForOp: AffineForOp,\n",
    "                            srcStats: LoopNestStats,\n",
    "                            dstForOp: AffineForOp,\n",
    "                            dstStats: LoopNestStats,\n",
    "                            sliceState: ComputationSliceState,\n",
    "                            ) -> int:\n",
    "  computeCostMap: Dict[Operation, int] = {}\n",
    "\n",
    "  sliceTripCountMap = sliceState.GetSliceTripCountMap()\n",
    "\n",
    "  sliceIterationCount = 1\n",
    "  for c in sliceTripCountMap.values():\n",
    "    sliceIterationCount *= c\n",
    "  assert (sliceIterationCount > 0)\n",
    "  storeLoadFwdGuaranteed: bool = (sliceIterationCount == 1)\n",
    "  insertPointParent: AffineForOp = sliceState.dstLoops.forOps[sliceState.dstDepth]\n",
    "\n",
    "  if (storeLoadFwdGuaranteed):\n",
    "    storeCount = 0\n",
    "    storeMemrefs: Set[Value] = set()\n",
    "\n",
    "    def lamb(op: OpView) -> bool:\n",
    "      if isinstance(op, AffineStoreOp):\n",
    "        storeMemrefs.add(op.memref)\n",
    "      return True\n",
    "    walker = IrVisitor(afterOperation=lamb)\n",
    "    walker.visit(srcForOp)\n",
    "    if (storeCount > 0):\n",
    "      computeCostMap[insertPointParent] = -storeCount\n",
    "    for memref in storeMemrefs:\n",
    "      for user in memref.uses:\n",
    "        userOp: OpView = user.owner\n",
    "        if isinstance(userOp, AffineLoadOp):\n",
    "          loops: List[AffineForOp] = GetAffineForIVs(userOp)\n",
    "          if (loops.count(insertPointParent)):\n",
    "            parentOp = userOp.operation.parent.opview\n",
    "            if isinstance(parentOp, AffineForOp):\n",
    "              computeCostMap.setdefault(parentOp, 1)\n",
    "              computeCostMap[parentOp] -= 1\n",
    "\n",
    "  sliceComputeCost = GetLoopComputeCost(\n",
    "      srcForOp, srcStats, sliceTripCountMap, computeCostMap)\n",
    "\n",
    "  computeCostMap[insertPointParent] = sliceComputeCost\n",
    "\n",
    "  computeCost = GetLoopComputeCost(dstForOp, dstStats, None, computeCostMap)\n",
    "  return computeCost\n",
    "\n",
    "\n",
    "fusedComputeCostTest = GetFusedLoopComputeCost(srcLoopNest.forOps[0], srcLoopStats,\n",
    "                                               dstLoopNest.forOps[0], dstLoopStats, sliceStateTest)\n",
    "print(fusedComputeCostTest)\n",
    "additionalComputeCost = (fusedComputeCostTest / (srcLoopNestCost + dstLoopNestCost)) - 1\n",
    "print(f\"additional compute fraction: {additionalComputeCost * 100} %\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 执行合并\n",
    "最后我们要基于分析的结果在对应的位置执行Fuse操作, 这里分三步走, 第一步detach srcLoops并在dstLoop的对应位置插入它."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "module {\n",
      "  func.func @main(%arg0: memref<8x128x384xf32>, %arg1: memref<8x384x512xf32>, %arg2: memref<8x128x512xf32>, %arg3: memref<8x512x64xf32>, %arg4: memref<8x128x64xf32>) {\n",
      "    affine.for %arg5 = 0 to 8 {\n",
      "      affine.for %arg6 = 0 to 8 {\n",
      "        affine.for %arg7 = 0 to 128 {\n",
      "          affine.for %arg8 = 0 to 512 {\n",
      "            affine.for %arg9 = 0 to 384 {\n",
      "              %0 = affine.load %arg0[%arg6, %arg7, %arg9] : memref<8x128x384xf32>\n",
      "              %1 = affine.load %arg1[%arg6, %arg9, %arg8] : memref<8x384x512xf32>\n",
      "              %2 = affine.load %arg2[%arg6, %arg7, %arg8] : memref<8x128x512xf32>\n",
      "              %3 = arith.mulf %0, %1 : f32\n",
      "              %4 = arith.addf %2, %3 : f32\n",
      "              affine.store %4, %arg2[%arg6, %arg7, %arg8] : memref<8x128x512xf32>\n",
      "            }\n",
      "          }\n",
      "        }\n",
      "      }\n",
      "      affine.for %arg6 = 0 to 128 {\n",
      "        affine.for %arg7 = 0 to 64 {\n",
      "          affine.for %arg8 = 0 to 512 {\n",
      "            %0 = affine.load %arg2[%arg5, %arg6, %arg8] : memref<8x128x512xf32>\n",
      "            %1 = affine.load %arg3[%arg5, %arg8, %arg7] : memref<8x512x64xf32>\n",
      "            %2 = affine.load %arg4[%arg5, %arg6, %arg7] : memref<8x128x64xf32>\n",
      "            %3 = arith.mulf %0, %1 : f32\n",
      "            %4 = arith.addf %2, %3 : f32\n",
      "            affine.store %4, %arg4[%arg5, %arg6, %arg7] : memref<8x128x64xf32>\n",
      "          }\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "    return\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "def MoveSrcLoopsIntoDstLoops(srcLoops: PerfectLoopNest,\n",
    "                           dstLoops: PerfectLoopNest,\n",
    "                           sliceState: ComputationSliceState):\n",
    "  srcLoopRoot: Operation = srcLoops.forOps[0].operation\n",
    "\n",
    "  with InsertionPoint.at_block_begin(dstLoops.forOps[sliceState.dstDepth].region.blocks[0]) as ip, Location.unknown():\n",
    "    ip.insert(srcLoopRoot.detach_from_parent())\n",
    "\n",
    "MoveSrcLoopsIntoDstLoops(srcLoopNest, dstLoopNest, sliceStateTest)\n",
    "mod.dump()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "然后从DomainRelation中分析出srcLoops和dstLoops中迭代变量的映射关系, 这里我的做法支持Identity的映射分析, 基于所有的约束, 找到Range和Domain之间的Dim等价约束, 同时排出其他额外的约束. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{0: 0}\n"
     ]
    }
   ],
   "source": [
    "def AnalysisIvMapping(sliceState: ComputationSliceState):\n",
    "  eqMat = sliceState.sliceDomainRel.equalities_matrix(\n",
    "      isl.ISL_DIM_TYPE.IN,\n",
    "      isl.ISL_DIM_TYPE.OUT,\n",
    "      isl.ISL_DIM_TYPE.PARAM,\n",
    "      isl.ISL_DIM_TYPE.DIV,\n",
    "      isl.ISL_DIM_TYPE.CST)\n",
    "  domainVarMap: Dict[int, int] = {}\n",
    "  inRank = sliceState.sliceDomainRel.dim(isl.ISL_DIM_TYPE.IN)\n",
    "  outRank = sliceState.sliceDomainRel.dim(isl.ISL_DIM_TYPE.OUT)\n",
    "  cstRank = sliceState.sliceDomainRel.dim(isl.ISL_DIM_TYPE.CST)\n",
    "  for r in range(eqMat.rows()):\n",
    "    noCoff = True\n",
    "    for i in range(inRank + outRank, eqMat.cols()):\n",
    "      noCoff &= eqMat.element_val(r, i).is_zero()\n",
    "    if (not noCoff):\n",
    "      continue\n",
    "    for i in range(0, inRank):\n",
    "      inv = eqMat.element_val(r, i)\n",
    "      for j in range(inRank, inRank + outRank):\n",
    "        outv = eqMat.element_val(r, j)\n",
    "        if not inv.is_zero() and not outv.is_zero() and inv.add(outv).is_zero():\n",
    "          if domainVarMap.get(i, None) is None:\n",
    "            domainVarMap.setdefault(i, j - inRank)\n",
    "          else:\n",
    "            raise NotImplementedError(\"the same input dim can't equal to muli output dim\")\n",
    "\n",
    "  ineqMat = sliceState.sliceDomainRel.inequalities_matrix(\n",
    "      isl.ISL_DIM_TYPE.IN,\n",
    "      isl.ISL_DIM_TYPE.OUT,\n",
    "      isl.ISL_DIM_TYPE.PARAM,\n",
    "      isl.ISL_DIM_TYPE.DIV,\n",
    "      isl.ISL_DIM_TYPE.CST)\n",
    "\n",
    "  for (k, v) in domainVarMap.items():\n",
    "    for r in range(ineqMat.rows()):\n",
    "      if not ineqMat.element_val(r, k).is_zero():\n",
    "        noCoff = True\n",
    "        for i in range(0, ineqMat.cols() - cstRank):\n",
    "          if i == k:\n",
    "            continue\n",
    "          noCoff &= ineqMat.element_val(r, i).is_zero()\n",
    "        if not noCoff:\n",
    "          raise NotImplementedError(\"not support non identity mapping!\")\n",
    "  return domainVarMap\n",
    "\n",
    "\n",
    "ivMapTest = AnalysisIvMapping(sliceStateTest)\n",
    "print(ivMapTest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "基于分析的结果将srcLoop循环中的Load/Store所使用的迭代变量替换为dstLoops中的迭代变量, 并将被替换那个循环消除:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "replaced!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "module {\n",
      "  func.func @main(%arg0: memref<8x128x384xf32>, %arg1: memref<8x384x512xf32>, %arg2: memref<8x128x512xf32>, %arg3: memref<8x512x64xf32>, %arg4: memref<8x128x64xf32>) {\n",
      "    affine.for %arg5 = 0 to 8 {\n",
      "      affine.for %arg6 = 0 to 128 {\n",
      "        affine.for %arg7 = 0 to 512 {\n",
      "          affine.for %arg8 = 0 to 384 {\n",
      "            %0 = affine.load %arg0[%arg5, %arg6, %arg8] : memref<8x128x384xf32>\n",
      "            %1 = affine.load %arg1[%arg5, %arg8, %arg7] : memref<8x384x512xf32>\n",
      "            %2 = affine.load %arg2[%arg5, %arg6, %arg7] : memref<8x128x512xf32>\n",
      "            %3 = arith.mulf %0, %1 : f32\n",
      "            %4 = arith.addf %2, %3 : f32\n",
      "            affine.store %4, %arg2[%arg5, %arg6, %arg7] : memref<8x128x512xf32>\n",
      "          }\n",
      "        }\n",
      "      }\n",
      "      affine.for %arg6 = 0 to 128 {\n",
      "        affine.for %arg7 = 0 to 64 {\n",
      "          affine.for %arg8 = 0 to 512 {\n",
      "            %0 = affine.load %arg2[%arg5, %arg6, %arg8] : memref<8x128x512xf32>\n",
      "            %1 = affine.load %arg3[%arg5, %arg8, %arg7] : memref<8x512x64xf32>\n",
      "            %2 = affine.load %arg4[%arg5, %arg6, %arg7] : memref<8x128x64xf32>\n",
      "            %3 = arith.mulf %0, %1 : f32\n",
      "            %4 = arith.addf %2, %3 : f32\n",
      "            affine.store %4, %arg4[%arg5, %arg6, %arg7] : memref<8x128x64xf32>\n",
      "          }\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "    return\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "def ReplaceIVAndCleanUp(srcLoops: PerfectLoopNest,\n",
    "                        dstLoops: PerfectLoopNest,\n",
    "                        ivMap: Dict[int, int]):\n",
    "\n",
    "  argMap = {}\n",
    "  candidates = set()\n",
    "  for (k, v) in ivMap.items():\n",
    "    argMap[srcLoops.forOps[k].region.blocks[0].arguments[0]\n",
    "           ] = dstLoopNest.forOps[v].region.blocks[0].arguments[0]\n",
    "    candidates.add(srcLoops.forOps[k])\n",
    "\n",
    "  def replaceArgs(op: OpView):\n",
    "    if len(op.regions) == 0:\n",
    "      for value in op.operands:\n",
    "        if BlockArgument.isinstance(value) and argMap.get(value, None) is not None:\n",
    "          value.replace_all_uses_with(argMap[value])\n",
    "          print(\"replaced!\")\n",
    "    return True\n",
    "\n",
    "  v = IrVisitor(beforeOperation=replaceArgs)\n",
    "  v.visit(dstLoops.forOps[0])\n",
    "\n",
    "  # remove the candidates\n",
    "\n",
    "  def removeCandidates(op: OpView):\n",
    "    if isinstance(op, AffineForOp):\n",
    "      childBlock: Block = op.region.blocks[0]\n",
    "      if childBlock.operations[0] in candidates:\n",
    "        removeOp: OpView = childBlock.operations[0]\n",
    "        with InsertionPoint.at_block_begin(childBlock) as ip, Location.unknown():\n",
    "          ip.insert(removeOp.region.blocks[0].operations[0].detach_from_parent())\n",
    "        removeOp.detach_from_parent()\n",
    "        candidates.remove(removeOp)\n",
    "        return False\n",
    "\n",
    "    return True\n",
    "  while len(candidates):\n",
    "    v = IrVisitor(beforeOperation=removeCandidates)\n",
    "    v.visit(dstLoops.forOps[0])\n",
    "\n",
    "\n",
    "ReplaceIVAndCleanUp(srcLoopNest, dstLoopNest, ivMapTest)\n",
    "mod.dump()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 整体执行流程\n",
    "到这里一步我们已经走完了所有的流程, 最终我们将上述流程串起来则是整个的Affine Fusion优化Pass:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fused src Loops at dst Loops 3, got additional compute cost 5400.0 %\n",
      "Fused src Loops at dst Loops 2, got additional compute cost 5400.0 %\n",
      "Fused src Loops at dst Loops 1, got additional compute cost 0.0 %\n",
      "Fused src Loops at dst Loops 0, got additional compute cost 0.0 %\n",
      "replaced!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "module {\n",
      "  func.func @main(%arg0: memref<8x128x384xf32>, %arg1: memref<8x384x512xf32>, %arg2: memref<8x128x512xf32>, %arg3: memref<8x512x64xf32>, %arg4: memref<8x128x64xf32>) {\n",
      "    affine.for %arg5 = 0 to 8 {\n",
      "      affine.for %arg6 = 0 to 128 {\n",
      "        affine.for %arg7 = 0 to 128 {\n",
      "          affine.for %arg8 = 0 to 512 {\n",
      "            affine.for %arg9 = 0 to 384 {\n",
      "              %0 = affine.load %arg0[%arg5, %arg7, %arg9] : memref<8x128x384xf32>\n",
      "              %1 = affine.load %arg1[%arg5, %arg9, %arg8] : memref<8x384x512xf32>\n",
      "              %2 = affine.load %arg2[%arg5, %arg7, %arg8] : memref<8x128x512xf32>\n",
      "              %3 = arith.mulf %0, %1 : f32\n",
      "              %4 = arith.addf %2, %3 : f32\n",
      "              affine.store %4, %arg2[%arg5, %arg7, %arg8] : memref<8x128x512xf32>\n",
      "            }\n",
      "          }\n",
      "        }\n",
      "        affine.for %arg7 = 0 to 64 {\n",
      "          affine.for %arg8 = 0 to 512 {\n",
      "            %0 = affine.load %arg2[%arg5, %arg6, %arg8] : memref<8x128x512xf32>\n",
      "            %1 = affine.load %arg3[%arg5, %arg8, %arg7] : memref<8x512x64xf32>\n",
      "            %2 = affine.load %arg4[%arg5, %arg6, %arg7] : memref<8x128x64xf32>\n",
      "            %3 = arith.mulf %0, %1 : f32\n",
      "            %4 = arith.addf %2, %3 : f32\n",
      "            affine.store %4, %arg4[%arg5, %arg6, %arg7] : memref<8x128x64xf32>\n",
      "          }\n",
      "        }\n",
      "      }\n",
      "    }\n",
      "    return\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "with open(\"test1.mlir\") as f:\n",
    "  mod = Module.parse(f.read(), ctx)\n",
    "\n",
    "srcLoopNest, dstLoopNest = LoopNestPairCollector.collect(mod)\n",
    "producerConsumerMemrefs = GatherProducerConsumerMemrefs(srcLoopNest, dstLoopNest)\n",
    "dependentOpPair: Tuple[AffineStoreOp, AffineLoadOp] = GatherDependentOpPairs(srcLoopNest, dstLoopNest)[\n",
    "    0]\n",
    "srcMemAccess = MemRefAccess(dependentOpPair[0])\n",
    "dstMemAccess = MemRefAccess(dependentOpPair[1])\n",
    "srcAccessRel = GetAccessRelation(srcMemAccess)\n",
    "dstAccessRel = GetAccessRelation(dstMemAccess)\n",
    "dstSrcDomainRel = GetDstSrcDomainRelation(srcAccessRel, dstAccessRel)\n",
    "dstMemrefOps = FilterOps(dstLoopNest, producerConsumerMemrefs)\n",
    "InnermostLoopDepth = GetInnermostCommonLoopDepth(dstMemrefOps)\n",
    "srcLoopStats = LoopNestStats.collect(srcLoopNest.forOps[0])\n",
    "srcLoopNestCost = GetLoopComputeCost(srcLoopNest.forOps[0], srcLoopStats)\n",
    "dstLoopStats = LoopNestStats.collect(dstLoopNest.forOps[0])\n",
    "dstLoopNestCost = GetLoopComputeCost(dstLoopNest.forOps[0], dstLoopStats)\n",
    "\n",
    "sliceStates: ComputationSliceState = []\n",
    "for depth in range(0, InnermostLoopDepth):\n",
    "  sliceState = ComputationSliceState(srcLoopNest, dstLoopNest, dstSrcDomainRel, depth)\n",
    "  sliceStates.append(sliceState)\n",
    "\n",
    "bestAdditionalComputeCost = 0.30\n",
    "bestSliceState = None\n",
    "for sliceState in sliceStates[::-1]:\n",
    "  fusedCost = GetFusedLoopComputeCost(srcLoopNest.forOps[0], srcLoopStats,\n",
    "                                      dstLoopNest.forOps[0], dstLoopStats, sliceState)\n",
    "\n",
    "  additionalComputeCost = (fusedCost / (srcLoopNestCost + dstLoopNestCost)) - 1\n",
    "  print(f\"Fused src Loops at dst Loops {sliceState.dstDepth}, got additional compute cost {additionalComputeCost*100} %\")\n",
    "  if additionalComputeCost < bestAdditionalComputeCost:\n",
    "    bestAdditionalComputeCost = additionalComputeCost\n",
    "    bestSliceState = sliceState\n",
    "\n",
    "if bestSliceState is not None:\n",
    "  MoveSrcLoopsIntoDstLoops(srcLoopNest, dstLoopNest, bestSliceState)\n",
    "  ivMapTest = AnalysisIvMapping(sliceStateTest)\n",
    "  ReplaceIVAndCleanUp(srcLoopNest, dstLoopNest, ivMapTest)\n",
    "\n",
    "mod.dump()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 扩展\n",
    "\n",
    "如果读者看懂了上述操作, 考虑进一步扩展提升的话, 这里从易到难留一些问题:\n",
    "1. 如何支持更多的依赖对/更具复杂的上下文\n",
    "2. 如何支持采用复杂的AffineMap映射迭代变量的情况, 比如`A[i+j]`这种情况\n",
    "3. 除了贪心的合并最后的循环体, 是否有更好的策略\n",
    "4. 如何设计更好的cost model来评估收益"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ci",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}