[mlir] Allow index as element type of memref

mlir/docs/Rationale/Rationale.md

 
 For more context, block arguments were previously used in the Swift
 [SIL Intermediate Representation](https://github.com/apple/swift/blob/master/docs/SIL.rst),
 and described in
 [a talk on YouTube](https://www.youtube.com/watch?v=Ntj8ab-5cvE). The section of
 interest
 [starts here](https://www.google.com/url?q=https://youtu.be/Ntj8ab-5cvE?t%3D596&sa=D&ust=1529450150971000&usg=AFQjCNFQHEWL7m8q3eO-1DiKw9zqC2v24Q).
 
-### Index type disallowed in vector/memref types
+### Index type disallowed in vector types
 
-Index types are not allowed as elements of `vector` and `memref` types. Index
+Index types are not allowed as elements of `vector` types. Index
 types are intended to be used for platform-specific "size" values and may appear
 in subscripts, sizes of aggregate types and affine expressions. They are also
 tightly coupled with `affine.apply` and affine.load/store operations; having
 `index` type is a necessary precondition of a value to be acceptable by these
-operations. While it may be useful to have `memref<?xindex>` to express indirect
-accesses, e.g. sparse matrix manipulations or lookup tables, it creates problems
-MLIR is not ready to address yet. MLIR needs to internally store constants of
-aggregate types and emit code operating on values of those types, which are
-subject to target-specific size and alignment constraints. Since MLIR does not
-have a target description mechanism at the moment, it cannot reliably emit such
-code. Moreover, some platforms may not support vectors of type equivalent to
-`index`.
-
-Indirect access use cases can be alternatively supported by providing and
-`index_cast` instruction that allows for conversion between `index` and
-fixed-width integer types, at the SSA value level. It has an additional benefit
-of supporting smaller integer types, e.g. `i8` or `i16`, for small indices
-instead of (presumably larger) `index` type.
-
-Index types are allowed as element types of `tensor` types. The `tensor` type
-specifically abstracts the target-specific aspects that intersect with the
-code-generation-related/lowering-related concerns explained above. In fact, the
-`tensor` type even allows dialect-specific types as element types.
+operations.
+
+We allow `index` types in tensors and memrefs as a code generation strategy has
+to map `index` to an implementation type and hence needs to be able to
+materialize corresponding values. However, the target might lack support for
+`vector` values with the target specfic equivalent of the `index` type.
 
 ### Bit width of a non-primitive type and `index` is undefined
 
 The bit width of a compound type is not defined by MLIR, it may be defined by a
 specific lowering pass. In MLIR, bit width is a property of certain primitive
 _type_, in particular integers and floats. It is equal to the number that
 appears in the type definition, e.g. the bit width of `i32` is `32`, so is the
 bit width of `f32`. The bit width is not _necessarily_ related to the amount of

mlir/lib/IR/StandardTypes.cpp

 /// pass in an instance of UnknownLoc.
 MemRefType MemRefType::getImpl(ArrayRef<int64_t> shape, Type elementType,
                                ArrayRef<AffineMap> affineMapComposition,
                                unsigned memorySpace,
                                Optional<Location> location) {
   auto *context = elementType.getContext();
 
   // Check that memref is formed from allowed types.
-  if (!elementType.isIntOrFloat() &&
+  if (!elementType.isIntOrIndexOrFloat() &&
       !elementType.isa<VectorType, ComplexType>())
     return emitOptionalError(location, "invalid memref element type"),
            MemRefType();
 
   for (int64_t s : shape) {
     // Negative sizes are not allowed except for `-1` that means dynamic size.
     if (s < -1)
       return emitOptionalError(location, "invalid memref size"), MemRefType();

mlir/lib/Parser/TypeParser.cpp

Type Parser::parseMemRefType() {
 
   // Parse the element type.
   auto typeLoc = getToken().getLoc();
   auto elementType = parseType();
   if (!elementType)
     return nullptr;
 
   // Check that memref is formed from allowed types.
-  if (!elementType.isIntOrFloat() &&
+  if (!elementType.isIntOrIndexOrFloat() &&
       !elementType.isa<VectorType, ComplexType>())
     return emitError(typeLoc, "invalid memref element type"), nullptr;
 
   // Parse semi-affine-map-composition.
   SmallVector<AffineMap, 2> affineMapComposition;
   Optional<unsigned> memorySpace;
   unsigned numDims = dimensions.size();
 

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

 // -----
 
 // CHECK-LABEL: func @bfloat
 // CHECK-SAME: !llvm.bfloat) -> !llvm.bfloat
 func @bfloat(%arg0: bf16) -> bf16 {
   return %arg0 : bf16
 }
 // CHECK-NEXT: return %{{.*}} : !llvm.bfloat
+
+// -----
+
+// CHECK-LABEL: func @memref_index
+// CHECK-SAME: %arg0: !llvm<"i64*">, %arg1: !llvm<"i64*">,
+// CHECK-SAME: %arg2: !llvm.i64, %arg3: !llvm.i64, %arg4: !llvm.i64)
+// CHECK-SAME: -> !llvm<"{ i64*, i64*, i64, [1 x i64], [1 x i64] }">
+// CHECK32-LABEL: func @memref_index
+// CHECK32-SAME: %arg0: !llvm<"i32*">, %arg1: !llvm<"i32*">,
+// CHECK32-SAME: %arg2: !llvm.i32, %arg3: !llvm.i32, %arg4: !llvm.i32)
+// CHECK32-SAME: -> !llvm<"{ i32*, i32*, i32, [1 x i32], [1 x i32] }">
+func @memref_index(%arg0: memref<32xindex>) -> memref<32xindex> {
+  return %arg0 : memref<32xindex>
+}

mlir/test/IR/invalid.mlir

 
 func @nestedtensor(tensor<tensor<i8>>) -> () // expected-error {{invalid tensor element type}}
 
 // -----
 
 func @indexvector(vector<4 x index>) -> () // expected-error {{vector elements must be int or float type}}
 
 // -----
-
-func @indexmemref(memref<? x index>) -> () // expected-error {{invalid memref element type}}
-
-// -----
 // Test no map in memref type.
 func @memrefs(memref<2x4xi8, >) // expected-error {{expected list element}}
 
 // -----
 // Test non-existent map in memref type.
 func @memrefs(memref<2x4xi8, #map7>) // expected-error {{undefined symbol alias id 'map7'}}
 
 // -----

mlir/test/IR/parser.mlir

 func @memrefs_compose_with_id(memref<2x2xi8, affine_map<(d0, d1) -> (d0, d1)>,
                                              affine_map<(d0, d1) -> (d1, d0)>>)
 
 
 // CHECK: func @complex_types(complex<i1>) -> complex<f32>
 func @complex_types(complex<i1>) -> complex<f32>
 
 
+// CHECK: func @memref_with_index_elems(memref<1x?xindex>)
+func @memref_with_index_elems(memref<1x?xindex>)
+
 // CHECK: func @memref_with_complex_elems(memref<1x?xcomplex<f32>>)
 func @memref_with_complex_elems(memref<1x?xcomplex<f32>>)
 
 // CHECK: func @memref_with_vector_elems(memref<1x?xvector<10xf32>>)
 func @memref_with_vector_elems(memref<1x?xvector<10xf32>>)
 
 // CHECK: func @unranked_memref_with_complex_elems(memref<*xcomplex<f32>>)
 func @unranked_memref_with_complex_elems(memref<*xcomplex<f32>>)