Reorder edsc python tests - NFC
authorNicolas Vasilache <ntv@google.com>
Fri, 10 May 2019 14:51:07 +0000 (07:51 -0700)
committerMehdi Amini <joker.eph@gmail.com>
Sat, 11 May 2019 02:28:21 +0000 (19:28 -0700)
    This CL orders the python tests to:
    1. allow introspecting on the EdscTest class and avoid the error-prone process of having to add the test call by hand;
    2. account for differences in the order of `dir(edscTest)` between python2, <= python3.5 and >= python 3.6

--

PiperOrigin-RevId: 247609687

mlir/bindings/python/test/test_py2and3.py

index bc6582f..0832c02 100644 (file)
@@ -35,72 +35,17 @@ class EdscTest:
     self.f32Type = self.module.make_scalar_type("f32")
     self.indexType = self.module.make_index_type()
 
-  def testFunctionContext(self):
-    self.setUp()
-    with self.module.function_context("foo", [], []):
-      pass
-      printWithCurrentFunctionName(self.module.get_function("foo"))
-    # CHECK-LABEL: testFunctionContext
-    #       CHECK: func @foo() {
-
-  def testMultipleFunctions(self):
-    self.setUp()
-    with self.module.function_context("foo", [], []):
-      pass
-    with self.module.function_context("foo", [], []):
-      E.constant_index(0)
-    printWithCurrentFunctionName(str(self.module))
-    # CHECK-LABEL: testMultipleFunctions
-    #       CHECK: func @foo()
-    #       CHECK: func @foo_0()
-    #       CHECK: %c0 = constant 0 : index
-
-  def testFunctionArgs(self):
-    self.setUp()
-    with self.module.function_context("foo", [self.f32Type, self.f32Type],
-                                      [self.indexType]) as fun:
-      pass
-      printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testFunctionArgs
-    #       CHECK: func @foo(%arg0: f32, %arg1: f32) -> index
-
-  def testLoopContext(self):
+  def testBlockArguments(self):
     self.setUp()
     with self.module.function_context("foo", [], []) as fun:
-      lhs = E.constant_index(0)
-      rhs = E.constant_index(42)
-      with E.LoopContext(lhs, rhs, 1) as i:
-        lhs + rhs + i
-        with E.LoopContext(rhs, rhs + rhs, 2) as j:
-          x = i + j
+      E.constant_index(42)
+      with E.BlockContext([self.f32Type, self.f32Type]) as b:
+        b.arg(0) + b.arg(1)
       printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testLoopContext
-    #       CHECK: "affine.for"() (
-    #       CHECK:   ^bb1(%i0: index):
-    #       CHECK: "affine.for"(%c42, %2) (
-    #       CHECK: ^bb2(%i1: index):
-    #       CHECK: "affine.apply"(%i0, %i1) {map: (d0, d1) -> (d0 + d1)} : (index, index) -> index
-    #       CHECK: {lower_bound: (d0) -> (d0), step: 2 : index, upper_bound: (d0) -> (d0)} : (index, index) -> ()
-    #       CHECK: {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (42)}
-
-  def testLoopNestContext(self):
-    self.setUp()
-    with self.module.function_context("foo", [], []) as fun:
-      lbs = [E.constant_index(i) for i in range(4)]
-      ubs = [E.constant_index(10 * i + 5) for i in range(4)]
-      with E.LoopNestContext(lbs, ubs, [1, 3, 5, 7]) as (i, j, k, l):
-        i + j + k + l
-    printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testLoopNestContext
-    #       CHECK: "affine.for"() (
-    #       CHECK: ^bb1(%i0: index):
-    #       CHECK: "affine.for"() (
-    #       CHECK: ^bb2(%i1: index):
-    #       CHECK: "affine.for"() (
-    #       CHECK: ^bb3(%i2: index):
-    #       CHECK: "affine.for"() (
-    #       CHECK: ^bb4(%i3: index):
-    #       CHECK: %2 = "affine.apply"(%i0, %i1, %i2, %i3) {map: (d0, d1, d2, d3) -> (d0 + d1 + d2 + d3)} : (index, index, index, index) -> index
+    # CHECK-LABEL: testBlockArguments
+    #       CHECK: %c42 = constant 42 : index
+    #       CHECK: ^bb1(%0: f32, %1: f32):
+    #       CHECK:   %2 = addf %0, %1 : f32
 
   def testBlockContext(self):
     self.setUp()
@@ -157,17 +102,28 @@ class EdscTest:
     #       CHECK: %c56 = constant 56 : index
     #       CHECK: %c57 = constant 57 : index
 
-  def testBlockArguments(self):
+  def testBooleanOps(self):
     self.setUp()
-    with self.module.function_context("foo", [], []) as fun:
-      E.constant_index(42)
-      with E.BlockContext([self.f32Type, self.f32Type]) as b:
-        b.arg(0) + b.arg(1)
+    with self.module.function_context(
+        "booleans", [self.boolType for _ in range(4)], []) as fun:
+      i, j, k, l = (fun.arg(x) for x in range(4))
+      stmt1 = (i < j) & (j >= k)
+      stmt2 = ~(stmt1 | (k == l))
       printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testBlockArguments
-    #       CHECK: %c42 = constant 42 : index
-    #       CHECK: ^bb1(%0: f32, %1: f32):
-    #       CHECK:   %2 = addf %0, %1 : f32
+    # CHECK-LABEL: testBooleanOps
+    #       CHECK: %0 = cmpi "slt", %arg0, %arg1 : i1
+    #       CHECK: %1 = cmpi "sge", %arg1, %arg2 : i1
+    #       CHECK: %2 = muli %0, %1 : i1
+    #       CHECK: %3 = cmpi "eq", %arg2, %arg3 : i1
+    #       CHECK: %true = constant 1 : i1
+    #       CHECK: %4 = subi %true, %2 : i1
+    #       CHECK: %true_0 = constant 1 : i1
+    #       CHECK: %5 = subi %true_0, %3 : i1
+    #       CHECK: %6 = muli %4, %5 : i1
+    #       CHECK: %true_1 = constant 1 : i1
+    #       CHECK: %7 = subi %true_1, %6 : i1
+    #       CHECK: %true_2 = constant 1 : i1
+    #       CHECK: %8 = subi %true_2, %7 : i1
 
   def testBr(self):
     self.setUp()
@@ -182,19 +138,6 @@ class EdscTest:
     #       CHECK: ^bb1:
     #       CHECK:   return
 
-  def testBrDeclaration(self):
-    self.setUp()
-    with self.module.function_context("foo", [], []) as fun:
-      blk = E.BlockContext()
-      E.br(blk.handle())
-      with blk:
-        E.ret()
-      printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testBrDeclaration
-    #       CHECK:   br ^bb1
-    #       CHECK: ^bb1:
-    #       CHECK:   return
-
   def testBrArgs(self):
     self.setUp()
     with self.module.function_context("foo", [], []) as fun:
@@ -210,42 +153,18 @@ class EdscTest:
     #       CHECK: ^bb1(%0: index, %1: index):
     #       CHECK:   br ^bb1(%1, %0 : index, index)
 
-  def testCondBr(self):
-    self.setUp()
-    with self.module.function_context("foo", [self.boolType], []) as fun:
-      with E.BlockContext() as blk1:
-        E.ret([])
-      with E.BlockContext([self.indexType]) as blk2:
-        E.ret([])
-      cst = E.constant_index(0)
-      E.cond_br(fun.arg(0), blk1, [], blk2, [cst])
-      printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testCondBr
-    #       CHECK:   cond_br %arg0, ^bb1, ^bb2(%c0 : index)
-
-  def testRet(self):
-    self.setUp()
-    with self.module.function_context("foo", [],
-                                      [self.indexType, self.indexType]) as fun:
-      c42 = E.constant_index(42)
-      c0 = E.constant_index(0)
-      E.ret([c42, c0])
-      printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testRet
-    #       CHECK:    %c42 = constant 42 : index
-    #       CHECK:    %c0 = constant 0 : index
-    #       CHECK:    return %c42, %c0 : index, index
-
-  def testSelectOp(self):
+  def testBrDeclaration(self):
     self.setUp()
-    with self.module.function_context("foo", [self.boolType],
-                                      [self.i32Type]) as fun:
-      a = E.constant_int(42, 32)
-      b = E.constant_int(0, 32)
-      E.ret([E.select(fun.arg(0), a, b)])
+    with self.module.function_context("foo", [], []) as fun:
+      blk = E.BlockContext()
+      E.br(blk.handle())
+      with blk:
+        E.ret()
       printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testSelectOp
-    #       CHECK:  %0 = select %arg0, %c42_i32, %c0_i32 : i32
+    # CHECK-LABEL: testBrDeclaration
+    #       CHECK:   br ^bb1
+    #       CHECK: ^bb1:
+    #       CHECK:   return
 
   def testCallOp(self):
     self.setUp()
@@ -260,51 +179,18 @@ class EdscTest:
     #       CHECK:   %f = constant @sqrtf : (f32) -> f32
     #       CHECK:   %0 = call_indirect %f(%arg0) : (f32) -> f32
 
-  def testBooleanOps(self):
-    self.setUp()
-    with self.module.function_context(
-        "booleans", [self.boolType for _ in range(4)], []) as fun:
-      i, j, k, l = (fun.arg(x) for x in range(4))
-      stmt1 = (i < j) & (j >= k)
-      stmt2 = ~(stmt1 | (k == l))
-      printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testBooleanOps
-    #       CHECK: %0 = cmpi "slt", %arg0, %arg1 : i1
-    #       CHECK: %1 = cmpi "sge", %arg1, %arg2 : i1
-    #       CHECK: %2 = muli %0, %1 : i1
-    #       CHECK: %3 = cmpi "eq", %arg2, %arg3 : i1
-    #       CHECK: %true = constant 1 : i1
-    #       CHECK: %4 = subi %true, %2 : i1
-    #       CHECK: %true_0 = constant 1 : i1
-    #       CHECK: %5 = subi %true_0, %3 : i1
-    #       CHECK: %6 = muli %4, %5 : i1
-    #       CHECK: %true_1 = constant 1 : i1
-    #       CHECK: %7 = subi %true_1, %6 : i1
-    #       CHECK: %true_2 = constant 1 : i1
-    #       CHECK: %8 = subi %true_2, %7 : i1
-
-  def testDivisions(self):
-    self.setUp()
-    with self.module.function_context(
-        "division", [self.indexType, self.i32Type, self.i32Type], []) as fun:
-      # indices only support floor division
-      fun.arg(0) // E.constant_index(42)
-      # regular values only support regular division
-      fun.arg(1) / fun.arg(2)
-      printWithCurrentFunctionName(str(self.module))
-    # CHECK-LABEL: testDivisions
-    #       CHECK:  floordiv 42
-    #       CHECK:  divis %arg1, %arg2 : i32
-
-  def testCustom(self):
+  def testCondBr(self):
     self.setUp()
-    with self.module.function_context("custom", [self.indexType, self.f32Type],
-                                      []) as fun:
-      E.op("foo", [fun.arg(0)], [self.f32Type]) + fun.arg(1)
+    with self.module.function_context("foo", [self.boolType], []) as fun:
+      with E.BlockContext() as blk1:
+        E.ret([])
+      with E.BlockContext([self.indexType]) as blk2:
+        E.ret([])
+      cst = E.constant_index(0)
+      E.cond_br(fun.arg(0), blk1, [], blk2, [cst])
       printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testCustom
-    #       CHECK: %0 = "foo"(%arg0) : (index) -> f32
-    #       CHECK:  %1 = addf %0, %arg1 : f32
+    # CHECK-LABEL: testCondBr
+    #       CHECK:   cond_br %arg0, ^bb1, ^bb2(%c0 : index)
 
   def testConstants(self):
     self.setUp()
@@ -333,6 +219,84 @@ class EdscTest:
     #       CHECK:  constant 123 : index
     #       CHECK:  constant @constants : () -> ()
 
+  def testCustom(self):
+    self.setUp()
+    with self.module.function_context("custom", [self.indexType, self.f32Type],
+                                      []) as fun:
+      E.op("foo", [fun.arg(0)], [self.f32Type]) + fun.arg(1)
+      printWithCurrentFunctionName(str(fun))
+    # CHECK-LABEL: testCustom
+    #       CHECK: %0 = "foo"(%arg0) : (index) -> f32
+    #       CHECK:  %1 = addf %0, %arg1 : f32
+
+  # Create 'addi' using the generic Op interface.  We need an operation known
+  # to the execution engine so that the engine can compile it.
+  def testCustomOpCompilation(self):
+    self.setUp()
+    with self.module.function_context("adder", [self.i32Type], []) as f:
+      c1 = E.op(
+          "std.constant", [], [self.i32Type],
+          value=self.module.integerAttr(self.i32Type, 42))
+      E.op("std.addi", [c1, f.arg(0)], [self.i32Type])
+      E.ret([])
+    self.module.compile()
+    printWithCurrentFunctionName(str(self.module.get_engine_address() == 0))
+    # CHECK-LABEL: testCustomOpCompilation
+    #       CHECK: False
+
+  def testDivisions(self):
+    self.setUp()
+    with self.module.function_context(
+        "division", [self.indexType, self.i32Type, self.i32Type], []) as fun:
+      # indices only support floor division
+      fun.arg(0) // E.constant_index(42)
+      # regular values only support regular division
+      fun.arg(1) / fun.arg(2)
+      printWithCurrentFunctionName(str(self.module))
+    # CHECK-LABEL: testDivisions
+    #       CHECK:  floordiv 42
+    #       CHECK:  divis %arg1, %arg2 : i32
+
+  def testFunctionArgs(self):
+    self.setUp()
+    with self.module.function_context("foo", [self.f32Type, self.f32Type],
+                                      [self.indexType]) as fun:
+      pass
+      printWithCurrentFunctionName(str(fun))
+    # CHECK-LABEL: testFunctionArgs
+    #       CHECK: func @foo(%arg0: f32, %arg1: f32) -> index
+
+  def testFunctionContext(self):
+    self.setUp()
+    with self.module.function_context("foo", [], []):
+      pass
+      printWithCurrentFunctionName(self.module.get_function("foo"))
+    # CHECK-LABEL: testFunctionContext
+    #       CHECK: func @foo() {
+
+  def testFunctionDeclaration(self):
+    self.setUp()
+    boolAttr = self.module.boolAttr(True)
+    t = self.module.make_memref_type(self.f32Type, [10])
+    t_llvm_noalias = t({"llvm.noalias": boolAttr})
+    t_readonly = t({"readonly": boolAttr})
+    f = self.module.declare_function("foo", [t, t_llvm_noalias, t_readonly], [])
+    printWithCurrentFunctionName(str(self.module))
+    # CHECK-LABEL: testFunctionDeclaration
+    #       CHECK: func @foo(memref<10xf32>, memref<10xf32> {llvm.noalias: true}, memref<10xf32> {readonly: true})
+
+  def testFunctionMultiple(self):
+    self.setUp()
+    with self.module.function_context("foo", [], []):
+      pass
+    with self.module.function_context("foo", [], []):
+      E.constant_index(0)
+    printWithCurrentFunctionName(str(self.module))
+    # CHECK-LABEL: testFunctionMultiple
+    #       CHECK: func @foo()
+    #       CHECK: func @foo_0()
+    #       CHECK: %c0 = constant 0 : index
+
   def testIndexedValue(self):
     self.setUp()
     memrefType = self.module.make_memref_type(self.f32Type, [10, 42])
@@ -355,32 +319,77 @@ class EdscTest:
     #       CHECK: {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (42)}
     #       CHECK: {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (10)}
 
-  def testMatrixMultiply(self):
+  def testLoopContext(self):
     self.setUp()
-    memrefType = self.module.make_memref_type(self.f32Type, [32, 32])
-    with self.module.function_context(
-        "matmul", [memrefType, memrefType, memrefType], []) as fun:
-      A = E.IndexedValue(fun.arg(0))
-      B = E.IndexedValue(fun.arg(1))
-      C = E.IndexedValue(fun.arg(2))
-      c0 = E.constant_index(0)
-      c32 = E.constant_index(32)
-      with E.LoopNestContext([c0, c0, c0], [c32, c32, c32], [1, 1, 1]) as (i, j,
-                                                                           k):
-        C.store([i, j], A.load([i, k]) * B.load([k, j]))
-      E.ret([])
+    with self.module.function_context("foo", [], []) as fun:
+      lhs = E.constant_index(0)
+      rhs = E.constant_index(42)
+      with E.LoopContext(lhs, rhs, 1) as i:
+        lhs + rhs + i
+        with E.LoopContext(rhs, rhs + rhs, 2) as j:
+          x = i + j
       printWithCurrentFunctionName(str(fun))
-    # CHECK-LABEL: testMatrixMultiply
-    #       CHECK: "affine.for"()
-    #       CHECK: "affine.for"()
-    #       CHECK: "affine.for"()
-    #   CHECK-DAG:  %0 = load %arg0[%i0, %i2] : memref<32x32xf32>
-    #   CHECK-DAG:  %1 = load %arg1[%i2, %i1] : memref<32x32xf32>
-    #       CHECK:  %2 = mulf %0, %1 : f32
-    #       CHECK:  store %2, %arg2[%i0, %i1] : memref<32x32xf32>
-    #       CHECK: {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (32)} : () -> ()
-    #       CHECK: {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (32)} : () -> ()
-    #       CHECK: {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (32)} : () -> ()
+    # CHECK-LABEL: testLoopContext
+    #       CHECK: "affine.for"() (
+    #       CHECK:   ^bb1(%i0: index):
+    #       CHECK: "affine.for"(%c42, %2) (
+    #       CHECK: ^bb2(%i1: index):
+    #       CHECK: "affine.apply"(%i0, %i1) {map: (d0, d1) -> (d0 + d1)} : (index, index) -> index
+    #       CHECK: {lower_bound: (d0) -> (d0), step: 2 : index, upper_bound: (d0) -> (d0)} : (index, index) -> ()
+    #       CHECK: {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (42)}
+
+  def testLoopNestContext(self):
+    self.setUp()
+    with self.module.function_context("foo", [], []) as fun:
+      lbs = [E.constant_index(i) for i in range(4)]
+      ubs = [E.constant_index(10 * i + 5) for i in range(4)]
+      with E.LoopNestContext(lbs, ubs, [1, 3, 5, 7]) as (i, j, k, l):
+        i + j + k + l
+    printWithCurrentFunctionName(str(fun))
+    # CHECK-LABEL: testLoopNestContext
+    #       CHECK: "affine.for"() (
+    #       CHECK: ^bb1(%i0: index):
+    #       CHECK: "affine.for"() (
+    #       CHECK: ^bb2(%i1: index):
+    #       CHECK: "affine.for"() (
+    #       CHECK: ^bb3(%i2: index):
+    #       CHECK: "affine.for"() (
+    #       CHECK: ^bb4(%i3: index):
+    #       CHECK: %2 = "affine.apply"(%i0, %i1, %i2, %i3) {map: (d0, d1, d2, d3) -> (d0 + d1 + d2 + d3)} : (index, index, index, index) -> index
+
+  def testMLIRBooleanCompilation(self):
+    self.setUp()
+    m = self.module.make_memref_type(self.boolType, [10])  # i1 tensor
+    with self.module.function_context("mkbooltensor", [m, m], []) as f:
+      input = E.IndexedValue(f.arg(0))
+      output = E.IndexedValue(f.arg(1))
+      zero = E.constant_index(0)
+      ten = E.constant_index(10)
+      with E.LoopNestContext([zero] * 3, [ten] * 3, [1] * 3) as (i, j, k):
+        b1 = (i < j) & (j < k)
+        b2 = ~b1
+        b3 = b2 | (k < j)
+        output.store([i], input.load([i]) & b3)
+      E.ret([])
+    self.module.compile()
+    printWithCurrentFunctionName(str(self.module.get_engine_address() == 0))
+    # CHECK-LABEL: testMLIRBooleanCompilation
+    #       CHECK: False
+
+  def testMLIRFunctionCreation(self):
+    self.setUp()
+    module = E.MLIRModule()
+    t = module.make_scalar_type("f32")
+    m = module.make_memref_type(t, [3, 4, -1, 5])
+    printWithCurrentFunctionName(str(t))
+    print(str(m))
+    print(str(module.make_function("copy", [m, m], [])))
+    print(str(module.make_function("sqrtf", [t], [t])))
+    # CHECK-LABEL: testMLIRFunctionCreation
+    #       CHECK:  f32
+    #       CHECK:  memref<3x4x?x5xf32>
+    #       CHECK: func @copy(%arg0: memref<3x4x?x5xf32>, %arg1: memref<3x4x?x5xf32>) {
+    #       CHECK:  func @sqrtf(%arg0: f32) -> f32
 
   def testMLIRScalarTypes(self):
     self.setUp()
@@ -405,106 +414,69 @@ class EdscTest:
     #       CHECK:  i123
     #       CHECK:  index
 
-  def testMLIRFunctionCreation(self):
-    self.setUp()
-    module = E.MLIRModule()
-    t = module.make_scalar_type("f32")
-    m = module.make_memref_type(t, [3, 4, -1, 5])
-    printWithCurrentFunctionName(str(t))
-    print(str(m))
-    print(str(module.make_function("copy", [m, m], [])))
-    print(str(module.make_function("sqrtf", [t], [t])))
-    # CHECK-LABEL: testMLIRFunctionCreation
-    #       CHECK:  f32
-    #       CHECK:  memref<3x4x?x5xf32>
-    #       CHECK: func @copy(%arg0: memref<3x4x?x5xf32>, %arg1: memref<3x4x?x5xf32>) {
-    #       CHECK:  func @sqrtf(%arg0: f32) -> f32
-
-  def testFunctionDeclaration(self):
+  def testMatrixMultiply(self):
     self.setUp()
-    boolAttr = self.module.boolAttr(True)
-    t = self.module.make_memref_type(self.f32Type, [10])
-    t_llvm_noalias = t({"llvm.noalias": boolAttr})
-    t_readonly = t({"readonly": boolAttr})
-    f = self.module.declare_function("foo", [t, t_llvm_noalias, t_readonly], [])
-    printWithCurrentFunctionName(str(self.module))
-    # CHECK-LABEL: testFunctionDeclaration
-    #       CHECK: func @foo(memref<10xf32>, memref<10xf32> {llvm.noalias: true}, memref<10xf32> {readonly: true})
+    memrefType = self.module.make_memref_type(self.f32Type, [32, 32])
+    with self.module.function_context(
+        "matmul", [memrefType, memrefType, memrefType], []) as fun:
+      A = E.IndexedValue(fun.arg(0))
+      B = E.IndexedValue(fun.arg(1))
+      C = E.IndexedValue(fun.arg(2))
+      c0 = E.constant_index(0)
+      c32 = E.constant_index(32)
+      with E.LoopNestContext([c0, c0, c0], [c32, c32, c32], [1, 1, 1]) as (i, j,
+                                                                           k):
+        C.store([i, j], A.load([i, k]) * B.load([k, j]))
+      E.ret([])
+      printWithCurrentFunctionName(str(fun))
+    # CHECK-LABEL: testMatrixMultiply
+    #       CHECK: "affine.for"()
+    #       CHECK: "affine.for"()
+    #       CHECK: "affine.for"()
+    #   CHECK-DAG:  %0 = load %arg0[%i0, %i2] : memref<32x32xf32>
+    #   CHECK-DAG:  %1 = load %arg1[%i2, %i1] : memref<32x32xf32>
+    #       CHECK:  %2 = mulf %0, %1 : f32
+    #       CHECK:  store %2, %arg2[%i0, %i1] : memref<32x32xf32>
+    #       CHECK: {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (32)} : () -> ()
+    #       CHECK: {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (32)} : () -> ()
+    #       CHECK: {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (32)} : () -> ()
 
-  def testMLIRBooleanCompilation(self):
+  def testRet(self):
     self.setUp()
-    m = self.module.make_memref_type(self.boolType, [10])  # i1 tensor
-    with self.module.function_context("mkbooltensor", [m, m], []) as f:
-      input = E.IndexedValue(f.arg(0))
-      output = E.IndexedValue(f.arg(1))
-      zero = E.constant_index(0)
-      ten = E.constant_index(10)
-      with E.LoopNestContext([zero] * 3, [ten] * 3, [1] * 3) as (i, j, k):
-        b1 = (i < j) & (j < k)
-        b2 = ~b1
-        b3 = b2 | (k < j)
-        output.store([i], input.load([i]) & b3)
-      E.ret([])
-    self.module.compile()
-    printWithCurrentFunctionName(str(self.module.get_engine_address() == 0))
-    # CHECK-LABEL: testMLIRBooleanCompilation
-    #       CHECK: False
+    with self.module.function_context("foo", [],
+                                      [self.indexType, self.indexType]) as fun:
+      c42 = E.constant_index(42)
+      c0 = E.constant_index(0)
+      E.ret([c42, c0])
+      printWithCurrentFunctionName(str(fun))
+    # CHECK-LABEL: testRet
+    #       CHECK:    %c42 = constant 42 : index
+    #       CHECK:    %c0 = constant 0 : index
+    #       CHECK:    return %c42, %c0 : index, index
 
-  # Create 'addi' using the generic Op interface.  We need an operation known
-  # to the execution engine so that the engine can compile it.
-  def testCustomOpCompilation(self):
+  def testSelectOp(self):
     self.setUp()
-    with self.module.function_context("adder", [self.i32Type], []) as f:
-      c1 = E.op(
-          "std.constant", [], [self.i32Type],
-          value=self.module.integerAttr(self.i32Type, 42))
-      E.op("std.addi", [c1, f.arg(0)], [self.i32Type])
-      E.ret([])
-    self.module.compile()
-    printWithCurrentFunctionName(str(self.module.get_engine_address() == 0))
-    # CHECK-LABEL: testCustomOpCompilation
-    #       CHECK: False
+    with self.module.function_context("foo", [self.boolType],
+                                      [self.i32Type]) as fun:
+      a = E.constant_int(42, 32)
+      b = E.constant_int(0, 32)
+      E.ret([E.select(fun.arg(0), a, b)])
+      printWithCurrentFunctionName(str(fun))
+    # CHECK-LABEL: testSelectOp
+    #       CHECK:  %0 = select %arg0, %c42_i32, %c0_i32 : i32
 
 # Until python 3.6 this cannot be used because the order in the dict is not the
 # order of method declaration.
-def runTests(edscTest):
+def runTests():
   def isTest(attr):
-    return inspect.ismethod(attr) and "__init" not in str(attr)
+    return inspect.ismethod(attr) and "EdscTest.setUp " not in str(attr)
 
-  tests = filter(isTest, (getattr(edscTest, attr) for attr in dir(edscTest)))
+  edscTest = EdscTest()
+  tests = sorted(filter(isTest,
+                        (getattr(edscTest, attr) for attr in dir(edscTest))),
+                 key = lambda x : str(x))
   for test in tests:
     test()
 
-# So instead one must list the functions in order of their Filecheck appearance.
-def main():
-  edscTest = EdscTest()
-  edscTest.testFunctionContext()
-  edscTest.testMultipleFunctions()
-  edscTest.testFunctionArgs()
-  edscTest.testLoopContext()
-  edscTest.testLoopNestContext()
-  edscTest.testBlockContext()
-  edscTest.testBlockContextAppend()
-  edscTest.testBlockContextStandalone()
-  edscTest.testBlockArguments()
-  edscTest.testBr()
-  edscTest.testBrDeclaration()
-  edscTest.testBrArgs()
-  edscTest.testCondBr()
-  edscTest.testRet()
-  edscTest.testSelectOp()
-  edscTest.testCallOp()
-  edscTest.testBooleanOps()
-  edscTest.testDivisions()
-  edscTest.testCustom()
-  edscTest.testConstants()
-  edscTest.testIndexedValue()
-  edscTest.testMatrixMultiply()
-  edscTest.testMLIRScalarTypes()
-  edscTest.testMLIRFunctionCreation()
-  edscTest.testFunctionDeclaration()
-  edscTest.testMLIRBooleanCompilation()
-  edscTest.testCustomOpCompilation()
-
 if __name__ == '__main__':
-    main()
+    runTests()