Changed ff_dim to ff_dim_t, added in nonnegative_int type

flexflow · Oct 31, 2024 · f2b6d2a · f2b6d2a
1 parent 1d5140d
commit f2b6d2a
Show file tree

Hide file tree

Showing 32 changed files with 576 additions and 43 deletions.
diff --git a/lib/kernels/include/kernels/legion_dim.h b/lib/kernels/include/kernels/legion_dim.h
@@ -8,7 +8,7 @@ namespace FlexFlow {
 
 legion_dim_t add_to_legion_dim(legion_dim_t legion_dim, int value);
 
-legion_dim_t legion_dim_from_ff_dim(ff_dim_t, int num_dimensions);
+legion_dim_t legion_dim_from_ff_dim_t(ff_dim_t, int num_dimensions);
 
 template <typename T>
 using LegionOrdered = DimOrdered<legion_dim_t, T>;

diff --git a/lib/kernels/src/array_shape.cc b/lib/kernels/src/array_shape.cc
@@ -47,7 +47,7 @@ std::size_t ArrayShape::at(legion_dim_t idx) const {
 }
 
 std::size_t ArrayShape::at(ff_dim_t idx) const {
-  return dims.at(legion_dim_from_ff_dim(idx, this->num_dims()));
+  return dims.at(legion_dim_from_ff_dim_t(idx, this->num_dims()));
 }
 
 ArrayShape ArrayShape::sub_shape(
@@ -65,7 +65,7 @@ std::optional<std::size_t> ArrayShape::at_maybe(legion_dim_t index) const {
 }
 
 std::optional<std::size_t> ArrayShape::at_maybe(ff_dim_t index) const {
-  return this->at_maybe(legion_dim_from_ff_dim(index, this->num_dims()));
+  return this->at_maybe(legion_dim_from_ff_dim_t(index, this->num_dims()));
 }
 
 size_t get_volume(ArrayShape const &shape) {

diff --git a/lib/kernels/src/cuda/ops/transpose_kernels.cu b/lib/kernels/src/cuda/ops/transpose_kernels.cu
@@ -36,7 +36,7 @@ TransposePerDeviceState init_kernel(int num_dim,
   std::vector<legion_dim_t> perm_vector;
   assert(length <= MAX_TENSOR_DIM);
   for (int i = 0; i < length; ++i) {
-    perm_vector.push_back(legion_dim_from_ff_dim(perm[i], num_dim));
+    perm_vector.push_back(legion_dim_from_ff_dim_t(perm[i], num_dim));
   }
 
   return {num_dim, perm_vector};

diff --git a/lib/kernels/src/legion_dim.cc b/lib/kernels/src/legion_dim.cc
@@ -6,8 +6,8 @@ legion_dim_t add_to_legion_dim(legion_dim_t legion_dim, int value) {
   return legion_dim_t(legion_dim.value + value);
 }
 
-legion_dim_t legion_dim_from_ff_dim(ff_dim_t ff_dim, int num_dimensions) {
-  return legion_dim_t(num_dimensions - ff_dim.value - 1);
+legion_dim_t legion_dim_from_ff_dim_t(ff_dim_t ff_dim_t, int num_dimensions) {
+  return legion_dim_t(num_dimensions - ff_dim_t.value - 1);
 }
 
 } // namespace FlexFlow
diff --git a/lib/local-execution/include/local-execution/legion_tensor_shape.h b/lib/local-execution/include/local-execution/legion_tensor_shape.h
@@ -3,7 +3,7 @@
 
 #include "kernels/legion_dim.h"
 #include "op-attrs/datatype.h"
-#include "op-attrs/ff_dim.h"
+#include "op-attrs/ff_dim_t.h"
 #include "op-attrs/tensor_shape.dtg.h"
 #include "utils/stack_vector.h"
 #include "utils/visitable.h"
@@ -30,10 +30,10 @@ struct LegionTensorShape : public use_visitable_cmp<LegionTensorShape>,
 };
 
 ff_dim_t to_ff(legion_dim_t, size_t num_dims);
-legion_dim_t legion_dim_from_ff_dim(ff_dim_t, size_t num_dims);
+legion_dim_t legion_dim_from_ff_dim_t(ff_dim_t, size_t num_dims);
 
 ff_dim_t to_ff(legion_dim_t, TensorShape const &);
-legion_dim_t legion_dim_from_ff_dim(ff_dim_t, TensorShape const &);
+legion_dim_t legion_dim_from_ff_dim_t(ff_dim_t, TensorShape const &);
 
 } // namespace FlexFlow
 

diff --git a/lib/local-execution/src/legion_tensor_shape.cc b/lib/local-execution/src/legion_tensor_shape.cc
@@ -3,12 +3,12 @@
 
 namespace FlexFlow {
 
-legion_dim_t legion_dim_from_ff_dim(ff_dim_t ff_dim, size_t num_dims) {
-  return legion_dim_t(num_dims - ff_dim.value - 1);
+legion_dim_t legion_dim_from_ff_dim_t(ff_dim_t ff_dim_t, size_t num_dims) {
+  return legion_dim_t(num_dims - ff_dim_t.value - 1);
 }
 
-legion_dim_t legion_dim_from_ff_dim(ff_dim_t ff_dim, TensorShape const &shape) {
-  return legion_dim_t(num_dims(shape) - ff_dim.value - 1);
+legion_dim_t legion_dim_from_ff_dim_t(ff_dim_t ff_dim_t, TensorShape const &shape) {
+  return legion_dim_t(num_dims(shape) - ff_dim_t.value - 1);
 }
 
 } // namespace FlexFlow
diff --git a/lib/local-execution/src/ops/gather.cc b/lib/local-execution/src/ops/gather.cc
@@ -67,7 +67,7 @@ static DeviceSpecificDeviceStates
   PerDeviceFFHandle handle = acc.get_argument<PerDeviceFFHandle>(HANDLE);
   auto const &attrs = acc.get_argument<GatherAttrs>(ATTRS);
   legion_dim_t legion_dim =
-      legion_dim_from_ff_dim(attrs.dim, input.shape.num_dims());
+      legion_dim_from_ff_dim_t(attrs.dim, input.shape.num_dims());
 
   assert(input.shape.get_dim() == index.shape.get_dim());
   assert(output.shape.get_dim() == index.shape.get_dim());

diff --git a/lib/local-execution/src/ops/layer_norm.cc b/lib/local-execution/src/ops/layer_norm.cc
@@ -123,7 +123,7 @@ static DeviceSpecificDeviceStates
   int64_t effective_batch_size, effective_num_elements;
   int M = 1;
   for (int i = 0; i < attrs.axes.size(); i++) {
-    legion_dim_t legion_dim = legion_dim_from_ff_dim(
+    legion_dim_t legion_dim = legion_dim_from_ff_dim_t(
         attrs.axes[i], get_tensor_shape(input.shape, input.data_type));
     M *= input.shape.at(legion_dim);
   }

diff --git a/lib/local-execution/src/ops/linear.cc b/lib/local-execution/src/ops/linear.cc
@@ -1,7 +1,7 @@
 #include "linear.h"
 #include "kernels/linear_kernels.h"
 #include "local-execution/task_argument_accessor.h"
-#include "op-attrs/ff_dim.h"
+#include "op-attrs/ff_dim_t.h"
 #include "op-attrs/get_output_shapes.h"
 #include "utils/exception.h"
 #include "utils/hash-utils.h"

diff --git a/lib/op-attrs/include/op-attrs/dim_ordered/dim_ordered.h b/lib/op-attrs/include/op-attrs/dim_ordered/dim_ordered.h
@@ -1,7 +1,7 @@
 #ifndef _FLEXFLOW_OPATTRS_INCLUDE_OPATTRS_FF_STACK_VECTOR_H
 #define _FLEXFLOW_OPATTRS_INCLUDE_OPATTRS_FF_STACK_VECTOR_H
 
-#include "op-attrs/ff_dim.dtg.h"
+#include "op-attrs/ff_dim_t.dtg.h"
 #include "utils/fmt/vector.h"
 #include "utils/stack_vector.h"
 #include <nlohmann/json.hpp>

diff --git a/lib/op-attrs/include/op-attrs/ff_dim_t.h b/lib/op-attrs/include/op-attrs/ff_dim_t.h
@@ -0,0 +1,17 @@
+#ifndef _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_FF_DIM_T_H
+#define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_FF_DIM_T_H
+
+#include "op-attrs/ff_dim_t.dtg.h"
+#include "rapidcheck.h"
+
+namespace rc {
+template <>
+struct Arbitrary<FlexFlow::ff_dim_t> {
+  static Gen<FlexFlow::ff_dim_t> arbitrary() {
+    return gen::construct<FlexFlow::ff_dim_t>(
+        gen::inRange<int>(0, MAX_TENSOR_DIM));
+  }
+};
+} // namespace rc
+
+#endif // _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_FF_DIM_T_H
diff --git a/lib/op-attrs/include/op-attrs/ff_dim_t.struct.toml b/lib/op-attrs/include/op-attrs/ff_dim_t.struct.toml
@@ -0,0 +1,14 @@
+namespace = "FlexFlow"
+name = "ff_dim_t"
+
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "fmt",
+]
+
+[[fields]]
+name = "value"
+type = "int"
diff --git a/lib/op-attrs/include/op-attrs/ops/combine_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/combine_attrs.struct.toml
@@ -10,8 +10,8 @@ features = [
 ]
 
 includes = [
-  "op-attrs/ff_dim.h",
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.h",
+  "op-attrs/ff_dim_t.dtg.h",
 ]
 
 [[fields]]

diff --git a/lib/op-attrs/include/op-attrs/ops/concat_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/concat_attrs.struct.toml
@@ -10,8 +10,8 @@ features = [
 ]
 
 includes = [
-  "op-attrs/ff_dim.h",
-  "op-attrs/ff_dim.dtg.h"
+  "op-attrs/ff_dim_t.h",
+  "op-attrs/ff_dim_t.dtg.h"
 ]
 
 [[fields]]

diff --git a/lib/op-attrs/include/op-attrs/ops/flat_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/flat_attrs.struct.toml
@@ -11,14 +11,14 @@ features = [
 
 includes = [
   "<optional>",
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.dtg.h",
 ]
 
 src_includes = [
   "utils/fmt/optional.h",
   "utils/json/optional.h",
   "utils/rapidcheck/optional.h",
-  "op-attrs/ff_dim.h",
+  "op-attrs/ff_dim_t.h",
 ]
 
 [[fields]]

diff --git a/lib/op-attrs/include/op-attrs/ops/gather_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/gather_attrs.struct.toml
@@ -10,8 +10,8 @@ features = [
 ]
 
 includes = [
-  "op-attrs/ff_dim.h",
-  "op-attrs/ff_dim.dtg.h"
+  "op-attrs/ff_dim_t.h",
+  "op-attrs/ff_dim_t.dtg.h"
 ]
 
 [[fields]]

diff --git a/lib/op-attrs/include/op-attrs/ops/layer_norm_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/layer_norm_attrs.struct.toml
@@ -10,8 +10,8 @@ features = [
 ]
 
 includes = [
-  "op-attrs/ff_dim.h",
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.h",
+  "op-attrs/ff_dim_t.dtg.h",
   "utils/stack_vector.h",
 ]
 

diff --git a/lib/op-attrs/include/op-attrs/ops/reduce_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/reduce_attrs.struct.toml
@@ -11,8 +11,8 @@ features = [
 
 includes = [
   "op-attrs/operator_type.dtg.h",
-  "op-attrs/ff_dim.h",
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.h",
+  "op-attrs/ff_dim_t.dtg.h",
   "utils/stack_vector.h",
 ]
 

diff --git a/lib/op-attrs/include/op-attrs/ops/repartition_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/repartition_attrs.struct.toml
@@ -10,8 +10,8 @@ features = [
 ]
 
 includes = [
-  "op-attrs/ff_dim.h",
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.h",
+  "op-attrs/ff_dim_t.dtg.h",
 ]
 
 [[fields]]

diff --git a/lib/op-attrs/include/op-attrs/ops/reverse_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/reverse_attrs.struct.toml
@@ -10,8 +10,8 @@ features = [
 ]
 
 includes = [
-  "op-attrs/ff_dim.h",
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.h",
+  "op-attrs/ff_dim_t.dtg.h",
 ]
 
 [[fields]]

diff --git a/lib/op-attrs/include/op-attrs/ops/softmax_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/softmax_attrs.struct.toml
@@ -10,8 +10,8 @@ features = [
 ]
 
 includes = [
-  "op-attrs/ff_dim.h",
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.h",
+  "op-attrs/ff_dim_t.dtg.h",
 ]
 
 [[fields]]

diff --git a/lib/op-attrs/include/op-attrs/ops/split_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/split_attrs.struct.toml
@@ -11,8 +11,8 @@ features = [
 
 includes = [
   "utils/stack_vector.h",
-  "op-attrs/ff_dim.h",
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.h",
+  "op-attrs/ff_dim_t.dtg.h",
 ]
 
 [[fields]]

diff --git a/lib/op-attrs/include/op-attrs/ops/transpose_attrs.struct.toml b/lib/op-attrs/include/op-attrs/ops/transpose_attrs.struct.toml
@@ -10,8 +10,8 @@ features = [
 ]
 
 includes = [
-  "op-attrs/ff_dim.h",
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.h",
+  "op-attrs/ff_dim_t.dtg.h",
   "op-attrs/dim_ordered/dim_ordered.h",
 ]
 

diff --git a/lib/op-attrs/include/op-attrs/parallel_tensor_dim_idx_t.variant.toml b/lib/op-attrs/include/op-attrs/parallel_tensor_dim_idx_t.variant.toml
@@ -9,7 +9,7 @@ features = [
 ]
 
 includes = [
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.dtg.h",
   "op-attrs/replica_type.dtg.h",
 ]
 

diff --git a/lib/op-attrs/include/op-attrs/relative_ff_dim_t.h b/lib/op-attrs/include/op-attrs/relative_ff_dim_t.h
@@ -0,0 +1,17 @@
+#ifndef _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_RELATIVE_FF_DIM_T_H
+#define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_RELATIVE_FF_DIM_T_H
+
+#include "op-attrs/relative_ff_dim_t.dtg.h"
+#include "rapidcheck.h"
+
+namespace rc {
+template <>
+struct Arbitrary<FlexFlow::relative_ff_dim_t> {
+  static Gen<FlexFlow::relative_ff_dim_t> arbitrary() {
+    return gen::construct<FlexFlow::relative_ff_dim_t>(
+        gen::inRange<int>(0, MAX_TENSOR_DIM));
+  }
+};
+} // namespace rc
+
+#endif // _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_RELATIVE_FF_DIM_T_H
diff --git a/lib/op-attrs/include/op-attrs/relative_ff_dim_t.struct.toml b/lib/op-attrs/include/op-attrs/relative_ff_dim_t.struct.toml
@@ -0,0 +1,14 @@
+namespace = "FlexFlow"
+name = "relative_ff_dim_t"
+
+features = [
+  "eq",
+  "ord",
+  "hash",
+  "json",
+  "fmt",
+]
+
+[[fields]]
+name = "value"
+type = "int"
diff --git a/lib/op-attrs/test/src/op-attrs/dim_ordered/zip.cc b/lib/op-attrs/test/src/op-attrs/dim_ordered/zip.cc
@@ -1,5 +1,5 @@
 #include "op-attrs/dim_ordered/zip.h"
-#include "op-attrs/ff_dim.dtg.h"
+#include "op-attrs/ff_dim_t.dtg.h"
 #include "test/utils/doctest/fmt/pair.h"
 #include <doctest/doctest.h>
 

diff --git a/lib/runtime/src/parallel_op_info.h b/lib/runtime/src/parallel_op_info.h
@@ -1,7 +1,7 @@
 #ifndef _FLEXFLOW_PARALLEL_OPS_PARALLEL_OP_INFO_H
 #define _FLEXFLOW_PARALLEL_OPS_PARALLEL_OP_INFO_H
 
-#include "op-attrs/ff_dim.h"
+#include "op-attrs/ff_dim_t.h"
 #include "op-attrs/operator_type.h"
 #include "utils/visitable.h"
 #include <functional>

diff --git a/...ubstitutions/include/substitutions/operator_pattern/operator_attribute_value.variant.toml b/...ubstitutions/include/substitutions/operator_pattern/operator_attribute_value.variant.toml
@@ -13,7 +13,7 @@ includes = [
   "<vector>",
   "<optional>",
   "op-attrs/operator_type.dtg.h",
-  "op-attrs/ff_dim.dtg.h",
+  "op-attrs/ff_dim_t.dtg.h",
   "op-attrs/activation.dtg.h",
   "op-attrs/aggregate_op.dtg.h",
   "op-attrs/regularizer_attrs.dtg.h",