Adding new optype

Signed-off-by: Alan Li <[email protected]>

compiler/src/iree/compiler/Codegen/Common/CPU/test/lower_to_ukernel_ops.mlir

@@ -511,7 +511,7 @@ func.func @unpack_f32f32_transpose_inner_and_outer(%arg0 : tensor<?x?x7x8xf32>,
 func.func @query_tile_sizes_2d() -> (index, index)  attributes {
   hal.executable.target = #hal.executable.target<"vmvx", "vmvx-bytecode-fb", {ukernels = "all"}>
 } {
-  %result:2 = iree_codegen.query_tile_sizes tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
+  %result:2 = iree_codegen.query_tile_sizes tensor<?x?xf32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [f32, f32, f32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
   return %result#0, %result#1 : index, index
 }
 

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

@@ -83,8 +83,10 @@ static RankedTensorType transposeIfNarrowNResult(RankedTensorType tensorType) {
   auto elemType = tensorType.getElementType();
   OpBuilder builder(context);
 
+  auto opTypeAttr = IREE::Encoding::EncodingOpTypeAttr::get(
+      context, IREE::Encoding::EncodingOpType::matmul);
   auto newEncoding = IREE::Encoding::EncodingAttr::get(
-      context, newIndex, encoding.getElementTypes(),
+      context, newIndex, opTypeAttr, encoding.getElementTypes(),
       TypeAttr::get(RankedTensorType::get(newOriginalShape, elemType)),
       encoding.getMatmulNarrow_N(), encoding.getMatmulNarrow_M(),
       newIndexingMaps, DenseI64ArrayAttr::get(context, newRoundDimsTo));

compiler/src/iree/compiler/Codegen/Common/test/convert_to_destination_passing_style.mlir

@@ -754,7 +754,7 @@ func.func @non_perfect_tiling_unpack() {
   %c512 = arith.constant 512 : index
   %c0 = arith.constant 0 : index
   %c16 = arith.constant 16 : index
-  %0:2 = iree_codegen.query_tile_sizes tensor<16x16xi32, #iree_encoding.encoding<operand_index = 2, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
+  %0:2 = iree_codegen.query_tile_sizes tensor<16x16xi32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
   %1 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%0#0]
   %2 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%0#1]
   %3 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c512) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi32>>{%1, %2, %0#0, %0#1}

compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_to_workgroups.mlir

@@ -2232,7 +2232,7 @@ hal.executable private @dynamic_unpack_fusion {
         %cst = arith.constant dense<[-918, -4433, 87, -234, -21393, 7738, 529, -8835, -16817, -375, -199, 572, 5082, 15569, -186, 4955]> : tensor<16xi32>
         %c12544 = arith.constant 12544 : index
         %c16 = arith.constant 16 : index
-        %0:2 = iree_codegen.query_tile_sizes tensor<12544x16xi32, #iree_encoding.encoding<operand_index = 2, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
+        %0:2 = iree_codegen.query_tile_sizes tensor<12544x16xi32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
         %1 = affine.apply affine_map<()[s0] -> (12544 ceildiv s0)>()[%0#0]
         %2 = affine.apply affine_map<()[s0] -> (16 ceildiv s0)>()[%0#1]
         %3 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c200960) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi32>>{%1, %2, %0#0, %0#1}

compiler/src/iree/compiler/Codegen/VMVX/test/pipeline.mlir

@@ -10,15 +10,15 @@ func.func @mmt4d_i8() attributes {hal.executable.target = #executable_target_vmv
   %c256 = arith.constant 256 : index
   %c512 = arith.constant 512 : index
   %c16 = arith.constant 16 : index
-  %0:2 = iree_codegen.query_tile_sizes tensor<16x16xi8, #iree_encoding.encoding<operand_index = 0, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
+  %0:2 = iree_codegen.query_tile_sizes tensor<16x16xi8, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
   %1 = affine.apply #map3()[%0#0]
   %2 = affine.apply #map3()[%0#1]
   %3 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%1, %2, %0#0, %0#1}
-  %4:2 = iree_codegen.query_tile_sizes tensor<16x16xi8, #iree_encoding.encoding<operand_index = 1, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
+  %4:2 = iree_codegen.query_tile_sizes tensor<16x16xi8, #iree_encoding.encoding<operand_index = 1, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
   %5 = affine.apply #map3()[%4#0]
   %6 = affine.apply #map3()[%4#1]
   %7 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c256) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?x?x?x?xi8>>{%5, %6, %4#0, %4#1}
-  %8:2 = iree_codegen.query_tile_sizes tensor<16x16xi32, #iree_encoding.encoding<operand_index = 2, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
+  %8:2 = iree_codegen.query_tile_sizes tensor<16x16xi32, #iree_encoding.encoding<operand_index = 2, op_type = matmul, element_types = [i8, i8, i32], user_indexing_maps = [#map, #map1, #map2]>> -> index, index
   %9 = affine.apply #map3()[%8#0]
   %10 = affine.apply #map3()[%8#1]
   %11 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c512) : !flow.dispatch.tensor<readwrite:tensor<?x?x?x?xi32>>{%9, %10, %8#0, %8#1}

compiler/src/iree/compiler/Codegen/VMVX/test/select_lowering_strategy.mlir

@@ -168,7 +168,7 @@ func.func @elem_pack_ukernels() attributes {hal.executable.target = #executable_
   %cst = arith.constant 0.000000e+00 : f32
   %c0 = arith.constant 0 : index
   %0 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<1024x2048xf32>>
-  %1:2 = iree_codegen.query_tile_sizes tensor<1024x2048xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32], original_type = tensor<1024x2048xf32>>> -> index, index
+  %1:2 = iree_codegen.query_tile_sizes tensor<1024x2048xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], original_type = tensor<1024x2048xf32>>> -> index, index
   %2 = affine.apply #map()[%1#0]
   %3 = affine.apply #map1()[%1#1]
   %4 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<?x?x?x?xf32>>{%2, %3, %1#0, %1#1}
@@ -179,12 +179,12 @@ func.func @elem_pack_ukernels() attributes {hal.executable.target = #executable_
     %15 = arith.addf %in, %in : f32
     linalg.yield %15 : f32
   } -> tensor<1024x2048xf32>
-  %8:2 = iree_codegen.query_tile_sizes tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32], original_type = tensor<1024x2048xf32>>> -> index, index
+  %8:2 = iree_codegen.query_tile_sizes tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], original_type = tensor<1024x2048xf32>>> -> index, index
   %9 = affine.apply #map()[%8#0]
   %10 = affine.apply #map1()[%8#1]
   %11 = tensor.empty(%9, %10, %8#0, %8#1) : tensor<?x?x?x?xf32>
   %pack = tensor.pack %7 padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [%8#0, %8#1] into %11 : tensor<1024x2048xf32> -> tensor<?x?x?x?xf32>
-  %12:2 = iree_codegen.query_tile_sizes tensor<1024x2048xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32], original_type = tensor<1024x2048xf32>>> -> index, index
+  %12:2 = iree_codegen.query_tile_sizes tensor<1024x2048xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], original_type = tensor<1024x2048xf32>>> -> index, index
   %13 = affine.apply #map()[%12#0]
   %14 = affine.apply #map1()[%12#1]
   flow.dispatch.tensor.store %pack, %4, offsets = [0, 0, 0, 0], sizes = [%13, %14, %12#0, %12#1], strides = [1, 1, 1, 1] : tensor<?x?x?x?xf32> -> !flow.dispatch.tensor<writeonly:tensor<?x?x?x?xf32>>{%13, %14, %12#0, %12#1}

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

@@ -44,6 +44,19 @@ class IREEEncoding_I32EnumAttr<string name, string summary, list<I32EnumAttrCase
 class IREEEncoding_EnumAttr<EnumAttrInfo enumInfo, string name = "">
   : EnumAttr<IREEEncoding_Dialect, enumInfo, name>;
 
+// Enums for tagging operand operation in an EncodingAttr
+def MATMUL : I32EnumAttrCase<"matmul", 0>;
+def CONV   : I32EnumAttrCase<"conv", 1>;
+
+def EncodingOpType : IREEEncoding_I32EnumAttr<"EncodingOpType",
+    "Tracks the type of operation of the operand.", [
+      MATMUL,
+      CONV,
+    ]>;
+
+def EncodingOpTypeAttr:
+  IREEEncoding_EnumAttr<EncodingOpType, "optype">;
+
 def EncodingAttr :
     IREEEncoding_Attr<"Encoding"> {
   let mnemonic = "encoding";
@@ -63,6 +76,7 @@ def EncodingAttr :
 
   let parameters = (ins
     AttrParameter<"IntegerAttr", "this tensor operand's index in the parameter list">:$operand_index,
+    AttrParameter<"EncodingOpTypeAttr", "operand type">:$op_type,
     AttrParameter<"ArrayAttr", "element types of the user's operands">:$element_types,
     OptionalParameter<"TypeAttr", "type of the original tensor type before padding">:$original_type,
     // TODO(#15466): generalize matmul_narrow_{M,N} into a list?
@@ -75,6 +89,7 @@ def EncodingAttr :
 
   let builders = [
     AttrBuilder<(ins "int64_t":$operandIndex,
+        "EncodingOpType":$opType,
         "ArrayRef<Type>":$elemTypes, "Type":$origType,
         CArg<"std::optional<int64_t>", "{}">:$matmulNarrowM,
         CArg<"std::optional<int64_t>", "{}">:$matmulNarrowN,

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

@@ -97,7 +97,8 @@ LogicalResult UnsetEncodingOp::reifyResultShapes(
 //===----------------------------------------------------------------------===//
 
 EncodingAttr EncodingAttr::get(MLIRContext *ctx, int64_t operandIndex,
-                               ArrayRef<Type> elemTypes, Type origType,
+                               EncodingOpType opType, ArrayRef<Type> elemTypes,
+                               Type origType,
                                std::optional<int64_t> matmulNarrowM,
                                std::optional<int64_t> matmulNarrowN,
                                ArrayRef<AffineMap> maps,
@@ -106,14 +107,15 @@ EncodingAttr EncodingAttr::get(MLIRContext *ctx, int64_t operandIndex,
   auto optionalToAttr = [&](std::optional<int64_t> x) {
     return x ? b.getIndexAttr(*x) : IntegerAttr();
   };
+  auto opTypeAttr = EncodingOpTypeAttr::get(ctx, opType);
   auto origTypeAttr = origType ? TypeAttr::get(origType) : TypeAttr();
   auto roundDimsToAttr = roundDimsTo.empty()
                              ? DenseI64ArrayAttr()
                              : b.getDenseI64ArrayAttr(roundDimsTo);
-  return get(ctx, b.getIndexAttr(operandIndex), b.getTypeArrayAttr(elemTypes),
-             origTypeAttr, optionalToAttr(matmulNarrowM),
-             optionalToAttr(matmulNarrowN), b.getAffineMapArrayAttr(maps),
-             roundDimsToAttr);
+  return get(ctx, b.getIndexAttr(operandIndex), opTypeAttr,
+             b.getTypeArrayAttr(elemTypes), origTypeAttr,
+             optionalToAttr(matmulNarrowM), optionalToAttr(matmulNarrowN),
+             b.getAffineMapArrayAttr(maps), roundDimsToAttr);
 }
 
 AffineMap EncodingAttr::getMapForOperandIndex() {

compiler/src/iree/compiler/Dialect/Encoding/IR/test/invalid.mlir

@@ -8,9 +8,9 @@ func.func @illegal_set_encoding_op_with_no_result_encoding(%arg0 : tensor<?x?xf3
 
 // -----
 
-func.func @illegal_set_encoding_op_with_source_encoding(%arg0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>>) -> tensor<?x?xf32> {
+func.func @illegal_set_encoding_op_with_source_encoding(%arg0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>) -> tensor<?x?xf32> {
   // expected-error @+1 {{source of set_encoding op cannot have a tensor encoding}}
-  %0 = iree_encoding.set_encoding %arg0: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>> -> tensor<?x?xf32>
+  %0 = iree_encoding.set_encoding %arg0: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }
 
@@ -24,18 +24,18 @@ func.func @illegal_set_encoding_op_with_unknown_encoding(%arg0 : tensor<?x?xf32>
 
 // -----
 
-func.func @illegal_set_encoding_op_with_rank_change(%arg0 : tensor<?x?xf32>) -> tensor<?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>> {
+func.func @illegal_set_encoding_op_with_rank_change(%arg0 : tensor<?x?xf32>) -> tensor<?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>> {
   // expected-error @+1 {{cannot change the rank of the tensor}}
-  %0 = iree_encoding.set_encoding %arg0: tensor<?x?xf32> -> tensor<?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>>
-  return %0 : tensor<?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>>
+  %0 = iree_encoding.set_encoding %arg0: tensor<?x?xf32> -> tensor<?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
+  return %0 : tensor<?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
 }
 
 // -----
 
-func.func @illegal_set_encoding_op_with_shape_change(%arg0 : tensor<10x20xf32>) -> tensor<20x30xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>> {
+func.func @illegal_set_encoding_op_with_shape_change(%arg0 : tensor<10x20xf32>) -> tensor<20x30xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>> {
   // expected-error @+1 {{expected to preserve the logical shape of the tensor}}
-  %0 = iree_encoding.set_encoding %arg0: tensor<10x20xf32> -> tensor<20x30xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>>
-  return %0 : tensor<20x30xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>>
+  %0 = iree_encoding.set_encoding %arg0: tensor<10x20xf32> -> tensor<20x30xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
+  return %0 : tensor<20x30xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
 }
 
 // -----
@@ -48,10 +48,10 @@ func.func @illegal_unset_encoding_op_with_no_source_encoding(%arg0 : tensor<?x?x
 
 // -----
 
-func.func @illegal_unset_encoding_op_with_result_encoding(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>> {
+func.func @illegal_unset_encoding_op_with_result_encoding(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>> {
   // expected-error @+1 {{result of unset_encoding op cannot have a tensor encoding}}
-  %0 = iree_encoding.unset_encoding %arg0: tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>>
-  return %0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>>
+  %0 = iree_encoding.unset_encoding %arg0: tensor<?x?xf32> -> tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
+  return %0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>
 }
 
 // -----
@@ -64,16 +64,16 @@ func.func @illegal_unset_encoding_op_with_unknown_encoding(%arg0 : tensor<?x?xf3
 
 // -----
 
-func.func @illegal_unset_encoding_op_with_rank_change(%arg0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>>) -> tensor<?xf32> {
+func.func @illegal_unset_encoding_op_with_rank_change(%arg0 : tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>) -> tensor<?xf32> {
   // expected-error @+1 {{cannot change the rank of the tensor}}
-  %0 = iree_encoding.unset_encoding %arg0: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>> -> tensor<?xf32>
+  %0 = iree_encoding.unset_encoding %arg0: tensor<?x?xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<?xf32>
   return %0 : tensor<?xf32>
 }
 
 // -----
 
-func.func @illegal_unset_encoding_op_with_shape_change(%arg0 : tensor<20x30xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>>) -> tensor<10x20xf32> {
+func.func @illegal_unset_encoding_op_with_shape_change(%arg0 : tensor<20x30xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>>) -> tensor<10x20xf32> {
   // expected-error @+1 {{expected to preserve the logical shape of the tensor}}
-  %0 = iree_encoding.unset_encoding %arg0: tensor<20x30xf32, #iree_encoding.encoding<operand_index = 0, element_types = [f32, f32, f32]>> -> tensor<10x20xf32>
+  %0 = iree_encoding.unset_encoding %arg0: tensor<20x30xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32]>> -> tensor<10x20xf32>
   return %0 : tensor<10x20xf32>
 }