Canonicalize std.dim into flow.dispatch.shape when relevant.

iree/compiler/Dialect/Flow/IR/BUILD

@@ -55,6 +55,7 @@ cc_library(
         "//iree/compiler/Dialect/Shape/IR",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InferTypeOpInterface",
         "@llvm-project//mlir:Parser",
         "@llvm-project//mlir:SideEffects",
         "@llvm-project//mlir:StandardOps",
@@ -77,6 +78,7 @@ gentbl(
         "//iree/compiler/Dialect/Shape/IR:td_files",
         "@llvm-project//mlir:OpBaseTdFiles",
         "@llvm-project//mlir:StdOpsTdFiles",
+        "@llvm-project//mlir:include/mlir/Interfaces/InferTypeOpInterface.td",
     ],
 )
 
@@ -94,6 +96,7 @@ gentbl(
         "//iree/compiler/Dialect/Shape/IR:td_files",
         "@llvm-project//mlir:OpBaseTdFiles",
         "@llvm-project//mlir:StdOpsTdFiles",
+        "@llvm-project//mlir:include/mlir/Interfaces/InferTypeOpInterface.td",
     ],
 )
 
@@ -112,6 +115,7 @@ gentbl(
         "@llvm-project//mlir:OpBaseTdFiles",
         "@llvm-project//mlir:SideEffectTdFiles",
         "@llvm-project//mlir:StdOpsTdFiles",
+        "@llvm-project//mlir:include/mlir/Interfaces/InferTypeOpInterface.td",
     ],
 )
 
@@ -129,5 +133,6 @@ iree_tablegen_doc(
         "@llvm-project//mlir:OpBaseTdFiles",
         "@llvm-project//mlir:SideEffectTdFiles",
         "@llvm-project//mlir:StdOpsTdFiles",
+        "@llvm-project//mlir:include/mlir/Interfaces/InferTypeOpInterface.td",
     ],
 )

iree/compiler/Dialect/Flow/IR/CMakeLists.txt

@@ -38,6 +38,7 @@ iree_cc_library(
   DEPS
     LLVMSupport
     MLIRIR
+    MLIRInferTypeOpInterface
     MLIRParser
     MLIRSideEffectInterfaces
     MLIRStandard

iree/compiler/Dialect/Flow/IR/FlowDialect.cpp

@@ -54,6 +54,7 @@ FlowDialect::FlowDialect(MLIRContext *context)
   addOperations<
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.cpp.inc"
       >();
+  context->getOrLoadDialect("shapex");
 }
 
 Operation *FlowDialect::materializeConstant(OpBuilder &builder, Attribute value,

iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp

@@ -289,6 +289,53 @@ void DispatchWorkgroupsOp::getCanonicalizationPatterns(
   // results.insert<DceDispatchWorkgroups>(context);
 }
 
+//===----------------------------------------------------------------------===//
+// flow.dispatch.input.load
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+// Some linalg patterns, due to being upstream, tend to introduce `dim` ops.
+// These generally fold with upstream patterns when tensors are involved, but
+// when DispatchInputLoadOp's are involved (with dispatch tensor types),
+// then this starts to break down, which causes the `dim` ops to survive
+// arbitrarily late into the pipeline. Often, they keep alive
+// DispatchInputLoadOp's that would otherwise be dead!
+//
+// To fix this, we convert the `std.dim` ops to `flow.dispatch.shape` ops.
+// ```
+// dim(flow.dispatch.input.load(%x), %const)
+// ->
+// shapex.ranked_dim(flow.dispatch.shape(%x), %const)
+// ``
+struct ConvertDimOfDispatchInputLoadToDispatchShape
+    : public OpRewritePattern<DimOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DimOp op,
+                                PatternRewriter &rewriter) const override {
+    auto dispatchInputLoad =
+        op.memrefOrTensor().getDefiningOp<DispatchInputLoadOp>();
+    if (!dispatchInputLoad) return failure();
+
+    Optional<int64_t> constantIndex = op.getConstantIndex();
+    if (!constantIndex.hasValue()) return failure();
+
+    auto rankedShape = rewriter.create<DispatchShapeOp>(
+        op.getLoc(), dispatchInputLoad.source());
+    rewriter.replaceOpWithNewOp<Shape::RankedDimOp>(op, rankedShape,
+                                                    *constantIndex);
+    return success();
+  }
+};
+
+}  // namespace
+
+void DispatchInputLoadOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<ConvertDimOfDispatchInputLoadToDispatchShape>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // flow.dispatch.workgroup.*
 //===----------------------------------------------------------------------===//

iree/compiler/Dialect/Flow/IR/FlowOps.cpp

@@ -735,6 +735,17 @@ void DispatchShapeOp::getAsmResultNames(
   setNameFn(result(), "shape");
 }
 
+LogicalResult DispatchShapeOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  auto dispatchTensorType = operands[0].getType().cast<DispatchTensorType>();
+  auto shape = dispatchTensorType.getShape();
+  auto rankedShapeType = Shape::RankedShapeType::get(shape, context);
+  inferredReturnTypes.assign({rankedShapeType});
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // flow.executable
 //===----------------------------------------------------------------------===//

iree/compiler/Dialect/Flow/IR/FlowOps.h

@@ -28,6 +28,7 @@
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/SymbolTable.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #define GET_OP_CLASSES

iree/compiler/Dialect/Flow/IR/FlowOps.td

@@ -18,6 +18,7 @@
 include "iree/compiler/Dialect/Flow/IR/FlowBase.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 class FLOW_PureOp<string mnemonic, list<OpTrait> traits = []> :
@@ -450,6 +451,7 @@ def FLOW_DispatchWorkgroupSizeOp : FLOW_PureOp<"dispatch.workgroup.size", [
 
 def FLOW_DispatchShapeOp : FLOW_PureOp<"dispatch.shape", [
     DeclareOpInterfaceMethods<OpAsmOpInterface>,
+    DeclareOpInterfaceMethods<InferTypeOpInterface>
   ]> {
   let summary = [{returns the shape of a dispatch region input/output tensor}];
   let description = [{
@@ -522,6 +524,8 @@ def FLOW_DispatchInputLoadOp : FLOW_PureOp<"dispatch.input.load", [
     ( `,` `strides` `=` `[` $strides^ `]` )?
     `:` type($source) `->` type($result) attr-dict-with-keyword
   }];
+
+  let hasCanonicalizer = 1;
 }
 
 def FLOW_DispatchOutputStoreOp : FLOW_Op<"dispatch.output.store", [

iree/compiler/Dialect/Flow/IR/test/dispatch_workgroups_folding.mlir

@@ -53,3 +53,18 @@ func @workgroupRankFolding(%arg0 : tensor<?x4xf32>) -> tensor<4x?xf32> {
   }
   return %0 : tensor<4x?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @convertDimOfDispatchInputLoadToDispatchShape
+// CHECK-SAME:    %[[ARG:.*]]: !flow.dispatch.input<?xf32>) {
+func @convertDimOfDispatchInputLoadToDispatchShape(%arg0: !flow.dispatch.input<?xf32>) {
+  // CHECK-NEXT: %[[RANKED_SHAPE:.*]] = flow.dispatch.shape %[[ARG]]
+  // CHECK-NEXT: %[[DIM:.*]] = shapex.ranked_dim %[[RANKED_SHAPE]][0]
+  // CHECK-NEXT: "test.sink"(%[[DIM]]) : (index) -> ()
+  %tensor = flow.dispatch.input.load %arg0 : !flow.dispatch.input<?xf32> -> tensor<?xf32>
+  %c0 = constant 0 : index
+  %dim = dim %tensor, %c0 : tensor<?xf32>
+  "test.sink"(%dim) : (index) -> ()
+  return
+}