updates

compiler/src/iree/compiler/Codegen/Common/test/CMakeLists.txt

@@ -64,6 +64,7 @@ iree_lit_test_suite(
     "repeated_matcher_use.mlir"
     "replace_slow_min_max_ops.mlir"
     "strip_compilation_info.mlir"
+    "subbyte_vectorize.mlir"
     "test_partitionable_loops_interface.mlir"
     "tile_and_distribute_to_workgroups.mlir"
     "tile_and_distribute_to_workgroups_func_scope.mlir"

compiler/src/iree/compiler/Dialect/Stream/Transforms/test/BUILD.bazel

@@ -27,6 +27,7 @@ iree_lit_test_suite(
             "emplace_allocations.mlir",
             "encode_device_tensors.mlir",
             "encode_device_tensors_packing.mlir",
+            "encode_device_tensors_packing_i1.mlir",
             "encode_host_tensors.mlir",
             "encode_host_tensors_packing.mlir",
             "fold_globals.mlir",

compiler/src/iree/compiler/Dialect/Stream/Transforms/test/CMakeLists.txt

@@ -25,6 +25,7 @@ iree_lit_test_suite(
     "emplace_allocations.mlir"
     "encode_device_tensors.mlir"
     "encode_device_tensors_packing.mlir"
+    "encode_device_tensors_packing_i1.mlir"
     "encode_host_tensors.mlir"
     "encode_host_tensors_packing.mlir"
     "fold_globals.mlir"

compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing.mlir

@@ -42,6 +42,7 @@ util.func public @denseTensorConstantI4() -> !stream.resource<constant> {
 // Checks that non-byte-aligned total size is not supported for constant.
 
 util.func public @denseTensorConstantI4() -> !stream.resource<constant> {
+  // expected-error @+1 {{failed to calculate total byte count: 'tensor<5xi4>' does not have integral number of total bytes}}
   %0 = stream.tensor.constant : tensor<5xi4> in !stream.resource<constant> = dense<[5, 15, 0, 3, 1]> : tensor<5xi4>
   util.return %0 : !stream.resource<constant>
 }
@@ -72,6 +73,7 @@ util.func public @denseTensorSizeOfStatic() -> index {
 // Checks that non-byte-aligned total size is not supported for sizeof.
 
 util.func public @denseTensorSizeOfStatic() -> index {
+  // expected-error @+1 {{failed to calculate total byte count: 'tensor<11xi4>' does not have integral number of total bytes}}
   %0 = stream.tensor.sizeof tensor<11xi4> : index
   util.return %0 : index
 }

compiler/src/iree/compiler/Dialect/Stream/Transforms/test/encode_host_tensors_packing_i1.mlir

@@ -0,0 +1,20 @@
+// RUN: iree-opt --split-input-file --iree-stream-encode-host-tensors --verify-diagnostics --iree-experimental-packed-i1-storage %s | FileCheck %s
+
+func.func @unaligned_i1_size() -> index {
+  %0 = stream.tensor.sizeof tensor<12xi1> : index
+  return %0 : index
+}
+// CHECK: func @unaligned_i1_size() -> index {
+// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+// CHECK: return %[[C2]] : index
+
+// -----
+
+func.func @aligned_i1_size() -> index {
+  %0 = stream.tensor.sizeof tensor<24xi1> : index
+  return %0 : index
+}
+
+// CHECK: func @aligned_i1_size() -> index {
+// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
+// CHECK: return %[[C3]] : index

compiler/src/iree/compiler/Utils/ElementPackingUtils.cpp

@@ -131,11 +131,11 @@ Value calculateStorageElementCountInBytes(Location loc,
       if (!clEnableI1Support) {
         return nullptr;
       }
+      // in case of subbyte element types, do ceil division as we need to pack
+      // at the back
       value = builder.createOrFold<arith::CeilDivUIOp>(loc, value, divisor);
-    }
-    // in case of subbyte element types, do ceil division as we need to pack at
-    // the back
-    value = builder.createOrFold<arith::DivUIOp>(loc, value, divisor);
+    } else
+      value = builder.createOrFold<arith::DivUIOp>(loc, value, divisor);
   }
 
   return value;