Legalize newly added compiler hints ops. (#18689)

This was missed in the prior patch adding the util.assume.* ops. Adds tests, including for the untested util.optimization_barrier op. Signed-off-by: Stella Laurenzo <[email protected]>
iree-org · Oct 3, 2024 · 79e979f · 79e979f
1 parent a9f5628
commit 79e979f
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 2 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Util/Conversion/ConversionPatterns.cpp b/compiler/src/iree/compiler/Dialect/Util/Conversion/ConversionPatterns.cpp
@@ -26,12 +26,35 @@ namespace mlir::iree_compiler {
 // Utilities
 //===----------------------------------------------------------------------===//
 
+namespace {
+
+struct ConvertAssumeNarrowOp
+    : public OpConversionPattern<IREE::Util::AssumeNarrowOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(IREE::Util::AssumeNarrowOp assumeOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // The op is a pass-through but we leave the narrow type unconverted since
+    // it is a hint and should not itself be subject to type conversion
+    // (narrowing, widening, etc).
+    rewriter.replaceOpWithNewOp<IREE::Util::AssumeNarrowOp>(
+        assumeOp, adaptor.getOperand().getType(), adaptor.getOperand(),
+        assumeOp.getNarrowType());
+    return success();
+  }
+};
+
+} // namespace
+
 void populateUtilConversionPatterns(MLIRContext *context,
                                     TypeConverter &typeConverter,
                                     RewritePatternSet &patterns) {
   patterns
-      .insert<GenericConvertTypesPattern<IREE::Util::OptimizationBarrierOp>>(
+      .insert<GenericConvertTypesPattern<IREE::Util::AssumeDivisibleOp>,
+              GenericConvertTypesPattern<IREE::Util::AssumeRangeOp>,
+              GenericConvertTypesPattern<IREE::Util::OptimizationBarrierOp>>(
           typeConverter, context);
+  patterns.insert<ConvertAssumeNarrowOp>(typeConverter, context);
 
   typeConverter.addConversion([&](IREE::Util::PtrType type,
                                   SmallVectorImpl<Type> &results) {
@@ -60,6 +83,11 @@ void populateUtilConversionPatterns(MLIRContext *context,
                                     ConversionTarget &conversionTarget,
                                     TypeConverter &typeConverter,
                                     RewritePatternSet &patterns) {
+  addGenericLegalOp<IREE::Util::AssumeDivisibleOp>(conversionTarget,
+                                                   typeConverter);
+  addGenericLegalOp<IREE::Util::AssumeNarrowOp>(conversionTarget,
+                                                typeConverter);
+  addGenericLegalOp<IREE::Util::AssumeRangeOp>(conversionTarget, typeConverter);
   addGenericLegalOp<IREE::Util::OptimizationBarrierOp>(conversionTarget,
                                                        typeConverter);
   addGenericLegalOp<IREE::Util::ListCreateOp>(conversionTarget, typeConverter);

diff --git a/compiler/src/iree/compiler/Dialect/Util/Conversion/test/BUILD.bazel b/compiler/src/iree/compiler/Dialect/Util/Conversion/test/BUILD.bazel
@@ -15,7 +15,10 @@ package(
 iree_lit_test_suite(
     name = "lit",
     srcs = enforce_glob(
-        ["structural_ops.mlir"],
+        [
+            "compiler_hints.mlir",
+            "structural_ops.mlir",
+        ],
         include = ["*.mlir"],
     ),
     cfg = "//compiler:lit.cfg.py",

diff --git a/compiler/src/iree/compiler/Dialect/Util/Conversion/test/CMakeLists.txt b/compiler/src/iree/compiler/Dialect/Util/Conversion/test/CMakeLists.txt
@@ -14,6 +14,7 @@ iree_lit_test_suite(
   NAME
     lit
   SRCS
+    "compiler_hints.mlir"
     "structural_ops.mlir"
   TOOLS
     FileCheck

diff --git a/compiler/src/iree/compiler/Dialect/Util/Conversion/test/compiler_hints.mlir b/compiler/src/iree/compiler/Dialect/Util/Conversion/test/compiler_hints.mlir
@@ -0,0 +1,32 @@
+// RUN: iree-opt --split-input-file --pass-pipeline='builtin.module(iree-util-test-conversion{widen-integers})' %s | FileCheck %s
+
+// CHECK-LABEL: @assumeDivisibleOp
+util.func public @assumeDivisibleOp(%arg0 : i16) -> i16 {
+  // CHECK: util.assume.divisible {{.*}} by 4 : i32
+  %0 = util.assume.divisible %arg0 by 4 : i16
+  util.return %0 : i16
+}
+
+// -----
+// CHECK-LABEL: @assumeNarrowOp
+util.func public @assumeNarrowOp(%arg0 : i16) -> i16 {
+  // CHECK: util.assume.narrow %arg0 : i32 to i8
+  %0 = util.assume.narrow %arg0 : i16 to i8
+  util.return %0 : i16
+}
+
+// -----
+// CHECK-LABEL: @assumeRangeOp
+util.func public @assumeRangeOp(%arg0 : i16) -> i16 {
+  // CHECK: util.assume.range %arg0 in [4, 12] : i32
+  %0 = util.assume.range %arg0 in [4, 12] : i16
+  util.return %0 : i16
+}
+
+// -----
+// CHECK-LABEL: @optimizationBarrier
+util.func public @optimizationBarrier(%arg0 : i16) -> i16 {
+  // CHECK: util.optimization_barrier %arg0 : i32
+  %0 = util.optimization_barrier %arg0 : i16
+  util.return %0 : i16
+}