已开启
微重构Clone #3
SorryNaCN创建于  2月6日
SorryNaCN
SorryNaCN成员
2月6日 创建


class SplitCloneHelper :  public dvm::CloneHelper {
  IntArrayRef *GetClone(IntArrayRef *shape) override {
    auto it = ref_map_.find(shape);
    return it != ref_map_.end() ? static_cast<IntArrayRef *>(it->second) : shape;
  }
  ScalarRef *GetClone(ScalarRef *scalar) override {
    auto it = ref_map_.find(scalar);
    return it != ref_map_.end() ? static_cast<ScalarRef *>(it->second) : scalar;
  }
  NDObject *GetClone(NDObject *op) override {
    auto it = op_map_.find(op);
    return it != op_map_.end() ? it->second : nullptr;
  }
  void SetClone(NDObject *op, NDObject *clone) override {
    op_map_[op] = clone;
  }
  std::unordered_map<NDObject *, NDObject *> op_map_;
  std::unordered_map<void *, void *> ref_map_;
};

class ExeGraphSplitKernelPy : public GraphSplitKernelPy {
 std::mutex executor_mutex_;
 std::vector<GraphSplitKernelPy *> executors_;
 
  void Setup() override {
    SetupRelocs();
    kernel_.Infer(); 
    executors_->push_back(this);    //当前节点作为第一个执行实例
  }
 
  GraphSplitKernelPy *Clone() {   // 基于当前节点clone其它执行实例
    SplitCloneHelper helper;
    auto kernel = new GraphSplitKernelPy();
    for (auto ref : shapes_) {
      helper[ref] = new ShapeWithRef(ref->size);
    }
    kernel->Clone(this, helper);
    for (auto load: loads_) {
      kernel->loads_.push_back(helper.GetClone(load));
    }
    for (auto store: stores_) {
      kernel->stores_.push_back(helper.GetClone(store));
    }
    kernel->contiguity_flags_ = contiguity_flags_;
    kernel->Setup();
    return kernel;
  }

  py::object Call(py::args inputs) override {
   GraphSplitKernelPy *executor = nullptr;
   {
     std::lock_guard<executor_mutex_> lock;
     if (!executors_.empty()) {
       executor = executors_.back();
       executors_.pop_back();
     }
   }
   if (executor == nullptr) {
     executor = Clone();
   }
   ...
   auto launch_call = [executor]() {
     executor->GraphSplitBase::Launch(...);
     {
       std::lock_guard<executor_mutex_> lock;
       executors_.push_back(executor);
     }
   };
   RunOpApiV2(op_name_, launch_call);
  }
};



likedislike
SorryNaCN
SorryNaCN成员
4月23日 评论:

@arg_mark(plat_marks=['platform_ascend910b'], level_mark='level0', card_mark='onecard', essential_mark='essential')
def test_spec_fused_mul_native_layer_norm_native_layer_norm_backward_19():
t = Tester("vector:spec", use_pass_opt=True)
shape = [32, 77, 512]
mean_shape = [32, 77, 1]

np.random.seed(0)
arg0_data = np.random.normal(0, 1, mean_shape).astype(np.float32)
arg1_storage = np.random.normal(0, 1, [77, 32, 512]).astype(np.float32)
arg1_storage += np.linspace(-0.1, 0.1, 512, dtype=np.float32).reshape(1, 1, 512)
arg2_data = np.random.normal(0, 1, 512).astype(np.float32)
arg3_data = np.random.normal(0, 1, 512).astype(np.float32)

arg0_1 = t.load(arg0_data)
arg1_1 = t.view_load(shape, [512, 16384, 1], arg1_storage)
arg2_1 = t.load(arg2_data)
arg3_1 = t.load(arg3_data)
div = t.div(arg0_1, 512.0)
expand = t.broadcast(div, shape)
sub = t.sub(arg1_1, expand)
pow_1 = t.pow(sub, 2)
sum_1 = t.sum(pow_1, [2], True)
t.spec_next()
expand_1 = t.broadcast(div, shape)
sub_1 = t.sub(arg1_1, expand_1)
div_1 = t.div(sum_1, 512.0)
add = t.add(div_1, 1e-05)
rsqrt = t.div(1, t.sqrt(add))
expand_2 = t.broadcast(rsqrt, shape)
mul = t.mul(sub_1, expand_2)
expand_3 = t.broadcast(arg2_1, shape)
mul_1 = t.mul(mul, expand_3)
expand_4 = t.broadcast(arg3_1, shape)
add_1 = t.add(mul_1, expand_4)
mul_2 = t.mul(add_1, 1)
div_2 = t.div(rsqrt, 512)

arg1_view = np.transpose(arg1_storage, (1, 0, 2))
mean = arg0_data / 512.0
centered = arg1_view - mean
variance_sum = np.sum(np.power(centered, 2), axis=2, keepdims=True)
inv_std = 1 / np.sqrt(variance_sum / 512.0 + 1e-05)
normalized = centered * inv_std
expected_mul_2 = normalized * arg2_data.reshape(1, 1, 512) + arg3_data.reshape(1, 1, 512)
expected_div_2 = inv_std / 512.0

t.store_expect(mul, normalized, 1e-4)
t.store_expect(mul_2, expected_mul_2, 1e-4)
t.store_expect(div_2, expected_div_2, 1e-4)
t.run(True)
assert (t.run_check())
likedislike