[Codegen] Implement ValueBoundsOpInterface and fully migrate vectorization for ToLayoutOp.

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/BUILD.bazel

@@ -38,6 +38,7 @@ iree_td_library(
         "@llvm-project//mlir:InferTypeOpInterfaceTdFiles",
         "@llvm-project//mlir:OpBaseTdFiles",
         "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+        "@llvm-project//mlir:ValueBoundsOpInterfaceTdFiles",
         "@llvm-project//mlir:VectorInterfacesTdFiles",
     ],
 )
@@ -82,6 +83,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:SideEffectInterfaces",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:ValueBoundsOpInterface",
         "@llvm-project//mlir:VectorDialect",
         "@llvm-project//mlir:VectorInterfaces",
     ],

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/CMakeLists.txt

@@ -45,6 +45,7 @@ iree_cc_library(
     MLIRSideEffectInterfaces
     MLIRSupport
     MLIRTensorDialect
+    MLIRValueBoundsOpInterface
     MLIRVectorDialect
     MLIRVectorInterfaces
     iree::compiler::Codegen::Utils::VectorOpUtils

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtDialect.h

@@ -14,6 +14,7 @@
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "mlir/Interfaces/VectorInterfaces.h"
 
 // clang-format off: must be included after all LLVM/MLIR headers

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.cpp

@@ -11,6 +11,7 @@
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 using namespace mlir;
 using namespace mlir::iree_compiler::IREE::VectorExt;
@@ -26,6 +27,11 @@ LogicalResult ToLayoutOp::verify() {
   return getLayout().isValidLayout(getInput().getType(), getLoc());
 }
 
+void ToLayoutOp::populateBoundsForShapedValueDim(
+    Value value, int64_t dim, ValueBoundsConstraintSet &cstr) {
+  cstr.bound(value)[dim] == cstr.getExpr(getInput(), dim);
+}
+
 // to_simd -> to_simt
 OpFoldResult ToSIMDOp::fold(FoldAdaptor) {
   if (auto simtOp = getOperand().getDefiningOp<ToSIMTOp>()) {

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtOps.td

@@ -14,6 +14,7 @@ include "iree/compiler/Codegen/Dialect/VectorExt/IR/VectorExtInterfaces.td"
 include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/ValueBoundsOpInterface.td"
 
 //===----------------------------------------------------------------------===//
 // Base class.
@@ -29,7 +30,9 @@ class IREEVectorExt_PureOp<string mnemonic, list<Trait> traits = []> :
 
 def IREEVectorExt_ToLayoutOp : IREEVectorExt_PureOp<"to_layout", [
   Pure,
-  AllTypesMatch<["input", "output"]>
+  AllTypesMatch<["input", "output"]>,
+  DeclareOpInterfaceMethods<ValueBoundsOpInterface,
+    ["populateBoundsForShapedValueDim"]>
   ]> {
   let summary = [{Layout conversion operator.}];
   let description = [{

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/BUILD.bazel

@@ -40,7 +40,6 @@ iree_compiler_cc_library(
         "LowerTransferGatherOps.cpp",
         "Passes.cpp",
         "VectorExtFoldUnitExtentDims.cpp",
-        "VectorizeIREEVectorExtOps.cpp",
     ],
     hdrs = [
         "BufferizationInterfaces.h",
@@ -52,7 +51,6 @@ iree_compiler_cc_library(
     deps = [
         ":PassesIncGen",
         "//compiler/src/iree/compiler/Codegen/Dialect/VectorExt/IR:IREEVectorExtDialect",
-        "//compiler/src/iree/compiler/Codegen/Interfaces:VectorizableOpInterface",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:BufferizationDialect",

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/CMakeLists.txt

@@ -34,7 +34,6 @@ iree_cc_library(
     "LowerTransferGatherOps.cpp"
     "Passes.cpp"
     "VectorExtFoldUnitExtentDims.cpp"
-    "VectorizeIREEVectorExtOps.cpp"
   DEPS
     ::PassesIncGen
     LLVMSupport
@@ -52,7 +51,6 @@ iree_cc_library(
     MLIRVectorDialect
     MLIRVectorUtils
     iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect
-    iree::compiler::Codegen::Interfaces::VectorizableOpInterface
   PUBLIC
 )
 

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/Passes.td

@@ -9,14 +9,6 @@
 
 include "mlir/Pass/PassBase.td"
 
-def VectorizeIREEVectorExtOpsPass :
-    Pass<"iree-vector-ext-vectorize-ops", ""> {
-  let summary = "Vectorizes then lowers a few iree_vector_ext ops before vectorization.";
-  let dependentDialects = [
-    "::mlir::iree_compiler::IREE::VectorExt::IREEVectorExtDialect"
-  ];
-}
-
 def VectorExtFoldUnitExtentDimsPass :
     Pass<"iree-vector-ext-fold-unit-extent-dims", ""> {
   let summary = "Folds unit dims for iree_vector_ext ops";

compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms/test/vectorize_vector_ext_ops.mlir

@@ -1,4 +1,4 @@
-// RUN: iree-opt %s -pass-pipeline='builtin.module(func.func(iree-vector-ext-vectorize-ops, iree-codegen-generic-vectorization{enable-vector-masking=true}),canonicalize,cse,func.func(iree-codegen-optimize-tensor-insert-extract-slices),canonicalize)' --split-input-file --mlir-print-local-scope | FileCheck %s
+// RUN: iree-opt %s -pass-pipeline='builtin.module(func.func(iree-codegen-generic-vectorization{enable-vector-masking=true}),canonicalize,cse,func.func(iree-codegen-optimize-tensor-insert-extract-slices),canonicalize)' --split-input-file --mlir-print-local-scope | FileCheck %s
 
 #layout = #iree_vector_ext.nested_layout<
   subgroup_tile = [1, 1],

compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp

@@ -14,7 +14,6 @@
 #include "iree/compiler/Codegen/Common/Passes.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
 #include "iree/compiler/Codegen/Dialect/GPU/Transforms/Passes.h"
-#include "iree/compiler/Codegen/Dialect/VectorExt/Transforms/Passes.h"
 #include "iree/compiler/Codegen/LLVMGPU/Passes.h"
 #include "iree/compiler/Codegen/LLVMGPU/ROCDLPasses.h"
 #include "iree/compiler/Codegen/Utils/GPUUtils.h"
@@ -274,8 +273,6 @@ static void addGPUVectorizationPasses(OpPassManager &funcPassManager,
   funcPassManager.addPass(IREE::LinalgExt::createDecomposeIm2colPass());
   funcPassManager.addPass(createCanonicalizerPass());
   funcPassManager.addPass(createCSEPass());
-  funcPassManager.addPass(
-      IREE::VectorExt::createVectorizeIREEVectorExtOpsPass());
   funcPassManager.addPass(IREE::GPU::createVectorizeIREEGPUOpsPass());
   // Vectorize.
   GenericVectorizationPassOptions options;