diff --git a/README.md b/README.md index 95ae0f9dc..cc871b18f 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ export LDFLAGS=-fuse-ld=$(which ld.lld-$LLVM_VERSION) export LLVM_SRC_DIR=/path/to/llvm-project # Check out last known good commit. -(cd $LLVM_SRC_DIR && git checkout 26777ad7a0916ad7853aa9229bb8ec0346c68a61) +(cd $LLVM_SRC_DIR && git checkout 0c4aab27b3da05dd1b0c0c39472525325fda5e23) ./tools/install_mlir.sh ./tools/cmake_configure.sh diff --git a/python/npcomp/CMakeLists.txt b/python/npcomp/CMakeLists.txt index e73beb493..dca1e178c 100644 --- a/python/npcomp/CMakeLists.txt +++ b/python/npcomp/CMakeLists.txt @@ -26,7 +26,6 @@ set(NPCOMP_PYEXT_LIBADD ${PYTHON_LIBRARIES}) set(extension_target NPCOMPNativePyExt) set(extension_pybind_sources native.cpp - mlir_edsc.cpp mlir_init.cpp mlir_ir.cpp pybind_utils.cpp diff --git a/python/npcomp/edsc_test.py b/python/npcomp/edsc_test.py deleted file mode 100644 index 773f43e77..000000000 --- a/python/npcomp/edsc_test.py +++ /dev/null @@ -1,536 +0,0 @@ -# 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) diff --git a/python/npcomp/mlir_edsc.cpp b/python/npcomp/mlir_edsc.cpp deleted file mode 100644 index 76a943b7e..000000000 --- a/python/npcomp/mlir_edsc.cpp +++ /dev/null @@ -1,1176 +0,0 @@ -//===- mlir_edsc.cpp - MLIR EDSC bindings ---------------------------------===// -// -// 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 -// -//===----------------------------------------------------------------------===// - -#include -#include - -#include -#include -#include - -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/raw_ostream.h" -#include "mlir-c/Core.h" -#include "mlir/Conversion/LoopToStandard/ConvertLoopToStandard.h" -#include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h" -#include "mlir/Dialect/Affine/EDSC/Intrinsics.h" -#include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h" -#include "mlir/EDSC/Builders.h" -#include "mlir/EDSC/Intrinsics.h" -#include "mlir/IR/AffineExpr.h" -#include "mlir/IR/AffineMap.h" -#include "mlir/IR/Attributes.h" -#include "mlir/IR/Function.h" -#include "mlir/IR/Module.h" -#include "mlir/IR/Types.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Target/LLVMIR.h" -#include "mlir/Transforms/Passes.h" - -namespace py = pybind11; -using mlir::edsc::intrinsics::AffineIndexedValue; - -using namespace mlir; -using namespace mlir::edsc; - -namespace mlir { -namespace edsc { - -struct CustomValueBuilder { - CustomValueBuilder(StringRef name, ArrayRef operands, - ArrayRef resultTypes, - ArrayRef attributes = {}) { - OperationState state(ScopedContext::getLocation(), name); - SmallVector ops(operands.begin(), operands.end()); - state.addOperands(ops); - state.addTypes(resultTypes); - for (const auto &attr : attributes) - state.addAttribute(attr.first, attr.second); - Operation *op = ScopedContext::getBuilder().createOperation(state); - if (op->getNumResults() != 1) - llvm_unreachable("unsupported operation, use an OperationHandle instead"); - value = op->getResult(0); - } - operator Value() { return value; } - Value value; -}; - -struct PythonAttribute; -struct PythonAttributedType; -struct PythonBindable; -struct PythonExpr; -struct PythonFunctionContext; -struct PythonStmt; -struct PythonBlock; -struct PythonAffineExpr; -struct PythonAffineMap; - -struct PythonType { - PythonType() : type{nullptr} {} - PythonType(mlir_type_t t) : type{t} {} - - operator mlir_type_t() const { return type; } - - PythonAttributedType attachAttributeDict( - const std::unordered_map &attrs) const; - - std::string str() { - mlir::Type f = mlir::Type::getFromOpaquePointer(type); - std::string res; - llvm::raw_string_ostream os(res); - f.print(os); - return res; - } - - mlir_type_t type; -}; - -struct PythonValueHandle { - PythonValueHandle(const PythonValueHandle &other) = default; - PythonValueHandle(const mlir::Value &other) : value(other) {} - PythonValueHandle(mlir::edsc::CustomValueBuilder vb) : value(vb) {} - template - PythonValueHandle(mlir::edsc::ValueBuilder vb) : value(vb) {} - operator mlir::Value() const { return value; } - operator mlir::Value &() { return value; } - - std::string str() const { - return std::to_string( - reinterpret_cast(value.getAsOpaquePointer())); - } - - PythonValueHandle call(const std::vector &args) { - std::vector argValues; - argValues.reserve(args.size()); - for (auto arg : args) argValues.push_back(arg.value); - return mlir::edsc::ValueBuilder(value, argValues); - } - - PythonType type() const { - return PythonType(value.getType().getAsOpaquePointer()); - } - - mlir::Value value; -}; - -struct PythonFunction { - PythonFunction() : function{nullptr} {} - PythonFunction(mlir_func_t f) : function{f} {} - PythonFunction(mlir::FuncOp f) - : function(const_cast(f.getAsOpaquePointer())) {} - operator mlir_func_t() { return function; } - std::string str() { - mlir::FuncOp f = mlir::FuncOp::getFromOpaquePointer(function); - std::string res; - llvm::raw_string_ostream os(res); - f.print(os); - return res; - } - - // If the function does not yet have an entry block, i.e. if it is a function - // declaration, add the entry block, transforming the declaration into a - // definition. Return true if the block was added, false otherwise. - bool define() { - auto f = mlir::FuncOp::getFromOpaquePointer(function); - if (!f.getBlocks().empty()) return false; - - f.addEntryBlock(); - return true; - } - - PythonValueHandle arg(unsigned index) { - auto f = mlir::FuncOp::getFromOpaquePointer(function); - assert(index < f.getNumArguments() && "argument index out of bounds"); - return PythonValueHandle(f.getArgument(index)); - } - - mlir_func_t function; -}; - -/// Trivial C++ wrappers make use of the EDSC C API. -struct PythonMLIRModule { - PythonMLIRModule() - : mlirContext(), - module(mlir::ModuleOp::create(mlir::UnknownLoc::get(&mlirContext))), - symbolTable(*module) {} - - PythonType makeMemRefType(PythonType elemType, std::vector sizes) { - return ::makeMemRefType(mlir_context_t{&mlirContext}, elemType, - int64_list_t{sizes.data(), sizes.size()}); - } - PythonType makeIndexType() { - return ::makeIndexType(mlir_context_t{&mlirContext}); - } - PythonType makeType(const std::string &type) { - return ::mlirParseType(type.c_str(), mlir_context_t{&mlirContext}, nullptr); - } - - // Declare a function with the given name, input types and their attributes, - // output types, and function attributes, but do not define it. - PythonFunction declareFunction(const std::string &name, - const py::list &inputs, - const std::vector &outputTypes, - const py::kwargs &funcAttributes); - - // Declare a function with the given name, input types and their attributes, - // output types, and function attributes. - PythonFunction makeFunction(const std::string &name, const py::list &inputs, - const std::vector &outputTypes, - const py::kwargs &funcAttributes) { - auto declaration = - declareFunction(name, inputs, outputTypes, funcAttributes); - declaration.define(); - return declaration; - } - - // Create a custom op given its name and arguments. - PythonExpr op(const std::string &name, PythonType type, - const py::list &arguments, const py::list &successors, - py::kwargs attributes); - - // Creates an integer attribute. - PythonAttribute integerAttr(PythonType type, int64_t value); - - // Creates a boolean attribute. - PythonAttribute boolAttr(bool value); - - // Creates a float attribute. - PythonAttribute floatAttr(float value); - - // Creates a string atrribute. - PythonAttribute stringAttr(const std::string &value); - - // Creates an Array attribute. - PythonAttribute arrayAttr(const std::vector &values); - - // Creates an AffineMap attribute. - PythonAttribute affineMapAttr(PythonAffineMap value); - - // Creates a dictionary attribute. - PythonAttribute dictionaryAttr( - const std::unordered_map &attrs); - - // Creates an affine constant expression. - PythonAffineExpr affineConstantExpr(int64_t value); - - // Creates an affine symbol expression. - PythonAffineExpr affineSymbolExpr(unsigned position); - - // Creates an affine dimension expression. - PythonAffineExpr affineDimExpr(unsigned position); - - // Creates a single constant result affine map. - PythonAffineMap affineConstantMap(int64_t value); - - // Creates an affine map. - PythonAffineMap affineMap(unsigned dimCount, unsigned symbolCount, - const std::vector &results); - - std::string getIR() { - std::string res; - llvm::raw_string_ostream os(res); - module->print(os); - return res; - } - - PythonFunction getNamedFunction(const std::string &name) { - return symbolTable.lookup(name); - } - - PythonFunctionContext makeNewFunctionContext( - const std::string &name, const py::list &inputs, - const std::vector &outputs, const py::kwargs &attributes); - - PythonFunctionContext makeFunctionContext( - PythonFunction f); - - private: - mlir::MLIRContext mlirContext; - // One single module in a python-exposed MLIRContext for now. - mlir::OwningModuleRef module; - mlir::SymbolTable symbolTable; -}; - -struct PythonFunctionContext { - PythonFunctionContext(PythonFunction f) : function(f) {} - PythonFunctionContext(PythonMLIRModule &module, const std::string &name, - const py::list &inputs, - const std::vector &outputs, - const py::kwargs &attributes) { - auto function = module.declareFunction(name, inputs, outputs, attributes); - function.define(); - } - - PythonFunction enter() { - assert(function.function && "function is not set up"); - auto mlirFunc = mlir::FuncOp::getFromOpaquePointer(function.function); - contextBuilder.emplace(mlirFunc.getBody()); - context = new mlir::edsc::ScopedContext(*contextBuilder, mlirFunc.getLoc()); - return function; - } - - void exit(py::object, py::object, py::object) { - delete context; - context = nullptr; - contextBuilder.reset(); - } - - PythonFunction function; - mlir::edsc::ScopedContext *context; - llvm::Optional contextBuilder; -}; - -PythonFunctionContext PythonMLIRModule::makeNewFunctionContext( - const std::string &name, const py::list &inputs, - const std::vector &outputs, const py::kwargs &attributes) { - auto func = declareFunction(name, inputs, outputs, attributes); - func.define(); - return PythonFunctionContext(func); -} - -PythonFunctionContext PythonMLIRModule::makeFunctionContext( - PythonFunction f) { - return PythonFunctionContext(f); -} - -struct PythonBlockHandle { - PythonBlockHandle() : value(nullptr) {} - PythonBlockHandle(const PythonBlockHandle &other) = default; - PythonBlockHandle(const mlir::edsc::BlockHandle &other) : value(other) {} - operator mlir::edsc::BlockHandle() const { return value; } - - PythonType type(int index) { return types[index].getAsOpaquePointer(); } - PythonValueHandle arg(int index) { return arguments[index]; } - - std::string str() { - std::string s; - llvm::raw_string_ostream os(s); - value.getBlock()->print(os); - return os.str(); - } - - mlir::edsc::BlockHandle value; - std::vector types; - std::vector arguments; -}; - -struct PythonOperationHandle { - PythonOperationHandle() : value(nullptr) {} - PythonOperationHandle(const PythonOperationHandle &other) = default; - PythonOperationHandle(const mlir::edsc::OperationHandle &other) - : value(other) {} - operator OperationHandle() const { return value; } - operator OperationHandle &() { return value; } - PythonValueHandle getResult(int index) const { - return value.getOperation()->getResult(index); - } - - std::string str() const { - std::string s; - llvm::raw_string_ostream os(s); - value.getOperation()->print(os); - return os.str(); - } - - PythonOperationHandle call(const std::vector &args) { - std::vector argValues; - argValues.reserve(args.size()); - for (auto arg : args) argValues.push_back(arg.value); - return OperationHandle::create(argValues[0], argValues); - } - - mlir::edsc::OperationHandle value; -}; - -struct PythonLoopContext { - PythonLoopContext(PythonValueHandle lb, PythonValueHandle ub, int64_t step) - : lb(lb), ub(ub), step(step) {} - PythonLoopContext(const PythonLoopContext &) = delete; - PythonLoopContext(PythonLoopContext &&) = default; - PythonLoopContext &operator=(const PythonLoopContext &) = delete; - PythonLoopContext &operator=(PythonLoopContext &&) = default; - ~PythonLoopContext() { assert(!builder && "did not exit from the context"); } - - PythonValueHandle enter() { - mlir::Value iv; - builder = new AffineLoopNestBuilder(&iv, lb.value, ub.value, step); - return iv; - } - - void exit(py::object, py::object, py::object) { - (*builder)({}); // exit from the builder's scope. - delete builder; - builder = nullptr; - } - - PythonValueHandle lb, ub; - int64_t step; - AffineLoopNestBuilder *builder = nullptr; -}; - -struct PythonLoopNestContext { - PythonLoopNestContext(const std::vector &lbs, - const std::vector &ubs, - const std::vector steps) - : lbs(lbs), ubs(ubs), steps(steps) { - assert(lbs.size() == ubs.size() && lbs.size() == steps.size() && - "expected the same number of lower, upper bounds, and steps"); - } - PythonLoopNestContext(const PythonLoopNestContext &) = delete; - PythonLoopNestContext(PythonLoopNestContext &&) = default; - PythonLoopNestContext &operator=(const PythonLoopNestContext &) = delete; - PythonLoopNestContext &operator=(PythonLoopNestContext &&) = default; - ~PythonLoopNestContext() { - assert(!builder && "did not exit from the context"); - } - - std::vector enter() { - if (steps.empty()) return {}; - - std::vector handles(steps.size()); - builder = new AffineLoopNestBuilder( - handles, std::vector(lbs.begin(), lbs.end()), - std::vector(ubs.begin(), ubs.end()), steps); - return std::vector(handles.begin(), handles.end()); - } - - void exit(py::object, py::object, py::object) { - (*builder)({}); // exit from the builder's scope. - delete builder; - builder = nullptr; - } - - std::vector lbs; - std::vector ubs; - std::vector steps; - AffineLoopNestBuilder *builder = nullptr; -}; - -struct PythonBlockAppender { - PythonBlockAppender(const PythonBlockHandle &handle) : handle(handle) {} - PythonBlockHandle handle; -}; - -struct PythonBlockContext { - public: - PythonBlockContext() { - createBlockBuilder(); - clearBuilder(); - } - PythonBlockContext(const std::vector &argTypes) { - handle.arguments.reserve(argTypes.size()); - for (const auto &t : argTypes) { - auto type = - Type::getFromOpaquePointer(reinterpret_cast(t.type)); - handle.types.emplace_back(type); - handle.arguments.push_back(Value()); - } - createBlockBuilder(); - clearBuilder(); - } - PythonBlockContext(const PythonBlockAppender &a) : handle(a.handle) {} - PythonBlockContext(const PythonBlockContext &) = delete; - PythonBlockContext(PythonBlockContext &&) = default; - PythonBlockContext &operator=(const PythonBlockContext &) = delete; - PythonBlockContext &operator=(PythonBlockContext &&) = default; - ~PythonBlockContext() { - assert(!builder && "did not exit from the block context"); - } - - // EDSC maintain an implicit stack of builders (mostly for keeping track of - // insertion points); every operation gets inserted using the top-of-the-stack - // builder. Creating a new EDSC Builder automatically puts it on the stack, - // effectively entering the block for it. - void createBlockBuilder() { - if (handle.value.getBlock()) - builder = new BlockBuilder(handle.value, mlir::edsc::Append()); - else - builder = new BlockBuilder(&handle.value, handle.types, handle.arguments); - } - - PythonBlockHandle enter() { - createBlockBuilder(); - return handle; - } - - void exit(py::object, py::object, py::object) { clearBuilder(); } - - PythonBlockHandle getHandle() { return handle; } - - // EDSC maintain an implicit stack of builders (mostly for keeping track of - // insertion points); every operation gets inserted using the top-of-the-stack - // builder. Calling operator() on a builder pops the builder from the stack, - // effectively resetting the insertion point to its position before we entered - // the block. - void clearBuilder() { - (*builder)({}); // exit from the builder's scope. - delete builder; - builder = nullptr; - } - - PythonBlockHandle handle; - BlockBuilder *builder = nullptr; -}; - -struct PythonAttribute { - PythonAttribute() : attr(nullptr) {} - PythonAttribute(const mlir_attr_t &a) : attr(a) {} - PythonAttribute(const PythonAttribute &other) = default; - operator mlir_attr_t() { return attr; } - - operator Attribute() const { return Attribute::getFromOpaquePointer(attr); } - - std::string str() const { - if (!attr) return "##null attr##"; - - std::string res; - llvm::raw_string_ostream os(res); - Attribute().print(os); - return res; - } - - mlir_attr_t attr; -}; - -struct PythonAttributedType { - PythonAttributedType() : type(nullptr) {} - PythonAttributedType(mlir_type_t t) : type(t) {} - PythonAttributedType( - PythonType t, - const std::unordered_map &attributes = - std::unordered_map()) - : type(t), attrs(attributes) {} - - operator mlir_type_t() const { return type.type; } - operator PythonType() const { return type; } - - // Return a vector of named attribute descriptors. The vector owns the - // mlir_named_attr_t objects it contains, but not the names and attributes - // those objects point to (names and opaque pointers to attributes are owned - // by `this`). - std::vector getNamedAttrs() const { - std::vector result; - result.reserve(attrs.size()); - for (const auto &namedAttr : attrs) - result.push_back({namedAttr.first.c_str(), namedAttr.second.attr}); - return result; - } - - std::string str() { - mlir::Type t = mlir::Type::getFromOpaquePointer(type); - std::string res; - llvm::raw_string_ostream os(res); - t.print(os); - if (attrs.empty()) return os.str(); - - os << '{'; - bool first = true; - for (const auto &namedAttr : attrs) { - if (first) - first = false; - else - os << ", "; - os << namedAttr.first << ": " << namedAttr.second.str(); - } - os << '}'; - - return os.str(); - } - - private: - PythonType type; - std::unordered_map attrs; -}; - -// Wraps mlir::AffineExpr. -struct PythonAffineExpr { - PythonAffineExpr() : affine_expr() {} - PythonAffineExpr(const AffineExpr &a) : affine_expr(a) {} - PythonAffineExpr(const PythonAffineExpr &other) = default; - - operator AffineExpr() const { return affine_expr; } - operator AffineExpr &() { return affine_expr; } - - AffineExpr get() const { return affine_expr; } - - std::string str() const { - std::string res; - llvm::raw_string_ostream os(res); - affine_expr.print(os); - return res; - } - - private: - AffineExpr affine_expr; -}; - -// Wraps mlir::AffineMap. -struct PythonAffineMap { - PythonAffineMap() : affine_map() {} - PythonAffineMap(const AffineMap &a) : affine_map(a) {} - PythonAffineMap(const PythonAffineMap &other) = default; - - operator AffineMap() const { return affine_map; } - operator AffineMap &() { return affine_map; } - - std::string str() const { - std::string res; - llvm::raw_string_ostream os(res); - affine_map.print(os); - return res; - } - - private: - AffineMap affine_map; -}; - -struct PythonIndexedValue { - explicit PythonIndexedValue(const AffineIndexedValue &other) - : indexed(other) {} - PythonIndexedValue(PythonValueHandle handle) : indexed(handle.value) {} - PythonIndexedValue(const PythonIndexedValue &other) = default; - - // Create a new indexed value with the same base as this one but with indices - // provided as arguments. - PythonIndexedValue index(const std::vector &indices) { - std::vector handles(indices.begin(), indices.end()); - return PythonIndexedValue(AffineIndexedValue(indexed(handles))); - } - - void store(const std::vector &indices, - PythonValueHandle value) { - // Uses the overloaded `operator=` to emit a store. - index(indices).indexed = value.value; - } - - PythonValueHandle load(const std::vector &indices) { - // Uses the overloaded cast to `mlir::Value` to emit a load. - return static_cast(index(indices).indexed); - } - - AffineIndexedValue indexed; -}; - -template -ListTy makeCList(SmallVectorImpl &owning, const py::list &list) { - for (auto &inp : list) { - owning.push_back(Ty{inp.cast()}); - } - return ListTy{owning.data(), owning.size()}; -} - -PythonFunction PythonMLIRModule::declareFunction( - const std::string &name, const py::list &inputs, - const std::vector &outputTypes, - const py::kwargs &funcAttributes) { - std::vector attributedInputs; - attributedInputs.reserve(inputs.size()); - for (const auto &in : inputs) { - std::string className = py::str(in.get_type()); - if (className.find(".Type'") != std::string::npos) - attributedInputs.emplace_back(in.cast()); - else - attributedInputs.push_back(in.cast()); - } - - // Create the function type. - std::vector ins(attributedInputs.begin(), - attributedInputs.end()); - std::vector outs(outputTypes.begin(), outputTypes.end()); - auto funcType = ::makeFunctionType( - mlir_context_t{&mlirContext}, mlir_type_list_t{ins.data(), ins.size()}, - mlir_type_list_t{outs.data(), outs.size()}); - - // Build the list of function attributes. - std::vector attrs; - attrs.reserve(funcAttributes.size()); - for (const auto &named : funcAttributes) - attrs.emplace_back( - Identifier::get(std::string(py::str(named.first)), &mlirContext), - mlir::Attribute::getFromOpaquePointer(reinterpret_cast( - named.second.cast().attr))); - - // Build the list of lists of function argument attributes. - std::vector inputAttrs; - inputAttrs.reserve(attributedInputs.size()); - for (const auto &in : attributedInputs) { - std::vector inAttrs; - for (const auto &named : in.getNamedAttrs()) - inAttrs.emplace_back(Identifier::get(named.name, &mlirContext), - mlir::Attribute::getFromOpaquePointer( - reinterpret_cast(named.value))); - inputAttrs.emplace_back(inAttrs); - } - - // Create the function itself. - auto func = mlir::FuncOp::create( - UnknownLoc::get(&mlirContext), name, - mlir::Type::getFromOpaquePointer(funcType).cast(), attrs, - inputAttrs); - symbolTable.insert(func); - return func; -} - -PythonAttributedType PythonType::attachAttributeDict( - const std::unordered_map &attrs) const { - return PythonAttributedType(*this, attrs); -} - -PythonAttribute PythonMLIRModule::integerAttr(PythonType type, int64_t value) { - return PythonAttribute(::makeIntegerAttr(type, value)); -} - -PythonAttribute PythonMLIRModule::boolAttr(bool value) { - return PythonAttribute(::makeBoolAttr(&mlirContext, value)); -} - -PythonAttribute PythonMLIRModule::floatAttr(float value) { - return PythonAttribute(::makeFloatAttr(&mlirContext, value)); -} - -PythonAttribute PythonMLIRModule::stringAttr(const std::string &value) { - return PythonAttribute(::makeStringAttr(&mlirContext, value.c_str())); -} - -PythonAttribute PythonMLIRModule::arrayAttr( - const std::vector &values) { - std::vector mlir_attributes(values.begin(), values.end()); - auto array_attr = ArrayAttr::get( - llvm::ArrayRef(mlir_attributes), &mlirContext); - return PythonAttribute(array_attr.getAsOpaquePointer()); -} - -PythonAttribute PythonMLIRModule::affineMapAttr(PythonAffineMap value) { - return PythonAttribute(AffineMapAttr::get(value).getAsOpaquePointer()); -} - -PythonAttribute PythonMLIRModule::dictionaryAttr( - const std::unordered_map &attribute_map) { - std::vector attributes; - attributes.reserve(attribute_map.size()); - for (const auto &pair : attribute_map) { - attributes.emplace_back(Identifier::get(pair.first, &mlirContext), - Attribute::getFromOpaquePointer(pair.second.attr)); - } - auto dictionaryAttribute = DictionaryAttr::get(attributes, &mlirContext); - return PythonAttribute(dictionaryAttribute.getAsOpaquePointer()); -} - -PythonAffineExpr PythonMLIRModule::affineConstantExpr(int64_t value) { - return PythonAffineExpr(getAffineConstantExpr(value, &mlirContext)); -} - -PythonAffineExpr PythonMLIRModule::affineSymbolExpr(unsigned position) { - return PythonAffineExpr(getAffineSymbolExpr(position, &mlirContext)); -} - -PythonAffineExpr PythonMLIRModule::affineDimExpr(unsigned position) { - return PythonAffineExpr(getAffineDimExpr(position, &mlirContext)); -} - -PythonAffineMap PythonMLIRModule::affineConstantMap(int64_t value) { - return PythonAffineMap(AffineMap::getConstantMap(value, &mlirContext)); -} - -PythonAffineMap PythonMLIRModule::affineMap( - unsigned dimCount, unsigned SymbolCount, - const std::vector &results) { - std::vector mlir_results(results.begin(), results.end()); - return PythonAffineMap( - AffineMap::get(dimCount, SymbolCount, - llvm::ArrayRef(mlir_results), &mlirContext)); -} - -void defineMlirEdscModule(py::module m) { - m.doc() = - "Python bindings for MLIR Embedded Domain-Specific Components (EDSCs)"; - m.def("version", []() { return "EDSC Python extensions v1.0"; }); - - py::class_( - m, "LoopContext", "A context for building the body of a 'for' loop") - .def(py::init()) - .def("__enter__", &PythonLoopContext::enter) - .def("__exit__", &PythonLoopContext::exit); - - py::class_(m, "LoopNestContext", - "A context for building the body of a the " - "innermost loop in a nest of 'for' loops") - .def(py::init &, - const std::vector &, - const std::vector &>()) - .def("__enter__", &PythonLoopNestContext::enter) - .def("__exit__", &PythonLoopNestContext::exit); - - m.def("constant_index", [](int64_t val) -> PythonValueHandle { - return mlir::edsc::ValueBuilder(val); - }); - m.def("constant_int", [](int64_t val, int width) -> PythonValueHandle { - return mlir::edsc::ValueBuilder(val, width); - }); - m.def("constant_float", [](double val, PythonType type) -> PythonValueHandle { - FloatType floatType = - Type::getFromOpaquePointer(type.type).cast(); - assert(floatType); - auto value = APFloat(val); - bool lostPrecision; - value.convert(floatType.getFloatSemantics(), APFloat::rmNearestTiesToEven, - &lostPrecision); - return mlir::edsc::ValueBuilder(value, floatType); - }); - m.def("constant_function", [](PythonFunction func) -> PythonValueHandle { - auto function = FuncOp::getFromOpaquePointer(func.function); - auto attr = SymbolRefAttr::get(function.getName(), function.getContext()); - return mlir::edsc::ValueBuilder(function.getType(), attr); - }); - m.def("appendTo", [](const PythonBlockHandle &handle) { - return PythonBlockAppender(handle); - }); - m.def( - "ret", - [](const std::vector &args) { - std::vector values(args.begin(), args.end()); - (intrinsics::std_ret(ArrayRef{values})); // vexing parse - return PythonValueHandle(mlir::Value()); - }, - py::arg("args") = std::vector()); - m.def( - "br", - [](const PythonBlockHandle &dest, - const std::vector &args) { - std::vector values(args.begin(), args.end()); - intrinsics::std_br(dest, values); - return PythonValueHandle(mlir::Value()); - }, - py::arg("dest"), py::arg("args") = std::vector()); - m.def( - "cond_br", - [](PythonValueHandle condition, const PythonBlockHandle &trueDest, - const std::vector &trueArgs, - const PythonBlockHandle &falseDest, - const std::vector &falseArgs) -> PythonValueHandle { - std::vector trueArguments(trueArgs.begin(), - trueArgs.end()); - std::vector falseArguments(falseArgs.begin(), - falseArgs.end()); - intrinsics::std_cond_br(condition, trueDest, trueArguments, falseDest, - falseArguments); - return PythonValueHandle(mlir::Value()); - }); - m.def("index_cast", - [](PythonValueHandle element, PythonType type) -> PythonValueHandle { - return mlir::edsc::ValueBuilder( - element.value, Type::getFromOpaquePointer(type.type)); - }); - m.def("select", - [](PythonValueHandle condition, PythonValueHandle trueValue, - PythonValueHandle falseValue) -> PythonValueHandle { - return mlir::edsc::ValueBuilder( - condition.value, trueValue.value, falseValue.value); - }); - m.def("op", - [](const std::string &name, - const std::vector &operands, - const std::vector &resultTypes, - const py::kwargs &attributes) -> PythonValueHandle { - std::vector operandHandles(operands.begin(), - operands.end()); - std::vector types; - types.reserve(resultTypes.size()); - for (auto t : resultTypes) - types.push_back(Type::getFromOpaquePointer(t.type)); - - std::vector attrs; - attrs.reserve(attributes.size()); - for (const auto &a : attributes) { - std::string name = py::str(a.first); - auto pyAttr = a.second.cast(); - auto cppAttr = Attribute::getFromOpaquePointer(pyAttr.attr); - auto identifier = - Identifier::get(name, ScopedContext::getContext()); - attrs.emplace_back(identifier, cppAttr); - } - return mlir::edsc::CustomValueBuilder(name, operandHandles, types, - attrs); - }); - m.def("multi_result_op", - [](const std::string &name, - const std::vector &operands, - const std::vector &resultTypes, - const py::kwargs &attributes) -> PythonOperationHandle { - std::vector operandHandles(operands.begin(), - operands.end()); - std::vector types; - types.reserve(resultTypes.size()); - for (auto t : resultTypes) - types.push_back(Type::getFromOpaquePointer(t.type)); - - std::vector attrs; - attrs.reserve(attributes.size()); - for (const auto &a : attributes) { - std::string name = py::str(a.first); - auto pyAttr = a.second.cast(); - auto cppAttr = Attribute::getFromOpaquePointer(pyAttr.attr); - auto identifier = - Identifier::get(name, ScopedContext::getContext()); - attrs.emplace_back(identifier, cppAttr); - } - return OperationHandle::create(name, operandHandles, types, attrs); - }); - - py::class_(m, "Function", "Wrapping class for mlir::FuncOp.") - .def(py::init()) - .def("__str__", &PythonFunction::str) - .def("define", &PythonFunction::define, - "Adds a body to the function if it does not already have one. " - "Returns true if the body was added") - .def("arg", &PythonFunction::arg, - "Get the mlir::Value to the indexed argument of the function"); - - py::class_(m, "Attribute", - "Wrapping class for mlir::Attribute") - .def(py::init()) - .def("__str__", &PythonAttribute::str); - - py::class_(m, "Type", "Wrapping class for mlir::Type.") - .def(py::init()) - .def("__call__", &PythonType::attachAttributeDict, - "Attach the attributes to these type, making it suitable for " - "constructing functions with argument attributes") - .def("__str__", &PythonType::str); - - py::class_( - m, "AttributedType", - "A class containing a wrapped mlir::Type and a wrapped " - "mlir::NamedAttributeList that are used together, e.g. in function " - "argument declaration") - .def(py::init()) - .def("__str__", &PythonAttributedType::str); - - py::class_( - m, "MLIRModule", - "An MLIRModule is the abstraction that owns the allocations to support " - "compilation of a single mlir::ModuleOp into an ExecutionEngine backed " - "by " - "the LLVM ORC JIT. A typical flow consists in creating an MLIRModule, " - "adding functions, compiling the module to obtain an ExecutionEngine on " - "which named functions may be called. For now the only means to retrieve " - "the ExecutionEngine is by calling `get_engine_address`. This mode of " - "execution is limited to passing the pointer to C++ where the function " - "is called. Extending the API to allow calling JIT compiled functions " - "directly require integration with a tensor library (e.g. numpy). This " - "is left as the prerogative of libraries and frameworks for now.") - .def(py::init<>()) - .def("boolAttr", &PythonMLIRModule::boolAttr, - "Creates an mlir::BoolAttr with the given value") - .def( - "integerAttr", &PythonMLIRModule::integerAttr, - "Creates an mlir::IntegerAttr of the given type with the given value " - "in the context associated with this MLIR module.") - .def("floatAttr", &PythonMLIRModule::floatAttr, - "Creates an mlir::FloatAttr with the given value") - .def("stringAttr", &PythonMLIRModule::stringAttr, - "Creates an mlir::StringAttr with the given value") - .def("arrayAttr", &PythonMLIRModule::arrayAttr, - "Creates an mlir::ArrayAttr of the given type with the given values " - "in the context associated with this MLIR module.") - .def("dictionaryAttr", &PythonMLIRModule::dictionaryAttr, - "Creates a dictonary attribute with the given map.") - .def("affineMapAttr", &PythonMLIRModule::affineMapAttr, - "Creates an mlir::AffineMapAttr of the given type with the given " - "value in the context associated with this MLIR module.") - .def("declare_function", &PythonMLIRModule::declareFunction, - "Declares a new mlir::FuncOp in the current mlir::ModuleOp. The " - "function arguments can have attributes. The function has no " - "definition and can be linked to an external library.") - .def("make_function", &PythonMLIRModule::makeFunction, - "Defines a new mlir::FuncOp in the current mlir::ModuleOp.") - .def("new_function_context", &PythonMLIRModule::makeNewFunctionContext, - "Defines a new mlir::FuncOp in the mlir::ModuleOp and creates the " - "function context for building the body of the function.") - .def("function_context", &PythonMLIRModule::makeFunctionContext, - "Creates a function context for building the body of an existing " - "function.") - .def("get_function", &PythonMLIRModule::getNamedFunction, - "Looks up the function with the given name in the module.") - .def("make_memref_type", &PythonMLIRModule::makeMemRefType, - "Returns an mlir::MemRefType of an elemental scalar. -1 is used to " - "denote symbolic dimensions in the resulting memref shape.") - .def("make_index_type", &PythonMLIRModule::makeIndexType, - "Returns an mlir::IndexType") - .def("make_type", &PythonMLIRModule::makeType, - "Returns an mlir::Type defined by the IR passed in as the argument.") - .def("get_ir", &PythonMLIRModule::getIR, - "Returns a dump of the MLIR representation of the module. This is " - "used for serde to support out-of-process execution as well as " - "debugging purposes.") - .def("affine_constant_expr", &PythonMLIRModule::affineConstantExpr, - "Returns an affine constant expression.") - .def("affine_symbol_expr", &PythonMLIRModule::affineSymbolExpr, - "Returns an affine symbol expression.") - .def("affine_dim_expr", &PythonMLIRModule::affineDimExpr, - "Returns an affine dim expression.") - .def("affine_constant_map", &PythonMLIRModule::affineConstantMap, - "Returns an affine map with single constant result.") - .def("affine_map", &PythonMLIRModule::affineMap, "Returns an affine map.", - py::arg("dimCount"), py::arg("symbolCount"), py::arg("results")) - .def("__str__", &PythonMLIRModule::getIR, - "Get the string representation of the module"); - - py::class_( - m, "FunctionContext", "A wrapper around mlir::edsc::ScopedContext") - .def(py::init()) - .def("__enter__", &PythonFunctionContext::enter) - .def("__exit__", &PythonFunctionContext::exit); - - { - using namespace mlir::edsc::op; - py::class_(m, "Value", "Wraps mlir::Value") - .def(py::init()) - .def("__add__", - [](PythonValueHandle lhs, PythonValueHandle rhs) - -> PythonValueHandle { return lhs.value + rhs.value; }) - .def("__sub__", - [](PythonValueHandle lhs, PythonValueHandle rhs) - -> PythonValueHandle { return lhs.value - rhs.value; }) - .def("__mul__", - [](PythonValueHandle lhs, PythonValueHandle rhs) - -> PythonValueHandle { return lhs.value * rhs.value; }) - .def("__div__", - [](PythonValueHandle lhs, PythonValueHandle rhs) - -> PythonValueHandle { return lhs.value / rhs.value; }) - .def("__truediv__", - [](PythonValueHandle lhs, PythonValueHandle rhs) - -> PythonValueHandle { return lhs.value / rhs.value; }) - .def("__floordiv__", - [](PythonValueHandle lhs, PythonValueHandle rhs) - -> PythonValueHandle { return floorDiv(lhs, rhs); }) - .def("__mod__", - [](PythonValueHandle lhs, PythonValueHandle rhs) - -> PythonValueHandle { return lhs.value % rhs.value; }) - .def("__lt__", - [](PythonValueHandle lhs, - PythonValueHandle rhs) -> PythonValueHandle { - return mlir::edsc::ValueBuilder(CmpIPredicate::slt, - lhs.value, rhs.value); - }) - .def("__le__", - [](PythonValueHandle lhs, - PythonValueHandle rhs) -> PythonValueHandle { - return mlir::edsc::ValueBuilder(CmpIPredicate::sle, - lhs.value, rhs.value); - }) - .def("__gt__", - [](PythonValueHandle lhs, - PythonValueHandle rhs) -> PythonValueHandle { - return mlir::edsc::ValueBuilder(CmpIPredicate::sgt, - lhs.value, rhs.value); - }) - .def("__ge__", - [](PythonValueHandle lhs, - PythonValueHandle rhs) -> PythonValueHandle { - return mlir::edsc::ValueBuilder(CmpIPredicate::sge, - lhs.value, rhs.value); - }) - .def("__eq__", - [](PythonValueHandle lhs, - PythonValueHandle rhs) -> PythonValueHandle { - return mlir::edsc::ValueBuilder(CmpIPredicate::eq, - lhs.value, rhs.value); - }) - .def("__ne__", - [](PythonValueHandle lhs, - PythonValueHandle rhs) -> PythonValueHandle { - return mlir::edsc::ValueBuilder(CmpIPredicate::ne, - lhs.value, rhs.value); - }) - .def("__invert__", - [](PythonValueHandle handle) -> PythonValueHandle { - return negate(handle.value); - }) - .def("__and__", - [](PythonValueHandle lhs, PythonValueHandle rhs) - -> PythonValueHandle { return lhs.value && rhs.value; }) - .def("__or__", - [](PythonValueHandle lhs, PythonValueHandle rhs) - -> PythonValueHandle { return lhs.value || rhs.value; }) - .def("__call__", &PythonValueHandle::call) - .def("type", &PythonValueHandle::type); - } - - py::class_( - m, "OperationHandle", "A wrapper around mlir::edsc::OperationHandle") - .def(py::init()) - .def("__call__", &PythonOperationHandle::call) - .def("result", &PythonOperationHandle::getResult); - - py::class_( - m, "BlockAppender", - "A dummy class signaling BlockContext to append IR to the given " - "block " - "instead of creating a new block") - .def(py::init()); - py::class_(m, "BlockHandle", - "A wrapper around mlir::edsc::BlockHandle") - .def(py::init()) - .def("arg", &PythonBlockHandle::arg); - - py::class_(m, "BlockContext", - "A wrapper around mlir::edsc::BlockBuilder") - .def(py::init<>()) - .def(py::init &>()) - .def(py::init()) - .def("__enter__", &PythonBlockContext::enter) - .def("__exit__", &PythonBlockContext::exit) - .def("handle", &PythonBlockContext::getHandle); - - py::class_(m, "IndexedValue", - "A wrapper around mlir::edsc::IndexedValue") - .def(py::init()) - .def("load", &PythonIndexedValue::load) - .def("store", &PythonIndexedValue::store); - - py::class_(m, "AffineExpr", - "A wrapper around mlir::AffineExpr") - .def(py::init()) - .def("__add__", - [](PythonAffineExpr lhs, int64_t rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get() + rhs); - }) - .def("__add__", - [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get() + rhs.get()); - }) - .def("__neg__", - [](PythonAffineExpr lhs) -> PythonAffineExpr { - return PythonAffineExpr(-lhs.get()); - }) - .def("__sub__", - [](PythonAffineExpr lhs, int64_t rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get() - rhs); - }) - .def("__sub__", - [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get() - rhs.get()); - }) - .def("__mul__", - [](PythonAffineExpr lhs, int64_t rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get() * rhs); - }) - .def("__mul__", - [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get() * rhs.get()); - }) - .def("__floordiv__", - [](PythonAffineExpr lhs, uint64_t rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get().floorDiv(rhs)); - }) - .def("__floordiv__", - [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get().floorDiv(rhs.get())); - }) - .def("ceildiv", - [](PythonAffineExpr lhs, uint64_t rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get().ceilDiv(rhs)); - }) - .def("ceildiv", - [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get().ceilDiv(rhs.get())); - }) - .def("__mod__", - [](PythonAffineExpr lhs, uint64_t rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get() % rhs); - }) - .def("__mod__", - [](PythonAffineExpr lhs, PythonAffineExpr rhs) -> PythonAffineExpr { - return PythonAffineExpr(lhs.get() % rhs.get()); - }) - .def("compose", - [](PythonAffineExpr self, PythonAffineMap map) -> PythonAffineExpr { - return PythonAffineExpr(self.get().compose(map)); - }) - .def( - "get_constant_value", - [](PythonAffineExpr self) -> py::object { - auto const_expr = self.get().dyn_cast(); - if (const_expr) return py::cast(const_expr.getValue()); - return py::none(); - }, - "Returns the constant value for the affine expression if any, or " - "returns None.") - .def("__str__", &PythonAffineExpr::str); - - py::class_(m, "AffineMap", - "A wrapper around mlir::AffineMap") - .def(py::init()) - .def("__str__", &PythonAffineMap::str); -} - -} // namespace edsc -} // namespace mlir diff --git a/python/npcomp/mlir_ir.cpp b/python/npcomp/mlir_ir.cpp index d3e166649..ea5e3c63e 100644 --- a/python/npcomp/mlir_ir.cpp +++ b/python/npcomp/mlir_ir.cpp @@ -146,14 +146,14 @@ public: pyResultTypes.end()); SmallVector operands(pyOperands.begin(), pyOperands.end()); - NamedAttributeList attrList; + MutableDictionaryAttr attrList; if (attrs) { auto dictAttrs = attrs->attr.dyn_cast(); if (!dictAttrs) { throw py::raiseValueError( "Expected `attrs` to be a DictionaryAttr"); } - attrList = NamedAttributeList(dictAttrs); + attrList = MutableDictionaryAttr(dictAttrs); } Operation *op = Operation::create(loc, opName, types, operands, attrList); diff --git a/python/npcomp/native.cpp b/python/npcomp/native.cpp index 5372b33ae..a6c618bd9 100644 --- a/python/npcomp/native.cpp +++ b/python/npcomp/native.cpp @@ -15,9 +15,6 @@ namespace mlir { void defineMlirIrModule(py::module m); -namespace edsc { -void defineMlirEdscModule(py::module m); -} // namespace edsc namespace npcomp { namespace python { @@ -67,8 +64,6 @@ PYBIND11_MODULE(native, m) { defineLLVMModule(llvm_m); auto mlir_m = m.def_submodule("mlir", "MLIR interop"); - auto mlir_edsc_m = mlir_m.def_submodule("edsc"); - edsc::defineMlirEdscModule(mlir_edsc_m); auto mlir_ir_m = mlir_m.def_submodule("ir"); defineMlirIrModule(mlir_ir_m); }