f345f7e30b
Allow pointwise operations to take rank zero input tensors similarly to scalar inputs. Use an empty indexing map to broadcast rank zero tensors to the iteration domain of the operation. Depends On D120734 Reviewed By: nicolasvasilache Differential Revision: https://reviews.llvm.org/D120807
612 lines
20 KiB
Python
612 lines
20 KiB
Python
# RUN: %PYTHON %s 2>&1 | FileCheck %s
|
|
|
|
import ctypes
|
|
import sys
|
|
from mlir.ir import *
|
|
from mlir.dialects import builtin
|
|
from mlir.dialects import func
|
|
from mlir.dialects import linalg
|
|
from mlir.passmanager import *
|
|
from mlir.execution_engine import *
|
|
|
|
from mlir.dialects.linalg.opdsl.lang import *
|
|
|
|
|
|
# Log everything to stderr and flush so that we have a unified stream to match
|
|
# errors/info emitted by MLIR to stderr.
|
|
def log(*args):
|
|
print(*args, file=sys.stderr)
|
|
sys.stderr.flush()
|
|
|
|
|
|
elemwise_boiler = """
|
|
func @main() -> f32 attributes {llvm.emit_c_interface} {
|
|
%v0 = arith.constant 0.0 : f32
|
|
%v1 = arith.constant 1.0 : f32
|
|
%v2 = arith.constant 2.0 : f32
|
|
|
|
%lhs = memref.alloc() : memref<f32>
|
|
%rhs = memref.alloc() : memref<4x8xf32>
|
|
%O0 = memref.alloc() : memref<4x8xf32>
|
|
%O1 = memref.alloc() : memref<4x8xf32>
|
|
linalg.fill(%v1, %lhs) : f32, memref<f32>
|
|
linalg.fill(%v2, %rhs) : f32, memref<4x8xf32>
|
|
linalg.fill(%v0, %O0) : f32, memref<4x8xf32>
|
|
linalg.fill(%v0, %O1) : f32, memref<4x8xf32>
|
|
|
|
call @elemwise_exp_add_on_buffers(%lhs, %rhs, %O0) :
|
|
(memref<f32>, memref<4x8xf32>, memref<4x8xf32>) -> ()
|
|
call @elemwise_log_mul_on_buffers(%lhs, %rhs, %O1) :
|
|
(memref<f32>, memref<4x8xf32>, memref<4x8xf32>) -> ()
|
|
|
|
%c0 = arith.constant 0 : index
|
|
%res0 = memref.load %O0[%c0, %c0] : memref<4x8xf32>
|
|
%res1 = memref.load %O1[%c0, %c0] : memref<4x8xf32>
|
|
|
|
%0 = arith.addf %res0, %res1 : f32
|
|
|
|
// TODO: FFI-based solution to allow testing and printing with python code.
|
|
return %0 : f32
|
|
}
|
|
"""
|
|
|
|
matmul_boiler = """
|
|
func @main() -> f32 attributes {llvm.emit_c_interface} {
|
|
%v0 = arith.constant 0.0 : f32
|
|
%v1 = arith.constant -1 : i8
|
|
%v2 = arith.constant 2.0 : f32
|
|
|
|
%A = memref.alloc() : memref<4x16xi8>
|
|
%B = memref.alloc() : memref<16x8xf32>
|
|
%C0 = memref.alloc() : memref<4x8xf32>
|
|
%C1 = memref.alloc() : memref<4x8xf32>
|
|
linalg.fill(%v1, %A) : i8, memref<4x16xi8>
|
|
linalg.fill(%v2, %B) : f32, memref<16x8xf32>
|
|
linalg.fill(%v0, %C0) : f32, memref<4x8xf32>
|
|
linalg.fill(%v0, %C1) : f32, memref<4x8xf32>
|
|
|
|
call @matmul_signed_on_buffers(%A, %B, %C0) :
|
|
(memref<4x16xi8>, memref<16x8xf32>, memref<4x8xf32>) -> ()
|
|
call @matmul_unsigned_on_buffers(%A, %B, %C1) :
|
|
(memref<4x16xi8>, memref<16x8xf32>, memref<4x8xf32>) -> ()
|
|
|
|
%c0 = arith.constant 0 : index
|
|
%res0 = memref.load %C0[%c0, %c0] : memref<4x8xf32>
|
|
%res1 = memref.load %C1[%c0, %c0] : memref<4x8xf32>
|
|
|
|
%0 = arith.addf %res0, %res1 : f32
|
|
|
|
// TODO: FFI-based solution to allow testing and printing with python code.
|
|
return %0 : f32
|
|
}
|
|
"""
|
|
|
|
fill_boiler = """
|
|
func @main() -> i32 attributes {llvm.emit_c_interface} {
|
|
%O0 = memref.alloc() : memref<i32>
|
|
%O1 = memref.alloc() : memref<16xi32>
|
|
%O2 = memref.alloc() : memref<4x16xi32>
|
|
|
|
%val0 = arith.constant 1.0 : f32
|
|
%val1 = arith.constant 2.0 : f32
|
|
%val2 = arith.constant 3.0 : f32
|
|
|
|
call @fill_0d_on_buffers(%val0, %O0) : (f32, memref<i32>) -> ()
|
|
call @fill_1d_on_buffers(%val1, %O1) : (f32, memref<16xi32>) -> ()
|
|
call @fill_2d_on_buffers(%val2, %O2) : (f32, memref<4x16xi32>) -> ()
|
|
|
|
%c0 = arith.constant 0 : index
|
|
%res0 = memref.load %O0[] : memref<i32>
|
|
%c8 = arith.constant 8 : index
|
|
%res1 = memref.load %O1[%c8] : memref<16xi32>
|
|
%c2 = arith.constant 2 : index
|
|
%res2 = memref.load %O2[%c2, %c8] : memref<4x16xi32>
|
|
|
|
%0 = arith.addi %res0, %res1 : i32
|
|
%1 = arith.addi %0, %res2 : i32
|
|
|
|
// TODO: FFI-based solution to allow testing and printing with python code.
|
|
return %1 : i32
|
|
}
|
|
"""
|
|
|
|
fill_rng_boiler = """
|
|
func @main() -> i32 attributes {llvm.emit_c_interface} {
|
|
%O = memref.alloc() : memref<4x16xi32>
|
|
%min = arith.constant -1000.0 : f64
|
|
%max = arith.constant 1000.0 : f64
|
|
%seed = arith.constant 42 : i32
|
|
|
|
call @fill_rng_on_buffers(%min, %max, %seed, %O) :
|
|
(f64, f64, i32, memref<4x16xi32>) -> ()
|
|
|
|
%c0 = arith.constant 0 : index
|
|
%0 = memref.load %O[%c0, %c0] : memref<4x16xi32>
|
|
|
|
// TODO: FFI-based solution to allow testing and printing with python code.
|
|
return %0 : i32
|
|
}
|
|
"""
|
|
|
|
conv_boiler = """
|
|
func @main() -> i32 attributes {llvm.emit_c_interface} {
|
|
%v0 = arith.constant 0 : i32
|
|
%v1 = arith.constant 1.0 : f64
|
|
%v2 = arith.constant 2.0 : f64
|
|
|
|
%input = memref.alloc() : memref<1x4x16x1xf64>
|
|
%filter = memref.alloc() : memref<2x2x1xf64>
|
|
%output = memref.alloc() : memref<1x2x4x1xi32>
|
|
linalg.fill(%v1, %input) : f64, memref<1x4x16x1xf64>
|
|
linalg.fill(%v2, %filter) : f64, memref<2x2x1xf64>
|
|
linalg.fill(%v0, %output) : i32, memref<1x2x4x1xi32>
|
|
|
|
call @conv_on_buffers(%input, %filter, %output) :
|
|
(memref<1x4x16x1xf64>, memref<2x2x1xf64>, memref<1x2x4x1xi32>) -> ()
|
|
|
|
%c0 = arith.constant 0 : index
|
|
%0 = memref.load %output[%c0, %c0, %c0, %c0] : memref<1x2x4x1xi32>
|
|
|
|
// TODO: FFI-based solution to allow testing and printing with python code.
|
|
return %0 : i32
|
|
}
|
|
"""
|
|
|
|
pooling_boiler = """
|
|
func @main() -> i32 attributes {llvm.emit_c_interface} {
|
|
%v0 = arith.constant 0 : i32
|
|
%v42 = arith.constant 42.0 : f64
|
|
%v77 = arith.constant 77.0 : f64
|
|
%v-13 = arith.constant -13.0 : f64
|
|
%v1 = arith.constant 1.0 : f64
|
|
|
|
%input = memref.alloc() : memref<1x4x16x1xf64>
|
|
%shape = memref.alloc() : memref<2x2xf64>
|
|
%output = memref.alloc() : memref<1x2x4x1xi32>
|
|
linalg.fill(%v1, %input) : f64, memref<1x4x16x1xf64>
|
|
linalg.fill(%v1, %shape) : f64, memref<2x2xf64>
|
|
linalg.fill(%v0, %output) : i32, memref<1x2x4x1xi32>
|
|
|
|
%c0 = arith.constant 0 : index
|
|
%c1 = arith.constant 1 : index
|
|
%c2 = arith.constant 2 : index
|
|
memref.store %v42, %input[%c0, %c0, %c0, %c0] : memref<1x4x16x1xf64>
|
|
memref.store %v77, %input[%c0, %c0, %c1, %c0] : memref<1x4x16x1xf64>
|
|
memref.store %v-13, %input[%c0, %c1, %c0, %c0] : memref<1x4x16x1xf64>
|
|
|
|
call @pooling_on_buffers(%input, %shape, %output) :
|
|
(memref<1x4x16x1xf64>, memref<2x2xf64>, memref<1x2x4x1xi32>) -> ()
|
|
|
|
%0 = memref.load %output[%c0, %c0, %c0, %c0] : memref<1x2x4x1xi32>
|
|
|
|
// TODO: FFI-based solution to allow testing and printing with python code.
|
|
return %0 : i32
|
|
}
|
|
"""
|
|
|
|
|
|
def transform(module, boilerplate):
|
|
import mlir.conversions
|
|
import mlir.all_passes_registration
|
|
import mlir.transforms
|
|
|
|
# TODO: Allow cloning functions from one module to another.
|
|
# Atm we have to resort to string concatenation.
|
|
ops = module.operation.regions[0].blocks[0].operations
|
|
mod = Module.parse("\n".join([str(op) for op in ops]) + boilerplate)
|
|
|
|
pm = PassManager.parse(
|
|
"builtin.func(convert-linalg-to-loops, lower-affine, " +
|
|
"convert-math-to-llvm, convert-scf-to-cf, arith-expand, memref-expand), "
|
|
+ "convert-vector-to-llvm, convert-memref-to-llvm, convert-func-to-llvm," +
|
|
"reconcile-unrealized-casts")
|
|
pm.run(mod)
|
|
return mod
|
|
|
|
|
|
def test_elemwise_builtin():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
i8 = IntegerType.get_signless(8)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((), f32), MemRefType.get((4, 8), f32),
|
|
MemRefType.get((4, 8), f32))
|
|
def elemwise_exp_add_on_buffers(lhs, rhs, out):
|
|
linalg.elemwise_unary(lhs, outs=[out])
|
|
linalg.elemwise_binary(out, rhs, outs=[out])
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((), f32), MemRefType.get((4, 8), f32),
|
|
MemRefType.get((4, 8), f32))
|
|
def elemwise_log_mul_on_buffers(lhs, rhs, out):
|
|
linalg.elemwise_unary(lhs, outs=[out], fun=UnaryFn.log)
|
|
linalg.elemwise_binary(out, rhs, outs=[out], fun=BinaryFn.mul)
|
|
|
|
execution_engine = ExecutionEngine(transform(module, elemwise_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result f32.
|
|
# Arguments must be passed as pointers.
|
|
c_float_p = ctypes.c_float * 1
|
|
res = c_float_p(-1.)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# elemwise_exp_add_on_buffers: exp(1.0) + 2.0 = 4.71828182846
|
|
# elemwise_log_mul_on_buffers: log(1.0) * 2.0 = 0.0
|
|
# CHECK: RESULT: 4.71828
|
|
|
|
|
|
test_elemwise_builtin()
|
|
|
|
|
|
def test_elemwise_generic():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
i8 = IntegerType.get_signless(8)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((), f32), MemRefType.get((4, 8), f32),
|
|
MemRefType.get((4, 8), f32))
|
|
def elemwise_exp_add_on_buffers(lhs, rhs, out):
|
|
linalg.elemwise_unary(lhs, outs=[out], emit_generic=True)
|
|
linalg.elemwise_binary(out, rhs, outs=[out], emit_generic=True)
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((), f32), MemRefType.get((4, 8), f32),
|
|
MemRefType.get((4, 8), f32))
|
|
def elemwise_log_mul_on_buffers(lhs, rhs, out):
|
|
linalg.elemwise_unary(
|
|
lhs, outs=[out], fun=UnaryFn.log, emit_generic=True)
|
|
linalg.elemwise_binary(
|
|
out, rhs, outs=[out], fun=BinaryFn.mul, emit_generic=True)
|
|
|
|
execution_engine = ExecutionEngine(transform(module, elemwise_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result f32.
|
|
# Arguments must be passed as pointers.
|
|
c_float_p = ctypes.c_float * 1
|
|
res = c_float_p(-1.)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# elemwise_exp_add_on_buffers: exp(1.0) + 2.0 = 4.71828182846
|
|
# elemwise_log_mul_on_buffers: log(1.0) * 2.0 = 0.0
|
|
# CHECK: RESULT: 4.71828
|
|
|
|
|
|
test_elemwise_generic()
|
|
|
|
|
|
def test_matmul_builtin():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
i8 = IntegerType.get_signless(8)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((4, 16), i8), MemRefType.get((16, 8), f32),
|
|
MemRefType.get((4, 8), f32))
|
|
def matmul_signed_on_buffers(lhs, rhs, out):
|
|
linalg.matmul(lhs, rhs, outs=[out])
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((4, 16), i8), MemRefType.get((16, 8), f32),
|
|
MemRefType.get((4, 8), f32))
|
|
def matmul_unsigned_on_buffers(lhs, rhs, out):
|
|
linalg.matmul(lhs, rhs, outs=[out], cast=TypeFn.cast_unsigned)
|
|
|
|
execution_engine = ExecutionEngine(transform(module, matmul_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result f32.
|
|
# Arguments must be passed as pointers.
|
|
c_float_p = ctypes.c_float * 1
|
|
res = c_float_p(-1.)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# matmul_signed_on_buffers: -1 * 2.0 * 16 = -32
|
|
# matmul_unsigned_on_buffers: (2^8-1) * 2.0 * 16 = 8160
|
|
# CHECK: RESULT: 8128
|
|
|
|
|
|
test_matmul_builtin()
|
|
|
|
|
|
def test_matmul_generic():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
i8 = IntegerType.get_signless(8)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((4, 16), i8), MemRefType.get((16, 8), f32),
|
|
MemRefType.get((4, 8), f32))
|
|
def matmul_signed_on_buffers(lhs, rhs, out):
|
|
linalg.matmul(lhs, rhs, outs=[out], emit_generic=True)
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((4, 16), i8), MemRefType.get((16, 8), f32),
|
|
MemRefType.get((4, 8), f32))
|
|
def matmul_unsigned_on_buffers(lhs, rhs, out):
|
|
linalg.matmul(
|
|
lhs, rhs, outs=[out], cast=TypeFn.cast_unsigned, emit_generic=True)
|
|
|
|
execution_engine = ExecutionEngine(transform(module, matmul_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result f32.
|
|
# Arguments must be passed as pointers.
|
|
c_float_p = ctypes.c_float * 1
|
|
res = c_float_p(-1.)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# matmul_signed_on_buffers = -1 * 2.0 * 16 = -32
|
|
# matmul_unsigned_on_buffers = (2^8-1) * 2.0 * 16 = 8160
|
|
# CHECK: RESULT: 8128
|
|
|
|
|
|
test_matmul_generic()
|
|
|
|
|
|
def test_fill_builtin():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
i32 = IntegerType.get_signless(32)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(f32, MemRefType.get([], i32))
|
|
def fill_0d_on_buffers(value, out):
|
|
linalg.fill_tensor(value, outs=[out])
|
|
|
|
@builtin.FuncOp.from_py_func(f32, MemRefType.get([16], i32))
|
|
def fill_1d_on_buffers(value, out):
|
|
linalg.fill_tensor(value, outs=[out])
|
|
|
|
@builtin.FuncOp.from_py_func(f32, MemRefType.get([4, 16], i32))
|
|
def fill_2d_on_buffers(value, out):
|
|
linalg.fill_tensor(value, outs=[out])
|
|
|
|
execution_engine = ExecutionEngine(transform(module, fill_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result i32.
|
|
# Arguments must be passed as pointers.
|
|
c_int_p = ctypes.c_int * 1
|
|
res = c_int_p(-1)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# CHECK: RESULT: 6
|
|
|
|
|
|
test_fill_builtin()
|
|
|
|
|
|
def test_fill_generic():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
i32 = IntegerType.get_signless(32)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(f32, MemRefType.get([], i32))
|
|
def fill_0d_on_buffers(value, out):
|
|
linalg.fill_tensor(value, outs=[out], emit_generic=True)
|
|
|
|
@builtin.FuncOp.from_py_func(f32, MemRefType.get([16], i32))
|
|
def fill_1d_on_buffers(value, out):
|
|
linalg.fill_tensor(value, outs=[out], emit_generic=True)
|
|
|
|
@builtin.FuncOp.from_py_func(f32, MemRefType.get([4, 16], i32))
|
|
def fill_2d_on_buffers(value, out):
|
|
linalg.fill_tensor(value, outs=[out], emit_generic=True)
|
|
|
|
execution_engine = ExecutionEngine(transform(module, fill_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result i32.
|
|
# Arguments must be passed as pointers.
|
|
c_int_p = ctypes.c_int * 1
|
|
res = c_int_p(-1)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# CHECK: RESULT: 6
|
|
|
|
|
|
test_fill_generic()
|
|
|
|
|
|
def test_fill_rng_builtin():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f64 = F64Type.get()
|
|
i32 = IntegerType.get_signless(32)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(f64, f64, i32, MemRefType.get((4, 16), i32))
|
|
def fill_rng_on_buffers(min, max, seed, out):
|
|
linalg.fill_rng_2d(min, max, seed, outs=[out])
|
|
|
|
execution_engine = ExecutionEngine(transform(module, fill_rng_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result i32.
|
|
# Arguments must be passed as pointers.
|
|
c_int_p = ctypes.c_int * 1
|
|
res = c_int_p(-1)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# CHECK: RESULT: -480
|
|
|
|
|
|
test_fill_rng_builtin()
|
|
|
|
|
|
def test_fill_rng_generic():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f64 = F64Type.get()
|
|
i32 = IntegerType.get_signless(32)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(f64, f64, i32, MemRefType.get((4, 16), i32))
|
|
def fill_rng_on_buffers(min, max, seed, out):
|
|
linalg.fill_rng_2d(min, max, seed, outs=[out], emit_generic=True)
|
|
|
|
execution_engine = ExecutionEngine(transform(module, fill_rng_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result i32.
|
|
# Arguments must be passed as pointers.
|
|
c_int_p = ctypes.c_int * 1
|
|
res = c_int_p(-1)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# CHECK: RESULT: -480
|
|
|
|
|
|
test_fill_rng_generic()
|
|
|
|
|
|
def test_max_pooling_builtin():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f64 = F64Type.get()
|
|
i32 = IntegerType.get_signless(32)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((1, 4, 16, 1), f64), MemRefType.get((2, 2), f64),
|
|
MemRefType.get((1, 2, 4, 1), i32))
|
|
def pooling_on_buffers(input, shape, output):
|
|
linalg.pooling_nhwc_max(
|
|
input, shape, outs=[output], strides=[2, 4], dilations=[1, 2])
|
|
|
|
execution_engine = ExecutionEngine(transform(module, pooling_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result i32.
|
|
# Arguments must be passed as pointers.
|
|
c_int_p = ctypes.c_int * 1
|
|
res = c_int_p(-1)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# 77 is not selected due to the dilation 2 in the second dimension.
|
|
# CHECK: RESULT: 42
|
|
|
|
|
|
test_max_pooling_builtin()
|
|
|
|
|
|
def test_max_pooling_generic():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f64 = F64Type.get()
|
|
i32 = IntegerType.get_signless(32)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((1, 4, 16, 1), f64), MemRefType.get((2, 2), f64),
|
|
MemRefType.get((1, 2, 4, 1), i32))
|
|
def pooling_on_buffers(input, shape, output):
|
|
linalg.pooling_nhwc_max(
|
|
input,
|
|
shape,
|
|
outs=[output],
|
|
strides=[2, 4],
|
|
dilations=[1, 2],
|
|
emit_generic=True)
|
|
|
|
execution_engine = ExecutionEngine(transform(module, pooling_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result i32.
|
|
# Arguments must be passed as pointers.
|
|
c_int_p = ctypes.c_int * 1
|
|
res = c_int_p(-1)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# 77 is not selected due to the dilation 2 in the second dimension.
|
|
# CHECK: RESULT: 42
|
|
|
|
|
|
test_max_pooling_generic()
|
|
|
|
|
|
def test_min_pooling_builtin():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f64 = F64Type.get()
|
|
i32 = IntegerType.get_signless(32)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((1, 4, 16, 1), f64), MemRefType.get((2, 2), f64),
|
|
MemRefType.get((1, 2, 4, 1), i32))
|
|
# Set the strides and use the default dilations.
|
|
def pooling_on_buffers(input, shape, output):
|
|
linalg.pooling_nhwc_min(input, shape, outs=[output], strides=[2, 4])
|
|
|
|
execution_engine = ExecutionEngine(transform(module, pooling_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result i32.
|
|
# Arguments must be passed as pointers.
|
|
c_int_p = ctypes.c_int * 1
|
|
res = c_int_p(-1)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# CHECK: RESULT: -13
|
|
|
|
|
|
test_min_pooling_builtin()
|
|
|
|
|
|
def test_min_pooling_generic():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f64 = F64Type.get()
|
|
i32 = IntegerType.get_signless(32)
|
|
with InsertionPoint(module.body):
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
MemRefType.get((1, 4, 16, 1), f64), MemRefType.get((2, 2), f64),
|
|
MemRefType.get((1, 2, 4, 1), i32))
|
|
# Set the strides and use the default dilations.
|
|
def pooling_on_buffers(input, shape, output):
|
|
linalg.pooling_nhwc_min(
|
|
input, shape, outs=[output], strides=[2, 4], emit_generic=True)
|
|
|
|
execution_engine = ExecutionEngine(transform(module, pooling_boiler))
|
|
|
|
# TODO: FFI-based solution to allow testing and printing with python code.
|
|
# Prepare arguments: one result i32.
|
|
# Arguments must be passed as pointers.
|
|
c_int_p = ctypes.c_int * 1
|
|
res = c_int_p(-1)
|
|
execution_engine.invoke("main", res)
|
|
|
|
log("RESULT: ", res[0])
|
|
# CHECK: RESULT: -13
|
|
|
|
|
|
test_min_pooling_generic()
|