[python] Add walk_with_filter to walk subset of IR

[uenoku]

This adds walk_with_filter python method to invoke callbacks only for subset of operations. We require GIL to call python function so this could improve performance of the tools based on Python API. This functionality could be better to be in upstream but for now we decided to add to CIRCT repo.

In the following benchmark there are 5*10^5 hw.constant ops and single hw.output op and it walks only hw.output op. The runtime is substantially faster.

import circt
from circt.dialects import hw
from circt.support import walk_with_filter

from circt.ir import (Context, Location, InsertionPoint, IntegerType,
                      IntegerAttr, Module, WalkOrder,
                      WalkResult)

n = 500000
with Context() as ctx, Location.unknown():
  circt.register_dialects(ctx)

  i32 = IntegerType.get_signless(32)

  m = Module.create()
  with InsertionPoint(m.body):

    def build(module):
      elems = []
      for i in range(n):
          elems.append(hw.ConstantOp(IntegerAttr.get(i32, i)))
    hw.HWModuleOp(name="test", body_builder=build)
  cnt = 0
  def callback(op):
    if isinstance(op.opview, hw.OutputOp):
      global cnt
      cnt += 1
    return WalkResult.ADVANCE

  from time import time
  # Benchmark walk_with_filter
  print("Starting walk_with_filter")
  start_time = time()
  walk_with_filter(m.operation, [hw.OutputOp], callback, WalkOrder.POST_ORDER)
  end_time = time()
  print(f"walk_with_filter elapsed time: {end_time - start_time:.6f} seconds cnt={cnt}")

  # Benchmark operation.walk
  print("Starting operation.walk")
  start_time = time()
  m.operation.walk(callback, WalkOrder.POST_ORDER)
  end_time = time()
  print(f"operation.walk elapsed time: {end_time - start_time:.6f} seconds cnt={cnt}")

Starting walk_with_filter
walk_with_filter elapsed time: 0.005462 seconds cnt=1
Starting operation.walk
operation.walk elapsed time: 1.061360 seconds cnt=2

[mikeurbach]

Awesome, I think this will be really helpful. Looks good to me, just one minor comment about the test code.

[dtzSiFive]

I'm not a deep python expert. Does this compose with itself re:changing exception type thrown from what's caught? Presumably throwing up to outer walk will be translated to python exception and back again, that sort of thing?

Anyway this looks super useful and a huge win, great work!

[uenoku]

I'm not a deep python expert. Does this compose with itself re:changing exception type thrown from what's caught? Presumably throwing up to outer walk will be translated to python exception and back again, that sort of thing?

Good point. Exceptions raised in a python callback is caught in IRCore.cpp and it raises another exception std::runtime_error which will be propagated to ptyhon side as RuntimeError. I followed upstream implementation for normal walk here: https://github.com/llvm/llvm-project/blob/9b67c99dc5b0d669a28b714e10e257d3988d1066/mlir/lib/Bindings/Python/IRCore.cpp#L1296-L1313