def create_matmul_basic(target, m, n, k):
m, n, k = [ir.Expr(_) for _ in (m, n, k)]
a = lang.Placeholder("float32", "A", [m, k])
b = lang.Placeholder("float32", "B", [k, n])
c_init = lang.compute([m, n], lambda v: ir.Expr(0.), "c_init")
k1 = ir.Var(k.as_int32(), "k1")
c = lang.compute([m, n], lambda v: lang.sum(
a(v[0], k1.to_expr_mutable()) * b(k1.to_expr_mutable(), v[1])), "c",
[k1])
stages = create_stages([c_init, c])
c_stage = stages[c]
stages[c_init].share_buffer_with(c_stage)
stages[c].ctrl_depend(c_init)
builder = lang.Module.Builder("matmul", target)
ts = [a.to_tensor(), b.to_tensor(), c_init, c]
func = lang.lower("matmul", stages, ts)
print('func', func)
builder.add_function(func)
return builder.build()
def reduction_tester(self, fn_name, cinn_fn, np_fn, axes, keep_dims,
initial):
m, n = [ir.Expr(_) for _ in (
self.m,
self.n,
)]
axes_expr = [ir.Expr(_) for _ in axes]
x = lang.Placeholder("float32", "x", [m, n])
func_name = "test_" + fn_name
stages = create_stages([x.to_tensor()])
if initial:
y, y_init = cinn_fn(x.to_tensor(), stages, axes_expr, keep_dims,
ir.Expr(initial))
func = lang.lower(func_name, stages, [x.to_tensor(), y, y_init])
else:
y = cinn_fn(x.to_tensor(), stages, axes_expr, keep_dims)
func = lang.lower(func_name, stages, [x.to_tensor(), y])
builder = lang.Module.Builder("reduction_module", self.target)
builder.add_function(func)
print(func)
module = builder.build()
self.compiler.build(module)
fn = self.compiler.lookup(func_name)
x_data, x_buf, out_buf, *args = self.create_data(axes, keep_dims)
fn(args)
self.assertTrue(
np.allclose(out_buf.numpy(),
self.create_target_data(x_data, np_fn, axes, keep_dims,
initial),
atol=1e-4))
def test_cxx_register(self):
add_int = ir.Registry.get("test_add_int64")
self.assertEqual(add_int(2, 3), 5)
add_expr = ir.Registry.get("test_add_expr")
x = ir.Expr(1)
y = ir.Expr(2)
z = x + y
r = add_expr(x, y)
self.assertEqual(r.node_type(), z.node_type())
mul_float = ir.Registry.get("test_mul_float")
self.assertTrue(isclose(mul_float(2.4, 2.5), 6.0, abs_tol=1e-5))
def union_tester(self, fn_name, cinn_fn, np_fn, dtype="float32"):
m, n = [ir.Expr(_) for _ in (
self.m,
self.n,
)]
x = lang.Placeholder(dtype, "x", [m, n])
y = cinn_fn(x.to_tensor())
func_name = "test_" + fn_name
stages = create_stages([x.to_tensor(), y])
func = lang.lower(func_name, stages, [x.to_tensor(), y])
builder = lang.Module.Builder("elementwise_module", self.target)
builder.add_function(func)
module = builder.build()
self.compiler.build(module)
fn = self.compiler.lookup(func_name)
x_data, x_buf, out_buf, *args = self.create_data(dtype)
fn(args)
self.assertTrue(
np.allclose(
out_buf.numpy(),
self.create_target_data(x_data, np_fn),
atol=1e-4))
def to_test_op(self, input_shapes, output_shape, op_name, attrs):
'''
Test the operator.
'''
self.compiler = cinn.Compiler.create(self.target)
inputs = []
inputs_data = []
for i_shape in input_shapes:
expr_shape = []
inputs_data.append(
np.around(np.random.random(i_shape).astype("float32"), 3))
for dim_shape in i_shape:
expr_shape.append(ir.Expr(dim_shape))
inputs.append(
lang.Placeholder("float32", self.__gen_var_name(),
expr_shape).to_tensor())
args = []
temp_inputs = []
for in_data in inputs_data:
temp_inputs.append(
runtime.cinn_buffer_t(in_data, runtime.cinn_x86_device))
for in_data in temp_inputs:
args.append(runtime.cinn_pod_value_t(in_data))
if output_shape == None:
correct_result, output_shape = self.create_target_data(
inputs_data, attrs)
else:
correct_result = self.create_target_data(inputs_data, attrs)
module = self.__codegen(op_name, inputs, attrs)
self.compiler.build(module)
fn = self.compiler.lookup(op_name)
out = []
for out_shape in output_shape:
out.append(
runtime.cinn_buffer_t(
np.zeros(out_shape).astype("float32"),
runtime.cinn_x86_device))
for out_data in out:
args.append(runtime.cinn_pod_value_t(out_data))
fn(args)
out_result = out[len(out) - 1].numpy()
self.assertTrue(np.allclose(out_result, correct_result, atol=1e-4))
def transform_matmul_tester(self, fn_name, cinn_fn, np_fn, trans_a,
trans_b, x_num_col_dims, y_num_col_dims):
m, n, k = [ir.Expr(_) for _ in (
self.m,
self.n,
self.k,
)]
x_shape_expr = [k, m] if trans_a else [m, k]
y_shape_expr = [n, k] if trans_b else [k, n]
x = lang.Placeholder("float32", "x", x_shape_expr)
y = lang.Placeholder("float32", "y", y_shape_expr)
func_name = "test_" + fn_name
z = cinn_fn(x.to_tensor(), y.to_tensor(), trans_a, trans_b,
x_num_col_dims, y_num_col_dims)
stages = create_stages([z])
func = lang.lower(func_name, stages, [x.to_tensor(), y.to_tensor(), z])
print(func)
builder = lang.Module.Builder("transform_module", self.target)
builder.add_function(func)
module = builder.build()
self.compiler.build(module)
fn = self.compiler.lookup(func_name)
x_data, y_data, x_buf, y_buf, out_buf, *args = self.create_data(
(self.m, self.n), trans_a, trans_b)
fn(args)
self.assertTrue(
np.allclose(out_buf.numpy(),
self.create_target_data(np_fn, x_data, y_data, trans_a,
trans_b),
atol=1e-4))