[ConvertToLLVM] Don't choke on alloc of memref of index (#14002)

Prior to this patch the compiler would crash on alloc of memrefs of
index on shared memory, because we would try to get the bitwidth of the
index type as if it was an integer or a floating point type.

The bitwidth of the index type is actually carried by the datalayout of
the module.

Update the code to query the datalayout to get the bitwidth.

Note: This particular pattern is called as part of convert-to-nvvm. In
that pass we actually lower the index type to an integer type, but we do
it after having lowered the `alloc`s from shared memory, i.e., where the
code is choking.

compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp
index 13bcd11..2ad79f4 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp
@@ -150,7 +150,11 @@
       if (auto shapeType = llvm::dyn_cast<ShapedType>(elType))
         alignement =
             shapeType.getNumElements() * shapeType.getElementTypeBitWidth() / 8;
-      else
+      else if (elType.isIndex()) {
+        auto mod = allocOp->getParentOfType<ModuleOp>();
+        LowerToLLVMOptions options(mod.getContext(), DataLayout(mod));
+        alignement = options.getIndexBitwidth() / 8;
+      } else
         alignement = elType.getIntOrFloatBitWidth() / 8;
     }
     // In CUDA workgroup memory is represented by a global variable.

compiler/src/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir
index 07c9073..4737034 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir
@@ -383,3 +383,31 @@
 //   CHECK-NOT: unrealized
 //       CHECK: llvm.return
 
+// -----
+
+// Check that we don't choke on memref of index.
+#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
+  #hal.descriptor_set.layout<0, bindings = [
+    #hal.descriptor_set.binding<0, storage_buffer>
+  ]>
+]>
+hal.executable @shared_memory_lowering_index {
+  hal.executable.variant @cuda, target = <"cuda", "cuda-nvptx-fb"> {
+    hal.executable.export @shared_memory_lowering_index layout(#pipeline_layout)
+    builtin.module {
+      func.func @shared_memory_lowering_index() {
+        %c0 = arith.constant 0 : index
+        %cst = arith.constant dense<0> : vector<4xindex>
+        %0 = memref.alloc() : memref<1x16x32xindex, #gpu.address_space<workgroup>>
+        vector.store %cst, %0[%c0, %c0, %c0] : memref<1x16x32xindex, #gpu.address_space<workgroup>>, vector<4xindex>
+        return
+      }
+    }
+  }
+}
+//       CHECK: llvm.mlir.global external @__dynamic_shared_memory__() {addr_space = 3 : i32, alignment = 16 : i64} : !llvm.array<0 x i8>
+// CHECK-LABEL: llvm.func @shared_memory_lowering_index() {
+//       CHECK: %{{.*}} = llvm.mlir.addressof @__dynamic_shared_memory__ : !llvm.ptr<array<0 x i8>, 3>
+//  CHECK-NEXT: %{{.*}} = llvm.mlir.constant(0 : i64) : i64
+//  CHECK-NEXT: %{{.*}} = llvm.mlir.constant(0 : i64) : i64
+//  CHECK-NEXT: %{{.*}} = llvm.getelementptr %{{.*}} : (!llvm.ptr<array<0 x i8>, 3>, i64, i64) -> !llvm.ptr<array<0 x i8>, 3>