Add cdist op

docs/source/oneflow.rst

@@ -387,6 +387,7 @@ Other Ops
  adaptive_avg_pool3d
  broadcast_like 
  cast
+ cdist
  cumprod 
  cumsum 
  diag 

oneflow/core/autograd/gradient_funcs/cdist.cpp

@@ -0,0 +1,96 @@
+/*
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include "oneflow/core/framework/attr_map.h"
+#include "oneflow/core/framework/op_expr_grad_function.h"
+#include "oneflow/core/framework/op_builder.h"
+#include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
+#include "oneflow/core/framework/op_expr.h"
+#include "oneflow/core/functional/functional.h"
+
+namespace oneflow {
+namespace one {
+
+namespace {
+
+struct CDistCaptureState : public AutoGradCaptureState {
+ bool requires_grad = false;
+ size_t x1_index = 0;
+ size_t x2_index = 0;
+ size_t out_index = 0;
+ double p = 0.0;
+};
+
+class CDistGrad : public OpExprGradFunction<CDistCaptureState> {
+ public:
+ virtual ~CDistGrad() = default;
+
+ using OpExprGradFunction<CDistCaptureState>::Init;
+
+ Maybe<void> Init(const OpExpr& op) override;
+ Maybe<void> Capture(CDistCaptureState* ctx, const TensorTuple& inputs, const TensorTuple& outputs,
+ const AttrMap& attrs) const override;
+ Maybe<void> Apply(const CDistCaptureState* ctx, const TensorTuple& out_grads,
+ TensorTuple* in_grads) const override;
+
+ private:
+ AttrMap base_attrs_;
+};
+
+Maybe<void> CDistGrad::Init(const OpExpr& op) {
+ const auto* fw_op_expr = dynamic_cast<const UserOpExpr*>(&op);
+ CHECK_NOTNULL_OR_RETURN(fw_op_expr); // NOLINT(maybe-need-error-msg)
+ base_attrs_ = MakeAttrMapFromUserOpConf(fw_op_expr->proto());
+ return Maybe<void>::Ok();
+}
+
+Maybe<void> CDistGrad::Capture(CDistCaptureState* ctx, const TensorTuple& inputs,
+ const TensorTuple& outputs, const AttrMap& attrs) const {
+ ctx->requires_grad = inputs.at(0)->requires_grad();
+ if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
+
+ ctx->x1_index = ctx->SaveTensorForBackward(inputs.at(0));
+ ctx->x2_index = ctx->SaveTensorForBackward(inputs.at(1));
+ ctx->out_index = ctx->SaveTensorForBackward(outputs.at(0));
+ ComposedAttrMap composed_attrs(attrs, base_attrs_);
+ ctx->p = JUST(composed_attrs.GetAttr<double>("p"));
+
+ return Maybe<void>::Ok();
+}
+
+Maybe<void> CDistGrad::Apply(const CDistCaptureState* ctx, const TensorTuple& out_grads,
+ TensorTuple* in_grads) const {
+ if (!ctx->requires_grad) { return Maybe<void>::Ok(); }
+ CHECK_LE_OR_RETURN(out_grads.size(), 1); // NOLINT(maybe-need-error-msg)
+
+ const auto& x1 = ctx->SavedTensors().at(ctx->x1_index);
+ const auto& x2 = ctx->SavedTensors().at(ctx->x2_index);
+ const auto& out = ctx->SavedTensors().at(ctx->out_index);
+ const double p = ctx->p;
+
+ in_grads->resize(2);
+ auto results = JUST(functional::CDistGrad(x1, x2, out, out_grads.at(0), p));
+ (*in_grads)[0] = results->at(0);
+ (*in_grads)[1] = results->at(1);
+ return Maybe<void>::Ok();
+}
+
+} // namespace
+
+REGISTER_OP_EXPR_GRAD_FUNCTION("cdist", CDistGrad);
+
+} // namespace one
+} // namespace oneflow

oneflow/core/functional/functional_api.yaml

@@ -2612,6 +2612,14 @@
  signature: "Tensor (Tensor x, Tensor y, Int32 dim=1, Double eps=1e-8) => CosineSimilarity"
  bind_python: True
 
+- name: "cdist"
+ signature: 'Tensor (Tensor x1, Tensor x2, Double p=2.0, String compute_mode=None) => CDist'
+ bind_python: True
+
+- name: "cdist_grad"
+ signature: "TensorTuple (Tensor x1, Tensor x2, Tensor out, Tensor dy, Double p=2.0) => CDistGrad"
+ bind_python: False
+
 - name: "normalize"
  signature: "Tensor (Tensor input, Float p=2.0, Int32 dim=1, Float eps=1e-12, Bool use_l2_norm_kernel=True) => Normalize"
  bind_python: True

oneflow/core/functional/impl/array_functor.cpp

@@ -3801,6 +3801,7 @@ class BroadcastTensorsFunctor {
  return outputs;
  }
 };
+
 class BinCountFunctor {
  public:
  BinCountFunctor() {

oneflow/core/functional/impl/nn_functor.cpp

@@ -3431,6 +3431,83 @@ class CosineSimilarityFunctor {
  }
 };
 
+class CDistFunctor {
+ public:
+ CDistFunctor() {
+ op_ = CHECK_JUST(OpBuilder("cdist").Input("x1").Input("x2").Output("out").Build());
+ }
+ Maybe<Tensor> euclidean_dist(const std::shared_ptr<Tensor>& x1,
+ const std::shared_ptr<Tensor>& x2) const {
+ const auto& x1_norm = JUST(ReduceSum(JUST(ScalarPow(x1, 2, false)), {-1}, true));
+ const auto& x2_norm = JUST(ReduceSum(JUST(ScalarPow(x2, 2, false)), {-1}, true));
+ const auto& x1_ones = JUST(OnesLike(x1_norm));
+ const auto& x2_ones = JUST(OnesLike(x2_norm));
+ const auto& x1_cat = JUST(Concat({JUST(ScalarMul(x1, -2, false)), x1_norm, x1_ones}, -1));
+ const auto& x2_cat = JUST(Concat({x2, x2_ones, x2_norm}, -1));
+ const auto& result =
+ JUST(MatMul(x1_cat, JUST(Transpose2dim(x2_cat, -1, -2)), false, false, 1.0));
+ return Sqrt(JUST(ClampMin(result, 0.0)));
+ };
+
+ Maybe<Tensor> operator()(const std::shared_ptr<Tensor>& x1, const std::shared_ptr<Tensor>& x2,
+ const double& p, const Optional<std::string>& compute_mode) const {
+ const int64_t x1_ndim = x1->ndim();
+ const int64_t x2_ndim = x2->ndim();
+ CHECK_OR_RETURN(x1_ndim >= 2) << "cdist only supports at least 2D tensors, X1 got: "
+ << x1->ndim() << "D";
+ CHECK_OR_RETURN(x2_ndim >= 2) << "cdist only supports at least 2D tensors, X2 got: "
+ << x2->ndim() << "D";
+ CHECK_OR_RETURN(x1->dim(x1_ndim - 1) == x2->dim(x2_ndim - 1))
+ << "X1 and X2 must have the same number of columns. X1: " << x1->dim(x1_ndim - 1)
+ << " X2: " << x2->dim(x2_ndim - 1);
+ CHECK_OR_RETURN(p >= 0) << "cdist only supports non-negative p values, got " << p;
+
+ if (compute_mode.has_value()) {
+ OF_LOG_ONCE(LOG(WARNING)
+ << "'compute_mode' argument is not supported yet, cdist "
+ "will not use matrix multiplication approach to calculate euclidean distance");
+ }
+
+ int64_t r1 = x1->dim(x1_ndim - 2);
+ int64_t r2 = x2->dim(x2_ndim - 2);
+ int64_t d = x1->dim(x1_ndim - 1);
+
+ std::vector<Shape> shape_vector = {
+ Shape(DimVector({x1->shape()->begin(), x1->shape()->end() - 2})),
+ Shape(DimVector({x2->shape()->begin(), x2->shape()->end() - 2})),
+ };
+ auto broadcasted_shape = JUST(BroadcastShapes(shape_vector));
+ Shape x1_expand_shape(*broadcasted_shape);
+ x1_expand_shape.emplace_back(r1);
+ x1_expand_shape.emplace_back(d);
+ broadcasted_shape->emplace_back(r2);
+ broadcasted_shape->emplace_back(d);
+
+ const auto x1_expand = JUST(Expand(x1, x1_expand_shape));
+ const auto x2_expand = JUST(Expand(x2, *broadcasted_shape));
+
+ // mm_for_euclid_dist has accuracy issue
+ // if (p == 2 && (mode == 1 || (mode == 0 && (r1 > 25 || r2 > 25)))) {
+ // shape output_shape(max_batch_shape);
+ // output_shape.emplace_back(r1);
+ // output_shape.emplace_back(r2);
+ // return JUST(Reshape(JUST(euclidean_dist(x1_expand, x2_expand)), output_shape));
+ // }
+
+ TensorProcessor tensor_processor;
+ JUST(tensor_processor.PromoteInputsToCommonDtype(true)
+ .AddInputs({x1_expand, x2_expand})
+ .Apply());
+
+ auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("p");
+ attrs.SetAllAttrs(p);
+ return OpInterpUtil::Dispatch<Tensor>(*op_, {x1, x2}, attrs);
+ }
+
+ private:
+ std::shared_ptr<OpExpr> op_;
+};
+
 class L2NormalizeFunctor {
  public:
  L2NormalizeFunctor() {
@@ -5501,6 +5578,7 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
  m.add_functor<impl::PairwiseDistanceFunctor>("PairwiseDistance");
  m.add_functor<impl::CosineSimilarityFunctor>("CosineSimilarity");
  m.add_functor<impl::NormalizeFunctor>("Normalize");
+ m.add_functor<impl::CDistFunctor>("CDist");
  m.add_functor<impl::L2NormalizeFunctor>("L2Normalize");
  m.add_functor<impl::L2NormalizeGradFunctor>("L2NormalizeGrad");
  m.add_functor<impl::FusedBiasAddGeluFunctor>("FusedBiasAddGelu");

oneflow/core/functional/impl/nn_grad_functor.cpp

@@ -659,6 +659,31 @@ class BinaryCrossEntropyWithLogitsReduceMeanLossTargetGradFunctor {
  }
 };
 
+class CDistGradFunctor {
+ public:
+ CDistGradFunctor() {
+ op_ = CHECK_JUST(one::OpBuilder("cdist_grad")
+ .Input("x1")
+ .Input("x2")
+ .Input("out")
+ .Input("dy")
+ .Output("dx1")
+ .Output("dx2")
+ .Build());
+ }
+ Maybe<TensorTuple> operator()(const std::shared_ptr<Tensor>& x1,
+ const std::shared_ptr<Tensor>& x2,
+ const std::shared_ptr<Tensor>& out,
+ const std::shared_ptr<Tensor>& dy, const double& p) const {
+ auto& attrs = THREAD_CACHED_MUTABLE_ATTR_MAP("p");
+ attrs.SetAllAttrs(p);
+ return OpInterpUtil::Dispatch<TensorTuple>(*op_, {x1, x2, out, dy}, attrs);
+ }
+
+ private:
+ std::shared_ptr<OpExpr> op_;
+};
+
 class CombinedMarginLossGradFunctor {
  public:
  CombinedMarginLossGradFunctor() {
@@ -1692,6 +1717,7 @@ ONEFLOW_FUNCTION_LIBRARY(m) {
  m.add_functor<impl::SparseSoftmaxCrossEntropyGrad>("SparseSoftmaxCrossEntropyGrad");
  m.add_functor<impl::SparseSoftmaxCrossEntropyMsGrad>("SparseSoftmaxCrossEntropyMsGrad");
  m.add_functor<impl::SmoothL1LossGradFunctor>("SmoothL1LossGrad");
+ m.add_functor<impl::CDistGradFunctor>("CDistGrad");
  m.add_functor<impl::CombinedMarginLossGradFunctor>("CombinedMarginLossGrad");
  m.add_functor<impl::AffineGridGradFunctor>("AffineGridGrad");
  m.add_functor<impl::GridSampleGradFunctor>("GridSampleGrad");

oneflow/ir/include/OneFlow/OneFlowUserOps.td

@@ -4483,6 +4483,43 @@ def OneFlow_AbsGradOp : OneFlow_BaseOp<"abs_grad", [NoMemoryEffect, DeclareOpInt
  let has_data_type_infer_fn = 1;
 }
 
+def OneFlow_CDistOp : OneFlow_BaseOp<"cdist", [NoMemoryEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
+ let input = (ins
+ OneFlow_Tensor:$x1,
+ OneFlow_Tensor:$x2
+ );
+ let output = (outs
+ OneFlow_Tensor:$out
+ );
+ let attrs = (ins
+ DefaultValuedAttr<F64Attr, "2.0">:$p
+ );
+ let has_logical_tensor_desc_infer_fn = 1;
+ let has_physical_tensor_desc_infer_fn = 1;
+ let has_get_sbp_fn = 1;
+ let has_data_type_infer_fn = 1;
+}
+
+def OneFlow_CDistGradOp : OneFlow_BaseOp<"cdist_grad", [NoMemoryEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
+ let input = (ins
+ OneFlow_Tensor:$x1,
+ OneFlow_Tensor:$x2,
+ OneFlow_Tensor:$out,
+ OneFlow_Tensor:$dy
+ );
+ let output = (outs
+ OneFlow_Tensor:$dx1,
+ OneFlow_Tensor:$dx2
+ );
+ let attrs = (ins
+ DefaultValuedAttr<F64Attr, "2.0">:$p
+ );
+ let has_logical_tensor_desc_infer_fn = 1;
+ let has_physical_tensor_desc_infer_fn = 1;
+ let has_get_sbp_fn = 1;
+ let has_data_type_infer_fn = 1;
+}
+
 def OneFlow_ErfOp : OneFlow_BaseOp<"erf", [NoMemoryEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
  let input = (ins
  OneFlow_Tensor:$x