[XTile] Introduce to/from tensor ops.

PiperOrigin-RevId: 827514977

Author: WillFroom

xla/codegen/xtile/ir/tests/to_from_tensor.mlir

@@ -0,0 +1,17 @@
+// RUN: emitters_opt %s  -canonicalize | FileCheck %s
+
+// CHECK-LABEL: @to_scalar_roundtrip
+func.func @to_scalar_roundtrip(%arg0: tensor<i32>) -> tensor<i32> {
+  %0 = xtile.to_scalar %arg0 : tensor<i32>
+  %1 = xtile.to_tensor %0 : i32
+  // CHECK: return %arg0 : tensor<i32>
+  return %1 : tensor<i32>
+}
+
+// CHECK-LABEL: @to_tensor_roundtrip
+func.func @to_tensor_roundtrip(%arg0: i32) -> i32 {
+  %0 = xtile.to_tensor %arg0 : i32
+  %1 = xtile.to_scalar %0 : tensor<i32>
+  // CHECK: return %arg0 : i32
+  return %1 : i32
+}

xla/codegen/xtile/ir/xtile_ops.cc

@@ -196,4 +196,58 @@ mlir::TypedValue<mlir::RankedTensorType> InsertTileOp::getTile() {
 
 mlir::LogicalResult InsertTileOp::verify() { return VerifyBufferOp(*this); }
 
+mlir::LogicalResult ToScalarOp::inferReturnTypes(
+    mlir::MLIRContext* context, ::std::optional<mlir::Location> location,
+    mlir::ValueRange operands, mlir::DictionaryAttr attributes,
+    mlir::OpaqueProperties properties, mlir::RegionRange regions,
+    ::llvm::SmallVectorImpl<mlir::Type>& inferredReturnTypes) {
+  if (operands.size() != 1) {
+    return mlir::failure();
+  }
+
+  auto tensor_type =
+      mlir::dyn_cast<mlir::RankedTensorType>(operands[0].getType());
+  if (!tensor_type) {
+    return mlir::failure();
+  }
+
+  if (tensor_type.getRank() != 0) {
+    return mlir::failure();
+  }
+
+  inferredReturnTypes.push_back(tensor_type.getElementType());
+  return mlir::success();
+}
+
+mlir::OpFoldResult ToScalarOp::fold(FoldAdaptor adaptor) {
+  if (auto to_tensor = getOperand().getDefiningOp<ToTensorOp>()) {
+    // to_scalar(to_tensor(x)) -> x
+    return to_tensor.getOperand();
+  }
+
+  return {};
+}
+
+mlir::LogicalResult ToTensorOp::inferReturnTypes(
+    mlir::MLIRContext* context, ::std::optional<mlir::Location> location,
+    mlir::ValueRange operands, mlir::DictionaryAttr attributes,
+    mlir::OpaqueProperties properties, mlir::RegionRange regions,
+    ::llvm::SmallVectorImpl<mlir::Type>& inferredReturnTypes) {
+  if (operands.size() != 1) {
+    return mlir::failure();
+  }
+  inferredReturnTypes.push_back(
+      mlir::RankedTensorType::get({}, operands[0].getType()));
+  return mlir::success();
+}
+
+mlir::OpFoldResult ToTensorOp::fold(FoldAdaptor adaptor) {
+  if (auto to_scalar = getOperand().getDefiningOp<ToScalarOp>()) {
+    // to_tensor(to_scalar(x)) -> x
+    return to_scalar.getOperand();
+  }
+
+  return {};
+}
+
 }  // namespace xla::xtile

xla/codegen/xtile/ir/xtile_ops.td

@@ -202,5 +202,27 @@ def InsertTileOp : XTile_Op<"insert", [TiledBufferInterface]> {
   let hasVerifier = 1;
 }
 
+// TODO(willfroom): Revisit if/where these are needed after the migration of the
+// triton emitter to support 0D tensors is complete.
+def ToScalarOp : XTile_Op<"to_scalar", [Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  let summary = "Converts a 0D tensor to a scalar";
+  let arguments = (ins 0DTensorOf<[AnyType]>:$input);
+  let results = (outs AnyType:$output);
+
+  let assemblyFormat = "$input `:` type($input) attr-dict";
+
+  let hasFolder = 1;
+}
+
+def ToTensorOp : XTile_Op<"to_tensor", [Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  let summary = "Converts a scalar to a 0D tensor";
+  let arguments = (ins AnyType:$input);
+  let results = (outs 0DTensorOf<[AnyType]>:$output);
+
+  let assemblyFormat = "$input `:` type($input) attr-dict";
+
+  let hasFolder = 1;
+}
+
 #endif // XLA_CODEGEN_XTILE_IR_XTILE_OPS