From bb5c72301200ddf801f9e9b7e4e3889c579ef43d Mon Sep 17 00:00:00 2001 From: yuanfz98 Date: Fri, 5 Aug 2022 16:25:59 +0200 Subject: [PATCH 1/4] format black --- python/tvm/relay/op/strategy/x86.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/op/strategy/x86.py b/python/tvm/relay/op/strategy/x86.py index 2b063b8bbec3..a56a789c51b0 100644 --- a/python/tvm/relay/op/strategy/x86.py +++ b/python/tvm/relay/op/strategy/x86.py @@ -837,6 +837,7 @@ def concatenate_strategy_cpu(attrs, inputs, out_type, target): ) return strategy + @batch_norm_strategy.register(["cpu"]) def batch_norm_strategy_cpu(attrs, inputs, out_type, target): """batch_norm x86 strategy""" @@ -846,4 +847,4 @@ def batch_norm_strategy_cpu(attrs, inputs, out_type, target): wrap_topi_schedule(topi.x86.schedule_batch_norm), name="batch_norm.cpu", ) - return strategy \ No newline at end of file + return strategy From 46ee230e0f75ed8eb349a4edbcbbee47bf10b3dd Mon Sep 17 00:00:00 2001 From: yuanfz98 Date: Fri, 5 Aug 2022 17:27:52 +0200 Subject: [PATCH 2/4] format black --- python/tvm/topi/x86/nn.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/tvm/topi/x86/nn.py b/python/tvm/topi/x86/nn.py index 53037e39f015..52c1f204fe80 100644 --- a/python/tvm/topi/x86/nn.py +++ b/python/tvm/topi/x86/nn.py @@ -108,18 +108,19 @@ def _callback(op): traverse_inline(s, outs[0].op, _callback) return s + def schedule_batch_norm(outs): s = te.create_schedule([x.op for x in outs]) # only parallelize outer dimensions up to axis - output_op=outs[0].op + output_op = outs[0].op axis = output_op.axis - outer_axes = [output_op.axis[i] for i in range(0, len(axis)-1)] + outer_axes = [output_op.axis[i] for i in range(0, len(axis) - 1)] fused_outer_axes = s[output_op].fuse(*outer_axes) s[output_op].parallel(fused_outer_axes) # when scale or center is enabled - if 'divide' not in output_op.name: + if "divide" not in output_op.name: div = output_op.input_tensors[0] substract = s[div].op.input_tensors[0] s[div].compute_inline() s[substract].compute_inline() - return s \ No newline at end of file + return s From 081fe56cf4c09d767d6e2ebd14eebd0187faff87 Mon Sep 17 00:00:00 2001 From: yuanfz98 Date: Fri, 5 Aug 2022 17:38:27 +0200 Subject: [PATCH 3/4] docstring --- python/tvm/topi/x86/nn.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/python/tvm/topi/x86/nn.py b/python/tvm/topi/x86/nn.py index 52c1f204fe80..7c3942fabc90 100644 --- a/python/tvm/topi/x86/nn.py +++ b/python/tvm/topi/x86/nn.py @@ -110,6 +110,19 @@ def _callback(op): def schedule_batch_norm(outs): + """Schedule for batch_norm + + Parameters + ---------- + outs: Array of Tensor + The computation graph description of batch_norm + in the format of an array of tensors. + + Returns + ------- + sch: Schedule + The computation schedule for the op. + """ s = te.create_schedule([x.op for x in outs]) # only parallelize outer dimensions up to axis output_op = outs[0].op From 4408218bde98ec6fb9a6932fef7b851b37505759 Mon Sep 17 00:00:00 2001 From: yuanfz98 Date: Sat, 6 Aug 2022 14:30:42 +0200 Subject: [PATCH 4/4] typo --- python/tvm/relay/op/strategy/x86.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/op/strategy/x86.py b/python/tvm/relay/op/strategy/x86.py index a56a789c51b0..6575e0f5c5a2 100644 --- a/python/tvm/relay/op/strategy/x86.py +++ b/python/tvm/relay/op/strategy/x86.py @@ -843,7 +843,7 @@ def batch_norm_strategy_cpu(attrs, inputs, out_type, target): """batch_norm x86 strategy""" strategy = _op.OpStrategy() strategy.add_implementation( - wrap_compute_concat(topi.nn.batch_norm), + wrap_compute_batch_norm(topi.nn.batch_norm), wrap_topi_schedule(topi.x86.schedule_batch_norm), name="batch_norm.cpu", )