repo-sync-2024-01-12T11:42:53+0800

.vscode/settings.json

@@ -40,5 +40,6 @@
     "git.ignoreLimitWarning": true,
     "[python]": {
         "editor.defaultFormatter": "ms-python.black-formatter"
-    }
+    },
+    "mlir.server_path": "bazel-bin/libspu/compiler/tools/mlir-pphlo-lsp"
 }

CHANGELOG.md

@@ -10,7 +10,9 @@
 >
 > please add your unreleased change here.
 
-## TBD
+- [Improvement] Optimize one-time setup for yacl ot
+
+## 20240105
 
 - [Feature] Add Odd-Even Merge Sort to replace the bitonic sort
 - [Feature] Add radix sort support for ABY3

bazel/repositories.bzl

@@ -140,8 +140,8 @@ def _com_github_xtensor_xtl():
     )
 
 def _com_github_openxla_xla():
-    OPENXLA_COMMIT = "d5791b01aa7541e3400224ac0a2985cc0f6940cb"
-    OPENXLA_SHA256 = "82dd50e6f51d79e8da69f109a234e33b8036f7b8798e41a03831b19c0c64d6e5"
+    OPENXLA_COMMIT = "fa9331a7e557b4ec1381f84cbbf7401a8f41ac66"
+    OPENXLA_SHA256 = "d19c570d434002b7b0490327d407fc7cf2b18633f4a2d3b1bb44f3f0e4b36533"
 
     maybe(
         http_archive,

examples/python/ml/flax_llama7b/README.md

@@ -22,6 +22,7 @@ This example demonstrates how to use SPU to run secure inference on a pre-traine
 
     Since EasyLM have an issue，so we have to make a samll change to support the option "streaming=false".
     Open and edit "convert_hf_to_easylm.py", chang this:
+
     ```python
        parser.add_argument("--streaming", action="store_true", default=True, help="whether is model weight saved stream format",)
     ```

libspu/compiler/passes/BUILD.bazel

@@ -92,7 +92,6 @@ spu_cc_library(
     deps = [
         ":map_stablehlo_to_pphlo_op",
         ":pass_details",
-        ":utils",
         ":visibility_inference",
         "//libspu/compiler/common:compilation_context",
         "//libspu/core:prelude",
@@ -204,7 +203,6 @@ spu_cc_library(
     hdrs = ["passes.h"],
     deps = [
         ":pass_details",
-        ":utils",
         "//libspu/dialect:pphlo_dialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:TransformUtils",
@@ -259,15 +257,6 @@ spu_cc_library(
     ],
 )
 
-spu_cc_library(
-    name = "utils",
-    srcs = ["utils.cc"],
-    hdrs = ["utils.h"],
-    deps = [
-        "@llvm-project//mlir:IR",
-    ],
-)
-
 spu_cc_library(
     name = "convert_push_down",
     srcs = ["convert_push_down.cc"],

libspu/compiler/passes/expand_secret_gather.cc

@@ -20,15 +20,14 @@
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "libspu/compiler/passes/pass_details.h"
-#include "libspu/compiler/passes/utils.h"
 #include "libspu/dialect/pphlo_ops.h"
 
 namespace mlir::pphlo {
 
 namespace {
 
 bool GatherIsBroadcast(GatherOp &op) {
-  auto gather_slice_size = op.getSliceSizes().getValues<int64_t>();
+  auto gather_slice_size = op.getSliceSizes();
   auto op_shape = op.getOperand().getType().getShape();
   return (gather_slice_size.size() == op_shape.size()) &&
          (std::equal(gather_slice_size.begin(), gather_slice_size.end(),
@@ -113,7 +112,7 @@ TransposeIndexVectorDimToLast(TypedValue<RankedTensorType> &start_indices,
       start_indices.getLoc(),
       RankedTensorType::get(result_shape,
                             start_indices.getType().getElementType()),
-      start_indices, ConvertDimensions(&builder, permutation));
+      start_indices, permutation);
 
   return transpose.getResult();
 }
@@ -211,17 +210,16 @@ CanonicalizeGatherIndices(TypedValue<RankedTensorType> &start_indices,
 }
 
 TypedValue<RankedTensorType> CreateGatherLoopAccumulatorInitValue(
-    GatherOp op, Type element_type, DenseIntElementsAttr slice_sizes,
+    GatherOp op, Type element_type, llvm::ArrayRef<int64_t> slice_sizes,
     int64_t gather_loop_trip_count,
     const GatherDimensionNumbersAttr &dim_numbers) {
   std::vector<int64_t> accumulator_state_shape_dims;
-  auto array_slice_size = slice_sizes.getValues<int64_t>();
-  accumulator_state_shape_dims.reserve(1 + array_slice_size.size());
+  accumulator_state_shape_dims.reserve(1 + slice_sizes.size());
   accumulator_state_shape_dims.push_back(gather_loop_trip_count);
-  for (int64_t i = 0; i < static_cast<int64_t>(array_slice_size.size()); i++) {
+  for (int64_t i = 0; i < static_cast<int64_t>(slice_sizes.size()); i++) {
     if (!std::binary_search(dim_numbers.getCollapsedSliceDims().begin(),
                             dim_numbers.getCollapsedSliceDims().end(), i)) {
-      accumulator_state_shape_dims.emplace_back(array_slice_size[i]);
+      accumulator_state_shape_dims.emplace_back(slice_sizes[i]);
     }
   }
 
@@ -392,9 +390,12 @@ llvm::SmallVector<Value> ExpandIndexVectorIntoOperandSpace(
       auto component_to_concat = builder->create<SliceOp>(
           index_vector.getLoc(),
           RankedTensorType::get({1}, index_vector.getType().getElementType()),
-          index_vector, ConvertDimensions(builder, {index_vector_dim_index}),
-          ConvertDimensions(builder, {index_vector_dim_index + 1}),
-          ConvertDimensions(builder, {1}));
+          index_vector,
+          DenseI64ArrayAttr::get(builder->getContext(),
+                                 {index_vector_dim_index}),
+          DenseI64ArrayAttr::get(builder->getContext(),
+                                 {index_vector_dim_index + 1}),
+          DenseI64ArrayAttr::get(builder->getContext(), {1}));
       auto reshaped = builder->create<ReshapeOp>(
           index_vector.getLoc(),
           RankedTensorType::get({}, index_vector.getType().getElementType()),
@@ -450,9 +451,9 @@ void GatherLoopBody(GatherOp gather, Region &body,
   if (has_scalar_indices) {
     // In this case start_indices has rank 1 and induction_var_as_vector (of
     // shape {1}) is an index into this rank 1 tensor.
-    auto ds = builder.create<DynamicSliceOp>(gather->getLoc(), start_indices,
-                                             ValueRange{induction_var},
-                                             ConvertDimensions(&builder, {1}));
+    auto ds = builder.create<DynamicSliceOp>(
+        gather->getLoc(), start_indices, ValueRange{induction_var},
+        DenseI64ArrayAttr::get(builder.getContext(), {1}));
     index_vector = ds.getResult();
   } else {
     // In this case start_indices has rank 2 and induction_var_as_vector (of
@@ -463,7 +464,7 @@ void GatherLoopBody(GatherOp gather, Region &body,
 
     auto index_vector_2d = builder.create<DynamicSliceOp>(
         gather->getLoc(), start_indices, ValueRange{induction_var, index_zero},
-        ConvertDimensions(&builder, {1, index_vector_size}));
+        DenseI64ArrayAttr::get(builder.getContext(), {1, index_vector_size}));
 
     index_vector = ElideDegenerateDims(&builder, index_vector_2d, {0});
   }
@@ -523,8 +524,8 @@ struct GatherConverter : public OpRewritePattern<GatherOp> {
           op->getLoc(), reshaped_type, op.getOperand());
       rewriter.replaceOpWithNewOp<BroadcastOp>(
           op, op->getResults().getType(), broadcast_operand,
-          ConvertDimensions(&builder,
-                            op.getDimensionNumbers().getOffsetDims()));
+          DenseI64ArrayAttr::get(builder.getContext(),
+                                 op.getDimensionNumbers().getOffsetDims()));
       return success();
     }
 
@@ -612,7 +613,7 @@ struct GatherConverter : public OpRewritePattern<GatherOp> {
     rewriter.replaceOpWithNewOp<TransposeOp>(
         op, op.getResult().getType(),
         accumulator_with_batch_dims_decanonicalized,
-        ConvertDimensions(&builder, permutation));
+        DenseI64ArrayAttr::get(builder.getContext(), permutation));
 
     return success();
   }

libspu/compiler/passes/hlo_legalize_to_pphlo.cc

@@ -28,7 +28,6 @@
 
 #include "libspu/compiler/passes/map_stablehlo_to_pphlo_op.h"
 #include "libspu/compiler/passes/pass_details.h"
-#include "libspu/compiler/passes/utils.h"
 #include "libspu/compiler/passes/value_visibility_map.h"
 #include "libspu/compiler/passes/visibility_inference.h"
 #include "libspu/core/prelude.h"
@@ -40,6 +39,11 @@
 namespace mlir::pphlo {
 namespace {
 
+DenseI64ArrayAttr ConvertDenseIntElementAttr(const DenseIntElementsAttr &attr) {
+  llvm::SmallVector<int64_t, 2> array(attr.getValues<int64_t>());
+  return DenseI64ArrayAttr::get(attr.getContext(), array);
+}
+
 ValueVisibilityMap
 VisibilityDiscovery(const llvm::ArrayRef<std::string> input_vis_list,
                     ModuleOp op) {
@@ -348,9 +352,14 @@ struct ReduceOpConverter : public OpConversionPattern<stablehlo::ReduceOp> {
       sig_conversion.addInputs(arg.getArgNumber(), lower_t);
     }
 
+    mlir::NamedAttribute dimAttr(
+        StringAttr::get(op->getContext(), "dimensions"),
+        ConvertDenseIntElementAttr(op.getDimensions()));
+
     auto new_op =
         rewriter.replaceOpWithNewOp<pphlo::HloToPPHloOp<stablehlo::ReduceOp>>(
-            op, result_types, materialized_operands, op->getAttrs());
+            op, result_types, materialized_operands,
+            llvm::SmallVector<mlir::NamedAttribute, 1>{dimAttr});
 
     // Copy over the operations inside the region.
     rewriter.inlineRegionBefore(op.getBody(), new_op.getBody(),
@@ -466,27 +475,30 @@ struct ReduceWindowOpConverter
           materialized_operands[idx] = rewriter.create<pphlo::PadOp>(
               op->getLoc(), materialized_operands[idx],
               materialized_operands[idx + num_results],
-              builder.getI64TensorAttr(padding_low),
-              builder.getI64TensorAttr(padding_high),
-              builder.getI64TensorAttr(interior_padding));
+              DenseI64ArrayAttr::get(op->getContext(), padding_low),
+              DenseI64ArrayAttr::get(op->getContext(), padding_high),
+              DenseI64ArrayAttr::get(op->getContext(), interior_padding));
         }
       }
     }
 
     llvm::SmallVector<mlir::NamedAttribute> attrs;
     {
       // I64ElementsAttr:$window_dimensions,
-      attrs.push_back({builder.getStringAttr("window_dimensions"),
-                       op.getWindowDimensionsAttr()});
+      attrs.push_back(
+          {builder.getStringAttr("window_dimensions"),
+           ConvertDenseIntElementAttr(op.getWindowDimensionsAttr())});
       // OptionalAttr<I64ElementsAttr>:$window_strides,
       if (op.getWindowStrides().has_value()) {
-        attrs.push_back({builder.getStringAttr("window_strides"),
-                         op.getWindowStridesAttr()});
+        attrs.push_back(
+            {builder.getStringAttr("window_strides"),
+             ConvertDenseIntElementAttr(op.getWindowStridesAttr())});
       }
       // OptionalAttr<I64ElementsAttr>:$window_dilations,
       if (op.getWindowDilations().has_value()) {
-        attrs.push_back({builder.getStringAttr("window_dilations"),
-                         op.getWindowDilationsAttr()});
+        attrs.push_back(
+            {builder.getStringAttr("window_dilations"),
+             ConvertDenseIntElementAttr(op.getWindowDilationsAttr())});
       }
     }
 
@@ -755,6 +767,40 @@ class HloToPPHloOpConverter : public OpConversionPattern<HloOpTy> {
   }
 };
 
+template <>
+class HloToPPHloOpConverter<stablehlo::BroadcastInDimOp>
+    : public OpConversionPattern<stablehlo::BroadcastInDimOp> {
+private:
+  const ValueVisibilityMap &vis_;
+
+public:
+  HloToPPHloOpConverter(TypeConverter &type_converter, MLIRContext *context,
+                        const ValueVisibilityMap &vis)
+      : OpConversionPattern<stablehlo::BroadcastInDimOp>(type_converter,
+                                                         context),
+        vis_(vis) {}
+
+  LogicalResult
+  matchAndRewrite(stablehlo::BroadcastInDimOp hlo_op,
+                  stablehlo::BroadcastInDimOpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto result_vis = vis_.getValueVisibility(hlo_op.getResult());
+
+    Type resultType = HloToPPHloTypeConverter::getTypeWithVisibility(
+        this->getTypeConverter()->convertType(hlo_op.getType()), result_vis);
+
+    mlir::NamedAttribute dim(
+        StringAttr::get(hlo_op.getContext(), "broadcast_dimensions"),
+        ConvertDenseIntElementAttr(hlo_op.getBroadcastDimensions()));
+
+    rewriter
+        .replaceOpWithNewOp<pphlo::HloToPPHloOp<stablehlo::BroadcastInDimOp>>(
+            hlo_op, resultType, adaptor.getOperands(), dim);
+
+    return success();
+  }
+};
+
 template <>
 class HloToPPHloOpConverter<stablehlo::ConstantOp>
     : public OpConversionPattern<stablehlo::ConstantOp> {
@@ -1003,14 +1049,15 @@ struct HloToPPHloOpConverter<stablehlo::SelectAndScatterOp>
 
       materialized_operand = rewriter.create<pphlo::PadOp>(
           op->getLoc(), materialized_operand, materialized_init_value,
-          builder.getI64TensorAttr(padding_low),
-          builder.getI64TensorAttr(padding_high),
-          builder.getI64TensorAttr(padding_interior));
+          DenseI64ArrayAttr::get(op->getContext(), padding_low),
+          DenseI64ArrayAttr::get(op->getContext(), padding_high),
+          DenseI64ArrayAttr::get(op->getContext(), padding_interior));
 
       new_op = rewriter.create<pphlo::SelectAndScatterOp>(
           op->getLoc(), materialized_operand.getType(), materialized_operand,
           adaptor.getSource(), materialized_init_value,
-          op.getWindowDimensionsAttr(), op.getWindowStridesAttr());
+          ConvertDenseIntElementAttr(op.getWindowDimensionsAttr()),
+          ConvertDenseIntElementAttr(op.getWindowStridesAttr()));
 
       llvm::SmallVector<int64_t, 2> slice_end(
           new_op.getType().dyn_cast<RankedTensorType>().getShape().begin(),
@@ -1022,15 +1069,18 @@ struct HloToPPHloOpConverter<stablehlo::SelectAndScatterOp>
 
       // Slice back
       rewriter.replaceOpWithNewOp<pphlo::SliceOp>(
-          op, result_type, new_op, ConvertDimensions(&builder, padding_low),
-          ConvertDimensions(&builder, slice_end),
-          ConvertDimensions(&builder,
-                            llvm::SmallVector<int64_t>(slice_end.size(), 1)));
+          op, result_type, new_op,
+          DenseI64ArrayAttr::get(builder.getContext(), padding_low),
+          DenseI64ArrayAttr::get(builder.getContext(), slice_end),
+          DenseI64ArrayAttr::get(
+              builder.getContext(),
+              llvm::SmallVector<int64_t>(slice_end.size(), 1)));
     } else {
       new_op = rewriter.replaceOpWithNewOp<pphlo::SelectAndScatterOp>(
           op, result_type, materialized_operand, adaptor.getSource(),
-          materialized_init_value, op.getWindowDimensionsAttr(),
-          op.getWindowStridesAttr());
+          materialized_init_value,
+          ConvertDenseIntElementAttr(op.getWindowDimensionsAttr()),
+          ConvertDenseIntElementAttr(op.getWindowStridesAttr()));
     }
 
     // Convert the region signature.
@@ -1204,7 +1254,8 @@ class HloToPPHloOpConverter<stablehlo::GatherOp>
 
     rewriter.replaceOpWithNewOp<pphlo::GatherOp>(
         op, resultType, adaptor.getOperands()[0], adaptor.getOperands()[1],
-        attr, op.getSliceSizes(), op.getIndicesAreSorted());
+        attr, ConvertDenseIntElementAttr(op.getSliceSizes()),
+        op.getIndicesAreSorted());
 
     return success();
   }
@@ -1262,17 +1313,15 @@ class HloToPPHloOpConverter<stablehlo::ConvolutionOp>
     }
 
     TypeTools type_tools;
-    auto indexType = rewriter.getIntegerType(64);
-    auto attrType = RankedTensorType::get({rank}, indexType);
     Value zero = rewriter.create<pphlo::ConstantOp>(
         loc, rewriter.getZeroAttr(RankedTensorType::get(
                  {}, type_tools.getExpressedType(inputType.getElementType()))));
     zero = rewriter.create<pphlo::ConvertOp>(
         loc, RankedTensorType::get({}, inputType.getElementType()), zero);
     return rewriter.create<pphlo::PadOp>(
-        loc, input, zero, DenseIntElementsAttr::get(attrType, padLow),
-        DenseIntElementsAttr::get(attrType, padHigh),
-        DenseIntElementsAttr::get(attrType, padInterior));
+        loc, input, zero, DenseI64ArrayAttr::get(loc.getContext(), padLow),
+        DenseI64ArrayAttr::get(loc.getContext(), padHigh),
+        DenseI64ArrayAttr::get(loc.getContext(), padInterior));
   }
 
 public:
@@ -1338,16 +1387,13 @@ class HloToPPHloOpConverter<stablehlo::ConvolutionOp>
 
       modifiedRhs = rewriter.create<pphlo::ReverseOp>(
           op.getLoc(), modifiedRhs,
-          mlir::DenseIntElementsAttr::get(
-              RankedTensorType::get(reversedDims.size(),
-                                    rewriter.getIntegerType(64)),
-              reversedDims));
+          DenseI64ArrayAttr::get(op->getContext(), reversedDims));
     }
 
     rewriter.replaceOpWithNewOp<pphlo::ConvolutionOp>(
         op, resultType, modifiedLhs, modifiedRhs,
-        op.getWindowStrides().value_or(nullptr), attr,
-        op.getFeatureGroupCount(), op.getBatchGroupCount());
+        ConvertDenseIntElementAttr(op.getWindowStrides().value_or(nullptr)),
+        attr, op.getFeatureGroupCount(), op.getBatchGroupCount());
 
     return success();
   }