Merge pull request iree-org#5083 from google/benvanik-flow-passes

hanhanW · Mar 12, 2021 · 1ef76ba · 1ef76ba
2 parents 0668b55 + bb819df
commit 1ef76ba
Show file tree

Hide file tree

Showing 20 changed files with 343 additions and 1,195 deletions.
diff --git a/iree/compiler/Conversion/LinalgToSPIRV/VectorizeMemref.cpp b/iree/compiler/Conversion/LinalgToSPIRV/VectorizeMemref.cpp
@@ -50,7 +50,6 @@ static bool getUsesIfAllTransferOp(Value v,
   return true;
 }
 
-
 /// Returns the bitwidth of a scalar or vector type.
 static Optional<unsigned> getBitWidth(Type type) {
   if (type.isIntOrFloat()) {
@@ -370,7 +369,7 @@ class ProcessPlaceHolder final
     auto vecMemRef = getVectorizedMemRefType(rewriter, placeholder.getResult());
     if (!vecMemRef) return failure();
     rewriter.replaceOpWithNewOp<IREE::PlaceholderOp>(
-        placeholder, *vecMemRef, ValueRange(), placeholder.getAttrs());
+        placeholder, *vecMemRef, ValueRange(), placeholder->getAttrs());
     return success();
   }
 };

diff --git a/iree/compiler/Dialect/Flow/Conversion/TypeConverter.cpp b/iree/compiler/Dialect/Flow/Conversion/TypeConverter.cpp
@@ -23,40 +23,43 @@ namespace iree_compiler {
 FlowTypeConverter::FlowTypeConverter() {
   // Allow types through by default.
   addConversion([](Type type) { return type; });
-  // addConversion([](IndexType type) {
-  //   // Always treat as 32-bit.
-  //   return IntegerType::get(32, type.getContext());
-  // });
+
+  addConversion([this](RankedTensorType tensorType) -> Optional<Type> {
+    auto convertedElementType = convertType(tensorType.getElementType());
+    if (!convertedElementType) {
+      return llvm::None;
+    }
+    return RankedTensorType::get(tensorType.getShape(), convertedElementType);
+  });
+
+  addConversion([](UnrankedTensorType tensorType) {
+    // We only support ranked tensors. We could convert unranked to ranked
+    // here for certain cases (such as * on the LHS).
+    return Type();
+  });
+
+  // UNSAFE: narrow 64-bit integers to 32-bit ones.
+  // This is a workaround for lower levels of the stack not always supporting
+  // 64-bit types natively.
+  // TODO(benvanik): make whether to narrow integers an option.
   addConversion([](IntegerType integerType) -> Optional<Type> {
     if (integerType.isSignlessInteger() && integerType.getWidth() > 32) {
       // Don't support 64-bit types in general. Rewrite to i32 (if desired).
-      // TODO(benvanik): split to i32+i32? allow and use availability?
-      // TODO(benvanik): make an option.
       return IntegerType::get(integerType.getContext(), 32);
     }
     return llvm::None;
   });
+
+  // UNSAFE: narrow 64-bit floats to 32-bit ones.
+  // This is a workaround for lower levels of the stack not always supporting
+  // 64-bit types natively.
+  // TODO(benvanik): make whether to narrow floats an option.
   addConversion([](FloatType floatType) -> Optional<Type> {
     if (floatType.getWidth() > 32) {
-      // Don't support 64-bit types in general. Rewrite to f32 (if desired).
-      // TODO(benvanik): make an option.
       return FloatType::getF32(floatType.getContext());
     }
     return llvm::None;
   });
-  addConversion([this](RankedTensorType tensorType) -> Optional<Type> {
-    auto convertedElementType = convertType(tensorType.getElementType());
-    if (!convertedElementType) {
-      return llvm::None;
-    }
-    return RankedTensorType::get(tensorType.getShape(), convertedElementType);
-  });
-  addConversion([](UnrankedTensorType tensorType) {
-    // We only support ranked tensors. We could convert unranked to ranked
-    // here for certain cases (such as * on the LHS).
-    return Type();
-  });
-  // TODO(b/145876978): add conversion materializer
 }
 
 }  // namespace iree_compiler

diff --git a/iree/compiler/Dialect/Flow/Transforms/BUILD b/iree/compiler/Dialect/Flow/Transforms/BUILD
@@ -35,12 +35,10 @@ cc_library(
         "FormStreams.cpp",
         "HLOToHLOPreprocessing.cpp",
         "HoistUnstreamableOps.cpp",
-        "IdentifyDispatchRegions.cpp",
         "IdentifyDispatchRegions2.cpp",
         "InjectDispatchTracing.cpp",
         "LegalizeInputTypes.cpp",
-        "MaterializeExportedReflection.cpp",
-        "MergeExportedReflection.cpp",
+        "MaterializeReflectionAttrs.cpp",
         "OutlineDispatchRegions.cpp",
         "OutlineDispatchRegions2.cpp",
         "OutlineLargeConstants.cpp",

diff --git a/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt b/iree/compiler/Dialect/Flow/Transforms/CMakeLists.txt
@@ -30,12 +30,10 @@ iree_cc_library(
     "FormStreams.cpp"
     "HLOToHLOPreprocessing.cpp"
     "HoistUnstreamableOps.cpp"
-    "IdentifyDispatchRegions.cpp"
     "IdentifyDispatchRegions2.cpp"
     "InjectDispatchTracing.cpp"
     "LegalizeInputTypes.cpp"
-    "MaterializeExportedReflection.cpp"
-    "MergeExportedReflection.cpp"
+    "MaterializeReflectionAttrs.cpp"
     "OutlineDispatchRegions.cpp"
     "OutlineDispatchRegions2.cpp"
     "OutlineLargeConstants.cpp"