# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. # See https://llvm.org/LICENSE.txt for license information. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception """Test for the MLIR EDSC Python bindings""" import inspect import sys from npcomp.native.mlir import edsc as E from npcomp.utils import test_utils # Prints `str` prefixed by the current test function name so we can use it in # Filecheck label directives. # This is achieved by inspecting the stack and getting the parent name. def printWithCurrentFunctionName(str): print(inspect.stack()[1][3]) print(str) class EdscTest: def setUp(self): self.module = E.MLIRModule() self.boolType = self.module.make_type("i1") self.i32Type = self.module.make_type("i32") self.f32Type = self.module.make_type("f32") self.indexType = self.module.make_index_type() def testBlockArguments(self): self.setUp() with self.module.new_function_context("foo", [], []) as fun: E.constant_index(42) with E.BlockContext([self.f32Type, self.f32Type]) as b: b.arg(0) + b.arg(1) printWithCurrentFunctionName(str(fun)) # CHECK-LABEL: testBlockArguments # CHECK: %{{.*}} = constant 42 : index # CHECK: ^bb{{.*}}(%{{.*}}: f32, %{{.*}}: f32): # CHECK: %{{.*}} = addf %{{.*}}, %{{.*}} : f32 def testBlockContext(self): self.setUp() with self.module.new_function_context("foo", [], []) as fun: cst = E.constant_index(42) with E.BlockContext(): cst + cst printWithCurrentFunctionName(str(fun)) # CHECK-LABEL: testBlockContext # CHECK: %{{.*}} = constant 42 : index # CHECK: ^bb # CHECK: %{{.*}} = affine.apply affine_map<() -> (84)>() def testBlockContextAppend(self): self.setUp() with self.module.new_function_context("foo", [], []) as fun: E.constant_index(41) with E.BlockContext() as b: blk = b # save block handle for later E.constant_index(0) E.constant_index(42) with E.BlockContext(E.appendTo(blk)): E.constant_index(1) printWithCurrentFunctionName(str(fun)) # CHECK-LABEL: testBlockContextAppend # CHECK: %{{.*}} = constant 41 : index # CHECK: %{{.*}} = constant 42 : index # CHECK: ^bb # CHECK: %{{.*}} = constant 0 : index # CHECK: %{{.*}} = constant 1 : index def testBlockContextStandalone(self): self.setUp() with self.module.new_function_context("foo", [], []) as fun: blk1 = E.BlockContext() blk2 = E.BlockContext() with blk1: E.constant_index(0) with blk2: E.constant_index(56) E.constant_index(57) E.constant_index(41) with blk1: E.constant_index(1) E.constant_index(42) printWithCurrentFunctionName(str(fun)) # CHECK-LABEL: testBlockContextStandalone # CHECK: %{{.*}} = constant 41 : index # CHECK: %{{.*}} = constant 42 : index # CHECK: ^bb # CHECK: %{{.*}} = constant 0 : index # CHECK: %{{.*}} = constant 1 : index # CHECK: ^bb # CHECK: %{{.*}} = constant 56 : index # CHECK: %{{.*}} = constant 57 : index def testBooleanOps(self): self.setUp() with self.module.new_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: %{{.*}} = cmpi "slt", %{{.*}}, %{{.*}} : i1 # CHECK: %{{.*}} = cmpi "sge", %{{.*}}, %{{.*}} : i1 # CHECK: %{{.*}} = and %{{.*}}, %{{.*}} : i1 # CHECK: %{{.*}} = cmpi "eq", %{{.*}}, %{{.*}} : i1 # CHECK: %{{.*}} = or %{{.*}}, %{{.*}} : i1 # CHECK: %{{.*}} = constant 1 : i1 # CHECK: %{{.*}} = subi %{{.*}}, %{{.*}} : i1 def testBr(self): self.setUp() with self.module.new_function_context("foo", [], []) as fun: with E.BlockContext() as b: blk = b E.ret() E.br(blk) printWithCurrentFunctionName(str(fun)) # CHECK-LABEL: testBr # CHECK: br ^bb # CHECK: ^bb # CHECK: return def testBrArgs(self): self.setUp() with self.module.new_function_context("foo", [], []) as fun: # Create an infinite loop. with E.BlockContext([self.indexType, self.indexType]) as b: E.br(b, [b.arg(1), b.arg(0)]) E.br(b, [E.constant_index(0), E.constant_index(1)]) printWithCurrentFunctionName(str(fun)) # CHECK-LABEL: testBrArgs # CHECK: %{{.*}} = constant 0 : index # CHECK: %{{.*}} = constant 1 : index # CHECK: br ^bb{{.*}}(%{{.*}}, %{{.*}} : index, index) # CHECK: ^bb{{.*}}(%{{.*}}: index, %{{.*}}: index): # CHECK: br ^bb{{.*}}(%{{.*}}, %{{.*}} : index, index) def testBrDeclaration(self): self.setUp() with self.module.new_function_context("foo", [], []) as fun: blk = E.BlockContext() E.br(blk.handle()) with blk: E.ret() printWithCurrentFunctionName(str(fun)) # CHECK-LABEL: testBrDeclaration # CHECK: br ^bb # CHECK: ^bb # CHECK: return def testCallOp(self): self.setUp() callee = self.module.declare_function("sqrtf", [self.f32Type], [self.f32Type]) with self.module.new_function_context("call", [self.f32Type], []) as fun: funCst = E.constant_function(callee) funCst([fun.arg(0)]) + E.constant_float(42., self.f32Type) printWithCurrentFunctionName(str(self.module)) # CHECK-LABEL: testCallOp # CHECK: func @sqrtf(f32) -> f32 # CHECK: %{{.*}} = constant @sqrtf : (f32) -> f32 # CHECK: %{{.*}} = call_indirect %{{.*}}(%{{.*}}) : (f32) -> f32 def testCondBr(self): self.setUp() with self.module.new_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 %{{.*}}, ^bb{{.*}}, ^bb{{.*}}(%{{.*}} : index) def testConstantAffineExpr(self): self.setUp() with self.module.new_function_context("constant_affine", [], []) as fun: a1 = self.module.affine_dim_expr(0) a2 = self.module.affine_dim_expr(1) a3 = a1 + a2 + 3 composedExpr = a3.compose( self.module.affine_map(2, 0, [ self.module.affine_constant_expr(4), self.module.affine_constant_expr(7) ])) printWithCurrentFunctionName(str(fun)) print("constant value : %d" % composedExpr.get_constant_value()) # CHECK-LABEL: testConstantAffineExpr # CHECK: constant value : 14 def testConstants(self): self.setUp() with self.module.new_function_context("constants", [], []) as fun: E.constant_float(1.23, self.module.make_type("bf16")) E.constant_float(1.23, self.module.make_type("f16")) E.constant_float(1.23, self.module.make_type("f32")) E.constant_float(1.23, self.module.make_type("f64")) E.constant_int(1, 1) E.constant_int(123, 8) E.constant_int(123, 16) E.constant_int(123, 32) E.constant_int(123, 64) E.constant_index(123) E.constant_function(fun) printWithCurrentFunctionName(str(fun)) # CHECK-LABEL: testConstants # CHECK: constant 1.230000e+00 : bf16 # CHECK: constant 1.230470e+00 : f16 # CHECK: constant 1.230000e+00 : f32 # CHECK: constant 1.230000e+00 : f64 # CHECK: constant 1 : i1 # CHECK: constant 123 : i8 # CHECK: constant 123 : i16 # CHECK: constant 123 : i32 # CHECK: constant 123 : index # CHECK: constant @constants : () -> () def testCustom(self): self.setUp() with self.module.new_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: %{{.*}} = "foo"(%{{.*}}) : (index) -> f32 # CHECK: %{{.*}} = addf %{{.*}}, %{{.*}} : f32 def testDictionaryAttributes(self): self.setUp() dictionaryAttr = self.module.dictionaryAttr({ "int_0": self.module.integerAttr(self.i32Type, 43), "int_1": self.module.integerAttr(self.i32Type, 33), }) f = self.module.declare_function("foo", [], [], dict_attr=dictionaryAttr) printWithCurrentFunctionName(str(self.module)) # CHECK-LABEL: testDictionaryAttributes # CHECK: func @foo() attributes {dict_attr = {int_0 = 43 : i32, int_1 = 33 : i32}} def testDivisions(self): self.setUp() with self.module.new_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: divi_signed %{{.*}}, %{{.*}} : i32 def testFunctionArgs(self): self.setUp() with self.module.new_function_context("foo", [self.f32Type, self.f32Type], [self.indexType]) as fun: pass printWithCurrentFunctionName(str(fun)) # CHECK-LABEL: testFunctionArgs # CHECK: func @foo(%{{.*}}: f32, %{{.*}}: f32) -> index def testFunctionContext(self): self.setUp() with self.module.new_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 testFunctionDeclarationWithAffineAttr(self): self.setUp() a1 = self.module.affine_constant_expr(23) a2 = self.module.affine_constant_expr(44) a3 = self.module.affine_dim_expr(1) s0 = self.module.affine_symbol_expr(0) aMap1 = self.module.affine_map(2, 0, [a1, a2, s0]) aMap2 = self.module.affine_constant_map(42) aMap3 = self.module.affine_map( 2, 0, [a1 + a2 * a3, a1 // a3 % a2, a1.ceildiv(a2), a1 - 2, a2 * 2, -a3]) affineAttr1 = self.module.affineMapAttr(aMap1) affineAttr2 = self.module.affineMapAttr(aMap2) affineAttr3 = self.module.affineMapAttr(aMap3) t = self.module.make_memref_type(self.f32Type, [10]) t_with_attr = t({ "affine_attr_1": affineAttr1, "affine_attr_2": affineAttr2, "affine_attr_3": affineAttr3, }) f = self.module.declare_function("foo", [t, t_with_attr], []) printWithCurrentFunctionName(str(self.module)) # CHECK-LABEL: testFunctionDeclarationWithAffineAttr # CHECK: func @foo(memref<10xf32>, memref<10xf32> {affine_attr_1 = affine_map<(d0, d1) -> (23, 44, s0)>, affine_attr_2 = affine_map<() -> (42)>, affine_attr_3 = affine_map<(d0, d1) -> (d1 * 44 + 23, (23 floordiv d1) mod 44, 1, 21, 88, -d1)>}) def testFunctionDeclarationWithArrayAttr(self): self.setUp() arrayAttr = self.module.arrayAttr([ self.module.integerAttr(self.i32Type, 43), self.module.integerAttr(self.i32Type, 33), ]) t = self.module.make_memref_type(self.f32Type, [10]) t_with_attr = t({"array_attr": arrayAttr}) f = self.module.declare_function("foo", [t, t_with_attr], []) printWithCurrentFunctionName(str(self.module)) # CHECK-LABEL: testFunctionDeclarationWithArrayAttr # CHECK: func @foo(memref<10xf32>, memref<10xf32> {array_attr = [43 : i32, 33 : i32]}) def testFunctionDeclarationWithFloatAndStringAttr(self): self.setUp() float_attr = self.module.floatAttr(23.3) string_attr = self.module.stringAttr("TEST_STRING") f = self.module.declare_function( "foo", [], [], float_attr=float_attr, string_attr=string_attr) printWithCurrentFunctionName(str(self.module)) # CHECK-LABEL: testFunctionDeclarationWithFloatAndStringAttr # CHECK: func @foo() attributes {float_attr = 2.330000e+01 : f32, string_attr = "TEST_STRING"} def testFunctionMultiple(self): self.setUp() with self.module.new_function_context("foo", [], []): pass with self.module.new_function_context("foo", [], []): E.constant_index(0) printWithCurrentFunctionName(str(self.module)) # CHECK-LABEL: testFunctionMultiple # CHECK: func @foo() # CHECK: func @foo_0() # CHECK: %{{.*}} = constant 0 : index def testIndexCast(self): self.setUp() with self.module.new_function_context("testIndexCast", [], []): index = E.constant_index(0) E.index_cast(index, self.i32Type) printWithCurrentFunctionName(str(self.module)) # CHECK-LABEL: testIndexCast # CHECK: index_cast %{{.*}} : index to i32 def testIndexedValue(self): self.setUp() memrefType = self.module.make_memref_type(self.f32Type, [10, 42]) with self.module.new_function_context("indexed", [memrefType], [memrefType]) as fun: A = E.IndexedValue(fun.arg(0)) cst = E.constant_float(1., self.f32Type) with E.LoopNestContext( [E.constant_index(0), E.constant_index(0)], [E.constant_index(10), E.constant_index(42)], [1, 1]) as (i, j): A.store([i, j], A.load([i, j]) + cst) E.ret([fun.arg(0)]) printWithCurrentFunctionName(str(fun)) # CHECK-LABEL: testIndexedValue # CHECK: affine.for # CHECK: affine.for # CHECK: %{{.*}} affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x42xf32> # CHECK: %{{.*}} = addf %{{.*}}, %{{.*}} : f32 # CHECK: affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x42xf32> def testLoopContext(self): self.setUp() with self.module.new_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: testLoopContext # CHECK: affine.for # CHECK: {{.*}} = affine.apply affine_map<() -> (42)>() # CHECK: {{.*}} = affine.apply affine_map<(d0) -> (d0 + 42)>({{.*}}) # CHECK: {{.*}} = affine.apply affine_map<() -> (84)>() # CHECK: affine.for {{.*}} = affine_map<(d0) -> (d0)>(%c42) to affine_map<(d0) -> (d0)>({{.*}}) step 2 { # CHECK: {{.*}} = affine.apply affine_map<(d0, d1) -> (d0 + d1)>({{.*}}, {{.*}}) def testLoopNestContext(self): self.setUp() with self.module.new_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: affine.for # CHECK: affine.for # CHECK: affine.for # CHECK: {{.*}} = affine.apply affine_map<(d0, d1) -> (d0 + d1)>({{.*}}, {{.*}}) # CHECK: {{.*}} = affine.apply affine_map<(d0, d1, d2) -> (d0 + d1 + d2)>({{.*}}, {{.*}}, {{.*}}) # CHECK: {{.*}} = affine.apply affine_map<(d0, d1, d2, d3) -> (d0 + d1 + d2 + d3)>({{.*}}, {{.*}}, {{.*}}, {{.*}}) def testMLIRFunctionCreation(self): self.setUp() module = E.MLIRModule() t = module.make_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(%{{.*}}: memref<3x4x?x5xf32>, %{{.*}}: memref<3x4x?x5xf32>) { # CHECK: func @sqrtf(%{{.*}}: f32) -> f32 def testMLIRScalarTypes(self): self.setUp() module = E.MLIRModule() printWithCurrentFunctionName(str(module.make_type("bf16"))) print(str(module.make_type("f16"))) print(str(module.make_type("f32"))) print(str(module.make_type("f64"))) print(str(module.make_type("i1"))) print(str(module.make_type("i8"))) print(str(module.make_type("i32"))) print(str(module.make_type("i123"))) print(str(module.make_type("index"))) # CHECK-LABEL: testMLIRScalarTypes # CHECK: bf16 # CHECK: f16 # CHECK: f32 # CHECK: f64 # CHECK: i1 # CHECK: i8 # CHECK: i32 # CHECK: i123 # CHECK: index def testMatrixMultiply(self): self.setUp() memrefType = self.module.make_memref_type(self.f32Type, [32, 32]) with self.module.new_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: %{{.*}} = affine.load # CHECK-DAG: %{{.*}} = affine.load # CHECK: %{{.*}} = mulf %{{.*}}, %{{.*}} : f32 # CHECK: affine.store # CHECK-SAME: memref<32x32xf32> def testRet(self): self.setUp() with self.module.new_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: %{{.*}} = constant 42 : index # CHECK: %{{.*}} = constant 0 : index # CHECK: return %{{.*}}, %{{.*}} : index, index def testSelectOp(self): self.setUp() with self.module.new_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: %{{.*}} = select %{{.*}}, %{{.*}}, %{{.*}} : i32 def testType(self): self.setUp() printWithCurrentFunctionName("") with self.module.new_function_context( "foo", [self.module.make_memref_type(self.f32Type, [10])], []) as fun: c42 = E.constant_int(42, 32) print(str(c42.type())) print(str(fun.arg(0).type())) # CHECK-LABEL: testType # CHECK: i32 # CHECK: memref<10xf32> # Until python 3.6 this cannot be used because the order in the dict is not the # order of method declaration. def runTests(): def isTest(attr): return inspect.ismethod(attr) and "EdscTest.setUp " not in str(attr) edscTest = EdscTest() tests = sorted( filter(isTest, (getattr(edscTest, attr) for attr in dir(edscTest))), key=lambda x: str(x)) for test in tests: print("--> Running test:", test.__name__, file=sys.stderr) test() if __name__ == "__main__": test_utils.run_under_filecheck(__file__, runTests)