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 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 __codegen(self, op_name, inputs, attrs):
types = [common.Float(32)]
strategy_map = framework.Operator.get_op_attrs("CINNStrategy")
res = strategy_map.apply_strategy(op_name, attrs, inputs, types,
self.target)
stages = create_stages(res)
func = lang.lower(op_name, stages, res)
logging.warning('func:\n\n%s\n', func)
builder = lang.Module.Builder(op_name, self.target)
builder.add_function(func)
return builder.build()
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))