[mlir][sparse] Add layout to the memref for the indices buffers to prepare for the AOS storage optimization for COO regions.

Fix relevant FileCheck tests.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D140742

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

@@ -1122,10 +1122,24 @@ public:
     Type resType = op.getType();
     Type indType = resType.cast<ShapedType>().getElementType();
     SmallString<15> name{"sparseIndices", overheadTypeFunctionSuffix(indType)};
-    Value dim =
-        constantIndex(rewriter, op->getLoc(), op.getDimension().getZExtValue());
-    replaceOpWithFuncCall(rewriter, op, name, resType,
-                          {adaptor.getTensor(), dim}, EmitCInterface::On);
+    Location loc = op->getLoc();
+    Value dim = constantIndex(rewriter, loc, op.getDimension().getZExtValue());
+
+    // The function returns a MemRef without a layout.
+    MemRefType callRetType = get1DMemRefType(indType, false);
+    SmallVector<Value> operands{adaptor.getTensor(), dim};
+    auto fn = getFunc(op->getParentOfType<ModuleOp>(), name, callRetType,
+                      operands, EmitCInterface::On);
+    Value callRet =
+        rewriter.create<func::CallOp>(loc, callRetType, fn, operands)
+            .getResult(0);
+
+    // Cast the MemRef type to the type expected by the users, though these
+    // two types should be compatible at runtime.
+    if (resType != callRetType)
+      callRet = rewriter.create<memref::CastOp>(loc, resType, callRet);
+    rewriter.replaceOp(op, callRet);
+
     return success();
   }
 };