def testIndexedValue(self): memrefType = self.module.make_memref_type(self.f32Type, [10, 42]) with self.module.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)]) code = str(fun) self.assertIn('"affine.for"()', code) self.assertIn( "{lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (10)}", code) self.assertIn('"affine.for"()', code) self.assertIn( "{lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (42)}", code) self.assertIn("%0 = load %arg0[%i0, %i1] : memref<10x42xf32>", code) self.assertIn("%1 = addf %0, %cst : f32", code) self.assertIn("store %1, %arg0[%i0, %i1] : memref<10x42xf32>", code)
def testLoopNestContext(self): 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 code = str(fun) self.assertIn( ' "affine.for"() {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (5)} : () -> () {\n', code) self.assertIn(" ^bb1(%i0: index):", code) self.assertIn( ' "affine.for"() {lower_bound: () -> (1), step: 3 : index, upper_bound: () -> (15)} : () -> () {\n', code) self.assertIn(" ^bb2(%i1: index):", code) self.assertIn( ' "affine.for"() {lower_bound: () -> (2), step: 5 : index, upper_bound: () -> (25)} : () -> () {\n', code) self.assertIn(" ^bb3(%i2: index):", code) self.assertIn( ' "affine.for"() {lower_bound: () -> (3), step: 7 : index, upper_bound: () -> (35)} : () -> () {\n', code) self.assertIn(" ^bb4(%i3: index):", code) self.assertIn( ' %2 = "affine.apply"(%i0, %i1, %i2, %i3) {map: (d0, d1, d2, d3) -> (d0 + d1 + d2 + d3)} : (index, index, index, index) -> index', code)
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))
def testIndexedValue(self): self.setUp() memrefType = self.module.make_memref_type(self.f32Type, [10, 42]) with self.module.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))
def testMatrixMultiply(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([]) printWithCurrentFunctionName(str(fun))
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))
def testMatrixMultiply(self): 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([]) code = str(fun) self.assertIn( '"affine.for"() {lower_bound: () -> (0), step: 1 : index, upper_bound: () -> (32)} : () -> ()', code) self.assertIn("%0 = load %arg0[%i0, %i2] : memref<32x32xf32>", code) self.assertIn("%1 = load %arg1[%i2, %i1] : memref<32x32xf32>", code) self.assertIn("%2 = mulf %0, %1 : f32", code) self.assertIn("store %2, %arg2[%i0, %i1] : memref<32x32xf32>", code)