14_affine_fusion_en.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Abstract\n",
    "\n",
    "Currently, the Triton language is experiencing rapid growth in the field of AI compilers. In this context, compiler optimization techniques for the Triton language are of utmost importance. Within business circles, there is consensus that the key to optimizing the Triton language is to lower it to the MLIR Linalg dialect (e.g. Triton-Linalg/Triton-Shared) and then leverage the Affine form to implement further optimizations.\n",
    "The affine optimization methods not only improve the performance of compiled programs but also ensure the extensibility and maintainability of the compilation process. This article will provide a comprehensive overview of the Affine-based Loop Fusion pass, aiming to contribute to the advancement of AI compiler techniques.\n",
    "\n",
    "Note:\n",
    "1. This article mainly references the mlir built-in transform affine-loop-fusion.\n",
    "2. this article assumes many situations to simplify the process.\n",
    "3. this article's source code is in my [tutorial repo](https://github.com/zhen8838/isl_learn/blob/main/14_affine_fusion_en.ipynb).\n",
    "4. this tutorial requires LLVM 17.04 and applying these patches bindings.patch and python.patch which are provided in my repo.\n",
    "5. this tutorial requires the ISL library which I provided [here](https://github.com/zhen8838/pet)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implement Details\n",
    "\n",
    "## 1. Load and Parse IR\n",
    "First of all, we parse the mlir source file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": [
    "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\n",
    "\n",
    "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 Pass aims to fuse the adjacent two perfect loops, so we need to identify them first. we define the class LoopNestPairCollector to collect perfect loops that are adjacent to each other:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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. Analyze  and Extract Polyhedron Model\n",
    "We need to try to fusion the adjacent loops that we have just collected. To maintain the program's correctness, we need to ensure that the fusion operation does not break the dependence relation between the loops. Let's begin by gathering the dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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": [
    "In this tutorial, I simplify the problem by assuming only one pair of producer/consumer dependence and getting it directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": [
    "Based on the dependent pair, we can construct the polyhedron representation of the access relation. In this article, we use the Integer Set Library to manipulate polyhedron representation, and the MLIR uses a built-in [FPL](https://grosser.science/FPL/) library. The extraction process mainly involves analyzing the AffineForOp and AffineLoadOp/AffineStoreOp:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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": [
    "Now we have created access relation srcAccRel and dstAccRel. They represent two different loop domains read and write within the same buffer, and we can get the mapping relationship between the loop domains by composing them:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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. Compute the profitable\n",
    "\n",
    "so far we have enough information to execute fusion. but actually, we have many insert points in the dst loops to contain the src loops and they will not violate the dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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": [
    "The InnermostLoopDepth indicates that there are four insertion points in the destination loops [0:InnermostLoopDepth). To calculate the profits at each insertion point, we are selecting a starting point for computation.\n",
    "\n",
    "We have defined the ComputationSliceState class to keep track of the recomputation part from the source loops when they are merged into the destination loops. It is assumed that the position to be inserted is dstLoopDepthTest = 0 for the next steps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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": [
    "The following code implements a simple class for calculating the cost of nested loop  as a cost model to evaluate benefits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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": [
    "Computing the additional cost when the source loops were merged into dstLoopDepth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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. Perform Fusion\n",
    "\n",
    "Finally, we need to perform fusion on the selected inert point based on the analysis results. We will separate the process into three parts. The first part involves detaching source loops and inserting them into the point:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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": [
    "In the second part, we analyze the relationship of iterative variables in srcLoops and dstLoops from DomainRelation. In my implementation, only the identity relationship is supported. For each constraint, we need to find the equivalent dimension constraint between Range and Domain, and exclude any additional constraints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "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": [
    "During the last step, we replace the iteration variables used by the AffineLoadOp/AffineStoreOp in the srcLoops with new iteration variables from dstLoops according to the analysis results. Then, we detach the loop whose iteration variable is no longer used in the srcLoops:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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": [
    "# Complete Processes\n",
    "At this point, we have completed all the steps, and finally, we string together the above processes as the full Affine Fusion Transform:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "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": [
    "# Extended Thinking\n",
    "If readers understand the above optimization and want to go deeper, here are some questions from easy to difficult:\n",
    "1. How to support more dependence pairs in the context?\n",
    "2. How to support more complex access relations, e.g. A[i*2] or A[i+j].\n",
    "3. Besides the greedy fusion strategy, Does it have other better fusion strategy?\n",
    "4. How to design a more precise cost model for computing profit?"
   ]
  }
 ],
 "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
}