interface

Author: lialan

compiler/src/iree/compiler/Codegen/Common/EncodingUtils.cpp

@@ -6,6 +6,7 @@
 
 #include "iree/compiler/Codegen/Common/EncodingUtils.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/Utils/Utils.h"
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
@@ -83,7 +84,8 @@ static RankedTensorType transposeIfNarrowNResult(RankedTensorType tensorType) {
   }
   auto newEncoding = IREE::Encoding::EncodingAttr::get(
       context, operandIndex, encoding.getOpType().getValue(),
-      encoding.getElementTypesArray(), maps, newBcastMap, newRoundDimsTo);
+      encoding.getElementTypesArray(), maps, newBcastMap, newRoundDimsTo,
+      encoding.getLayouts());
   return RankedTensorType::get(newShape, elemType, newEncoding);
 }
 
@@ -133,7 +135,9 @@ MaterializeEncodingConversionTarget::MaterializeEncodingConversionTarget(
   markUnknownOpDynamicallyLegal([](Operation *op) {
     auto typeHasEncoding = [](Type t) -> bool {
       auto tensorType = dyn_cast<RankedTensorType>(t);
-      return tensorType && tensorType.getEncoding();
+      if (!(tensorType && tensorType.getEncoding()))
+        return false;
+      return !IREE::Encoding::hasPackedStorageAttr(tensorType);
     };
     auto valueHasEncoding = [=](Value v) -> bool {
       return typeHasEncoding(v.getType());

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.cpp

@@ -111,6 +111,12 @@ EncodingAttr EncodingAttr::clone(AffineMap bcastMap) {
              AffineMapAttr::get(bcastMap), getRoundDimsTo(), getLayouts());
 }
 
+bool EncodingAttr::i1PackedStorage() const {
+  return llvm::any_of(getLayouts(), [&](const Attribute &a) {
+    return llvm::isa<PackedStorageAttr>(a);
+  });
+}
+
 MatmulNarrowDim getMatmulNarrowDim(EncodingAttr encoding) {
   if (encoding.getOpType().getValue() != EncodingOpType::matmul) {
     return {};
@@ -142,6 +148,10 @@ EncodingAttr getEncodingAttr(RankedTensorType type) {
   return dyn_cast_or_null<EncodingAttr>(type.getEncoding());
 }
 
+bool hasPackedStorageAttr(RankedTensorType type) {
+  return dyn_cast_or_null<PackedStorageAttr>(type.getEncoding()) != nullptr;
+}
+
 FailureOr<linalg::ContractionDimensions>
 getEncodingContractionDims(EncodingAttr encoding) {
   auto indexingMapsAttr = encoding.getUserIndexingMaps();

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingAttrs.td

@@ -40,6 +40,17 @@ def EncodingOpType : IREEEncoding_I32EnumAttr<"EncodingOpType",
 def EncodingOpTypeAttr:
   IREEEncoding_EnumAttr<EncodingOpType, "optype">;
 
+
+def PackedStorageAttr : IREEEncoding_Attr<"PackedStorage"> {
+  let mnemonic = "packed_storage";
+  let summary = [{Indicates packed storage datatype.}];
+  let description = [{
+    This attribute indicates this is a back-to-back packed storage in memory.
+    This attribute takes no arguments.
+  }];
+  let genVerifyDecl = 0;
+}
+
 def EncodingAttr :
     IREEEncoding_Attr<"Encoding"> {
   let mnemonic = "encoding";
@@ -108,6 +119,9 @@ def EncodingAttr :
 
     /// Clones an encoding with a new bcast_map
     EncodingAttr clone(AffineMap bcastMap);
+
+    /// Returns whether the encoding indicates packed storage for i1 datatype.
+    bool i1PackedStorage() const;
   }];
 
   let genVerifyDecl = 0;

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingTypes.h

@@ -38,6 +38,9 @@ namespace mlir::iree_compiler::IREE::Encoding {
 /// Otherwise, returns null.
 EncodingAttr getEncodingAttr(RankedTensorType type);
 
+/// Returns true if the type contains packed_storage attribute.
+bool hasPackedStorageAttr(RankedTensorType type);
+
 /// Returns the ContractionDimensions for the encoding user_indexing_maps.
 FailureOr<linalg::ContractionDimensions>
 getEncodingContractionDims(EncodingAttr encoding);

compiler/src/iree/compiler/Dialect/Stream/Conversion/HALToStream/Patterns.cpp

@@ -12,6 +12,10 @@
 #include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 
+namespace mlir::iree_compiler::IREE::Encoding {
+bool hasPackedStorageAttr(mlir::RankedTensorType);
+} // namespace mlir::iree_compiler::IREE::Encoding
+
 namespace mlir::iree_compiler {
 
 namespace {
@@ -90,6 +94,11 @@ struct ConvertTensorImportOp
                                                RankedTensorType tensorType,
                                                ValueRange dynamicDims,
                                                OpBuilder &builder) {
+    // If the encoding attr is about packed storage then we don't need all this
+    if (IREE::Encoding::hasPackedStorageAttr(tensorType)) {
+      return success();
+    }
+
     auto expectedElementType = builder.create<IREE::HAL::ElementTypeOp>(
         loc, tensorType.getElementType());
     auto expectedEncodingType = builder.create<IREE::HAL::EncodingTypeOp>(

compiler/src/iree/compiler/Dialect/Stream/Transforms/EncodeTensors.cpp

@@ -6,6 +6,7 @@
 
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
 // TODO(benvanik): have a stream/upstream equivalent of the flow.dispatch.* ops.
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
@@ -46,7 +47,8 @@ static LogicalResult checkEncoding(Operation *op, RankedTensorType encodingType,
                                    ValueRange encodingDims,
                                    PatternRewriter &rewriter) {
   auto encoding = encodingType.getEncoding();
-  if (encoding && !llvm::isa<IREE::Encoding::EncodingAttr>(encoding)) {
+  if (encoding && !llvm::isa<IREE::Encoding::EncodingAttr,
+                             IREE::Encoding::PackedStorageAttr>(encoding)) {
     return rewriter.notifyMatchFailure(op, [=](Diagnostic &d) {
       d << "unsupported tensor encoding: " << encodingType;
     });

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

@@ -1,7 +1,8 @@
-// RUN: iree-opt --split-input-file --iree-stream-encode-host-tensors --iree-experimental-packed-i1-storage %s | FileCheck %s
+// RUN: iree-opt --split-input-file --iree-stream-encode-host-tensors %s | FileCheck %s
 
+#packed = #iree_encoding.packed_storage
 func.func @unaligned_i1_size() -> index {
-  %0 = stream.tensor.sizeof tensor<12xi1> : index
+  %0 = stream.tensor.sizeof tensor<12xi1, #packed> : index
   return %0 : index
 }
 // CHECK: func @unaligned_i1_size() -> index {
@@ -10,8 +11,9 @@ func.func @unaligned_i1_size() -> index {
 
 // -----
 
+#packed = #iree_encoding.packed_storage
 func.func @aligned_i1_size() -> index {
-  %0 = stream.tensor.sizeof tensor<24xi1> : index
+  %0 = stream.tensor.sizeof tensor<24xi1, #packed> : index
   return %0 : index
 }
 

compiler/src/iree/compiler/DispatchCreation/SetEncoding.cpp

@@ -12,6 +12,7 @@
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingDialect.h"
 #include "iree/compiler/Dialect/Encoding/IR/EncodingOps.h"
+#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowDialect.h"
 #include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
 #include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
@@ -25,18 +26,33 @@
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
+llvm::cl::opt<bool> clEnableI1Support(
+    "iree-experimental-packed-i1-storage",
+    llvm::cl::desc(
+        "Experimental feature: use packed storage for i1 tensors. This feature "
+        "can be dropped once the frontend can emit packed i1 tensors."),
+    llvm::cl::init(false));
+
 #define DEBUG_TYPE "iree-dispatch-creation-set-encoding"
 
 namespace mlir::iree_compiler::DispatchCreation {
 #define GEN_PASS_DEF_SETENCODINGPASS
 #include "iree/compiler/DispatchCreation/Passes.h.inc"
 
 using IREE::Encoding::EncodingAttr;
+using IREE::Encoding::PackedStorageAttr;
 
 //===---------------------------------------------------------------------===//
 // Utility functions
 //===---------------------------------------------------------------------===//
 
+static std::optional<Attribute> getI1PackedStorageAttr(MLIRContext *context) {
+  if (clEnableI1Support) {
+    return PackedStorageAttr::get(context);
+  }
+  return {};
+}
+
 Value setEncoding(OpBuilder &builder, Location loc, Value source,
                   EncodingAttr encodingAttr) {
   auto sourceType = cast<RankedTensorType>(source.getType());
@@ -216,9 +232,17 @@ class setContractionOpEncoding
       if (narrowDim.isN()) {
         roundDimsTo[1] = llvm::PowerOf2Ceil(narrowDim.size);
       }
-      auto encoding = EncodingAttr::get(linalgOp.getContext(), operandIndex,
-                                        opType, elemTypes, maps,
-                                        /*bcastMap=*/std::nullopt, roundDimsTo);
+      // This is a temporary solution that we can use
+      // `--iree-experimental-packed-i1-storage` to enable packed storage for i1
+      // tensors without specifying attributes.
+      auto optI1PackedStorageAttr =
+          getI1PackedStorageAttr(rewriter.getContext());
+      auto layouts = optI1PackedStorageAttr
+                         ? ArrayRef<Attribute>(*optI1PackedStorageAttr)
+                         : ArrayRef<Attribute>{};
+      auto encoding = EncodingAttr::get(
+          linalgOp.getContext(), operandIndex, opType, elemTypes, maps,
+          /*bcastMap=*/std::nullopt, roundDimsTo, layouts);
       return setEncoding(rewriter, loc, src, encoding);
     };
     Value encodedLhs = setEncodingWrapper(lhs, IREE::Encoding::MATMUL_LHS);

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

@@ -15,15 +15,10 @@
 
 namespace mlir::iree_compiler {
 
-llvm::cl::opt<bool> clEnableI1Support(
-    "iree-experimental-packed-i1-storage",
-    llvm::cl::desc(
-        "Experimental feature: enable i1 data type support in codegen"),
-    llvm::cl::init(false));
-
-bool needToPackSubByteElementBitWidth(unsigned bitWidth) {
+bool needToPackSubByteElementBitWidth(unsigned bitWidth,
+                                      bool isI1PackedStorage) {
   // Enable i1 support if requested.
-  if (clEnableI1Support && bitWidth == 1) {
+  if (isI1PackedStorage && bitWidth == 1) {
     return true;
   }
   // Require the original bit width to be some power of two for now to avoid
@@ -35,18 +30,20 @@ bool needToPackSubByteElementBitWidth(unsigned bitWidth) {
 
 bool needToPackSubByteElements(RankedTensorType shapedType) {
   unsigned bitWidth = IREE::Util::getTypeBitWidth(shapedType.getElementType());
-  return needToPackSubByteElementBitWidth(bitWidth);
+  auto encoding = IREE::Encoding::getEncodingAttr(shapedType);
+  bool isI1PackedStorage = encoding && encoding.i1PackedStorage();
+  return needToPackSubByteElementBitWidth(bitWidth, isI1PackedStorage);
 }
 
-Type legalizeStorageElementType(Type elementType) {
+Type legalizeStorageElementType(Type elementType, bool isI1PackedStorage) {
   // Only handle integers; floats in MLIR all have aligned widths (today).
   auto intType = dyn_cast<IntegerType>(elementType);
   if (!intType)
     return elementType;
 
   // For sub-byte elements, default to pack them into bytes.
   unsigned bitWidth = intType.getWidth();
-  if (needToPackSubByteElementBitWidth(bitWidth))
+  if (needToPackSubByteElementBitWidth(bitWidth, isI1PackedStorage))
     return elementType;
 
   // Otherwise, extend them to the next power-of-two bit width.
@@ -62,21 +59,22 @@ Value calculateStorageElementCountInBytes(Location loc,
                                           RankedTensorType shapedType,
                                           ValueRange dynamicDims,
                                           OpBuilder &builder) {
+  auto encoding = IREE::Encoding::getEncodingAttr(shapedType);
+  bool packedStorage = IREE::Encoding::hasPackedStorageAttr(shapedType);
   Type alignedElementType =
-      legalizeStorageElementType(shapedType.getElementType());
+      legalizeStorageElementType(shapedType.getElementType(), packedStorage);
   unsigned elementBits = IREE::Util::getTypeBitWidth(alignedElementType);
 
   // Calculate all static dims first, if any.
   int64_t staticCount = 1;
-  if (!needToPackSubByteElementBitWidth(elementBits)) {
+  if (!needToPackSubByteElementBitWidth(elementBits, packedStorage)) {
     staticCount *= IREE::Util::getRoundedElementByteWidth(alignedElementType);
   }
 
   // TODO: Do we use makeComposedFoldedAffineApply here, so the index
   // computation an be much simpler.
   SmallVector<int64_t> paddedShape(shapedType.getShape());
   SmallVector<Value> paddedDynamicDims(dynamicDims.begin(), dynamicDims.end());
-  auto encoding = IREE::Encoding::getEncodingAttr(shapedType);
   if (encoding && !encoding.getRoundDimsToArray().empty()) {
     auto roundDimsTo = encoding.getRoundDimsToArray();
     FailureOr<linalg::ContractionDimensions> cDims =
@@ -121,13 +119,13 @@ Value calculateStorageElementCountInBytes(Location loc,
     value = builder.createOrFold<arith::MulIOp>(loc, value, dim);
   }
   // Sub-byte packing requires putting multiple elements in the same byte.
-  if (needToPackSubByteElementBitWidth(elementBits)) {
+  if (needToPackSubByteElementBitWidth(elementBits, packedStorage)) {
     assert(8 % elementBits == 0);
     unsigned byteElements = 8 / elementBits;
     // TODO(antiagainst): We may want to emit runtime check to make sure this is
     // divisible.
     auto divisor = builder.create<arith::ConstantIndexOp>(loc, byteElements);
-    if (!clEnableI1Support && paddedDynamicDims.empty() &&
+    if (!packedStorage && paddedDynamicDims.empty() &&
         (staticCount * elementBits) % 8 != 0) {
       return nullptr;
     }
@@ -145,8 +143,11 @@ Value calculateStorageElementOffsetInBytes(Location loc,
       legalizeStorageElementType(originalType.getElementType());
   unsigned elementBits = IREE::Util::getTypeBitWidth(alignedElementType);
 
+  auto encoding = IREE::Encoding::getEncodingAttr(originalType);
+  bool isI1PackedStorage = encoding && encoding.i1PackedStorage();
+
   // Sub-byte packing requires putting multiple elements in the same byte.
-  if (needToPackSubByteElementBitWidth(elementBits)) {
+  if (needToPackSubByteElementBitWidth(elementBits, isI1PackedStorage)) {
     Value byteElements =
         builder.create<arith::ConstantIndexOp>(loc, 8 / elementBits);
     // TODO(antiagainst): We may want to emit runtime check to make sure this is

compiler/src/iree/compiler/Utils/ElementPackingUtils.h

@@ -15,7 +15,9 @@ namespace mlir::iree_compiler {
 
 /// Returns true if the given |bitWidth|, if appearing at runtime-kernel
 /// interface, is less than a byte that should be tightly packed together.
-bool needToPackSubByteElementBitWidth(unsigned bitWidth);
+bool needToPackSubByteElementBitWidth(unsigned bitWidth,
+                                      bool isI1PackedStorage = false);
+
 /// Returns true if the given |shapedType|, if appearing at runtime-kernel
 /// interface, has sub-byte element types that should be tightly packed
 /// together.
@@ -27,7 +29,8 @@ bool needToPackSubByteElements(RankedTensorType shapedType);
 /// runtime and kernel. For such cases, we perform tight packing for supported
 /// sub-byte elements, and expand to the next power-of-two bit width for other
 /// cases.
-Type legalizeStorageElementType(Type elementType);
+Type legalizeStorageElementType(Type elementType,
+                                bool isI1PackedStorage = false);
 
 /// Emits IR with the given |builder| to calculate the total number of bytes
 /// required for the given |shapedType| in storage. Returns the value for the