[StableHLO][CHLO] Add CHLO to stablehlo pipeline (iree-org#13835)

Unlike in the MHLO pipeline, run CHLO legalization as a separate pass to better controll required canonicalization patterns. This also decouples the CHLO conversion from the custom type converter used during lowering to linalg and IREE dialects. Also add a pattern to handle non-broadcasting constants. Issue: iree-org#13803
kuhar · May 29, 2023 · b218383 · b218383
1 parent 4c2c575
commit b218383
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 5 deletions.
diff --git a/compiler/src/iree/compiler/InputConversion/StableHLO/LegalizeCHLO.cpp b/compiler/src/iree/compiler/InputConversion/StableHLO/LegalizeCHLO.cpp
@@ -279,6 +279,17 @@ struct ConvertRankedDynamicBroadcastBinaryOp final
   }
 };
 
+struct ConvertConstantOp final : OpConversionPattern<mlir::chlo::ConstantOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      mlir::chlo::ConstantOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<mlir::stablehlo::ConstantOp>(op, op.getValue());
+    return success();
+  }
+};
+
 struct ConvertConstantLikeOp final
     : OpConversionPattern<mlir::chlo::ConstantLikeOp> {
   using OpConversionPattern::OpConversionPattern;
@@ -483,8 +494,7 @@ void populateLegalizeChloPatterns(MLIRContext* context,
       context, patterns, 10);
   populateForBroadcastingBinaryOp<ConvertRankedDynamicBroadcastBinaryOp>(
       context, patterns, 5);
-  patterns
-      ->add<ConvertConstantLikeOp, ConvertDynamicReshapeOp, ConvertSelectOp>(
-          context);
+  patterns->add<ConvertConstantOp, ConvertConstantLikeOp,
+                ConvertDynamicReshapeOp, ConvertSelectOp>(context);
 }
 }  // namespace mlir::iree_compiler::stablehlo
diff --git a/compiler/src/iree/compiler/InputConversion/StableHLO/Passes.cpp b/compiler/src/iree/compiler/InputConversion/StableHLO/Passes.cpp
@@ -110,6 +110,7 @@ void buildStableHLOInputConversionPassPipelineImpl(OpPassManager &passManager,
       stablehlo::createLegalizeShapeComputations());
   passManager.addNestedPass<func::FuncOp>(
       stablehlo::createConvertStableHloToLinalgExt());
+  passManager.addNestedPass<func::FuncOp>(stablehlo::createLegalizeChlo());
   passManager.addPass(createConvertStableHloToIreeInputDialects());
   // Ensure conversion completed.
   passManager.addPass(createReconcileUnrealizedCastsPass());

diff --git a/compiler/src/iree/compiler/InputConversion/StableHLO/StableHLOToIREEInputDialects.cpp b/compiler/src/iree/compiler/InputConversion/StableHLO/StableHLOToIREEInputDialects.cpp
@@ -432,8 +432,6 @@ struct ConvertStableHloToIreeInputDialects final
     // expensive expansions.
     populateCanonicalizationPatterns(context, &patterns, /*benefit=*/1024);
 
-    // TODO(#12678): Handle chlo lowering.
-
     populateStableHloToLinalgOnTensorsConversionPatterns(
         context, *typeConverter, &patterns);
     populateStableHloCollectivesConversionPatterns(context, *typeConverter,

diff --git a/compiler/src/iree/compiler/InputConversion/StableHLO/test/legalize_chlo_no_broadcast.mlir b/compiler/src/iree/compiler/InputConversion/StableHLO/test/legalize_chlo_no_broadcast.mlir
@@ -4,6 +4,18 @@
 // Check the non-broadcast case for each registered op, then just check a
 // representative op for detailed broadcast semantics.
 
+// CHECK-LABEL: @constants
+func.func @constants() -> (tensor<4xi32>, tensor<2x2xf32>) {
+  %0 = chlo.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+  %1 = chlo.constant dense<0.0> : tensor<2x2xf32>
+
+  // CHECK-DAG: stablehlo.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+  // CHECK-DAG: stablehlo.constant dense<0.000000e+00> : tensor<2x2xf32>
+  func.return %0, %1 : tensor<4xi32>, tensor<2x2xf32>
+}
+
+// -----
+
 // CHECK-LABEL: @addWithoutBroadcast
 func.func @addWithoutBroadcast(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
   // CHECK: stablehlo.add %arg0, %arg1