Check if type is int or float

Signed-off-by: Ian Wood <ianwood2024@u.northwestern.edu>
iree-org · Nov 16, 2024 · 2d3bd37 · 2d3bd37
1 parent 60cf4ab
commit 2d3bd37
Showing 1 changed file with 12 additions and 10 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -2922,16 +2922,18 @@ setLoweringConfigForComputeOps(mlir::FunctionOpInterface entryPointFn,
   // loads and stores will have a performance impact.
   auto resultTypes = rootOperation->getResultTypes();
   if (commonVecTileSizes.size() != 0 && !resultTypes.empty()) {
-    auto elementTypeSize =
-        cast<ShapedType>(rootOperation->getResultTypes().front())
-            .getElementType()
-            .getIntOrFloatBitWidth();
-    // for now just enable for i1
-    if (elementTypeSize == 1) {
-      auto innermostTileSize = commonVecTileSizes.back();
-      commonVecTileSizes.back() =
-          llvm::alignTo(innermostTileSize * elementTypeSize, 8) /
-          elementTypeSize;
+    auto resultType = cast<ShapedType>(rootOperation->getResultTypes().front())
+                          .getElementType();
+    if (resultType.isIntOrFloat()) {
+      auto elementTypeSize = resultType.getIntOrFloatBitWidth();
+
+      // for now just enable for i1
+      if (elementTypeSize == 1) {
+        auto innermostTileSize = commonVecTileSizes.back();
+        commonVecTileSizes.back() =
+            llvm::alignTo(innermostTileSize * elementTypeSize, 8) /
+            elementTypeSize;
+      }
     }
   }