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))