CHLO defns for a ragged dot that permits ragged batch and contraction.

PiperOrigin-RevId: 706767602

Author: pravnar

mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo_pass.cc

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <optional>
 #include <utility>
 
 #include "mhlo/IR/hlo_ops.h"
@@ -22,9 +23,12 @@ limitations under the License.
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "stablehlo/dialect/ChloOps.h"
@@ -56,7 +60,8 @@ struct ChloLegalizeToHighLevelMhloPass
     // Consider the mhlo dialect legal for tests. Also add helper dialects
     // that are needed by the patterns.
     conversionTarget.addLegalDialect<chlo::ChloDialect, mhlo::MhloDialect>();
-    conversionTarget.addIllegalOp<chlo::TopKOp, chlo::ErfOp>();
+    conversionTarget
+        .addIllegalOp<chlo::TopKOp, chlo::ErfOp, chlo::RaggedDotOp>();
 
     if (failed(applyPartialConversion(getOperation(), conversionTarget,
                                       std::move(conversionPatterns)))) {
@@ -93,6 +98,64 @@ struct ChloLegalizeToHloPass
   }
 };
 
+struct RaggedDotChloToMhlo : public OpRewritePattern<chlo::RaggedDotOp> {
+  using OpRewritePattern<chlo::RaggedDotOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(chlo::RaggedDotOp raggedDotOp,
+                                PatternRewriter &rewriter) const override {
+    auto moduleOp = raggedDotOp->getParentOfType<ModuleOp>();
+
+    OpBuilder builder(moduleOp.getBodyRegion());
+    builder.setInsertionPointToStart(&moduleOp.getBodyRegion().front());
+
+    auto chloRaggedDotDimNums = raggedDotOp.getRaggedDotDimensionNumbers();
+    auto dotDimNums = mhlo::DotDimensionNumbersAttr::get(
+        builder.getContext(), chloRaggedDotDimNums.getLhsBatchingDimensions(),
+        chloRaggedDotDimNums.getRhsBatchingDimensions(),
+        chloRaggedDotDimNums.getLhsContractingDimensions(),
+        chloRaggedDotDimNums.getRhsContractingDimensions());
+    auto raggedDotDimNums = mhlo::RaggedDotDimensionNumbersAttr::get(
+        builder.getContext(), dotDimNums,
+        chloRaggedDotDimNums.getLhsRaggedDimensions(),
+        chloRaggedDotDimNums.getRhsGroupDimensions());
+
+    auto mhloPrecision =
+        [](chlo::Precision precision) -> std::optional<mhlo::Precision> {
+      switch (precision) {
+        case chlo::Precision::DEFAULT:
+          return mhlo::Precision::DEFAULT;
+        case chlo::Precision::HIGH:
+          return mhlo::Precision::HIGH;
+        case chlo::Precision::HIGHEST:
+          return mhlo::Precision::HIGHEST;
+      }
+    };
+    ArrayAttr precisionConfig = rewriter.getArrayAttr({});
+    if (raggedDotOp.getPrecisionConfig().has_value()) {
+      SmallVector<Attribute> vector;
+      for (auto configValue : raggedDotOp.getPrecisionConfig()
+                                  .value()
+                                  .getAsRange<chlo::PrecisionAttr>()) {
+        vector.push_back(
+            PrecisionAttr::get(raggedDotOp.getContext(),
+                               mhloPrecision(configValue.getValue()).value()));
+      }
+      precisionConfig = rewriter.getArrayAttr(vector);
+    }
+
+    rewriter.replaceOp(
+        raggedDotOp,
+        rewriter
+            .create<mhlo::RaggedDotOp>(
+                raggedDotOp.getLoc(), raggedDotOp.getResult().getType(),
+                raggedDotOp.getLhs(), raggedDotOp.getRhs(),
+                raggedDotOp.getGroupSizes(), raggedDotDimNums, precisionConfig)
+            .getOperation());
+
+    return success();
+  }
+};
+
 }  // namespace
 
 }  // namespace mhlo
@@ -105,6 +168,7 @@ namespace {
 
 void populateChloToHighLevelMhloOpPatterns(MLIRContext *,
                                            RewritePatternSet *patterns) {
+  patterns->add<mhlo::RaggedDotChloToMhlo>(patterns->getContext());
   populateWithGenerated(*patterns);
 }
 

stablehlo/BUILD

@@ -315,6 +315,7 @@ cc_library(
         ":chlo_attrs_inc_gen",
         ":chlo_enums_inc_gen",
         ":chlo_ops_inc_gen",
+        ":stablehlo_assembly_format",
         ":stablehlo_type_inference",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:BytecodeOpInterface",

stablehlo/stablehlo/dialect/ChloEnums.td

@@ -70,4 +70,29 @@ def CHLO_ComparisonType : I32EnumAttr<"ComparisonType",
 
 def CHLO_ComparisonTypeAttr : EnumAttr<CHLO_Dialect, CHLO_ComparisonType, "comparison_type">;
 
+//===----------------------------------------------------------------------===//
+// Ragged dot op definitions.
+//===----------------------------------------------------------------------===//
+
+// These mirror the XLA PrecisionConfig proto enum.
+def CHLO_PRECISION_DEFAULT : I32EnumAttrCase<"DEFAULT", 0>;
+def CHLO_PRECISION_HIGH    : I32EnumAttrCase<"HIGH", 1>;
+def CHLO_PRECISION_HIGHEST : I32EnumAttrCase<"HIGHEST", 2>;
+
+def CHLO_Precision : I32EnumAttr<"Precision",
+    "XLA precision for an operand. Has backend specific meaning.",
+    [
+      CHLO_PRECISION_DEFAULT,
+      CHLO_PRECISION_HIGH,
+      CHLO_PRECISION_HIGHEST
+    ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::chlo";
+}
+
+def CHLO_PrecisionAttr : EnumAttr<CHLO_Dialect, CHLO_Precision, "precision">;
+
+def CHLO_PrecisionConfigAttr:
+  TypedArrayAttrBase<CHLO_PrecisionAttr, "Precision Config attribute">;
+
 #endif // STABLEHLO_DIALECT_CHLO_ENUMS