mlir/test/Bindings/Python/dialects.py

# RUN: %PYTHON %s | FileCheck %s

import gc
import mlir

def run(f):
  print("\nTEST:", f.__name__)
  f()
  gc.collect()
  assert mlir.ir.Context._get_live_count() == 0


# CHECK-LABEL: TEST: testDialectDescriptor
def testDialectDescriptor():
  ctx = mlir.ir.Context()
  d = ctx.get_dialect_descriptor("std")
  # CHECK: <DialectDescriptor std>
  print(d)
  # CHECK: std
  print(d.namespace)
  try:
    _ = ctx.get_dialect_descriptor("not_existing")
  except ValueError:
    pass
  else:
    assert False, "Expected exception"

run(testDialectDescriptor)


# CHECK-LABEL: TEST: testUserDialectClass
def testUserDialectClass():
  ctx = mlir.ir.Context()
  # Access using attribute.
  d = ctx.dialects.std
  # Note that the standard dialect namespace prints as ''. Others will print
  # as "<Dialect %namespace (..."
  # CHECK: <Dialect (class mlir.dialects.std._Dialect)>
  print(d)
  try:
    _ = ctx.dialects.not_existing
  except AttributeError:
    pass
  else:
    assert False, "Expected exception"

  # Access using index.
  d = ctx.dialects["std"]
  # CHECK: <Dialect (class mlir.dialects.std._Dialect)>
  print(d)
  try:
    _ = ctx.dialects["not_existing"]
  except IndexError:
    pass
  else:
    assert False, "Expected exception"

  # Using the 'd' alias.
  d = ctx.d["std"]
  # CHECK: <Dialect (class mlir.dialects.std._Dialect)>
  print(d)

run(testUserDialectClass)


# CHECK-LABEL: TEST: testCustomOpView
# This test uses the standard dialect AddFOp as an example of a user op.
# TODO: Op creation and access is still quite verbose: simplify this test as
# additional capabilities come online.
def testCustomOpView():
  ctx = mlir.ir.Context()
  ctx.allow_unregistered_dialects = True
  f32 = mlir.ir.F32Type.get(ctx)
  loc = ctx.get_unknown_location()
  m = ctx.create_module(loc)

  def createInput():
    op = ctx.create_operation("pytest_dummy.intinput", loc, results=[f32])
    # TODO: Auto result cast from operation
    return op.results[0]

  with mlir.ir.InsertionPoint.at_block_terminator(m.body):
    # Create via dialects context collection.
    input1 = createInput()
    input2 = createInput()
    op1 = ctx.dialects.std.AddFOp(loc, input1, input2)

    # Create via an import
    from mlir.dialects.std import AddFOp
    AddFOp(loc, input1, op1.result)

  # CHECK: %[[INPUT0:.*]] = "pytest_dummy.intinput"
  # CHECK: %[[INPUT1:.*]] = "pytest_dummy.intinput"
  # CHECK: %[[R0:.*]] = addf %[[INPUT0]], %[[INPUT1]] : f32
  # CHECK: %[[R1:.*]] = addf %[[INPUT0]], %[[R0]] : f32
  m.operation.print()

run(testCustomOpView)