From 79e979fa349f5376f430b5371afac49bd0d1ab01 Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident@gmail.com>
Date: Thu, 3 Oct 2024 16:19:04 -0700
Subject: [PATCH] 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 <stellaraccident@gmail.com>
---
 .../Util/Conversion/ConversionPatterns.cpp    | 30 ++++++++++++++++-
 .../Dialect/Util/Conversion/test/BUILD.bazel  |  5 ++-
 .../Util/Conversion/test/CMakeLists.txt       |  1 +
 .../Util/Conversion/test/compiler_hints.mlir  | 32 +++++++++++++++++++
 4 files changed, 66 insertions(+), 2 deletions(-)
 create mode 100644 compiler/src/iree/compiler/Dialect/Util/Conversion/test/compiler_hints.mlir

diff --git a/compiler/src/iree/compiler/Dialect/Util/Conversion/ConversionPatterns.cpp b/compiler/src/iree/compiler/Dialect/Util/Conversion/ConversionPatterns.cpp
index a07a41f5f903..ecb3fde54883 100644
--- 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
index c315d4fdfe6d..54a003dd2c21 100644
--- 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
index 9c1bbe91259a..70a5b311fa44 100644
--- 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
new file mode 100644
index 000000000000..5cbc52402eb6
--- /dev/null
+++ 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
+}