[mlir] Slow TopologicalSort for large functions

[Jokeren]

The TopologicalSort implementation adds every node it sees into the queue without filtering those have been visited. Therefore, the searching space grows rapidly, and immediately the size of the queue explodes.

TopologicalSort

Adding a single line can resolve the problem:

    if (!state->seen.insert(current).second)
      continue;

cc @ptillet, @ngimel

Tested with llvm-15.0.6

Reproduced with:

mlir-opt ./test.mlir -test-print-topological-sort

test.mlir:

func.func @region(
  %arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index,
  %arg4 : i32, %arg5 : i32, %arg6 : i32,
  %buffer : memref<i32>) {
  %0 = arith.addi %arg4, %arg5 {__test_sort_original_idx__ = 0} : i32 
  %idx = arith.addi %arg0, %arg1 {__test_sort_original_idx__ = 3} : index
  scf.for %arg7 = %idx to %arg2 step %arg3  {
    %2 = arith.addi %0, %arg5 : i32
    %3 = arith.subi %2, %arg6 {__test_sort_original_idx__ = 1} : i32                                                                                                                         
    %4 = arith.addi %3, %2: i32
    %5 = arith.addi %4, %3: i32
    %6 = arith.addi %5, %4: i32
    %7 = arith.addi %6, %5: i32
    %8 = arith.addi %7, %6: i32
    %9 = arith.addi %8, %7: i32
    %10 = arith.addi %9, %8: i32
    %11 = arith.addi %10, %9: i32
    %12 = arith.addi %11, %10: i32
    %13 = arith.addi %12, %11: i32
    %14 = arith.addi %13, %12: i32
    %15 = arith.addi %14, %13: i32
    %16 = arith.addi %15, %14: i32
    %17 = arith.addi %16, %15: i32
    %18 = arith.addi %17, %16: i32
    %19 = arith.addi %18, %17: i32
    %20 = arith.addi %19, %18: i32
    %21 = arith.addi %20, %19: i32
    %22 = arith.addi %21, %20: i32
    %23 = arith.addi %22, %21: i32
    %24 = arith.addi %23, %22: i32
    %25 = arith.addi %24, %23: i32
    %26 = arith.addi %25, %24: i32
    %27 = arith.addi %26, %25: i32
    %28 = arith.addi %27, %26: i32
    %29 = arith.addi %28, %27: i32
    %30 = arith.addi %29, %28: i32
    %31 = arith.addi %30, %29: i32
    %32 = arith.addi %31, %30: i32
    %33 = arith.addi %32, %31: i32
    %34 = arith.addi %33, %32: i32
    %35 = arith.addi %34, %33: i32
    %36 = arith.addi %35, %34: i32
    %37 = arith.addi %36, %35: i32
    %38 = arith.addi %37, %36: i32
    %39 = arith.addi %38, %37: i32
    %40 = arith.addi %39, %38: i32
    %41 = arith.addi %40, %39: i32
    %42 = arith.addi %41, %40: i32
    %43 = arith.addi %42, %41: i32
    %44 = arith.addi %43, %42: i32
    %45 = arith.addi %44, %43: i32
    %46 = arith.addi %45, %44: i32
    %47 = arith.addi %46, %45: i32
    %48 = arith.addi %47, %46: i32
    %49 = arith.addi %48, %47: i32
    %50 = arith.addi %49, %48: i32
    %51 = arith.addi %50, %49: i32
    %52 = arith.addi %51, %50: i32
    %53 = arith.addi %52, %51: i32
    %54 = arith.addi %53, %52: i32
    %55 = arith.addi %54, %53: i32
    %56 = arith.addi %55, %54: i32
    %57 = arith.addi %56, %55: i32
    %58 = arith.addi %57, %56: i32
    %59 = arith.addi %58, %57: i32
    %60 = arith.addi %59, %58: i32
    %61 = arith.addi %60, %59: i32
    %62 = arith.addi %61, %60: i32
    %63 = arith.addi %62, %61: i32
    %64 = arith.addi %63, %62: i32
    %65 = arith.addi %64, %63: i32
    %66 = arith.addi %65, %64: i32
    %67 = arith.addi %66, %65: i32
    %68 = arith.addi %67, %66: i32
    %69 = arith.addi %68, %67: i32
    %70 = arith.addi %69, %68: i32
    %71 = arith.addi %70, %69: i32
    %72 = arith.addi %71, %70: i32
    %73 = arith.addi %72, %71: i32
    %74 = arith.addi %73, %72: i32
    %75 = arith.addi %74, %73: i32
    %76 = arith.addi %75, %74: i32
    %77 = arith.addi %76, %75: i32
    %78 = arith.addi %77, %76: i32
    %79 = arith.addi %78, %77: i32
    %80 = arith.addi %79, %78: i32
    %81 = arith.addi %80, %79: i32
    %82 = arith.addi %81, %80: i32
    %83 = arith.addi %82, %81: i32
    %84 = arith.addi %83, %82: i32
    %85 = arith.addi %84, %83: i32
    %86 = arith.addi %85, %84: i32
    %87 = arith.addi %86, %85: i32
    %88 = arith.addi %87, %86: i32
    %89 = arith.addi %88, %87: i32
    %90 = arith.addi %89, %88: i32
    %91 = arith.addi %90, %89: i32
    %92 = arith.addi %91, %90: i32
    %93 = arith.addi %92, %91: i32
    %94 = arith.addi %93, %92: i32
    %95 = arith.addi %94, %93: i32
    %96 = arith.addi %95, %94: i32
    %97 = arith.addi %96, %95: i32
    %98 = arith.addi %97, %96: i32
    %99 = arith.addi %98, %97: i32
    %100 = arith.addi %99, %98: i32
    %101 = arith.addi %100, %99: i32
    %102 = arith.addi %101, %100: i32
    %103 = arith.addi %102, %101: i32
    %104 = arith.addi %103, %102: i32
    %105 = arith.addi %104, %103: i32
    %106 = arith.addi %105, %104: i32
    %107 = arith.addi %106, %105: i32
    %108 = arith.addi %107, %106: i32
    %109 = arith.addi %108, %107: i32
    %110 = arith.addi %109, %108: i32
    %111 = arith.addi %110, %109: i32
    %112 = arith.addi %111, %110: i32
    %113 = arith.addi %112, %111: i32
    %114 = arith.addi %113, %112: i32
    %115 = arith.addi %114, %113: i32
    %116 = arith.addi %115, %114: i32
    %117 = arith.addi %116, %115: i32
    %118 = arith.addi %117, %116: i32
    %119 = arith.addi %118, %117: i32
    %120 = arith.addi %119, %118: i32
    %121 = arith.addi %120, %119: i32
    %122 = arith.addi %121, %120: i32
    %123 = arith.addi %122, %121: i32
    %124 = arith.addi %123, %122: i32
    %125 = arith.addi %124, %123: i32
    %126 = arith.addi %125, %124: i32
    %127 = arith.addi %126, %125: i32
    %128 = arith.addi %127, %126: i32
    %129 = arith.addi %128, %127: i32
    %130 = arith.addi %129, %128: i32
    %131 = arith.addi %130, %129: i32
    %132 = arith.addi %131, %130: i32
    %133 = arith.addi %132, %131: i32
    %134 = arith.addi %133, %132: i32
    %135 = arith.addi %134, %133: i32
    %136 = arith.addi %135, %134: i32
    %137 = arith.addi %136, %135: i32
    %138 = arith.addi %137, %136: i32
    %139 = arith.addi %138, %137: i32
    %140 = arith.addi %139, %138: i32
    %141 = arith.addi %140, %139: i32
    %142 = arith.addi %141, %140: i32
    %143 = arith.addi %142, %141: i32
    %144 = arith.addi %143, %142: i32
    %145 = arith.addi %144, %143: i32
    %146 = arith.addi %145, %144: i32
    %147 = arith.addi %146, %145: i32
    %148 = arith.addi %147, %146: i32
    %149 = arith.addi %148, %147: i32
    %150 = arith.addi %149, %148: i32
    %151 = arith.addi %150, %149: i32
    %152 = arith.addi %151, %150: i32
    %153 = arith.addi %152, %151: i32
    %154 = arith.addi %153, %152: i32
    %155 = arith.addi %154, %153: i32
    %156 = arith.addi %155, %154: i32
    %157 = arith.addi %156, %155: i32
    %158 = arith.addi %157, %156: i32
    %159 = arith.addi %158, %157: i32
    %160 = arith.addi %159, %158: i32
    %161 = arith.addi %160, %159: i32
    %162 = arith.addi %161, %160: i32
    %163 = arith.addi %162, %161: i32
    %164 = arith.addi %163, %162: i32
    %165 = arith.addi %164, %163: i32
    %166 = arith.addi %165, %164: i32
    %167 = arith.addi %166, %165: i32
    %168 = arith.addi %167, %166: i32
    %169 = arith.addi %168, %167: i32
    %170 = arith.addi %169, %168: i32
    %171 = arith.addi %170, %169: i32
    %172 = arith.addi %171, %170: i32
    %173 = arith.addi %172, %171: i32
    %174 = arith.addi %173, %172: i32
    %175 = arith.addi %174, %173: i32
    %176 = arith.addi %175, %174: i32
    %177 = arith.addi %176, %175: i32
    %178 = arith.addi %177, %176: i32
    %179 = arith.addi %178, %177: i32
    %180 = arith.addi %179, %178: i32
    %181 = arith.addi %180, %179: i32
    %182 = arith.addi %181, %180: i32
    %183 = arith.addi %182, %181: i32
    %184 = arith.addi %183, %182: i32
    %185 = arith.addi %184, %183: i32
    %186 = arith.addi %185, %184: i32
    %187 = arith.addi %186, %185: i32
    %188 = arith.addi %187, %186: i32
    %189 = arith.addi %188, %187: i32
    %190 = arith.addi %189, %188: i32
    %191 = arith.addi %190, %189: i32
    %192 = arith.addi %191, %190: i32
    %193 = arith.addi %192, %191: i32
    %194 = arith.addi %193, %192: i32
    %195 = arith.addi %194, %193: i32
    %196 = arith.addi %195, %194: i32
    %197 = arith.addi %196, %195: i32
    %198 = arith.addi %197, %196: i32
    %199 = arith.addi %198, %197: i32
    %200 = arith.addi %199, %198: i32
    %201 = arith.addi %200, %199: i32
    %202 = arith.addi %201, %200: i32
    %203 = arith.addi %202, %201: i32
    %204 = arith.addi %203, %202: i32
    %205 = arith.addi %204, %203: i32
    %206 = arith.addi %205, %204: i32
    %207 = arith.addi %206, %205: i32
    %208 = arith.addi %207, %206: i32
    %209 = arith.addi %208, %207: i32
    %210 = arith.addi %209, %208: i32
    %211 = arith.addi %210, %209: i32
    %212 = arith.addi %211, %210: i32
    %213 = arith.addi %212, %211: i32
    %214 = arith.addi %213, %212: i32
    %215 = arith.addi %214, %213: i32
    %216 = arith.addi %215, %214: i32
    %217 = arith.addi %216, %215: i32
    %218 = arith.addi %217, %216: i32
    %219 = arith.addi %218, %217: i32
    %220 = arith.addi %219, %218: i32
    %221 = arith.addi %220, %219: i32
    %222 = arith.addi %221, %220: i32
    %223 = arith.addi %222, %221: i32
    %224 = arith.addi %223, %222: i32
    %225 = arith.addi %224, %223: i32
    %226 = arith.addi %225, %224: i32
    %227 = arith.addi %226, %225: i32
    %228 = arith.addi %227, %226: i32
    %229 = arith.addi %228, %227: i32
    %230 = arith.addi %229, %228: i32
    %231 = arith.addi %230, %229: i32
    %232 = arith.addi %231, %230: i32
    %233 = arith.addi %232, %231: i32
    %234 = arith.addi %233, %232: i32
    %235 = arith.addi %234, %233: i32
    %236 = arith.addi %235, %234: i32
    %237 = arith.addi %236, %235: i32
    %238 = arith.addi %237, %236: i32
    %239 = arith.addi %238, %237: i32
    %240 = arith.addi %239, %238: i32
    %241 = arith.addi %240, %239: i32
    %242 = arith.addi %241, %240: i32
    %243 = arith.addi %242, %241: i32
    %244 = arith.addi %243, %242: i32
    %245 = arith.addi %244, %243: i32
    %246 = arith.addi %245, %244: i32
    %247 = arith.addi %246, %245: i32
    %248 = arith.addi %247, %246: i32
    %249 = arith.addi %248, %247: i32
    %250 = arith.addi %249, %248: i32
    %251 = arith.addi %250, %249: i32
    %252 = arith.addi %251, %250: i32
    %253 = arith.addi %252, %251: i32
    %254 = arith.addi %253, %252: i32
    %255 = arith.addi %254, %253: i32
    %256 = arith.addi %255, %254: i32
    %257 = arith.addi %256, %255: i32
    %258 = arith.addi %257, %256: i32
    %259 = arith.addi %258, %257: i32
    %260 = arith.addi %259, %258: i32
    %261 = arith.addi %260, %259: i32
    %262 = arith.addi %261, %260: i32
    %263 = arith.addi %262, %261: i32
    %264 = arith.addi %263, %262: i32
    %265 = arith.addi %264, %263: i32
    %266 = arith.addi %265, %264: i32
    %267 = arith.addi %266, %265: i32
    %268 = arith.addi %267, %266: i32
    %269 = arith.addi %268, %267: i32
    %270 = arith.addi %269, %268: i32
    %271 = arith.addi %270, %269: i32
    %272 = arith.addi %271, %270: i32
    %273 = arith.addi %272, %271: i32
    %274 = arith.addi %273, %272: i32
    %275 = arith.addi %274, %273: i32
    %276 = arith.addi %275, %274: i32
    %277 = arith.addi %276, %275: i32
    %278 = arith.addi %277, %276: i32
    %279 = arith.addi %278, %277: i32
    %280 = arith.addi %279, %278: i32
    %281 = arith.addi %280, %279: i32
    %282 = arith.addi %281, %280: i32
    %283 = arith.addi %282, %281: i32
    %284 = arith.addi %283, %282: i32
    %285 = arith.addi %284, %283: i32
    %286 = arith.addi %285, %284: i32
    %287 = arith.addi %286, %285: i32
    %288 = arith.addi %287, %286: i32
    %289 = arith.addi %288, %287: i32
    %290 = arith.addi %289, %288: i32
    %291 = arith.addi %290, %289: i32
    %292 = arith.addi %291, %290: i32
    %293 = arith.addi %292, %291: i32
    %294 = arith.addi %293, %292: i32
    %295 = arith.addi %294, %293: i32
    %296 = arith.addi %295, %294: i32
    %297 = arith.addi %296, %295: i32
    %298 = arith.addi %297, %296: i32
    %299 = arith.addi %298, %297: i32
    %300 = arith.addi %299, %298: i32
    memref.store %300, %buffer[] : memref<i32>
  } {__test_sort_original_idx__ = 2}
  return
}

[llvmbot]

@llvm/issue-subscribers-mlir

[joker-eph]

Are you sending a patch? (since you found the issue) or do you want me to land this for you?

[Jokeren]

I'm not familiar with the patching process yet. Maybe I'll check tomorrow to see if I could submit a patch with a test included for this issue. If not, I'll seek help from your side. Many thanks!

[Jokeren]

Our solution: https://github.com/openai/triton/blob/ccd17d6bf9da4975b15442b0de879594301062b5/lib/Analysis/Utility.cpp#L232

Will try to prepare a patch